Energy curve plotter

I am not much experienced in Python, just write some small script. All my codes are procedural. They work fine, and I always check them with pep8.

One of them is:

#!/usr/bin/python3
import os
import sys
import numpy as np
import scipy as sp
import scipy.optimize
import matplotlib.pyplot as plt

def f(x, a, b, c, d):
return a + b*x + c*x**2 + d*x**3

data = []
fout = "tmpft"  # sys.argv + "_trial.dat"
with open(fout, "w") as of:
for subdir, dirs, files in os.walk(sys.argv):
for inp in files:
if inp.endswith("SCF.out"):
fsys = subdir + "/" + inp
# print(fsys)
with open(fsys) as finp:
for line in finp:
if "lattice constant  ALAT" in line:
lata = float(line.strip()[-7:])/1.88973
if " ERR " in line:
mom = line.strip()[61:70]
if "SCF - cycle converged" in line:
etot = float(line.lstrip()[10:25])
of.write("{:<10f} {:<15f} {:<15f}\n"
.format(lata,  etot, float(mom)))

ffout = sys.argv+".dat"
with open(ffout, "w") as ff:
ff.write("#lat(A)          ETOT          MOM\n")
with open(fout, "r") as uns:
for line in sorted(uns):
ff.write(line)
values = np.genfromtxt(fout, dtype=None, delimiter='   ', usecols=[0, 1, 2])
os.remove(fout)
etot = values[:, 1]
latp = values[:, 0]

popt, pcov = sp.optimize.curve_fit(f, latp, etot)
res = sp.optimize.minimize(lambda x: f(x, *popt), 2.8)
print('Function is minimized for {0}.'.format(float(res['x'])))
with open(ffout, 'a') as ff:
ff.write('#Function is minimized for {0}\n'.format(float(res['x'])))

# Plot data
pout = sys.argv+".png"
x = sp.linspace(2.80, 3.05, 100)
y = f(x, *popt)
plt.plot(x, y, lw=3)
plt.xlabel("Lattice Parameter")
plt.ylabel("Energy")
plt.title(sys.argv)
plt.scatter(latp, etot, s=60)
plt.scatter(res['x'], f(res['x'], *popt), color='red', s=80)
plt.savefig(pout)
plt.show()

This works fine, i.e. gives me proper/correct result. But I wonder if this is a good way of writing Python. If I take this as my standard example (most of my Python interaction is like this: read a file, fetch some data, post-process it and plot), how I can make it better? By "better", I mean, optimized, reusable and readable.

Like, as I already said, most of my python involves reading a line that matched a string. So, I would like to have it as a function (in C/fortran sense) etc.

• Note that functional-programming is about more than just code that has "functions" in the C/Fortran sense. It requires you to write functions in the mathematical sense of "function". I have therefore removed the tag from your question. Nov 26 '15 at 19:01

Your function f takes 4 generically named parameters. Obviously it's based on an equation so the parameters can't be more meaningful but it would be best if you renamed the function with something more descriptive and gave it a docstring:

def equation(x, a, b, c, d):
"""Returns the result of a + b*x + c*x**2 + d*x**3

This is also known as an ________ equation."""

return a + b*x + c*x**2 + d*x**3

equation still isn't a good name, the type of equation would be better there.

You have a lot of nesting when reading the files. One way you could reduce it is inside the for inp in files block is to invert the condition and use continue:

for inp in files:
if not inp.endswith("SCF.out"):
continue

fsys = subdir + "/" + inp
# print(fsys)

continue will skip to the next iteration of the loop, which does what you need without putting everything into another level of nesting. You could also abstract out the rest of this block so a function. Currently with everything bunched up together it's harder to read and parse, but if you just had a function that processed the files then it'd look much neater:

with open(fout, "w") as of:
for subdir, dirs, files in os.walk(sys.argv):
for inp in files:
if not inp.endswith("SCF.out"):
continue

fsys = subdir + "/" + inp
process_file(fsys)

