I just spent almost one day to figure out the following code, which analyzes the output data from material science programs called VASP and Bader. As a Python beginner, I earnestly ask some aces to help me make this code more pythonic, or efficient. Because I know that although the code can run smoothly, it definitely needs more modification from professionals. The code is listed below, with the data file ACF.dat after it.#!/usr/bin/pythonimport numpy as npdef charge(x): # function to calculate charge difference     if x<4:        return 4-x # atom1    else:        return 6-x # atom2# read from file ACF.datf = open('ACF.dat', 'r')lines = []for row in f.readlines():    tmp = row.split()    lines.append(tmp)lines = lines[2:-4] # remove unecessary rowslines = np.array(lines, dtype=float)lines = lines[:,0:5] # remove unecessary columns# calculate residual chargechg = np.vectorize(charge)lines[:,4] = chg(lines[:,4])# calculate interatomic distancesdist = []for i in range(0,len(lines)):    d = np.sqrt(np.sum((lines[i,1:4]-lines[0,1:4])**2))    dist.append(d)lines = np.column_stack((lines, dist)) # attach dist to lineslines = lines[lines[:,5].argsort()]for row in lines:     print '{0:<3.0f} {1:12.6f} {2:12.6f}'.format( row[0], row[5], row[4])f.close()ACF.datThe structure of the file is seen as the heading. What I need are the first five columns: the atomic index, the x, y, z coordinates of each atom, the atomic charge.#         X           Y           Z        CHARGE     MIN DIST    ATOMIC VOL--------------------------------------------------------------------------------1      7.5119      7.5119      7.5119      2.9875      1.1144     20.76922     18.0286     18.0286     18.0286      2.8514      1.3688     23.00953      6.0058     18.0205     18.0205      2.8500      1.3599     22.92654     12.0323     18.0342     18.0342      2.8480      1.3638     22.98165     18.0205      6.0058     18.0205      2.8500      1.3599     22.92656      5.9968      5.9968     17.9789      2.7979      1.2317     22.3582--------------------------------------------------------------------------------VACUUM CHARGE:               0.0000VACUUM VOLUME:               0.0000NUMBER OF ELECTRONS:      1085.9999What the code does is:Read in the content of ACF.dat, and removes the first two rows and last four rows. Then saves only the first five columns into array lines.Calculates the residual charge using function charge(x).Calculates the interatomic distances using the x, y, z coordinates of the atoms, and attach the array of distances to lines as its last column.Sorts lines according to the interatomic distances (column #6 of lines).Prints out the atomic index, interatomic distances, and atomic charges.I understand that this is a lengthy question, but I still would like to ask someone helping me, and other newbies, learn Python faster by optimizing our own code. Any suggestion is welcomed.
Pythonic way to manipulate arrays: a case study
python;numpy
Coding styleFor starters, read PEP 8, which is the Python style guide. This gives suggestions on whitespace around operators and commas, variable names, and the like. Your code isnt bad if it doesnt follow these conventions, but it will slow down more experienced Python programmers reading your code, and might make it harder for you later when you read other peoples Python.You should also read PEP 257, which explain the Python conventions for docstrings. Rather than including a description for the function on the same line as the def, we put it on a special line inside the function definition.With those two in mind, I might rewrite your charge function as:def residual_charge(x):    Returns the residual charge of an atom.    if x < 4:        return 4 - x # atom 1    else:        return 6 - x # atom 2The rest of the style changes are only minor, and you can find them yourself.Opening/closing the fileAs far as I can tell, the six lines after youve defined charge are the only lines where youre directly interacting with the file ACF.dat, but you dont issue a close() command until the end of the script. That can cause problems if the input file is large, or if you want to load multiple files.Id make three changes:Use Pythons with open() ... construction  this keeps all the file handling code in one place, and automatically closes it for you as well.Wrap this in a function: it keeps the code for processing the file input in one place, and means you can use it again if you want to feed it a different file.Go through the file a line at a time (you might not need this, but it saves loading the whole file into memory if ACF.dat is large).So Id write something like:def parse_dat_file(filename):    Returns a numpy array based on the contents of the input file.    # read the file contents    with open(filename, 'r') as f:        for line in f:            parsed_lines.append(line.split())    # remove unnecessary rows and columns    parsed_lines = parsed_lines[2:-4]    parsed_lines = np.array(parsed_lines, dtype=float)    parsed_lines = parsed_lines[:,0:5]    return parsed_linesDoing the necessary calculationsAgain, Id break this into a single function. It separates that block of work from the rest of the file, and makes it easier to organise. This means that if you get the data from another source, you can use the same code untouched.Some notes:In Python, a range() starts at 0 by default, so you dont need to add it. If no start is supplied, it just counts from 0 upwards. You can find out more by typing help(range) at a Python shell.If the data set is large, you should use xrange() instead. This is faster when youre using a large range of objects. See What is the difference between range and xrange? for more details.This is what I get at the end:def process_lines(lines):    Returns the lines with the residual charge and    interatomic distances computed and added.        # Calculate residual charge    v_charge = np.vectorize(residual_charge)    lines[:, 4] = v_charge(lines[:, 4])    # Calculate atomic distances    distances = []    for i in xrange(len(lines)):        line_sum = np.sum((lines[i, 1:4] - lines[0, 1:4]) ** 2)        min_dist = np.sqrt(line_sum)        distances.append(min_dist)    lines = np.column_stack((lines, dist))    lines = lines[lines[:, 5].argsort()]    returnPrinting the fileWithout wishing to sound like a broken record, put this in a function as well. This one is easy:def print_data(lines):    Prints the parsed and formatted data.    for row in lines:         print '{0:<3.0f} {1:12.6f} {2:12.6f}'.format( row[0], row[5], row[4])Script vs. moduleAt the end of the script, we can add the linesacf_lines = parse_dat_file('ACF.dat')print_data(process_lines(acf_lines))and it will do the printing we originally set out to do.But now the code has been broken into separate functions, we can use it in other scripts. Say this was atomic.py; then we could write from atomic import residual_charge in another file, and wed have access to the residual_charge function.However, when we import this file in another script, we dont want it to start reading and printing ACF.dat; we just want the definitions. If we put the particular code for ACF.dat in a special if statement, then it only gets run if the file is called directly:if __name__ == '__main__':    acf_lines = parse_dat_file('ACF.dat')    print_data(process_lines(acf_lines))Now, if we type python atomic.py at a command line, this code gets run and the data is printed. If we use import atomic elsewhere, then we just get the function definitions. You can find out more in What does if __name__ == __main__: do?SummaryIf I combine everything Ive done above and put it into a single file, then this is what Im left with:#!/usr/bin/pythonimport numpy as npdef residual_charge(x):    Returns the residual charge of an atom.    if x < 4:        return 4 - x # atom 1    else:        return 6 - x # atom 2def parse_dat_file(filename):    Returns a numpy array based on the contents of the input file.    # read the file contents    parsed_lines = []    with open(filename, 'r') as f:        for line in f:            parsed_lines.append(line.split())    # Remove unnecessary rows and columns    parsed_lines = parsed_lines[2:-4]    parsed_lines = np.array(parsed_lines, dtype=float)    parsed_lines = parsed_lines[:,0:5]    return parsed_linesdef process_lines(lines):    Returns the lines with the residual charge and    interatomic distances computed and added.           # Calculate residual charge    v_charge = np.vectorize(residual_charge)    lines[:, 4] = v_charge(lines[:, 4])    # Calculate atomic distances    distances = []    for i in xrange(len(lines)):        line_sum = np.sum((lines[i, 1:4] - lines[0, 1:4]) ** 2)        min_dist = np.sqrt(line_sum)        distances.append(min_dist)    lines = np.column_stack((lines, distances))    lines = lines[lines[:, 5].argsort()]    return linesdef print_data(lines):    Prints the parsed and formatted data.    for row in lines:         print '{0:<3.0f} {1:12.6f} {2:12.6f}'.format( row[0], row[5], row[4])if __name__ == '__main__':    acf_lines = parse_dat_file('ACF.dat')    print_data(process_lines(acf_lines))