add guess_elements function to workflows modules (base, dihedrals), w…

…ith test

mdpow/tests/test_workflows_base.py

@@ -7,7 +7,9 @@
 
 import pybol
 import pytest
+import numpy as np
 import pandas as pd
+import MDAnalysis as mda
 
 from . import RESOURCES, MANIFEST, STATES
 from pkg_resources import resource_filename
@@ -21,6 +23,19 @@ def molname_workflows_directory(tmp_path):
 
 class TestWorkflowsBase(object):
 
+    check_guessed_elements = np.array(['C', 'N', 'C', 'Cl', 'C'])
+
+    @pytest.fixture(scope='function')
+    def create_universe(self):
+        masses = np.array([12.011, 14.007, 12.011, 35.45, 12.011])
+        names = np.array(["C", "Nx", "C0S", "Cl123", "C0U"])
+
+        # build minimal test universe
+        u = mda.Universe.empty(n_atoms=len(names))
+        u.add_TopologyAttr("names", names)
+        u.add_TopologyAttr("masses", masses)
+        return u
+
     @pytest.fixture(scope='function')
     def SM_tmp_dir(self, molname_workflows_directory):
         dirname = molname_workflows_directory
@@ -43,6 +58,12 @@ def test_df_data(self):
     def project_paths_data(self, SM_tmp_dir):
         project_paths = base.project_paths(parent_directory=SM_tmp_dir)
         return project_paths
+
+    def test_guess_elements(self, create_universe):
+        u = create_universe
+        guessed_elements = base.guess_elements(u.atoms)
+
+        assert guessed_elements == self.check_guessed_elements
 
     def test_project_paths(self, test_df_data, project_paths_data):
         test_df = test_df_data

mdpow/workflows/base.py

@@ -17,14 +17,19 @@
 
 .. autofunction:: project_paths
 .. autofunction:: automated_project_analysis
+.. autofunction:: guess_elements
 
 """
 
 import os
 import re
+import math
 import logging
 
+import numpy as np
 import pandas as pd
+import MDAnalysis as mda
+from MDAnalysis.topology import guessers, tables
 
 logger = logging.getLogger('mdpow.workflows.base')
 
@@ -173,3 +178,34 @@ def automated_project_analysis(project_paths, ensemble_analysis, **kwargs):
 
     logger.info('all analyses completed')
     return
+
+def guess_elements(universe_atom_group):
+    # input from dihedrals module
+    # universe_atom_group = universe.atoms
+    # universe from MDPOW FEP project results
+    names = universe_atom_group.names
+    masses = universe_atom_group.masses
+
+    guessed_elements = guessers.guess_types(names)
+    guessed_masses = np.array([guessers.get_atom_mass(n) for n in guessed_elements])
+
+    problem_cases = []
+    i=0
+    while i < len(masses):
+        if masses[i] != 0:
+            value = math.isclose(guessed_masses[i], masses[i], rel_tol=0.001)
+            value = not(value)
+        else:
+            value = False
+        problem_cases.append(value)
+        i+=1
+    problems = np.array(problem_cases)
+
+    problem_elements = []
+    for m in masses[problems]:
+        element = [k for k,v in tables.masses.items() if  np.isclose(m, v, rtol=0.001)][0]
+        problem_elements.append(element)
+
+    guessed_elements[problems] = problem_elements
+
+    return guessed_elements

mdpow/workflows/dihedrals.py

@@ -27,7 +27,6 @@
     :annotation: = 190
 .. autofunction:: automated_dihedral_analysis
 .. autofunction:: build_universe
-.. autofunction:: guess_atoms
 .. autofunction:: rdkit_conversion
 .. autofunction:: get_atom_indices
 .. autofunction:: get_bond_indices
@@ -148,9 +147,6 @@ def build_universe(dirname, solvents=SOLVENTS_DEFAULT):
 
     return u
 
-def guess_atoms():
-    return
-
 def rdkit_conversion(u, resname):
     """Converts the solute, `resname`, of the
        :class:`~MDAnalysis.core.universe.Universe` to :class:`rdkit.Chem.rdchem.Mol` object
@@ -185,11 +181,16 @@ def rdkit_conversion(u, resname):
 
     """
 
+    from .base import guess_elements
+    # temporary
+    # keep here to avoid circular imports?
+
     try:
         solute = u.select_atoms(f'resname {resname}')
         mol = solute.convert_to('RDKIT')
     except AttributeError:
-        u.add_TopologyAttr("elements", guess_types(u.atoms.names))
+        guessed_elements = guess_elements(u.atoms)
+        u.add_TopologyAttr("elements", guessed_elements)
         solute = u.select_atoms(f'resname {resname}')
         mol = solute.convert_to('RDKIT')