announce py38 and add Mol fns

README.md

@@ -16,7 +16,7 @@ This project also contains a generator, validator, and translator for [Molecule
 
 ## ✨ Getting Started
 
-- Installation. QCElemental supports Python 3.7+.
+- Installation. QCElemental supports Python 3.7+. Starting with v0.50 (aka "next", aka "QCSchema v2 available"), Python 3.8+ will be supported.
 
   ```sh
   python -m pip install qcelemental

qcelemental/models/molecule.py

@@ -1,7 +1,7 @@
 """
 Molecule Object Model
 """
-
+import collections
 import hashlib
 import json
 import warnings
@@ -875,6 +875,10 @@
         >>> two_pentanol_radcat.get_molecular_formula(chgmult=True)
         2^C5H12O+
 
+        Notes
+        -----
+        This includes all atoms in the molecule, including ghost atoms. See :py:meth:`element_composition` to exclude.
+
         """
 
         from ..molutil import molecular_formula_from_symbols
@@ -1151,13 +1155,15 @@
         tensor[2][1] = tensor[1][2] = -1.0 * np.sum(weight * geom[:, 1] * geom[:, 2])
         return tensor
 
-    def nuclear_repulsion_energy(self, ifr: int = None) -> float:
+    def nuclear_repulsion_energy(self, ifr: int = None, real_only: bool = True) -> float:
         r"""Nuclear repulsion energy.
 
         Parameters
         ----------
         ifr
             If not `None`, only compute for the `ifr`-th (0-indexed) fragment.
+        real_only
+            Only include real atoms in the sum.
 
         Returns
         -------
@@ -1165,7 +1171,11 @@
             Nuclear repulsion energy in entire molecule or in fragment.
 
         """
         Zeff = [z * int(real) for z, real in zip(cast(Iterable[int], self.atomic_numbers), self.real)]
+        if real_only:
+            Zeff = [z * int(real) for z, real in zip(cast(Iterable[int], self.atomic_numbers), self.real)]
+        else:
+            Zeff = self.atomic_numbers
         atoms = list(range(self.geometry.shape[0]))
 
         if ifr is not None:
@@ -1178,21 +1188,26 @@
                 nre += Zeff[at1] * Zeff[at2] / dist
         return nre
 
-    def nelectrons(self, ifr: int = None) -> int:
+    def nelectrons(self, ifr: int = None, real_only: bool = True) -> int:
         r"""Number of electrons.
 
         Parameters
         ----------
         ifr
             If not `None`, only compute for the `ifr`-th (0-indexed) fragment.
+        real_only
+            Only include real atoms in the sum.
 
         Returns
         -------
         nelec : int
             Number of electrons in entire molecule or in fragment.
 
         """
-        Zeff = [z * int(real) for z, real in zip(cast(Iterable[int], self.atomic_numbers), self.real)]
+        if real_only:
+            Zeff = [z * int(real) for z, real in zip(cast(Iterable[int], self.atomic_numbers), self.real)]
+        else:
+            Zeff = self.atomic_numbers
 
         if ifr is None:
             nel = sum(Zeff) - self.molecular_charge
@@ -1202,6 +1217,69 @@
 
         return int(nel)
 
+    def molecular_weight(self, ifr: int = None, real_only: bool = True) -> float:
+        r"""Molecular weight in uamu.
+
+        Parameters
+        ----------
+        ifr
+            If not `None`, only compute for the `ifr`-th (0-indexed) fragment.
+        real_only
+            Only include real atoms in the sum.
+
+        Returns
+        -------
+        mw : float
+            Molecular weight in entire molecule or in fragment.
+
+        """
+        if real_only:
+            masses = [mas * int(real) for mas, real in zip(cast(Iterable[float], self.masses), self.real)]
+        else:
+            masses = self.masses
+
+        if ifr is None:
+            mw = sum(masses)
+
+        else:
+            mw = sum([mas for iat, mas in enumerate(masses) if iat in self.fragments[ifr]])
+
+        return mw
+
+    def element_composition(self, ifr: int = None, real_only: bool = True) -> Dict[str, int]:
+        r"""Atomic count map.
+
+        Parameters
+        ----------
+        ifr
+            If not `None`, only compute for the `ifr`-th (0-indexed) fragment.
+        real_only
+            Only include real atoms.
+
+        Returns
+        -------
+        composition : Dict[str, int]
+            Atomic count map.
+
+        Notes
+        -----
+        This excludes ghost atoms by default whereas get_molecular_formula always includes them.
+
+        """
+        if real_only:
+            symbols = [sym * int(real) for sym, real in zip(cast(Iterable[str], self.symbols), self.real)]
+        else:
+            symbols = self.symbols
+
+        if ifr is None:
+            count = collections.Counter(sym.title() for sym in symbols)
+
+        else:
+            count = collections.Counter(sym.title() for iat, sym in enumerate(symbols) if iat in self.fragments[ifr])
+
+        count.pop("", None)
+        return dict(count)
+
     def align(
         self,
         ref_mol: "Molecule",

qcelemental/tests/test_molecule.py

@@ -913,3 +913,39 @@
         qcel.models.Molecule(**mol_args)
 
     assert error in str(e.value)
+
+
+_one_helium_mass = 4.00260325413
+
+
+@pytest.mark.parametrize(
+    "mol_string,args,formula,formula_dict,molecular_weight,nelec, nre",
+    [
+        ("He 0 0 0", {}, "He", {"He": 1}, _one_helium_mass, 2, 0.0),
+        ("He 0 0 0\n--\nHe 0 0 5", {}, "He2", {"He": 2}, 2 * _one_helium_mass, 4, 0.4233417684),
+        ("He 0 0 0\n--\n@He 0 0 5", {}, "He2", {"He": 1}, _one_helium_mass, 2, 0.0),
+        ("He 0 0 0\n--\n@He 0 0 5", {"ifr": 0}, "He2", {"He": 1}, _one_helium_mass, 2, 0.0),
+        ("He 0 0 0\n--\n@He 0 0 5", {"ifr": 1}, "He2", {}, 0.0, 0, 0.0),
+        ("He 0 0 0\n--\n@He 0 0 5", {"real_only": False}, "He2", {"He": 2}, 2 * _one_helium_mass, 4, 0.4233417684),
+        ("He 0 0 0\n--\n@He 0 0 5", {"real_only": False, "ifr": 0}, "He2", {"He": 1}, _one_helium_mass, 2, 0.0),
+        ("He 0 0 0\n--\n@He 0 0 5", {"real_only": False, "ifr": 1}, "He2", {"He": 1}, _one_helium_mass, 2, 0.0),
+        ("4He 0 0 0", {}, "He", {"He": 1}, _one_helium_mass, 2, 0.0),
+        ("5He4 0 0 0", {}, "He", {"He": 1}, 5.012057, 2, 0.0),  # suffix-4 is label
+        ("[email protected] 0 0 0", {}, "He", {"He": 1}, 3.14, 2, 0.0),
+    ],
+)
+def test_molecular_weight(mol_string, args, formula, formula_dict, molecular_weight, nelec, nre):
+    mol = Molecule.from_data(mol_string)
+
+    assert mol.molecular_weight(**args) == molecular_weight, f"molecular_weight: ret != {molecular_weight}"
+    assert mol.nelectrons(**args) == nelec, f"nelectrons: ret != {nelec}"
+    assert (abs(mol.nuclear_repulsion_energy(**args) - nre) < 1.0e-5, f"nre: ret != {nre}"
+    assert mol.element_composition(**args) == formula_dict, f"element_composition: ret != {formula_dict}"
+    assert mol.get_molecular_formula() == formula, f"get_molecular_formula: ret != {formula}"
+
+    # after py38
+    # assert (ret := mol.molecular_weight(**args)) == molecular_weight, f"molecular_weight: {ret} != {molecular_weight}"
+    # assert (ret := mol.nelectrons(**args)) == nelec, f"nelectrons: {ret} != {nelec}"
+    # assert (abs(ret := mol.nuclear_repulsion_energy(**args)) - nre) < 1.0e-5, f"nre: {ret} != {nre}"
+    # assert (ret := mol.element_composition(**args)) == formula_dict, f"element_composition: {ret} != {formula_dict}"
+    # assert (ret := mol.get_molecular_formula()) == formula, f"get_molecular_formula: {ret} != {formula}"