Refactor AlphaList functionality to neucbot/alpha.py and add black

.github/workflows/test.yaml

@@ -25,3 +25,6 @@ jobs:
 
       - name: Run tests
         run: python -m pytest -vv
+
+      - name: Run linter
+        run: python -m black neucbot/ tests/ --check

neucbot.py

@@ -10,6 +10,7 @@
 from neucbot import ensdf
 
 from neucbot import elements
+from neucbot import alpha
 
 class constants:
     N_A = 6.0221409e+23
@@ -46,46 +47,6 @@ def parseIsotope(iso):
             A = A*10 + int(i)
     return [ele,A]
 
-def generateAlphaFileName(ele,A):
-    outdir = './AlphaLists/'
-    fName = outdir + ele.capitalize() + str(A) + 'Alphas.dat'
-    return fName
-
-def generateAlphaList(ele, A):
-    print('generateAlphaList(',ele,A,')',file=constants.ofile)
-    ensdf.main(['parseENSDF',ele,A])
-
-def loadAlphaList(fname):
-    f = open(fname)
-    tokens = [line.split() for line in f.readlines()]
-    alpha_list = []
-    for words in tokens:
-        if words[0][0] == '#' or len(words) < 2:
-            continue
-
-        alpha = []
-        for word in words:
-            alpha.append(float(word))
-        alpha_list.append(alpha)
-    f.close()
-    return alpha_list
-
-def getAlphaList(ele,A):
-    fname = generateAlphaFileName(ele,A)
-    return loadAlphaList(fname)
-
-def getAlphaListIfExists(ele,A):
-    fName = generateAlphaFileName(ele,A)
-    tries = 3
-    while not os.path.isfile(fName):
-        if tries < 0:
-            print('Cannot generate alpha list for ele = ', ele,  ' and A = ', A,file = constants.ofile)
-            return 0
-        print('generating alpha file ', fName, file = constants.ofile)
-        generateAlphaList(ele,A)
-        tries -= 1
-    return getAlphaList(ele,A)
-
 def loadChainAlphaList(fname):
     f = open(fname)
     tokens = [line.split() for line in f.readlines()]
@@ -99,7 +60,7 @@ def loadChainAlphaList(fname):
         [ele,A] = parseIsotope(iso)
 
         # Now get the isotope's alpha list and add it to the chain's list
-        aList_forIso = getAlphaListIfExists(ele,A)
+        aList_forIso = alpha.AlphaList(ele, A).load_or_fetch()
         if constants.print_alphas:
             print(iso, file = constants.ofile)
             print('\t', aList_forIso, file = constants.ofile)
@@ -472,7 +433,7 @@ def main():
         if arg == '-l':
             alphalist_file = sys.argv[sys.argv.index(arg)+1]
             print('load alpha list', alphalist_file, file = sys.stdout)
-            alpha_list = loadAlphaList(alphalist_file)
+            alpha_list = alpha.AlphaList.from_filepath(alphalist_file)
         if arg == '-c':
             chain_file = sys.argv[sys.argv.index(arg)+1]
             print('load alpha chain', chain_file, file = sys.stdout)
@@ -545,4 +506,3 @@ def main():
 
 if __name__ == '__main__':
     main()
-

neucbot/alpha.py

