Add ability to use number density when calculating target box lenghts for shrinking and NVT.

flowermd/base/simulation.py

@@ -13,10 +13,8 @@
     HOOMDThermostats,
     StdOutLogger,
     UpdateWalls,
-    calculate_box_length,
     validate_ref_value,
 )
-from flowermd.utils.exceptions import ReferenceUnitError
 
 
 class Simulation(hoomd.simulation.Simulation):
@@ -586,12 +584,11 @@ def remove_walls(self, wall_axis):
 
     def run_update_volume(
         self,
+        final_box_lengths,
         n_steps,
         period,
         kT,
         tau_kt,
-        final_box_lengths=None,
-        final_density=None,
         thermalize_particles=True,
         write_at_start=True,
     ):
@@ -604,6 +601,8 @@ def run_update_volume(
 
         Parameters
         ----------
+        final_box_lengths : np.ndarray, shape=(3,), dtype=float, default None
+            The final box edge lengths in (x, y, z) order.
         n_steps : int, required
             Number of steps to run during volume update.
         period : int, required
@@ -612,52 +611,24 @@ def run_update_volume(
             The temperature to use during volume update.
         tau_kt : float, required
             Thermostat coupling period (in simulation time units).
-        final_box_lengths : np.ndarray, shape=(3,), dtype=float, default None
-            The final box edge lengths in (x, y, z) order.
-        final_density : float, default None
-            The final density of the simulation in g/cm^3.
         write_at_start : bool, default True
             When set to True, triggers writers that evaluate to True
             for the initial step to execute before the next simulation
             time step.
 
         """
-        if final_box_lengths is None and final_density is None:
-            raise ValueError(
-                "Must provide either `final_box_lengths` or `final_density`"
-            )
-        if final_box_lengths is not None and final_density is not None:
-            raise ValueError(
-                "Cannot provide both `final_box_lengths` and `final_density`."
-            )
-        if final_box_lengths is not None:
-            final_box = hoomd.Box(
-                Lx=final_box_lengths[0],
-                Ly=final_box_lengths[1],
-                Lz=final_box_lengths[2],
-            )
+        if self.reference_length and hasattr(final_box_lengths, "to"):
+            ref_unit = self.reference_length.units
+            final_box_lengths = final_box_lengths.to(ref_unit)
+            final_box_lengths /= self.reference_length
         else:
-            if not self.reference_values:
-                raise ReferenceUnitError(
-                    "Missing simulation units. Please "
-                    "provide units for mass, length, and"
-                    " energy."
-                )
-
-            if isinstance(final_density, u.unyt_quantity):
-                density_quantity = final_density.to(u.g / u.cm**3)
-            else:
-                density_quantity = u.unyt_quantity(
-                    final_density, u.g / u.cm**3
-                )
-            mass_g = self.mass.to("g")
-            L = calculate_box_length(mass_g, density_quantity)
-            # convert L from cm to reference units
-            L = (
-                L.to(self.reference_length.units) / self.reference_length.value
-            ).value
-            final_box = hoomd.Box(Lx=L, Ly=L, Lz=L)
-
+            # TODO: Do we need to do anything here?
+            pass
+        final_box = hoomd.Box(
+            Lx=final_box_lengths[0],
+            Ly=final_box_lengths[1],
+            Lz=final_box_lengths[2],
+        )
         resize_trigger = hoomd.trigger.Periodic(period)
         box_ramp = hoomd.variant.Ramp(
             A=0, B=1, t_start=self.timestep, t_ramp=int(n_steps)

flowermd/base/system.py

@@ -14,8 +14,9 @@
 from flowermd.base.forcefield import BaseHOOMDForcefield, BaseXMLForcefield
 from flowermd.base.molecule import Molecule
 from flowermd.utils import (
-    calculate_box_length,
     check_return_iterable,
+    get_target_box_mass_density,
+    get_target_box_number_density,
     validate_ref_value,
 )
 from flowermd.utils.exceptions import ForceFieldError, MoleculeLoadError
@@ -43,11 +44,6 @@ class System(ABC):
     ----------
     molecules : flowermd.Molecule or a list of Molecule objects, required
         The molecules to be placed in the system.
-    density : float, required
-        The desired density of the system (g/cm^3). Used to set the
-        target_box attribute. Can be useful when initializing
-        systems at low density and running a shrink simulation
-        to achieve a target density.
     base_units : dict, default {}
         Dictionary of base units to use for scaling.
         Dictionary keys are "length", "mass", and "energy". Values should be an
@@ -71,18 +67,15 @@ class System(ABC):
     def __init__(
         self,
         molecules,
-        density: float,
         base_units=dict(),
     ):
         self._molecules = check_return_iterable(molecules)
-        self.density = density
         self.all_molecules = []
         self.gmso_system = None
         self._reference_values = base_units
         self._hoomd_snapshot = None
         self._hoomd_forcefield = []
         self._gmso_forcefields_dict = dict()
-        self._target_box = None
         # Reference values used when last writing snapshot and forcefields
         self._ff_refs = dict()
         self._snap_refs = dict()
@@ -147,17 +140,17 @@ def n_molecules(self):
     @property
     def n_particles(self):
         """Total number of particles in the system."""
-        return self.gmso_system.n_sites
+        if self.gmso_system:
+            return self.gmso_system.n_sites
+        else:
+            return sum(mol.n_particles for mol in self.all_molecules)
 
     @property
     def mass(self):
         """Total mass of the system in amu."""
         if self.gmso_system:
-            return sum(
-                float(site.mass.to("amu").value)
-                for site in self.gmso_system.sites
-            )
-        return sum(mol.mass for mol in self.all_molecules)
+            return sum(site.mass.to("amu") for site in self.gmso_system.sites)
+        return sum(mol.mass * u.Unit("amu") for mol in self.all_molecules)
 
     @property
     def net_charge(self):
@@ -325,24 +318,6 @@ def hoomd_forcefield(self):
             )
         return self._hoomd_forcefield
 
-    @property
-    def target_box(self):
-        """The target box size of the system in form of a numpy array.
-
-        If reference length is set, the target box is in reduced units.
-
-        Notes
-        -----
-        The `target_box` property can be passed to
-        `flowermd.base.Simulation.run_update_volume` method to reach the
-        target density.
-
-        """
-        if self.reference_length:
-            return self._target_box / self.reference_length.value
-        else:
-            return self._target_box
-
     def remove_hydrogens(self):
         """Call this method to remove hydrogen atoms from the system.
 
@@ -622,38 +597,6 @@ def apply_forcefield(
         )
         self._hoomd_snapshot = self._create_hoomd_snapshot()
 
-    def set_target_box(
-        self, x_constraint=None, y_constraint=None, z_constraint=None
-    ):
-        """Set the target box size of the system.
-
-        If no constraints are set, the target box is cubic.
-        Setting constraints will hold those box vectors
-        constant and adjust others to match the target density.
-
-        Parameters
-        ----------
-        x_constraint : float, optional, defualt=None
-            Fixes the box length (nm) along the x axis.
-        y_constraint : float, optional, default=None
-            Fixes the box length (nm) along the y axis.
-        z_constraint : float, optional, default=None
-            Fixes the box length (nm) along the z axis.
-
-        """
-        if not any([x_constraint, y_constraint, z_constraint]):
-            Lx = Ly = Lz = self._calculate_L()
-        else:
-            constraints = np.array([x_constraint, y_constraint, z_constraint])
-            fixed_L = constraints[np.not_equal(constraints, None).nonzero()]
-            # Conv from nm to cm for _calculate_L
-            fixed_L *= 1e-7
-            L = self._calculate_L(fixed_L=fixed_L)
-            constraints[np.equal(constraints, None).nonzero()] = L
-            Lx, Ly, Lz = constraints
-
-        self._target_box = np.array([Lx, Ly, Lz])
-
     def visualize(self):
         """Visualize the system."""
         if self.system:
@@ -663,31 +606,6 @@ def visualize(self):
                 "The initial configuraiton has not been created yet."
             )
 
-    def _calculate_L(self, fixed_L=None):
-        """Calculate the box length.
-
-        Calculate the required box length(s) given the mass of a system and
-        the target density.
-        Box edge length constraints can be set by set_target_box().
-        If constraints are set, this will solve for the required
-        lengths of the remaining non-constrained edges to match
-        the target density.
-
-        Parameters
-        ----------
-        fixed_L : np.array, optional, defualt=None
-            Array of fixed box lengths to be accounted for when solving for L.
-
-        """
-        mass_quantity = u.unyt_quantity(self.mass, u.g / u.mol).to("g")
-        density_quantity = u.unyt_quantity(self.density, u.g / u.cm**3)
-        if fixed_L is not None:
-            fixed_L = u.unyt_array(fixed_L, u.cm)
-        L = calculate_box_length(
-            mass_quantity, density_quantity, fixed_L=fixed_L
-        )
-        return L.to("nm").value
-
 
 class Pack(System):
     """Uses PACKMOL via mbuild.packing.fill_box.
@@ -697,6 +615,11 @@ class Pack(System):
 
     Parameters
     ----------
+    density : float, required
+        The desired density of the system (g/cm^3). Used to set the
+        target_box attribute. Can be useful when initializing
+        systems at low density and running a shrink simulation
+        to achieve a target density.
     packing_expand_factor : int, default 5
         The factor by which to expand the box for packing.
     edge : float, default 0.2
@@ -731,21 +654,42 @@ def __init__(
         edge=0.2,
         overlap=0.2,
     ):
+        if not isinstance(density, u.array.unyt_quantity):
+            self.density = density * u.Unit("g") / u.Unit("cm**3")
+            warnings.warn(
+                "Units for density were not given, assuming "
+                "units of g/cm**3."
+            )
+        else:
+            self.density = density
         self.packing_expand_factor = packing_expand_factor
         self.edge = edge
         self.overlap = overlap
-        super(Pack, self).__init__(
-            molecules=molecules,
-            density=density,
-            base_units=base_units,
-        )
+        super(Pack, self).__init__(molecules=molecules, base_units=base_units)
 
     def _build_system(self):
-        self.set_target_box()
+        mass_density = u.Unit("kg") / u.Unit("m**3")
+        number_density = u.Unit("m**-3")
+        if self.density.units.dimensions == mass_density.dimensions:
+            target_box = get_target_box_mass_density(
+                density=self.density, mass=self.mass
+            ).to("nm")
+        elif self.density.units.dimensions == number_density.dimensions:
+            target_box = get_target_box_number_density(
+                density=self.density, n_beads=self.n_particles
+            ).to("nm")
+        else:
+            raise ValueError(
+                f"Density dimensions of {self.density.units.dimensions} "
+                "were given, but only mass density "
+                f"({mass_density.dimensions}) and "
+                f"number density ({number_density.dimensions}) are supported."
+            )
+
         system = mb.packing.fill_box(
             compound=self.all_molecules,
             n_compounds=[1 for i in self.all_molecules],
-            box=list(self._target_box * self.packing_expand_factor),
+            box=list(target_box * self.packing_expand_factor),
             overlap=self.overlap,
             edge=self.edge,
         )
@@ -773,7 +717,6 @@ class Lattice(System):
     def __init__(
         self,
         molecules,
-        density: float,
         x: float,
         y: float,
         n: int,
@@ -786,7 +729,6 @@ def __init__(
         self.basis_vector = basis_vector
         super(Lattice, self).__init__(
             molecules=molecules,
-            density=density,
             base_units=base_units,
         )
 
@@ -820,7 +762,6 @@ def _build_system(self):
         x_len = bounding_box.lengths[0] + self.x
         y_len = bounding_box.lengths[1] + self.y
         z_len = bounding_box.lengths[2] + 0.2
-        self.set_target_box(x_constraint=x_len, y_constraint=y_len)
         # Center the lattice in its box
         system.box = mb.box.Box(np.array([x_len, y_len, z_len]))
         system.translate_to(

flowermd/library/surfaces.py

@@ -93,7 +93,6 @@ def __init__(
         surface_mol = Molecule(num_mols=1, compound=surface)
         super(Graphene, self).__init__(
             molecules=[surface_mol],
-            density=1.0,
             base_units=base_units,
         )
 

flowermd/modules/surface_wetting/surface_wetting.py

@@ -6,7 +6,7 @@
 
 from flowermd.base import Simulation
 from flowermd.modules.surface_wetting.utils import combine_forces
-from flowermd.utils import HOOMDThermostats
+from flowermd.utils import HOOMDThermostats, get_target_box_mass_density
 
 
 class DropletSimulation(Simulation):
@@ -95,21 +95,27 @@ def run_droplet(
 
         """
         # Shrink down to high density
+        target_box = get_target_box_mass_density(
+            density=shrink_density * (u.g / (u.cm**3)), mass=self.mass.to("g")
+        )
         self.run_update_volume(
             n_steps=shrink_steps,
             period=shrink_period,
             kT=shrink_kT,
             tau_kt=tau_kt,
-            final_density=shrink_density * (u.g / (u.cm**3)),
+            final_box_lengths=target_box,
             write_at_start=True,
         )
         # Expand back up to low density
+        target_box = get_target_box_mass_density(
+            density=final_density * (u.g / (u.cm**3)), mass=self.mass.to("g")
+        )
         self.run_update_volume(
             n_steps=expand_steps,
             period=expand_period,
             kT=expand_kT,
             tau_kt=tau_kt,
-            final_density=final_density * (u.g / (u.cm**3)),
+            final_box_lengths=target_box,
         )
         # Run at low density
         self.run_NVT(n_steps=hold_steps, kT=hold_kT, tau_kt=tau_kt)