fix kwargs passing, update test

cmelab · Jan 30, 2025 · a2514f1 · a2514f1
1 parent e0fa262
commit a2514f1
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 11 deletions.
diff --git a/flowermd/base/system.py b/flowermd/base/system.py
@@ -68,6 +68,7 @@ def __init__(
         self,
         molecules,
         base_units=dict(),
+        **kwargs,
     ):
         self._molecules = check_return_iterable(molecules)
         self.all_molecules = []
@@ -123,7 +124,7 @@ def __init__(
                 self.n_mol_types += 1
 
         # Create mBuild system
-        self.system = self._build_system()
+        self.system = self._build_system(**kwargs)
         # Create GMSO topology
         self.gmso_system = self._convert_to_gmso()
 
@@ -678,7 +679,7 @@ def __init__(
             molecules=molecules, base_units=base_units, **kwargs
         )
 
-    def _build_system(self):
+    def _build_system(self, **kwargs):
         mass_density = u.Unit("kg") / u.Unit("m**3")
         number_density = u.Unit("m**-3")
         if self.density.units.dimensions == mass_density.dimensions:
@@ -705,6 +706,7 @@ def _build_system(self):
             seed=self.seed,
             edge=self.edge,
             fix_orientation=self.fix_orientation,
+            **kwargs,
         )
         return system
 

diff --git a/flowermd/tests/base/test_system.py b/flowermd/tests/base/test_system.py
@@ -205,23 +205,20 @@ def test_pack_seed(self, benzene_molecule):
         benzene_mol = benzene_molecule(n_mols=3)
         default_seed = Pack(molecules=[benzene_mol], density=0.1)
         change_seed = Pack(molecules=[benzene_mol], density=0.1, seed=12340)
-        assert not np.array_equal(default_seed.xyz,change_seed.xyz)
+        assert not np.array_equal(
+            default_seed.system.xyz, change_seed.system.xyz
+        )
 
-    #adding test for kwargs argument in system.py Pack class
+    # adding test for kwargs argument in system.py Pack class
     def test_pack_kwargs_attr(self, polyethylene):
         polyethylene = polyethylene(lengths=5, num_mols=1)
-        system = Pack(
+        Pack(
             molecules=[polyethylene],
             density=1.0,
-            compound,
-            n_compounds=None,
-            box=None,
-            density=None,
             overlap=0.2,
             seed=12345,
-            sidemax=100.0, edge=0.2, compound_ratio=None, aspect_ratio=None, fix_orientation=False, temp_file=None, update_port_locations=False, packmol_args=None
+            sidemax=100.0,
         )
-        assert
 
     def test_mass(self, pps_molecule):
         pps_mol = pps_molecule(n_mols=20)