<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Numpy arrays. Data manipulation
In this notebook we are going to work with some numerical data that we need to re-format.
Profasi is a Monte Carlo code for protein simulation. It can run Parallel Tempering simulations where each replica runs in a processor and exchanges temperatures with the other replicas. It is done this way because it is more efficient to exchange temperatures than to exchange molecular coordiantes between processors. The problem is that in the output files the temperatures are mixed. Usually one needs the data for all the sampled molecules at a the same temperature.
The above explanation is just to say that we need to reorder the data from different files and generated re-ordered files. All the our input files are called rt and each resides in a directory called ni where i corresponds to the processor. If the simulation was run with 16 processors i will go from 0 to 15.
The rt file is a text file that contains just columns of numbers, separated by spaces. Here is how it looks like
Step1:  The first column indicate the iteration or cycle, the second the temperature. The remaining are energy components and other observables. The temperature is an integer. In a different file one an find the conversion to kelvin.
We start with some setup...
Step2:  Use glob.glob to get a list of all the input rt files.
Step3:  We now read all the rt files into memory. If they were too big, we should think of a more efficient way to process them (maybe using memmap or pytables). In our case they are small enough.
As these are numeric files, we use loadtxt to automatically generte an array. As different files will generate different arrays, we collect them in a list that we finally transform into an array.
Step4:  This is the shape of the resulting array
Step5:  Now we need to reshape it so that it contains 2 dimensions with all the raws and the 14 columns
Step6:  Alternatively, we could have concatenated all the rows into the first loaded array. Here is a way to do that
Step7:  Now we need to get the temperatures. We know that there are as many temperatures as nodes, so this would work
Step8:  We can also use np.unique to the the unique values of an array or part of it. A little bit more efficient... (check it with %timeit).
Step9:  Now we eneed to extract the energies from the array based on the temperature value, and create separate sub-arrays. Array elements can be selected with Boolean arrays. This is called fancy indexing.
We start by difining an empty array and then fill it in
Step10:  The last step is to keep only those energy values that are beyond the equilibration point. So we want only to keep data from, a certain point. Let's plot the energy vs. iteration to see how it looks line. We'll plot temperature 5 as this is the lowest temperature (Profasi order from high to low).
Step11:  Ugly, isn't it? The array is not ordered by iteration. The order can be seen here
Step12:  We have the structures at the correct temperature but still ordered from the 6 replicas that were running. Let's order them with respect to the first column. We cannot use sort here, because we want to use the order of the first column to order all the raw elements. We can get that order with argsort and then apply it to the array.
The problem is that the sizes of the sorting order do no agree with the sizes of ene_temp. To solve that we need to do a trick which I pesonally find very cumbersome. The reason is we need to broadcas correctly the dimensions of order into ene_temp. It's simplers to understand if you see that for the first temperatures we want
Step13:  Finally we can select and save the submatrix from iteration 2000 onwards
Step14:  Advanced Topic
Step15:  Now we loop though all the rows and put each row to its corrrect first axis dimension. We need to keep an array of the filled positions.
Step16:  We check we are still getting the same result, and we time it
Step17:  Good (ironically)! Two orders of magnitude slower than the first approach... The reason is that in the previous approach we were using numpy fast looping abilities, whereas now the loops are implemented in pure python and therefore are much slower.
This is the typical case where numba can increase the performance of such loops.
Step18:  We first write a function of our lines. We avoid creating arrays into that function as those cannot be optimized with numba. We test our approach and check that the timings are the same.
Step19:  Now we can pass this function to numba. The nopython option tell numba not to create object code which is as slow as python code. That is why we created the arrays outside the function. We also check the timings.
Step20:  Wow! Three orders of magnitude faster than the python and one order faster than our original numpy code (with only 6 temperatures!).
But having to declare all arrays outside is ugly. Is there a workaroud? Yes! Numba is clever enough to separate the loops from the array creation, and optimize the loops. This called loop-lifting or loop-jitting. We need to remove the nopython option as part of the code will be object like, but we see that it is as efficient as before.
Here we also use a decoratior instead of a function call. The results are exactly the same, it just gives a shorter syntax.
<ASSISTANT_TASK:>
Python Code:
!head ../../data/profasi/n0/rt
Explanation: Numpy arrays. Data manipulation
In this notebook we are going to work with some numerical data that we need to re-format.
Profasi is a Monte Carlo code for protein simulation. It can run Parallel Tempering simulations where each replica runs in a processor and exchanges temperatures with the other replicas. It is done this way because it is more efficient to exchange temperatures than to exchange molecular coordiantes between processors. The problem is that in the output files the temperatures are mixed. Usually one needs the data for all the sampled molecules at a the same temperature.
The above explanation is just to say that we need to reorder the data from different files and generated re-ordered files. All the our input files are called rt and each resides in a directory called ni where i corresponds to the processor. If the simulation was run with 16 processors i will go from 0 to 15.
The rt file is a text file that contains just columns of numbers, separated by spaces. Here is how it looks like:
End of explanation
import numpy as np
import glob
Explanation: The first column indicate the iteration or cycle, the second the temperature. The remaining are energy components and other observables. The temperature is an integer. In a different file one an find the conversion to kelvin.
We start with some setup...
End of explanation
files = glob.glob('../../data/profasi/n*/rt')
Explanation: Use glob.glob to get a list of all the input rt files.
End of explanation
all_enes = []
for filein in files:
    print("Reading..... ", filein)
    all_enes.append(np.loadtxt(filein))
