Merge pull request #33 from chrisjonesBSU/scale-charges

Fix scale charges, add unit tests

hoomd_organics/base/system.py

@@ -137,9 +137,7 @@ def n_molecules(self):
 
     @property
     def n_particles(self):
-        if self.gmso_system:
-            return self.gmso_system.n_sites
-        return sum([mol.n_particles for mol in self.all_molecules])
+        return self.gmso_system.n_sites
 
     @property
     def mass(self):
@@ -150,6 +148,12 @@ def mass(self):
             )
         return sum(mol.mass for mol in self.all_molecules)
 
+    @property
+    def net_charge(self):
+        return sum(
+            site.charge if site.charge else 0 for site in self.gmso_system.sites
+        )
+
     @property
     def box(self):
         return self.system.box
@@ -299,6 +303,22 @@ def _remove_hydrogens(self):
         if len(hydrogens) > 0:
             self.gmso_system = from_parmed(parmed_struc)
 
+    def _scale_charges(self):
+        """"""
+        charges = np.array(
+            [
+                site.charge if site.charge else 0
+                for site in self.gmso_system.sites
+            ]
+        )
+        net_charge = sum(charges)
+        abs_charge = sum(abs(charges))
+        if abs_charge != 0:
+            for site in self.gmso_system.sites:
+                site.charge -= abs(site.charge if site.charge else 0) * (
+                    net_charge / abs_charge
+                )
+
     def to_gsd(self, file_name):
         with gsd.hoomd.open(file_name, "wb") as traj:
             traj.append(self.hoomd_snapshot)
@@ -344,8 +364,7 @@ def _apply_forcefield(self):
             for site in self.gmso_system.sites:
                 site.charge = 0
         if self.scale_charges and not self.remove_charges:
-            pass
-            # TODO: Scale charges from self.gmso_system
+            self._scale_charges()
         epsilons = [
             s.atom_type.parameters["epsilon"] for s in self.gmso_system.sites
         ]

hoomd_organics/tests/base/test_system.py

@@ -549,3 +549,25 @@ def test_lattice_molecule(self, benzene_molecule):
         assert len(system.hoomd_forcefield) > 0
         assert system.n_particles == system.hoomd_snapshot.particles.N
         assert system.reference_values.keys() == {"energy", "length", "mass"}
+
+    def test_scale_charges(self, pps):
+        pps_mol = pps(num_mols=5, lengths=5)
+        no_scale = Pack(
+            molecules=pps_mol,
+            density=0.5,
+            r_cut=2.4,
+            force_field=OPLS_AA_PPS(),
+            auto_scale=True,
+            scale_charges=False,
+        )
+
+        with_scale = Pack(
+            molecules=pps_mol,
+            density=0.5,
+            r_cut=2.4,
+            force_field=OPLS_AA_PPS(),
+            auto_scale=True,
+            scale_charges=True,
+        )
+        assert abs(no_scale.net_charge.value) > abs(with_scale.net_charge.value)
+        assert np.allclose(0, with_scale.net_charge.value, atol=1e-30)

hoomd_organics/tests/base_test.py

@@ -159,6 +159,26 @@ def __init__(self, lengths, num_mols, **kwargs):
 
         return _PolyEthylene
 
+    @pytest.fixture()
+    def pps(self, pps_smiles):
+        class _PPS(Polymer):
+            def __init__(self, lengths, num_mols, **kwargs):
+                smiles = pps_smiles
+                bond_indices = [7, 10]
+                bond_length = 0.176
+                bond_orientation = [[0, 0, 1], [0, 0, -1]]
+                super().__init__(
+                    lengths=lengths,
+                    num_mols=num_mols,
+                    smiles=smiles,
+                    bond_indices=bond_indices,
+                    bond_length=bond_length,
+                    bond_orientation=bond_orientation,
+                    **kwargs
+                )
+
+        return _PPS
+
     @pytest.fixture()
     def polyDME(self, dimethylether_smiles):
         class _PolyDME(Polymer):

hoomd_organics/utils/__init__.py

@@ -1,4 +1,4 @@
 from .actions import *
 from .base_types import FF_Types
 from .ff_utils import xml_to_gmso_ff
-from .utils import check_return_iterable, scale_charges
+from .utils import check_return_iterable

hoomd_organics/utils/utils.py

@@ -8,13 +8,3 @@ def check_return_iterable(obj):
         return obj
     except:  # noqa: E722
         return [obj]
-
-
-def scale_charges(charges):
-    net_charge = sum(charges)
-    abs_charge = sum([abs(charge) for charge in charges])
-    scaled_charges = []
-    for idx, charge in enumerate(charges):
-        new_charge = charge - (abs(charge) * (net_charge / abs_charge))
-        scaled_charges.append(new_charge)
-    return scaled_charges