Can this if statement be simplified?

Question

I would like to know if the following script can be simplified? Thank you!

def charge(atom, r_12):
    #Defining the charge
    if atom == "C":
        E = energy(-4, r_12)
    if atom == "N":
        E = energy(-3, r_12)
    if atom == "Cs" or atom == "H" :
        E = energy(1, r_12)
    if atom == "Ge" or atom == "Pb" or atom == "Si" :
        E = energy(2, r_12)
    if atom == "Br" or atom == "I" :
        E = energy(-1, r_12)
    if atom == "X":
        E = energy(-2, r_12)
    return E

The energy(q_ion, r_12) is just another function I defined for calculating the interaction energy and m is the math package in python.

def energy(q_ion, r_12):
    #Calculating the energy
    Q = -1
    l = 6
    E = (Q * q_ion * m.erf(r_12 / l)) / r_12
    return E

In the future, please don't use the "snippet" feature for Python. It's intended for HTML/CSS/JavaScript. — Fred Larson, Oct 28 '21 at 18:14
Yes, Since the second argument is same in all the if condition you can declare the value of r_12 in energy and use dict with 'c', 'n' and all other as keys and q_ion as value — r_batra, Oct 28 '21 at 18:15

ti7 · Answer 1 · 2021-10-28T19:57:52.060

I'd simplify and improve this this by both

keeping the by-atom input in a dictionary, which both
- makes the collection of possible inputs more manageable
- avoids successively comparing after a result has already been found (originally, each if is discretely checked even when E has been named, though this can be fixed by using elif!)
caching the results of the calculations, which will speed up successive requests for the same (atom, r_12) pair

from functools import lru_cache

atomic_charge = {
    "C": -4,
    "N": -3,
    ...
}

@lru_cache
def energy(q_ion, r_12):
    #Calculating the energy
    Q = -1
    l = 6
    E = (Q * q_ion * m.erf(r_12 / l)) / r_12
    return E

def charge(atom, r_12):
    return energy(atomic_charge[atom], r_12)

Note KeyError will be raised for atoms not in your source dict (where your initial version would raise NameError: E not defined if no if tripped)

functools.lru_cache keeps the last N (defaults to 128) results in an internal mapping which can be returned faster if the same input is given for a small overhead, this is known as memoization

If repeated requests for the same (atom, r_12) pair is unlikely, you can skip the lru_cache as it's only adding a small overhead without benefit

Finally, if you have a huge number of atoms to process (say >10000, not Avogadro's-sized!), Scipy may have a more direct function you can use with some scientific array, which will be much faster after the initial setup pain

https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.energy_distance.html

score 1 · Answer 2 · answered Oct 28 '21 at 18:17

1

One simplification would be to use a dictionary instead of all the if statements. If you're going to be calling "charge" a lot, or if this same information is reused elsewhere, you may want to define the dictionary in a broader context (e.g. globally).

def charge(atom, r_12):
    charges = {"C": -4, "N": -3, ...}
    return energy(charges[atom], r_12)

answered Oct 28 '21 at 18:17

inteoryx

785
2
9

This is great and indeed just what they need, but I'd push the dict creation outside of the function so it's not re-created each call! – ti7 Oct 28 '21 at 18:28
I agree. If this function is going to be called a lot, or if performance matters, then it would be better not to create this dictionary inside of the function. If those things aren't the case then it may just be a little simpler to create within the function. – inteoryx Oct 29 '21 at 00:25

score 0 · Answer 3 · answered Oct 28 '21 at 18:28

I would create a dictionary.

table = {
    "C": -4,
    "N": -3,
    "Cs": 1,
    "H": 1,
    "Ge": 2,
    "Pb": 2,
    "Si": 2,
    "Br": -1,
    "I": -1,
    "X": -2
}


def charge(atom, r_12):
    return energy(table[atom], r_12)

In fact, you can remove the charge function altogether and reference the dictionary directly in the energy function.

score 0 · Answer 4 · answered Oct 28 '21 at 18:32

if you update to Python 3.10 you will be able to use Switch statements, thatskes this kind of if chain much easier.
Other than that, the if statements where you are iterating over more than one value using "or", should be better if you add those values to a list or tuple (depending on your needs), and then iterate over the list. I've always heard that too many "or" statements are bad practice.
If you don't want the code to be bloated with the values you could add those to a list and use a for loop to iterate over it, it would look pretty much just as this one, but in case you need to change a value it's gonna be way easier.

Can this if statement be simplified?

4 Answers4