Energy curve plotter

Question

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[1] + "_trial.dat"
with open(fout, "w") as of:
    for subdir, dirs, files in os.walk(sys.argv[1]):
        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[1]+".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[1]+".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[1])
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.

Answer 1

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[1]):
        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)))

Answer 2

"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.

Answer 3

Readibility

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

   def readable_info(line)
             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])
                        return "{:<10f} {:<15f} {:<15f}\n"
                                 .format(lata,  etot, float(mom))

You may then call this foreach line in the file instead of inlining it. You also reduce nesting that is often considered good.