all_enes = np.asarray(all_enes)
Explanation: We now read all the rt files into memory. If they were too big, we should think of a more efficient way to process them (maybe using memmap or pytables). In our case they are small enough.
As these are numeric files, we use loadtxt to automatically generte an array. As different files will generate different arrays, we collect them in a list that we finally transform into an array.
End of explanation
all_enes.shape
Explanation: This is the shape of the resulting array:
End of explanation
all_enes=all_enes.reshape((-1,all_enes.shape[2]))
Explanation: Now we need to reshape it so that it contains 2 dimensions with all the raws and the 14 columns:
End of explanation
all_enes = None
for filein in files:
    print("Reading..... ", filein)
    if all_enes is not None:
        all_enes= np.r_[all_enes, np.loadtxt(filein)]
    else:
        all_enes= np.loadtxt(filein)
Explanation: Alternatively, we could have concatenated all the rows into the first loaded array. Here is a way to do that:
End of explanation
temperatures = set(all_enes[:,1])
Explanation: Now we need to get the temperatures. We know that there are as many temperatures as nodes, so this would work:
temperatures = np.arange(len(files))
However, image that for some purpose we did not process all the ni directories, but only a fraction of them. We can still get all the temperatures from the rt files. It corresponds to the 2nd column. A simple way to get the temperatures is with a set.
End of explanation
temperatures = np.unique(all_enes[:,1])
%timeit set(all_enes[:,1])
%timeit np.unique(all_enes[:,1])
Explanation: We can also use np.unique to the the unique values of an array or part of it. A little bit more efficient... (check it with %timeit).
End of explanation
ene_temp = np.zeros_like(all_enes)
ene_temp = ene_temp.reshape([len(temperatures), -1, all_enes.shape[1]])
for ti in temperatures.astype(int):
    ene_temp[ti] = all_enes[all_enes[:,1]==ti, :]