@@ -0,0 +1,63 @@
+import os
+from neucbot import ensdf
+
+ALPHA_LIST_DIR = "./AlphaLists"
+
+
+def _alphas_from_file_path(file_path):
+    file = open(file_path)
+
+    # Parse alphalist files:
+    # 1. Only parse lines that have 2+ tab-separated tokens
+    # 2. Ignore any lines starting with "#"
+    # 3. Return list of lists (where each sublist is a list of floats)
+    alphas = [
+        [float(token) for token in line.split()]  # Parse each token as float
+        for line in file.readlines()  # for each line in file
+        if line[0] != "#"
+        and len(line.split()) >= 2  # except for lines matching these conditions
+    ]
+
+    file.close()
+
+    return alphas
+
+
+class AlphaList:
+    def __init__(self, element, isotope):
+        self.element = element
+        self.isotope = isotope
+        self.file_path = f"{ALPHA_LIST_DIR}/{self.element}{self.isotope}Alphas.dat"
+        self.fetch_attempts = 3
+
+    @classmethod
+    def from_filepath(cls, file_path):
+        return _alphas_from_file_path(file_path)
+
+    def load_or_fetch(self):
+        while not os.path.isfile(self.file_path):
+            if self.fetch_attempts < 0:
+                raise RuntimeError(f"Unable to write alpha file to {self.file_path}")
+            self.write()
+            self.fetch_attempts -= 1
+
+        return self.load()
+
+    def load(self):
+        return _alphas_from_file_path(self.file_path)
+
+    def write(self):
+        if os.path.exists(self.file_path):
+            print(f"Alpha list file already exists at {self.file_path}")
+        else:
+            client = ensdf.Client(self.element, self.isotope)
+            decay_file_text = client.read_or_fetch_decay_file()
+            energyMaps = ensdf.Parser.parse(decay_file_text)
+            file = open(self.file_path, "w")
+
+            for energy, probability in energyMaps["alphas"].items():
+                file.write(f"{str(energy/1000)}\t{probability}\n")
+
+            file.close()
+
+        return True

neucbot/chemistry.py