def process_file(fsys):
with open(fsys) as finp:
for line in finp:
if "lattice constant  ALAT" in line:
lata = float(line.strip()[-7:])/1.88973
if " ERR " in line:
mom = line.strip()[61:70]
if "SCF - cycle converged" in line:
etot = float(line.lstrip()[10:25])
of.write("{:<10f} {:<15f} {:<15f}\n"
.format(lata,  etot, float(mom)))

process_file is a bad name, but I'm not entirely clear what this block does. It would actually be clearer to understand if it was now given a name and docstring so others could follow it easier.

Speaking of names, fsys, finp and inp are not clear. Why is each file called inp? What's an fsys? File system? System file? It's used as a file path, so maybe that would be a better name? The name should indicate the use or intent behind a variable. They should communicate to someone clearly. If you find a case where a good name is impossible, then at least explain what's happening with a comment, but that's a last resort.

Instead of looping over uns and writing each line, use str.join to make this faster. "".join(sorted(uns)) will create a single string from each line of uns concatenated together. Though are you aware that this down't add newlines to the end of each write? So instead of separating lines with \n you're just writing all the lines together (unless the lines themselves contain newline characters). If you did want the lines separated, you could just use "\n".join(sorted(uns)) to write out all the lines to a file with each on a new line. Also you can have two with open statements on the same line:

with open(ffout, "w") as ff, open(fout, "r") as uns:
ff.write("#lat(A)          ETOT          MOM\n")
ff.write("\n".join(sorted(uns)))
• f is a function even in the mathematical sense of the word. Equation is a different concept. Nov 27 '15 at 8:09

"Better" is obviously fairly subjective. After a quick glance I would like to make a few suggestions:

• If you already imported scipy, you don't need to import scipy.optimize separately. You can either limit the importing by just importing the (sub)modules you need, which is cleaner. Or you can import the entire module and just use what you need.
• Especially in terms of usability and readability for later (especially rereading your own code), I would suggest adding some docstring to describe what it does: https://www.python.org/dev/peps/pep-0257/#id16
• Similarly I would use a clearer naming convention. The program won't run faster if you call it "energy_curve_calculator_function_using_curve_fit" or "f", but when you look into your files, you will know the difference.
• Lastly, I would also advise you add a comment block (or one line) between each section of code explaining what it does.

Right now pep8 only checks that your lines aren't too long. I would concentrate a little more on making sure someone who didn't write the code (or you in a few months) can read what it's doing based on the flow, comments where needed and the variable names.

• Hi, Thanks for your comment. For the first point of importing scipy.optimize separately, if I just comment import scipy.optimize, and run, I am getting error: ./latres.py Rh Traceback (most recent call last): File "./latres.py", line 50, in <module> popt, pcov = sp.optimize.curve_fit(f, latp, etot) AttributeError: 'module' object has no attribute 'optimize' Nov 26 '15 at 15:37
• This is something special about how scipy manages its packages. It is indeed necessary to import scipy.optimize to be able to use scipy.optimize.curve_fit. It is not the most common behaviour for python packages in general though. Nov 26 '15 at 16:55
• Another thing...I just read up the docstring link. Does it do anything better than simple comment? I failed to see. Nov 27 '15 at 8:00
• @BaRud It does, it can be read programmatically. One example is if you call help(function) it will print the docstring from function if it has one. Comments on the other hand are ignored by the interpreter and not accessible. Nov 27 '15 at 9:27
• I wasn't aware of the scipy.optimize behavior, learned something new. As for the docstring, yes, a lot of IDEs have expanded that help function by showing the docstring when you want to run the function (for example if you import this from a package, you can type "package.function(" in your IDE and then it'll show you the docstring in a floating bubble. I find it very helpful because I make sure my dosctrings contain exactly what my function needs (deg vs rads, which order the params go in, etc). Another reason to give the variables a more descriptive name. Nov 27 '15 at 16:21

Isolating some pieces of code into functions makes the code easier to understand at a glance: