Repository: openforcefield/openff-evaluator

Issue:
Add easy way to create new substance for specific number of molecules
A Substance may be created with mole fractions of 0.5098 for one component and 0.4902 for another, but when actually creating a box of 100 molecules, this gets turned into actual mole fractions of 0.51 and 0.49 respectively. This will pose an issue when querying storage backends for existing equilibrated boxes, as what's in the storage backend (i.e. the latter, simpler mole fractions) won't match the exact substance in the query. Instead, there should be an easy way to create a new, simpler substance matching the total number of molecules for the query.

Discussion:

Patch:
diff --git a/openff/evaluator/substances/substances.py b/openff/evaluator/substances/substances.py
index 44afb8b..fc4bb80 100644
--- a/openff/evaluator/substances/substances.py
+++ b/openff/evaluator/substances/substances.py
@@ -220,6 +220,107 @@ class Substance(AttributeClass):
 
         return self.amounts[identifier]
 
+    def to_substance_n_molecules(
+        self,
+        maximum_molecules: int,
+        tolerance=None,
+        count_exact_amount=True,
+        truncate_n_molecules=True,
+    ) -> "Substance":
+        """
+        Returns a new substance given a maximum total number of molecules.
+
+        Parameters
+        ----------
+        maximum_molecules: int
+            The maximum number of molecules.
+        tolerance: float, optional
+            The tolerance within which this amount should be represented. As
+            an example, when converting a mole fraction into a number of molecules,
+            the total number of molecules may not be sufficiently large enough to
+            reproduce this amount.
+        count_exact_amount: bool
+            Whether components present in an exact amount (i.e. defined with an
+            ``ExactAmount``) should be considered when apply the maximum number
+             of molecules constraint. This may be set false, for example, when
+             building a separate solvated protein (n = 1) and solvated protein +
+             ligand complex (n = 2) system but wish for both systems to have the
+             same number of solvent molecules.
+        truncate_n_molecules: bool
+            Whether or not to attempt to truncate the number of molecules in the
+            substance if the total number is over the specified maximum. If False, an
+            exception will be raised in this case.
+
+            The truncation works by iteratively removing one molecule of the
+            predominant component up to a limit of removing a total number of molecules
+            equal to the number of components  in the substance (e.g. for a binary
+            substance a maximum of two molecules can be removed). An exception is
+            raised if the number of molecules cannot be sensibly truncated.
+        """
+
+        molecules_per_component = self.get_molecules_per_component(
+            maximum_molecules,
+            tolerance=tolerance,
+            count_exact_amount=count_exact_amount,
+            truncate_n_molecules=truncate_n_molecules,
+        )
+
+        new_amounts = defaultdict(list)
+        total_number_of_molecules = sum(molecules_per_component.values())
+
+        # Handle any exact amounts.
+        for component in self.components:
+            exact_amounts = [
+                amount
+                for amount in self.get_amounts(component)
+                if isinstance(amount, ExactAmount)
+            ]
+
+            if len(exact_amounts) == 0:
+                continue
+
+            total_number_of_molecules -= exact_amounts[0].value
+            new_amounts[component].append(exact_amounts[0])
+
+        # Recompute the mole fractions.
+        total_mole_fraction = 0.0
+        number_of_new_mole_fractions = 0
+
+        for component in self.components:
+            mole_fractions = [
+                amount
+                for amount in self.get_amounts(component)
+                if isinstance(amount, MoleFraction)
+            ]
+
+            if len(mole_fractions) == 0:
+                continue
+
+            molecule_count = molecules_per_component[component.identifier]
+
+            if component in new_amounts:
+                molecule_count -= new_amounts[component][0].value
+
+            new_mole_fraction = molecule_count / total_number_of_molecules
+            new_amounts[component].append(MoleFraction(new_mole_fraction))
+
+            total_mole_fraction += new_mole_fraction
+            number_of_new_mole_fractions += 1
+
+        if (
+            not np.isclose(total_mole_fraction, 1.0)
+            and number_of_new_mole_fractions > 0
+        ):
+            raise ValueError("The new mole fraction does not equal 1.0")
+
+        output_substance = Substance()
+
+        for component, amounts in new_amounts.items():
+            for amount in amounts:
+                output_substance.add_component(component, amount)
+
+        return output_substance
+
     def get_molecules_per_component(
         self,
         maximum_molecules,


Test patch:
diff --git a/openff/evaluator/_tests/test_substances.py b/openff/evaluator/_tests/test_substances.py
index 61f7237..55aaad9 100644
--- a/openff/evaluator/_tests/test_substances.py
+++ b/openff/evaluator/_tests/test_substances.py
@@ -98,3 +98,14 @@ def test_truncate_n_molecules():
 def test_substance_len():
     substance = Substance.from_components("C", "CC", "CCC", "CCC")
     assert len(substance) == 3
+
+
+def test_to_substance_n_molecules():
+    substance = Substance()
+    substance.add_component(Component("C"), MoleFraction(0.509800))
+    substance.add_component(Component("CC"), MoleFraction(0.490200))
+
+    new_substance = substance.to_substance_n_molecules(100)
+    assert new_substance.number_of_components == 2
+    assert new_substance.get_amounts(substance.components[0])[0].value == 0.51
+    assert new_substance.get_amounts(substance.components[1])[0].value == 0.49