Added dataclasses decorator to molecule class

src/qibochem/driver/molecule.py

@@ -2,6 +2,7 @@
 Driver for obtaining molecular integrals from either PySCF or PSI4
 """
 
+from dataclasses import dataclass, field
 from pathlib import Path
 
 import numpy as np
@@ -15,6 +16,7 @@
 )
 
 
+@dataclass
 class Molecule:
     """
     Class representing a single molecule
@@ -31,88 +33,83 @@
 
     """
 
-    def __init__(self, geometry=None, charge=0, multiplicity=1, basis=None, xyz_file=None, active=None):
-        # Basic properties
-        # Define using the function arguments if xyz_file not given
-        if xyz_file is None:
-            self.geometry = geometry
-            self.charge = charge
-            self.multiplicity = multiplicity
-        else:
-            # Check if xyz_file exists, then fill in the Molecule attributes
-            assert Path(f"{xyz_file}").exists(), f"{xyz_file} not found!"
-            self._process_xyz_file(xyz_file, charge, multiplicity)
-        if basis is None:
-            # Default bais is STO-3G
-            self.basis = "sto-3g"
-        else:
-            self.basis = basis
-
-        self.nelec = None  #: Total number of electrons for the molecule
-        self.norb = None  #: Number of molecular orbitals considered for the molecule
-        self.nso = None  #: Number of molecular spin-orbitals considered for the molecule
-        self.e_hf = None  #: Hartree-Fock energy
-        self.oei = None  #: One-electron integrals
-        self.tei = None  #: Two-electron integrals, order follows the second quantization notation
-
-        self.ca = None
-        self.pa = None
-        self.da = None
-        self.nalpha = None
-        self.nbeta = None
-        self.e_nuc = None
-        self.overlap = None
-        self.eps = None
-        self.fa = None
-        self.hcore = None
-        self.ja = None
-        self.ka = None
-        self.aoeri = None
-
-        # For HF embedding
-        self.active = active  #: Iterable of molecular orbitals included in the active space
-        self.frozen = None  #: Iterable representing the occupied molecular orbitals removed from the simulation
-
-        self.inactive_energy = None
-        self.embed_oei = None
-        self.embed_tei = None
-
-        self.n_active_e = None  #: Number of electrons included in the active space if HF embedding is used
-        self.n_active_orbs = None  #: Number of spin-orbitals in the active space if HF embedding is used
-
-    def _process_xyz_file(self, xyz_file, charge, multiplicity):
+    geometry: list = None
+    charge: int = None
+    multiplicity: int = 1
+    basis: str = "sto-3g"
+    xyz_file: str = None
+
+    # or HF embedding
+    active: None = None
+    frozen: None = field(default=None, init=False)
+
+    nelec: None = field(default=None, init=False)
+    norb: None = field(default=None, init=False)
+    nso: None = field(default=None, init=False)
+    e_hf: None = field(default=None, init=False)
+    oei: None = field(default=None, init=False)
+    tei: None = field(default=None, init=False)
+
+    ca: None = field(default=None, init=False)
+    pa: None = field(default=None, init=False)
+    da: None = field(default=None, init=False)
+    nalpha: None = field(default=None, init=False)
+    nbeta: None = field(default=None, init=False)
+    e_nuc: None = field(default=None, init=False)
+    overlap: None = field(default=None, init=False)
+    eps: None = field(default=None, init=False)
+    fa: None = field(default=None, init=False)
+    hcore: None = field(default=None, init=False)
+    ja: None = field(default=None, init=False)
+    ka: None = field(default=None, init=False)
+    aoeri: None = field(default=None, init=False)
+
+    inactive_energy: None = field(default=None, init=False)
+    embed_oei: None = field(default=None, init=False)
+    embed_tei: None = field(default=None, init=False)
+
+    n_active_e: None = field(default=None, init=False)
+    n_active_orbs: None = field(default=None, init=False)
+
+    # Runs after init, formerly the _process_xyz_file function
+    def __post_init__(self):
         """
-        Reads a .xyz file to obtain and set the molecular coordinates (in OpenFermion format),
-            charge, and multiplicity
-
-        Args:
-            xyz_file: .xyz file for molecule. Comment line should follow "{charge} {multiplicity}"
+        Initializes certain variables.
+        Inits multiplicit, basis, and certain parameters based on .xyz_file
+        Reads a .xyz_file to obtain and set the molecular coordinates (in OpenFermion format),
+        charge, and multiplicity.
         """
-        with open(xyz_file, encoding="utf-8") as file_handler:
-            # First two lines: # atoms and comment line (charge, multiplicity)
-            _n_atoms = int(file_handler.readline())  # Not needed/used
-
-            # Try to read charge and multiplicity from comment line
-            split_line = [int(_num) for _num in file_handler.readline().split()]
-            if len(split_line) == 2:
-                # Format of comment line matches (charge, multiplicity):
-                _charge, _multiplicity = split_line
-            else:
-                # Otherwise, use the default (from __init__) values of 0 and 1
-                _charge, _multiplicity = charge, multiplicity
-
-            # Start reading xyz coordinates from the 3rd line onwards
-            _geometry = []
-            for line in file_handler:
-                split_line = line.split()
-                # OpenFermion format: [('H', (0.0, 0.0, 0.0)), ('H', (0.0, 0.0, 0.7)), ...]
-                atom_xyz = [split_line[0], tuple(float(_xyz) for _xyz in split_line[1:4])]
-                _geometry.append(tuple(atom_xyz))
-
-        # Set the class attributes
-        self.charge = _charge
-        self.multiplicity = _multiplicity
-        self.geometry = _geometry
+        if self.multiplicity is None:
+            self.multiplicity = 1
+        if self.basis is None:
+            self.basis = "sto-3g"
+        if self.xyz_file is not None:
+            assert Path(f"{self.xyz_file}").exists(), f"{self.xyz_file} not found!"
+            with open(self.xyz_file, encoding="utf-8") as file_handler:
+                # First two lines: # atoms and comment line (charge, multiplicity)
+                _n_atoms = int(file_handler.readline())  # Not needed/used
+
+                # Try to read charge and multiplicity from comment line
+                split_line = [int(_num) for _num in file_handler.readline().split()]
+                if len(split_line) == 2:
+                    # Format of comment line matches (charge, multiplicity):
+                    _charge, _multiplicity = split_line
+                else:
+                    # Otherwise, use the default (from __init__) values of 0 and 1
+                    _charge, _multiplicity = self.charge, self.multiplicity
+
+                # Start reading xyz coordinates from the 3rd line onwards
+                _geometry = []
+                for line in file_handler:
+                    split_line = line.split()
+                    # OpenFermion format: [('H', (0.0, 0.0, 0.0)), ('H', (0.0, 0.0, 0.7)), ...]
+                    atom_xyz = [split_line[0], tuple(float(_xyz) for _xyz in split_line[1:4])]
+                    _geometry.append(tuple(atom_xyz))
+
+            # Set the class attributes
+            self.charge = _charge
+            self.multiplicity = _multiplicity
+            self.geometry = _geometry
 
     def run_pyscf(self, max_scf_cycles=50):
         """

tests/test_ucc.py

@@ -239,7 +239,7 @@ def test_ucc_ansatz_error_checks():
     """Test the checks for input validity"""
     mol = Molecule([("Li", (0.0, 0.0, 0.0)), ("H", (0.0, 0.0, 1.4))])
     # Define number of electrons and spin-obritals by hand
-    mol.nelec = (2, 2)
+    mol.nelec = 4
     mol.nso = 12
 
     # Excitation level check: Input is in ("S", "D", "T", "Q")