Cleaning Cosymlib class initialization

GrupEstructuraElectronicaSimetria · Feb 19, 2022 · 90cb5be · 90cb5be
1 parent 84d0405
commit 90cb5be
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 25 deletions.
diff --git a/cosymlib/__init__.py b/cosymlib/__init__.py
@@ -104,28 +104,14 @@ def __init__(self,
                  connectivity_thresh=None,
                  mode=0):
 
-        mode_list = [None,'EH', 'Dens']
+        mode_list = [None, 'EH', 'Dens']
         self._mode = mode_list[mode]
         self._molecules = []
-        if isinstance(structures, list):
+        if isinstance(structures, (list, tuple)):
             for structure in structures:
-                if isinstance(structure, Molecule):
-                    if self._mode is not None:
-                        structures._electronic_structure = self._mode
-                    self._molecules.append(structure)
-                elif isinstance(structure, Geometry):
-                    self._molecules.append(Molecule(structure,electronic_structure=self._mode))
-                else:
-                    raise AttributeError('Molecule object not found')
+                self._molecules.append(Molecule(structure, electronic_structure=self._mode))
         else:
-            if isinstance(structures, Molecule):
-                if self._mode is not None:
-                    structures._electronic_structure = self._mode
-                self._molecules.append(structures)
-            elif isinstance(structures, Geometry):
-                self._molecules.append(Molecule(structures,electronic_structure=self._mode))
-            else:
-                raise AttributeError('Molecule object not found')
+            self._molecules.append(Molecule(structures, electronic_structure=self._mode))
 
         for molecule in self._molecules:
             if ignore_atoms_labels:

diff --git a/cosymlib/molecule/__init__.py b/cosymlib/molecule/__init__.py
@@ -22,19 +22,22 @@ class Molecule:
     Main molecule class
 
     :param geometry: The geometry
-    :type geometry: Geometry
+    :type geometry: Geometry, Molecule
     :param electronic_structure: The electronic structure
-    :type electronic_structure: ElectronicStructure
+    :type electronic_structure: ElectronicStructure, str
     :param name: Molecule name
     :type name: str
 
     """
     def __init__(self, geometry,
-                 electronic_structure='HF',
+                 electronic_structure=None,
                  name=None):
 
-        if not geometry:
-            raise Exception('No geometry found in the input file, check out input file for possible errors')
+        # handle if geometry is already a molecule
+        if isinstance(geometry, Molecule):
+            if electronic_structure is None:
+                electronic_structure = geometry.electronic_structure
+            geometry = geometry.geometry
         if name is None:
             self._name = geometry.name
 
@@ -71,8 +74,13 @@ def electronic_structure(self):
         :return: The electronic structure
         :rtype: ElectronicStructure
         """
-        if self._electronic_structure == 'EH':
+
+        # if None default electronic structure is ExtendedHuckel
+        if self._electronic_structure is None:
+            self._electronic_structure = 'EH'
             warn('Electronic structure auto generated from Extended-Huckel calculation')
+
+        if self._electronic_structure == 'EH':
             eh = ExtendedHuckel(self.geometry)
 
             self._electronic_structure = ElectronicStructure(basis=eh.get_basis(),
@@ -82,8 +90,8 @@ def electronic_structure(self):
                                                              multiplicity=eh.get_multiplicity(),
                                                              alpha_occupancy=[1]*eh.get_alpha_electrons(),
                                                              beta_occupancy=[1]*eh.get_beta_electrons())
+
         elif self._electronic_structure == 'Dens':
-            warn('Electronic structure now generates electronic densities for quick symmetry measures')
             protodensity = ProtoElectronicDensity(self.geometry)
             self._electronic_structure = ProtoElectronicStructure(basis=protodensity.get_basis(),
                                                                   orbital_coefficients=[protodensity.get_mo_coefficients(),[]])