Merge pull request #169 from Autodesk/topology_consistency

Copyable wavefunctions and properties

moldesign/_tests/helpers.py

@@ -67,6 +67,12 @@ def native_str_buffer(*args, **kwargs):
         return io.StringIO(*args, **kwargs)
 
 
+if PY2:
+    PICKLE_PROTOCOLS = list(range(3))
+else:
+    PICKLE_PROTOCOLS = list(range(5))
+
+
 class ZeroEnergy(mdt.models.base.EnergyModelBase):
     """ All 0, all the time
     """

moldesign/_tests/molecule_fixtures.py

@@ -286,3 +286,16 @@ def cached_protein_with_default_amber_ff(cached_pdb1yu8):
 @typedfixture('hasmodel')
 def protein_default_amber_forcefield(cached_protein_with_default_amber_ff):
     return cached_protein_with_default_amber_ff.copy()
+
+
+@pytest.fixture(scope='session')
+def cached_h2_rhfwfn():
+    mol = h2()  # fixture is not cached, so just call it directly
+    mol.set_energy_model(mdt.models.PySCFPotential, basis='sto-3g', theory='rhf')
+    mol.calculate(requests=['forces'])
+    return mol
+
+
+@pytest.fixture
+def h2_rhfwfn(cached_h2_rhfwfn):
+    return cached_h2_rhfwfn.copy()

moldesign/_tests/object_fixtures.py

@@ -31,39 +31,38 @@ def fixture_wrapper(func):
 
 
 ######################################
-# Test the basic data structures
-
+# Various python objects
 TESTDICT = collections.OrderedDict((('a', 'b'),
                                     ('c', 3),
                                     ('d', 'e'),
                                     ('a', 1),
                                     (3, 35)))
 
 
-@typedfixture('object')
+@typedfixture('pickleable')
 def dotdict():
     dd = utils.DotDict(TESTDICT)
     return dd
 
 
 # Some objects with units
-@typedfixture('object')
+@typedfixture('pickleable')
 def list_of_units():
     return [1.0 * u.angstrom, 1.0 * u.nm, 1.0 * u.a0]
 
 
-@typedfixture('object', 'equality')
+@typedfixture('pickleable', 'equality')
 def simple_unit_array():
     return np.array([1.0, -2.0, 3.5]) * u.angstrom
 
 
-@typedfixture('object', 'equality')
+@typedfixture('pickleable', 'equality')
 def unit_number():
     return 391.23948 * u.ureg.kg * u.ang / u.alpha
 
 
 ######################################
-# Test underlying elements
+# Atom objects
 @typedfixture('atom')
 def carbon_atom():
     atom1 = mdt.Atom('C')
@@ -82,9 +81,7 @@ def carbon_copy(carbon_atom):
 
 
 ######################################
-# Tests around hydrogen
-
-
+# Hydrogen-related objects
 @typedfixture('molecule')
 def h2_harmonic(h2):
     mol = h2
@@ -96,6 +93,26 @@ def h2_harmonic(h2):
     return mol
 
 
+@typedfixture('pickleable')
+def atom_bond_graph(h2):
+    return h2.bond_graph[h2.atoms[0]]
+
+
+@typedfixture('pickleable')
+def mol_bond_graph(h2):
+    return h2.bond_graph
+
+
+@typedfixture('pickleable')
+def mol_wfn(h2_rhfwfn):
+    return h2_rhfwfn.copy().wfn
+
+
+@typedfixture('pickleable')
+def mol_properties(h2_rhfwfn):
+    return h2_rhfwfn.copy().properties
+
+
 @typedfixture('trajectory')
 def h2_trajectory(h2_harmonic):
     mol = h2_harmonic
@@ -135,4 +152,4 @@ def h2_harmonic_atoms(h2_harmonic):
                      registered_types['submolecule'] +
                      registered_types['trajectory'] +
                      registered_types['atom'])
-all_objects = registered_types['object'] + moldesign_objects
+pickleable = registered_types['pickleable'] + moldesign_objects

moldesign/_tests/test_atom_bond_computed_properties.py

@@ -5,8 +5,6 @@
 
 from .molecule_fixtures import *
 from . import helpers
-from .test_qm_xfaces import h2_rhfwfn
-
 
 
 def test_forcefield_atom_term_access(protein_default_amber_forcefield):

moldesign/_tests/test_copies.py

@@ -231,3 +231,52 @@ def test_constraints_copied_with_molecule(mol_with_zerocharge_params):
         if constraint.desc != 'hbonds':
             for atom in constraint.atoms:
                 assert atom.molecule is mcpy
