announce py38 and add Mol fns (#360)

* Update README.md

* Update README.md

* Update test_molecule.py

* Update test_molecule.py

* Update molecule.py

* Update test_molecule.py

* Update molecule.py

* Update changelog.rst

* Update test_molecule.py

* Update test_molecule.py

* Update Lint.yml

* Update test_molecule.py

Author: loriab

.github/workflows/Lint.yml

@@ -17,6 +17,8 @@ jobs:
           python-version: "3.7"
       - name: Install black
         run: pip install "black>=22.1.0,<23.0a0"
+      - name: Print code formatting with black (hints here if next step errors)
+        run: black --diff .
       - name: Check code formatting with black
         run: black --check .
 

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

docs/changelog.rst

@@ -31,6 +31,10 @@ Breaking Changes
 
 New Features
 ++++++++++++
+- (:pr:`360`) ``Molecule`` learned new functions ``element_composition`` and ``molecular_weight``.
+  The first gives a dictionary of element symbols and counts, while the second gives the weight in amu.
+  Both can access the whole molecule or per-fragment like the existing ``nelectrons`` and
+  ``nuclear_repulsion_energy``. All four can now select all atoms or exclude ghosts (default).
 
 Enhancements
 ++++++++++++

qcelemental/models/molecule.py

@@ -1,7 +1,7 @@
 """
 Molecule Object Model
 """
-
+import collections
 import hashlib
 import json
 import warnings
@@ -875,6 +875,10 @@ def get_molecular_formula(self, order: str = "alphabetical", chgmult: bool = Fal
         >>> 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,21 +1155,26 @@ def _inertial_tensor(geom, *, weight):
         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
         -------
         nre : float
             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 +1187,26 @@ def nuclear_repulsion_energy(self, ifr: int = None) -> float:
                 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 +1216,69 @@ def nelectrons(self, ifr: int = None) -> int:
 
         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 @@ def test_frag_multiplicity_types_errors(mult_in, validate, error):
         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}"