@@ -1,32 +1,130 @@
-elementsToZ = {'H'  : 1, 'He' : 2, 'Li' : 3, 'Be' : 4, 'B'  : 5,
-            'C'  : 6, 'N'  : 7, 'O'  : 8, 'F'  : 9, 'Ne' : 10,
-            'Na' : 11,'Mg' : 12,'Al' : 13,'Si' : 14,'P'  : 15,
-            'S'  : 16,'Cl' : 17,'Ar' : 18,'K'  : 19,'Ca' : 20,
-            'Sc' : 21,'Ti' : 22,'V'  : 23,'Cr' : 24,'Mn' : 25,
-            'Fe' : 26,'Co' : 27,'Ni' : 28,'Cu' : 29,'Zn' : 30,
-            'Ga' : 31,'Ge' : 32,'As' : 33,'Se' : 34,'Br' : 35,
-            'Kr' : 36,'Rb' : 37,'Sr' : 38,'Y'  : 39,'Zr' : 40,
-            'Nb' : 41,'Mo' : 42,'Tc' : 43,'Ru' : 44,'Rh' : 45,
-            'Pd' : 46,'Ag' : 47,'Cd' : 48,'In' : 49,'Sn' : 50,
-            'Sb' : 51,'Te' : 52,'I'  : 53,'Xe' : 54,'Cs' : 55,
-            'Ba' : 56,'La' : 57,'Ce' : 58,'Pr' : 59,'Nd' : 60,
-            'Pm' : 61,'Sm' : 62,'Eu' : 63,'Gd' : 64,'Tb' : 65,
-            'Dy' : 66,'Ho' : 67,'Er' : 68,'Tm' : 69,'Yb' : 70,
-            'Lu' : 71,'Hf' : 72,'Ta' : 73,'W'  : 74,'Re' : 75,
-            'Os' : 76,'Ir' : 77,'Pt' : 78,'Au' : 79,'Hg' : 80,
-            'Tl' : 81,'Pb' : 82,'Bi' : 83,'Po' : 84,'At' : 85,
-            'Rn' : 86,'Fr' : 87,'Ra' : 88,'Ac' : 89,'Th' : 90,
-            'Pa' : 91,'U'  : 92,'Np' : 93,'Pu' : 94,'Am' : 95,
-            'Cm' : 96,'Bk' : 97,'Cf' : 98,'Es' : 99,'Fm' :100,
-            'Md' :101,'No' :102,'Lr' :103,'Rf' :104,'Db' :105,
-            'Sg' :106,'Bh' :107,'Hs' :108,'Mt' :109,'Ds' :110,
-            'Rg' :111,'Cn' :112,'Uut':113,'Fl' :114,'Uup':115,
-            'Lv' :116,'Uus':117,'Uuo':118}
+elementsToZ = {
+    "H": 1,
+    "He": 2,
+    "Li": 3,
+    "Be": 4,
+    "B": 5,
+    "C": 6,
+    "N": 7,
+    "O": 8,
+    "F": 9,
+    "Ne": 10,
+    "Na": 11,
+    "Mg": 12,
+    "Al": 13,
+    "Si": 14,
+    "P": 15,
+    "S": 16,
+    "Cl": 17,
+    "Ar": 18,
+    "K": 19,
+    "Ca": 20,
+    "Sc": 21,
+    "Ti": 22,
+    "V": 23,
+    "Cr": 24,
+    "Mn": 25,
+    "Fe": 26,
+    "Co": 27,
+    "Ni": 28,
+    "Cu": 29,
+    "Zn": 30,
+    "Ga": 31,
+    "Ge": 32,
+    "As": 33,
+    "Se": 34,
+    "Br": 35,
+    "Kr": 36,
+    "Rb": 37,
+    "Sr": 38,
+    "Y": 39,
+    "Zr": 40,
+    "Nb": 41,
+    "Mo": 42,
+    "Tc": 43,
+    "Ru": 44,
+    "Rh": 45,
+    "Pd": 46,
+    "Ag": 47,
+    "Cd": 48,
+    "In": 49,
+    "Sn": 50,
+    "Sb": 51,
+    "Te": 52,
+    "I": 53,
+    "Xe": 54,
+    "Cs": 55,
+    "Ba": 56,
+    "La": 57,
+    "Ce": 58,
+    "Pr": 59,
+    "Nd": 60,
+    "Pm": 61,
+    "Sm": 62,
+    "Eu": 63,
+    "Gd": 64,
+    "Tb": 65,
+    "Dy": 66,
+    "Ho": 67,
+    "Er": 68,
+    "Tm": 69,
+    "Yb": 70,
+    "Lu": 71,
+    "Hf": 72,
+    "Ta": 73,
+    "W": 74,
+    "Re": 75,
+    "Os": 76,
+    "Ir": 77,
+    "Pt": 78,
+    "Au": 79,
+    "Hg": 80,
+    "Tl": 81,
+    "Pb": 82,
+    "Bi": 83,
+    "Po": 84,
+    "At": 85,
+    "Rn": 86,
+    "Fr": 87,
+    "Ra": 88,
+    "Ac": 89,
+    "Th": 90,
+    "Pa": 91,
+    "U": 92,
+    "Np": 93,
+    "Pu": 94,
+    "Am": 95,
+    "Cm": 96,
+    "Bk": 97,
+    "Cf": 98,
+    "Es": 99,
+    "Fm": 100,
+    "Md": 101,
+    "No": 102,
+    "Lr": 103,
+    "Rf": 104,
+    "Db": 105,
+    "Sg": 106,
+    "Bh": 107,
+    "Hs": 108,
+    "Mt": 109,
+    "Ds": 110,
+    "Rg": 111,
+    "Cn": 112,
+    "Uut": 113,
+    "Fl": 114,
+    "Uup": 115,
+    "Lv": 116,
+    "Uus": 117,
+    "Uuo": 118,
+}
 
 zToElements = {z: element for element, z in elementsToZ.items()}
 
+
 def getZ(ele):
     return elementsToZ.get(ele.capitalize())
 
+
 def getElement(z):
     return zToElements.get(z, "None")

neucbot/elements.py

@@ -30,6 +30,7 @@
 with open("./neucbot/elements.json", "r") as file:
     isotopesMap = json.load(file)
 
+
 class Element:
     def __init__(self, element):
         self.symbol = element.capitalize()