+
+
+def test_properties_copied_with_molecule(cached_h2_rhfwfn):
+    original = cached_h2_rhfwfn
+    assert original.potential_energy is not None  # sanity check
+
+    mol = cached_h2_rhfwfn.copy()
+
+    assert mol is not original
+    assert mol.properties is not original.properties
+
+    for prop, val in original.properties.items():
+        assert prop in mol.properties
+        if isinstance(val, (u.MdtQuantity, np.ndarray)) and getattr(val, 'shape', False):
+                assert (mol.properties[prop] == val).all()
+                assert mol.properties[prop] is not val
+
+        elif isinstance(val, (str, int, float, mdt.molecules.AtomicProperties)):
+            assert mol.properties[prop] == val
+
+        else:  # otherwise, just make sure it's not the original
+            assert mol.properties[prop] is not val
+
+
+def test_wfn_copied_with_molecule(cached_h2_rhfwfn):
+    original = cached_h2_rhfwfn
+    assert original.wfn is not None  # sanity check
+
+    mol = original.copy()
+
+    assert mol.wfn is not None
+
+    # should be completely equal
+    assert (mol.wfn.aobasis.fock == original.wfn.aobasis.fock).all()
+    # but different objects
+    assert mol.wfn.aobasis.fock is not original.wfn.aobasis.fock
+
+
+def test_wfn_copy(cached_h2_rhfwfn):
+    original = cached_h2_rhfwfn
+    wfn = original.wfn.copy()
+
+    assert wfn.mol is original
+    assert wfn is not original.wfn
+
+    # should be completely equal
+    assert (wfn.aobasis.fock == original.wfn.aobasis.fock).all()
+    # but different objects
+    assert wfn.aobasis.fock is not original.wfn.aobasis.fock

moldesign/_tests/test_molecules.py

@@ -15,6 +15,7 @@
 __PYTEST_MARK__ = 'internal'  # mark all tests in this module with this label (see ./conftest.py)
 
 
+
 def test_h2_protected_atom_arrays(h2):
     atom1, atom2 = h2.atoms
     with pytest.raises(TypeError):
@@ -225,15 +226,6 @@ def test_molecule_types(molkey, request):
         assert issubclass(type(residue), mdt.Residue)
 
 
-@pytest.mark.parametrize('objkey', all_objects)
-def test_pickling(objkey, request):
-    obj = request.getfixturevalue(objkey)
-    for iprotocol in (0,1,2):
-        x = pickle.dumps(obj, protocol=iprotocol)
-        y = pickle.loads(x)
-        assert type(y) == type(obj)
-
-
 def test_degrees_of_freedom(benzene):
     assert benzene.dynamic_dof == 36
 
@@ -242,7 +234,7 @@ def test_degrees_of_freedom(benzene):
 def test_pickled_equality(objkey, request):
     obj = request.getfixturevalue(objkey)
 
-    for iprotocol in (0,1,2):
+    for iprotocol in helpers.PICKLE_PROTOCOLS:
         x = pickle.dumps(obj, protocol=iprotocol)
         y = pickle.loads(x)
         assert type(y) == type(obj)

moldesign/_tests/test_objects.py

@@ -1,5 +1,8 @@
+import pickle
+
 from moldesign.utils import Alias
 from .object_fixtures import *
+from .molecule_fixtures import *
 
 
 __PYTEST_MARK__ = 'internal'  # mark all tests in this module with this label (see ./conftest.py)
@@ -12,6 +15,15 @@ def __init__(self):
         self.s = 'ABC'
 
 
+@pytest.mark.parametrize('objkey', pickleable)
+def test_pickling(objkey, request):
+    obj = request.getfixturevalue(objkey)
+    for iprotocol in (0,1,2):
+        x = pickle.dumps(obj, protocol=iprotocol)
+        y = pickle.loads(x)
+        assert type(y) == type(obj)
+
+
 def test_alias():
     t = ComposedClass()
     assert t.delegated() == 'abc'

moldesign/_tests/test_qm_xfaces.py

@@ -69,13 +69,6 @@ def nwchem_dft():
     return mdt.models.NWChemQM(basis='sto-3g', theory='rks', functional='b3lyp')
 
 
-@pytest.fixture
-def h2_rhfwfn(h2):
-    h2.set_energy_model(mdt.models.PySCFPotential, basis='sto-3g', theory='rhf')
-    h2.calculate(requests=['forces'])
-    return h2
-
-
 @pytest.mark.parametrize('objkey', TESTSET)
 def test_pyscf_rhf_sto3g_properties(objkey, request):
     mol = request.getfixturevalue(objkey)