Explanation: Now we eneed to extract the energies from the array based on the temperature value, and create separate sub-arrays. Array elements can be selected with Boolean arrays. This is called fancy indexing.
We start by difining an empty array and then fill it in:
End of explanation
import matplotlib.pyplot as plt
%matplotlib inline
plt.rcParams['figure.figsize'] = (10.0, 8.0)
plt.plot(ene_temp[5,:,0], ene_temp[5,:,2])
Explanation: The last step is to keep only those energy values that are beyond the equilibration point. So we want only to keep data from, a certain point. Let's plot the energy vs. iteration to see how it looks line. We'll plot temperature 5 as this is the lowest temperature (Profasi order from high to low).
End of explanation
plt.plot(ene_temp[5,:,0],'x')
Explanation: Ugly, isn't it? The array is not ordered by iteration. The order can be seen here:
End of explanation
order = ene_temp[:, :, 0].argsort()
ene_temp = ene_temp[np.arange(ene_temp.shape[0])[:, np.newaxis], order]
plt.subplot(2,1,1)
plt.plot(ene_temp[5,:,0],'x')
plt.subplot(2,1,2)
plt.plot(ene_temp[5,:,0], ene_temp[5,:,2])
Explanation: We have the structures at the correct temperature but still ordered from the 6 replicas that were running. Let's order them with respect to the first column. We cannot use sort here, because we want to use the order of the first column to order all the raw elements. We can get that order with argsort and then apply it to the array.
The problem is that the sizes of the sorting order do no agree with the sizes of ene_temp. To solve that we need to do a trick which I pesonally find very cumbersome. The reason is we need to broadcas correctly the dimensions of order into ene_temp. It's simplers to understand if you see that for the first temperatures we want:
ene_temp[0, order[0]]
ene_temp[1, order[1]]
and so on. It would seem that ene_temp[:, order[:]] but this performs the broadcasting in the wrong axis. Instead, we need to transpose the first axis, because what we are actually doing is:
ene_temp[[0], order[0]]
ene_temp[[1], order[1]]
And this can be done creating the vector [[0], [1], [2]... which is done with:  np.arange(ene_temp.shape[0])[:, np.newaxis].
End of explanation
np.save('energies_temperatures', ene_temp[ene_temp[:,:,0]>2000, :], )
Explanation: Finally we can select and save the submatrix from iteration 2000 onwards:
End of explanation
%%timeit
for ti in temperatures.astype(int):
    ene_temp[ti] = all_enes[all_enes[:,1]==ti, :]
Explanation: Advanced Topic: optimizing with numba
In the previous section, if there were $N$ temperatures we cycles the all_enes array $N$ times, which is not very efficient. We could potentially make if faster by running this in a single step. We create an empty arrray and fill it in with the correct values.
We first time our initial approach:
End of explanation
ene_temp2 = np.zeros_like(ene_temp)
filled = np.zeros(len(temperatures), np.int)
for row in all_enes:
    ti = int(row[1])
    ene_temp2[ti, filled[ti]] = row
    filled[ti] +=1
Explanation: Now we loop though all the rows and put each row to its corrrect first axis dimension. We need to keep an array of the filled positions.
End of explanation
np.all(ene_temp2==ene_temp)
%%timeit
filled = np.zeros(len(temperatures), np.int)
for row in all_enes:
    ti = int(row[1])
    ene_temp2[ti, filled[ti]] = row
    filled[ti] +=1
Explanation: We check we are still getting the same result, and we time it:
End of explanation
import numba
Explanation: Good (ironically)! Two orders of magnitude slower than the first approach... The reason is that in the previous approach we were using numpy fast looping abilities, whereas now the loops are implemented in pure python and therefore are much slower.
This is the typical case where numba can increase the performance of such loops.
End of explanation
def get_temperatures(array_in, array_out, filled):
    for r in range(array_in.shape[0]):
        ti = int(array_in[r,1])
        for j in range(array_in.shape[1]):
            array_out[ti, filled[ti], j] = array_in[r,j]
        filled[ti] +=1
    return array_out
%%timeit 
num_temp = len(temperatures)
m = all_enes.shape[0]
n = all_enes.shape[1]
m = m // num_temp
ene_temp = np.zeros((num_temp, m,n ))
filled = np.zeros(num_temp, np.int)
get_temperatures(all_enes, ene_temp, filled)
Explanation: We first write a function of our lines. We avoid creating arrays into that function as those cannot be optimized with numba. We test our approach and check that the timings are the same.
End of explanation
numba_get_temperatures = numba.jit(get_temperatures,nopython=True)
%%timeit 
num_temp = len(temperatures)
m = all_enes.shape[0]
n = all_enes.shape[1]
m = m // num_temp
ene_temp3 = np.zeros((num_temp, m,n ))
filled = np.zeros(num_temp, np.int)
numba_get_temperatures(all_enes, ene_temp3, filled)
Explanation: Now we can pass this function to numba. The nopython option tell numba not to create object code which is as slow as python code. That is why we created the arrays outside the function. We also check the timings.
End of explanation
@numba.jit 
def numba2_get_temperatures(array_in, num_temp):
    m = all_enes.shape[0]
    n = all_enes.shape[1]
    m = m // num_temp
    array_out = np.zeros((num_temp, m,n ))
    filled = np.zeros(num_temp, np.int)
    for r in range(array_in.shape[0]):
        ti = int(array_in[r,1])
        for j in range(array_in.shape[1]):
            array_out[ti, filled[ti], j] = array_in[r,j]
        filled[ti] +=1
    return array_out
%%timeit 
num_temp = len(temperatures)
ene_temp4 = numba2_get_temperatures(all_enes, num_temp)
np.all(numba2_get_temperatures(all_enes, num_temp)==ene_temp)
Explanation: Wow! Three orders of magnitude faster than the python and one order faster than our original numpy code (with only 6 temperatures!).
But having to declare all arrays outside is ugly. Is there a workaroud? Yes! Numba is clever enough to separate the loops from the array creation, and optimize the loops. This called loop-lifting or loop-jitting. We need to remove the nopython option as part of the code will be object like, but we see that it is as efficient as before.
Here we also use a decoratior instead of a function call. The results are exactly the same, it just gives a shorter syntax.
End of explanation
<END_TASK>