diff --git a/mbuild/compound.py b/mbuild/compound.py index b09569b3f..453ed1053 100644 --- a/mbuild/compound.py +++ b/mbuild/compound.py @@ -275,9 +275,10 @@ def _n_particles(self, include_ports=False): return sum(1 for _ in self._particles(include_ports)) def _contains_only_ports(self): - for part in self.children: - if not part.port_particle: - return False + if self.children: + for part in self.children: + if not part.port_particle: + return False return True def print_hierarchy(self, print_full=False, index=None, show_tree=True): @@ -547,7 +548,17 @@ def mass(self, value): @property def charge(self): - """Get the charge of the Compound.""" + """Return the total charge of a compound. + + If the compound contains children compouds, the total charge of all + children compounds is returned. + + If the charge of a particle has not been explicitly set + then the particle's charge is None, and are not used when + calculating the total charge. + """ + if self._contains_only_ports(): + return self._particle_charge(self) charges = [p._charge for p in self.particles()] if None in charges: warn( @@ -557,6 +568,11 @@ def charge(self): filtered_charges = [charge for charge in charges if charge is not None] return sum(filtered_charges) if filtered_charges else None + @staticmethod + def _particle_charge(particle): + """Return charge of a Compound with no children.""" + return particle._charge + @charge.setter def charge(self, value): if self._contains_only_ports(): diff --git a/mbuild/tests/test_compound.py b/mbuild/tests/test_compound.py index 99c19e1a5..bd80486d1 100644 --- a/mbuild/tests/test_compound.py +++ b/mbuild/tests/test_compound.py @@ -1702,16 +1702,15 @@ def test_charge_neutrality_warn(self, benzene): def test_none_charge(self): A = mb.Compound() - with pytest.warns(UserWarning): - A.charge + assert A.charge is None A.charge = 1 B = mb.Compound() container = mb.Compound(subcompounds=[A, B]) + assert A.charge == 1 + assert B.charge is None with pytest.warns(UserWarning): container_charge = container.charge - assert A.charge == 1 - assert B.charge == None assert container_charge == 1 @pytest.mark.skipif(not has_openbabel, reason="Open Babel not installed")