update aims

mushroom/aims/analyse.py

@@ -13,6 +13,7 @@
     "get_dimensions",
     "get_chemical_potential",
     "is_finished_aimsdir",
+    "get_atoms_from_geometry",
 ]
 
 _logger = loggers["aims"]
@@ -105,3 +106,21 @@ def is_finished_aimsdir(dirpath: Union[str, os.PathLike], aimsout_pat: str = "ai
             if stdout.is_finished():
                 return aimsout.name
     return None
+
+
+def get_atoms_from_geometry(path_geometry: str = "geometry.in"):
+    """Get the atoms in the geometry file
+
+    Args:
+        path_geometry (str)
+
+    Returns:
+        List of str
+    """
+    with open(path_geometry, 'r') as h:
+        lines = h.readlines()
+    atms = []
+    for l in lines:
+        if l.strip().startswith("atom"):
+            atms.append(l.split()[-1])
+    return atms

mushroom/aims/species.py

@@ -13,6 +13,7 @@
 
 __all__ = [
     "Species",
+    "readlines_species_defaults",
 ]
 _logger = loggers["aims"]
 
@@ -122,21 +123,20 @@ def get_species_defaults_directory():
     return species_defaults
 
 
-def get_specie_filename(elem: str,
-                        directory: str,
-                        species_defaults: str = None,
-                        error_dir_not_found: bool = True):
-    """get the name of specie file of a particular basis
+def get_species_filepaths(directory: str,
+                          elem: str, *elems: str,
+                          species_defaults: str = None,
+                          error_dir_not_found: bool = True):
+    """get the paths of species files of a particular basis
 
     Args:
-        elem (str)
         directory (str): directory of species category or its alias
+        elem and elems (str)
         species_defaults (str)
 
     Returns:
-        str
+        List of str
     """
-    elem_id = get_atomic_number(elem, False)
     if species_defaults is None:
         species_defaults = get_species_defaults_directory()
     directories_avail = search_basis_directories(species_defaults, error_dir_not_found=error_dir_not_found)
@@ -147,11 +147,30 @@ def get_specie_filename(elem: str,
     if error_dir_not_found and directory not in directories_avail:
         raise ValueError("{} is not found in available basis directories {}"
                          .format(directory, directories_avail))
-    pspecies = [species_defaults, directory, f"{elem_id:02d}_{elem}_default"]
-    pspecies = os.path.join(*pspecies)
+    elems = [elem, *elems]
+    pspecies = []
+    for e in elems:
+        elem_id = get_atomic_number(e, False)
+        pspecies.append(os.path.join(*[species_defaults, directory, f"{elem_id:02d}_{e}_default"]))
     return pspecies
 
 
+def readlines_species_defaults(directory: str, elem: str, *elems: str, species_defaults: str = None):
+    """read the species files specified by `directory` and elements `elem` and `elems`
+
+    Args:
+        directory (str): directory of species category or its alias
+        elem and elems (str)
+        species_defaults (str)
+    """
+    lines = []
+    for pspecie in get_species_filepaths(directory, elem, *elems,
+                                         species_defaults=species_defaults, error_dir_not_found=True):
+        with open(pspecie, 'r') as h:
+            lines.extend(h.readlines())
+    return "".join(lines)
+
+
 class NaoBasisPool:
     """"""
 
@@ -599,5 +618,5 @@ def read_default(cls, elem: str, directory: str = "defaults_2020/intermediate"):
                 e.g. "cc-pVDZ" and set category to "non-standard"
         """
         species_defaults = get_species_defaults_directory()
-        pspecies = get_specie_filename(elem, directory, species_defaults)
+        pspecies = get_species_filepaths(directory, elem, species_defaults)[0]
         return cls.read(pspecies)

mushroom/aims/test/test_species.py

@@ -12,7 +12,7 @@
 from mushroom.aims.species import get_num_obfs, get_l_nrad
 from mushroom.aims.species import _SPECIES_DEFAULTS_ENV, \
     search_basis_directories, get_species_defaults_directory, get_basis_directory_from_alias, \
-    get_specie_filename
+    get_species_filepaths
 
 
 def _check_if_species_defaults_set():
@@ -68,12 +68,13 @@ def test_alias(self):
         self.assertEqual(get_basis_directory_from_alias("nao-vcc-2z"), "NAO-VCC-nZ/NAO-VCC-2Z")
 
     def test_get_specie_filename(self):
-        self.assertRaises(ValueError, get_specie_filename, "Ne", "unknown", "path/to/species_defaults")
+        self.assertRaises(ValueError, get_species_filepaths,
+                          "unknown", "Ne", species_defaults="path/to/species_defaults")
         if _check_if_species_defaults_set():
             return
         d = get_species_defaults_directory()
-        self.assertEqual(os.path.join(d, "defaults_2010", "tight", "01_H_default"),
-                         get_specie_filename("H", "defaults_2010/tight"))
+        self.assertListEqual([os.path.join(d, "defaults_2010", "tight", "01_H_default"),],
+                             get_species_filepaths("defaults_2010/tight", "H"))
 
     def test_get_num_obfs(self):
         """Test get number of basis functions"""
@@ -157,7 +158,7 @@ def test_read_default(self):
 
     def test_read_multiple(self):
         try:
-            pspecies = get_specie_filename('H', 'defaults_2010/tight')
+            pspecies = get_species_filepaths('defaults_2010/tight', 'H')[0]
         except (ValueError, KeyError):
             return
         with open(pspecies, 'r') as h:

mushroom/aims/workflow.py

@@ -3,8 +3,5 @@
 import os
 
 
-class AimsInput:
-
-    def __init__(self, control, geometry):
-        self.geometry = geometry
-        self.control = control
+class Aims:
+    """class to define a FHI-aims calculation"""