moldesign/_tests/test_units.py

@@ -133,6 +133,32 @@ def test_array_unit_checks():
                                np.arange(10, 15))
 
 
+@pytest.mark.parametrize(['a1','a2'],
+                         ((np.ones(3), np.ones(3)*10/10),
+                          (np.ones(3), u.array(np.ones(3))),
+                          (np.ones(3)*u.nm, 10.0*np.ones(3)*u.angstrom),
+                          (np.ones((3,3)) * u.nm / u.angstrom, np.ones(3) * 10)),
+                         ids="numpy-numpy numpy-dimensionless nm-angstrom nm/angstrom-numpy".split()
+                         )
+def test_array_almost_equal_returns_true(a1, a2):
+    assert u.arrays_almost_equal(a1, a2)
+    assert u.arrays_almost_equal(a2, a1)
+
+
+@pytest.mark.parametrize(['a1','a2'],
+                         ((np.ones(3), np.ones(3)*10/10 * u.eV),
+                          (np.ones(3)*u.nm, 10.0*np.ones(3)*u.dalton),
+                          (u.array(np.ones(3)), np.ones(3) * u.kcalpermol)),
+                         ids="numpy-ev nm-dalton dimensionless-kcalpermole".split()
+                         )
+def test_array_almost_equal_raises_dimensionality_error(a1, a2):
+    with pytest.raises(u.DimensionalityError):
+        u.arrays_almost_equal(a1, a2)
+    with pytest.raises(u.DimensionalityError):
+        u.arrays_almost_equal(a2, a1)
+
+
+
 @pytest.mark.screening
 def test_default_unit_conversions():
     assert abs(10.0 - (1.0*units.nm).defunits_value()) < 1e-10

moldesign/interfaces/ambertools.py

@@ -26,6 +26,7 @@
 from .. import compute, utils
 from .. import units as u
 from .. import forcefields
+from ..molecules import AtomicProperties
 
 IMAGE = 'ambertools'
 
@@ -103,7 +104,7 @@ def _antechamber_calc_charges(mol, ambname, chargename, kwargs):
     def finish_job(job):
         """Callback to complete the job"""
         lines = iter(job.get_output('out.mol2').read().split('\n'))
-        charges = utils.DotDict(type='atomic')
+        charges = {}
 
         line = next(lines)
         while line.strip()[:len('@<TRIPOS>ATOM')] != '@<TRIPOS>ATOM':
@@ -117,7 +118,7 @@ def finish_job(job):
             charges[mol.atoms[idx]] = u.q_e*float(fields[-1])
             line = next(lines)
 
-        mol.properties[chargename] = charges
+        mol.properties[chargename] = AtomicProperties(charges)
         return charges
 
     job = pyccc.Job(image=mdt.compute.get_image_path(IMAGE),

moldesign/models/pyscf.py

@@ -28,6 +28,7 @@
 from ..interfaces.pyscf_interface import force_remote, mol_to_pyscf,  StatusLogger, SPHERICAL_NAMES
 from .base import QMBase
 from ..helpers import Logger
+from ..molecules import AtomicProperties
 
 if future.utils.PY2:
     from cStringIO import StringIO
@@ -231,8 +232,8 @@ def _get_properties(self, ref, kernel, grad):
         scf_matrices = self._get_scf_matrices(orb_calc, ao_matrices)
         if hasattr(orb_calc, 'mulliken_pop'):
             ao_pop, atom_pop = orb_calc.mulliken_pop(verbose=-1)
-            result['mulliken'] = utils.DotDict({a: p * u.q_e for a, p in zip(self.mol.atoms, atom_pop)})
-            result['mulliken'].type = 'atomic'
+            result['mulliken'] = AtomicProperties({a: p * u.q_e
+                                                   for a, p in zip(self.mol.atoms, atom_pop)})
 
         if hasattr(orb_calc, 'dip_moment'):
             result['dipole_moment'] = orb_calc.dip_moment() * u.debye

moldesign/molecules/__init__.py

@@ -4,6 +4,7 @@ def toplevel(o):
 __all__ = []
 
 
+from .properties import *
 from .bond_graph import *
 from .coord_arrays import *
 from .atomcollections import *

moldesign/molecules/atoms.py

@@ -73,7 +73,10 @@ def basis_functions(self):
         except mdt.exceptions.NotCalculatedError:
             return None
 
-        return wfn.aobasis.on_atom.get(self, [])
+        try:
+            return wfn.aobasis.get_basis_functions_on_atom(self)
+        except KeyError:
+            return None
 
     @property
     def properties(self):