From 4746169cd53add37a7228fd0074b6f62563e15f6 Mon Sep 17 00:00:00 2001
From: chrisjonesBSU <chrisjones4@u.boisestate.edu>
Date: Sat, 6 Jul 2024 18:14:09 -0600
Subject: [PATCH 01/10] add write and read simulation methods

---
 flowermd/base/simulation.py | 85 +++++++++++++++++++++++++++++++++++++
 1 file changed, 85 insertions(+)

diff --git a/flowermd/base/simulation.py b/flowermd/base/simulation.py
index 04ba2d44..ce845fbe 100644
--- a/flowermd/base/simulation.py
+++ b/flowermd/base/simulation.py
@@ -1,7 +1,9 @@
 """Base simulation class for flowerMD."""
 
 import inspect
+import os
 import pickle
+import tempfile
 import warnings
 from collections.abc import Iterable
 
@@ -150,6 +152,29 @@ def from_system(cls, system, **kwargs):
                 "or a system with a forcefield."
             )
 
+    @classmethod
+    def from_simulation_pickle(cls, file_path):
+        with open(file_path, "rb") as f:
+            string = f.read(len(b"FLOWERMD"))
+            if string != b"FLOWERMD":
+                raise ValueError(
+                    "It appears this pickle file "
+                    "was not created by flowermd.base.Simulation. "
+                    "See flowermd.base.Simulation.save_simulation()."
+                )
+            data = pickle.load(f)
+
+        state = data["state"]
+        forces = data["forcefield"]
+        ref_values = data["reference_values"]
+        sim_kwargs = data["sim_kwargs"]
+        return cls(
+            initial_state=state,
+            forcefield=forces,
+            reference_values=ref_values,
+            **sim_kwargs,
+        )
+
     @classmethod
     def from_snapshot_forces(cls, initial_state, forcefield, **kwargs):
         """Initialize a simulation from an initial state and HOOMD forces.
@@ -1085,6 +1110,66 @@ def save_restart_gsd(self, file_path="restart.gsd"):
         """
         hoomd.write.GSD.write(self.state, filename=file_path)
 
+    def save_simulation(self, file_path="simulation.pickle", save_walls=False):
+        """Save a pickle file with everything needed to retart a simulation.
+
+        This method is useful for saving the state of a simulation to a file
+        and reusing it for restarting a simulation or running a different
+        simulation.
+
+        Parameters
+        ----------
+        file_path : str, default "simulation.pickle"
+            The path to save the pickle file to.
+        save_walls : bool, default False
+            Determines if any wall forces are saved.
+
+        Notes
+        -----
+        This method creates a dictionary that contains the
+        simulation's forcefield, references values, and snapshot.
+        The key:value pairs are:
+            'reference_values': dict of str:float
+            'forcefield': list of hoomd forces
+            'state': hoomd.snapshot.Snapshot
+            'sim_kwargs': dict of flowermd.base.Simulation kwargs
+
+        """
+        if self._wall_forces and save_walls is False:
+            forces = []
+            for force in self._forcefield:
+                if not hasattr(force, "wall"):
+                    forces.append(force)
+        else:
+            forces = self._forcefield
+
+        # Make a temp restart gsd file
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            temp_file_path = os.path.join(tmp_dir, "temp.gsd")
+            self.save_restart_gsd(temp_file_path)
+            with gsd.hoomd.open(temp_file_path, "r") as traj:
+                snap = traj[0]
+                os.remove(temp_file_path)
+
+        sim_kwargs = {
+            "dt": self.dt,
+            "gsd_write_freq": self.gsd_write_freq,
+            "log_write_freq": self.log_write_freq,
+            "gsd_max_buffer_size": self.maximum_write_buffer_size,
+            "seed": self.seed,
+        }
+
+        sim_dict = {
+            "reference_values": self.reference_values,
+            "forcefield": self._forcefield,
+            "state": snap,
+            "sim_kwargs": sim_kwargs,
+        }
+        flower_string = b"FLOWERMD"
+        with open(file_path, "wb") as f:
+            f.write(flower_string)
+            pickle.dump(sim_dict, f)
+
     def flush_writers(self):
         """Flush all write buffers to file."""
         for writer in self.operations.writers:

From f6ad46706e9314b2d91faa3b6b05517086eba29f Mon Sep 17 00:00:00 2001
From: chrisjonesBSU <chrisjones4@u.boisestate.edu>
Date: Sat, 6 Jul 2024 18:17:57 -0600
Subject: [PATCH 02/10] add more doc strings

---
 flowermd/base/simulation.py | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/flowermd/base/simulation.py b/flowermd/base/simulation.py
index ce845fbe..6b20f5f5 100644
--- a/flowermd/base/simulation.py
+++ b/flowermd/base/simulation.py
@@ -1134,6 +1134,11 @@ def save_simulation(self, file_path="simulation.pickle", save_walls=False):
             'state': hoomd.snapshot.Snapshot
             'sim_kwargs': dict of flowermd.base.Simulation kwargs
 
+        Wall forces are not able to be reused when starting
+        a simulation. If your simulation has wall forces,
+        set `save_walls` to `False` and manually re-add them
+        in the new simulation if needed.
+
         """
         if self._wall_forces and save_walls is False:
             forces = []
@@ -1143,14 +1148,14 @@ def save_simulation(self, file_path="simulation.pickle", save_walls=False):
         else:
             forces = self._forcefield
 
-        # Make a temp restart gsd file
+        # Make a temp restart gsd file.
         with tempfile.TemporaryDirectory() as tmp_dir:
             temp_file_path = os.path.join(tmp_dir, "temp.gsd")
             self.save_restart_gsd(temp_file_path)
             with gsd.hoomd.open(temp_file_path, "r") as traj:
                 snap = traj[0]
                 os.remove(temp_file_path)
-
+        # Save dict of kwargs needed to restart simulation.
         sim_kwargs = {
             "dt": self.dt,
             "gsd_write_freq": self.gsd_write_freq,
@@ -1158,13 +1163,15 @@ def save_simulation(self, file_path="simulation.pickle", save_walls=False):
             "gsd_max_buffer_size": self.maximum_write_buffer_size,
             "seed": self.seed,
         }
-
+        # Create the final dict that holds everything.
         sim_dict = {
             "reference_values": self.reference_values,
             "forcefield": self._forcefield,
             "state": snap,
             "sim_kwargs": sim_kwargs,
         }
+        # Add a header to the pickle file.
+        # This will be checked in Simulation.from_simulation_pickle.
         flower_string = b"FLOWERMD"
         with open(file_path, "wb") as f:
             f.write(flower_string)

From 54311668382315f2bc43f3be1a2c191490f2c947 Mon Sep 17 00:00:00 2001
From: chrisjonesBSU <chrisjones4@u.boisestate.edu>
Date: Sat, 6 Jul 2024 18:53:13 -0600
Subject: [PATCH 03/10] add unit test

---
 flowermd/tests/base/test_simulation.py | 37 ++++++++++++++++++++++++++
 1 file changed, 37 insertions(+)

diff --git a/flowermd/tests/base/test_simulation.py b/flowermd/tests/base/test_simulation.py
index a756f808..692fe5a0 100644
--- a/flowermd/tests/base/test_simulation.py
+++ b/flowermd/tests/base/test_simulation.py
@@ -43,6 +43,43 @@ def test_initialize_from_state(self, benzene_system):
             reference_values=benzene_system.reference_values,
         )
 
+    def test_initialize_from_simulation_pickle(self, benzene_system):
+        sim = Simulation.from_snapshot_forces(
+            initial_state=benzene_system.hoomd_snapshot,
+            forcefield=benzene_system.hoomd_forcefield,
+            reference_values=benzene_system.reference_values,
+        )
+        sim.save_simulation("simulation.pickle")
+        new_sim = Simulation.from_simulation_pickle("simulation.pickle")
+        assert new_sim.dt == sim.dt
+        assert new_sim.gsd_write_freq == sim.gsd_write_freq
+        assert new_sim.log_write_freq == sim.log_write_freq
+        assert new_sim.seed == sim.seed
+        assert (
+            new_sim.maximum_write_buffer_size == sim.maximum_write_buffer_size
+        )
+        assert new_sim.volume_reduced == sim.volume_reduced
+        assert new_sim.mass_reduced == sim.mass_reduced
+        assert new_sim.reference_mass == sim.reference_mass
+        assert new_sim.reference_energy == sim.reference_energy
+        assert new_sim.reference_length == sim.reference_length
+        snap = sim.state.get_snapshot()
+        new_snap = new_sim.state.get_snapshot()
+        assert np.array_equal(
+            snap.particles.position, new_snap.particles.position
+        )
+        new_sim.run_NVT(n_steps=2, kT=1.0, tau_kt=0.001)
+
+    def test_initialize_from_bad_pickle(self, benzene_system):
+        sim = Simulation.from_snapshot_forces(
+            initial_state=benzene_system.hoomd_snapshot,
+            forcefield=benzene_system.hoomd_forcefield,
+            reference_values=benzene_system.reference_values,
+        )
+        sim.pickle_forcefield("forces.pickle")
+        with pytest.raises(ValueError):
+            Simulation.from_simulation_pickle("forces.pickle")
+
     def test_no_reference_values(self, benzene_system):
         sim = Simulation.from_snapshot_forces(
             initial_state=benzene_system.hoomd_snapshot,

From a1518cdcdf6ecead79b7f1506e073baa5351de92 Mon Sep 17 00:00:00 2001
From: chrisjonesBSU <chrisjones4@u.boisestate.edu>
Date: Sun, 7 Jul 2024 09:55:47 -0600
Subject: [PATCH 04/10] add walls check to pickle_forcefield()

---
 flowermd/base/simulation.py | 31 +++++++++++++++++++++++++------
 1 file changed, 25 insertions(+), 6 deletions(-)

diff --git a/flowermd/base/simulation.py b/flowermd/base/simulation.py
index 6b20f5f5..7e9219bc 100644
--- a/flowermd/base/simulation.py
+++ b/flowermd/base/simulation.py
@@ -1022,7 +1022,9 @@ def temperature_ramp(self, n_steps, kT_start, kT_final):
             A=kT_start, B=kT_final, t_start=self.timestep, t_ramp=int(n_steps)
         )
 
-    def pickle_forcefield(self, file_path="forcefield.pickle"):
+    def pickle_forcefield(
+        self, file_path="forcefield.pickle", save_walls=False
+    ):
         """Pickle the list of HOOMD forces.
 
         This method useful for saving the forcefield of a simulation to a file
@@ -1033,6 +1035,15 @@ def pickle_forcefield(self, file_path="forcefield.pickle"):
         ----------
         file_path : str, default "forcefield.pickle"
             The path to save the pickle file to.
+        save_walls : bool, default False
+            Determines if any wall forces are saved.
+
+        Notes
+        -----
+        Wall forces are not able to be reused when starting
+        a simulation. If your simulation has wall forces,
+        set `save_walls` to `False` and manually re-add them
+        in the new simulation if needed.
 
         Examples
         --------
@@ -1063,8 +1074,15 @@ def pickle_forcefield(self, file_path="forcefield.pickle"):
             tensile_sim.run_tensile(strain=0.05, kT=2.0, n_steps=1e3, period=10)
 
         """
+        if self._wall_forces and save_walls is False:
+            forces = []
+            for force in self._forcefield:
+                if not hasattr(force, "wall"):
+                    forces.append(force)
+        else:
+            forces = self._forcefield
         f = open(file_path, "wb")
-        pickle.dump(self._forcefield, f)
+        pickle.dump(forces, f)
 
     def save_restart_gsd(self, file_path="restart.gsd"):
         """Save a GSD file of the current simulation state.
@@ -1129,9 +1147,10 @@ def save_simulation(self, file_path="simulation.pickle", save_walls=False):
         This method creates a dictionary that contains the
         simulation's forcefield, references values, and snapshot.
         The key:value pairs are:
+
             'reference_values': dict of str:float
             'forcefield': list of hoomd forces
-            'state': hoomd.snapshot.Snapshot
+            'state': gsd.hoomd.Frame
             'sim_kwargs': dict of flowermd.base.Simulation kwargs
 
         Wall forces are not able to be reused when starting
@@ -1140,6 +1159,7 @@ def save_simulation(self, file_path="simulation.pickle", save_walls=False):
         in the new simulation if needed.
 
         """
+        # Make list of forces
         if self._wall_forces and save_walls is False:
             forces = []
             for force in self._forcefield:
@@ -1147,7 +1167,6 @@ def save_simulation(self, file_path="simulation.pickle", save_walls=False):
                     forces.append(force)
         else:
             forces = self._forcefield
-
         # Make a temp restart gsd file.
         with tempfile.TemporaryDirectory() as tmp_dir:
             temp_file_path = os.path.join(tmp_dir, "temp.gsd")
@@ -1166,7 +1185,7 @@ def save_simulation(self, file_path="simulation.pickle", save_walls=False):
         # Create the final dict that holds everything.
         sim_dict = {
             "reference_values": self.reference_values,
-            "forcefield": self._forcefield,
+            "forcefield": forces,
             "state": snap,
             "sim_kwargs": sim_kwargs,
         }
@@ -1234,7 +1253,7 @@ def _create_state(self, initial_state):
             self.create_state_from_gsd(initial_state)
         elif isinstance(initial_state, hoomd.snapshot.Snapshot):
             print(
-                "Initializing simulation state from a hoomd.snapshot.Snapshot"
+                "Initializing simulation state from a hoomd.snapshot.Snapshot."
             )
             self.create_state_from_snapshot(initial_state)
         elif isinstance(initial_state, gsd.hoomd.Frame):

From 4bc5f4f99d970b78d164fe8bad1eec3ad3349ad3 Mon Sep 17 00:00:00 2001
From: chrisjonesBSU <chrisjones4@u.boisestate.edu>
Date: Sun, 7 Jul 2024 10:30:17 -0600
Subject: [PATCH 05/10] update versions in actions

---
 .github/workflows/pytest.yml | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/pytest.yml b/.github/workflows/pytest.yml
index 284423a0..9e46e8b0 100644
--- a/.github/workflows/pytest.yml
+++ b/.github/workflows/pytest.yml
@@ -36,10 +36,10 @@ jobs:
 
     steps:
     - name: Check out repository
-      uses: actions/checkout@v2
+      uses: actions/checkout@v4
 
     - name: Build environment
-      uses: conda-incubator/setup-miniconda@v2
+      uses: conda-incubator/setup-miniconda@v3
       with:
         environment-file: environment-dev.yml
         python-version: ${{ matrix.python-version }}
@@ -56,7 +56,7 @@ jobs:
       run: python -m pytest -rs -v --cov=./ --cov-report=xml
 
     - name: Upload coverage to Codecov
-      uses: codecov/codecov-action@v3
+      uses: codecov/codecov-action@v4
       with:
         token: ${{ secrets.CODECOV_TOKEN }}
         file: ./coverage.xml

From 290d68e329321ab20c7baf953659265fb783903f Mon Sep 17 00:00:00 2001
From: chrisjonesBSU <chrisjones4@u.boisestate.edu>
Date: Sun, 7 Jul 2024 15:19:46 -0600
Subject: [PATCH 06/10] fix wall check logic, add test

---
 flowermd/base/simulation.py            |  4 ++--
 flowermd/tests/base/test_simulation.py | 30 ++++++++++++++++++++++++++
 2 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/flowermd/base/simulation.py b/flowermd/base/simulation.py
index 7e9219bc..52b767b1 100644
--- a/flowermd/base/simulation.py
+++ b/flowermd/base/simulation.py
@@ -1077,7 +1077,7 @@ def pickle_forcefield(
         if self._wall_forces and save_walls is False:
             forces = []
             for force in self._forcefield:
-                if not hasattr(force, "wall"):
+                if not isinstance(force, hoomd.md.external.wall.LJ):
                     forces.append(force)
         else:
             forces = self._forcefield
@@ -1163,7 +1163,7 @@ def save_simulation(self, file_path="simulation.pickle", save_walls=False):
         if self._wall_forces and save_walls is False:
             forces = []
             for force in self._forcefield:
-                if not hasattr(force, "wall"):
+                if not isinstance(force, hoomd.md.external.wall.LJ):
                     forces.append(force)
         else:
             forces = self._forcefield
diff --git a/flowermd/tests/base/test_simulation.py b/flowermd/tests/base/test_simulation.py
index 692fe5a0..145cf249 100644
--- a/flowermd/tests/base/test_simulation.py
+++ b/flowermd/tests/base/test_simulation.py
@@ -80,6 +80,36 @@ def test_initialize_from_bad_pickle(self, benzene_system):
         with pytest.raises(ValueError):
             Simulation.from_simulation_pickle("forces.pickle")
 
+    def test_save_forces_with_walls(self, benzene_system):
+        sim = Simulation.from_snapshot_forces(
+            initial_state=benzene_system.hoomd_snapshot,
+            forcefield=benzene_system.hoomd_forcefield,
+            reference_values=benzene_system.reference_values,
+        )
+        sim.add_walls(wall_axis=(1, 0, 0), sigma=1.0, epsilon=1.0, r_cut=2.0)
+        assert len(sim._wall_forces[(1, 0, 0)]) == 2
+
+        # Test without saving walls
+        sim.pickle_forcefield("forces_no_walls.pickle", save_walls=False)
+        found_wall_force = False
+        with open("forces_no_walls.pickle", "rb") as f:
+            forces = pickle.load(f)
+            for force in forces:
+                if isinstance(force, hoomd.md.external.wall.LJ):
+                    found_wall_force = True
+        assert found_wall_force is False
+        # Make sure wall force is still in sim object
+        assert len(sim._wall_forces[(1, 0, 0)]) == 2
+        # Test with saving walls
+        sim.pickle_forcefield("forces_walls.pickle", save_walls=True)
+        found_wall_force = False
+        with open("forces_walls.pickle", "rb") as f:
+            forces = pickle.load(f)
+            for force in forces:
+                if isinstance(force, hoomd.md.external.wall.LJ):
+                    found_wall_force = True
+        assert found_wall_force is True
+
     def test_no_reference_values(self, benzene_system):
         sim = Simulation.from_snapshot_forces(
             initial_state=benzene_system.hoomd_snapshot,

From bfca19d14e31c10576b59942557a88d719c4af33 Mon Sep 17 00:00:00 2001
From: chrisjonesBSU <chrisjones4@u.boisestate.edu>
Date: Tue, 9 Jul 2024 10:33:39 -0600
Subject: [PATCH 07/10] make forces into a list

---
 flowermd/base/simulation.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/flowermd/base/simulation.py b/flowermd/base/simulation.py
index 52b767b1..85475add 100644
--- a/flowermd/base/simulation.py
+++ b/flowermd/base/simulation.py
@@ -166,11 +166,14 @@ def from_simulation_pickle(cls, file_path):
 
         state = data["state"]
         forces = data["forcefield"]
+        for force in forces:
+            if isinstance(force, hoomd.md.external.wall.LJ):
+                pass
         ref_values = data["reference_values"]
         sim_kwargs = data["sim_kwargs"]
         return cls(
             initial_state=state,
-            forcefield=forces,
+            forcefield=list(forces),
             reference_values=ref_values,
             **sim_kwargs,
         )

From 2f4e6d1848d980b75d2af48b5395b37f142ef152 Mon Sep 17 00:00:00 2001
From: chrisjonesBSU <chrisjones4@u.boisestate.edu>
Date: Tue, 9 Jul 2024 15:24:07 -0600
Subject: [PATCH 08/10] possible work around for starting sim with walls

---
 flowermd/base/simulation.py | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/flowermd/base/simulation.py b/flowermd/base/simulation.py
index 85475add..a713fe81 100644
--- a/flowermd/base/simulation.py
+++ b/flowermd/base/simulation.py
@@ -168,7 +168,18 @@ def from_simulation_pickle(cls, file_path):
         forces = data["forcefield"]
         for force in forces:
             if isinstance(force, hoomd.md.external.wall.LJ):
-                pass
+                new_walls = []
+                for _wall in force.walls:
+                    new_walls.append(
+                        hoomd.wall.Plane(
+                            origin=_wall.origin, normal=_wall.normal
+                        )
+                    )
+                new_wall = hoomd.md.external.wall.LJ(walls=new_walls)
+                for param in force.params:
+                    new_wall.params[param] = force.params[param]
+                forces.remove(force)
+                forces.append(new_wall)
         ref_values = data["reference_values"]
         sim_kwargs = data["sim_kwargs"]
         return cls(

From 813ece4a84513c3cf2e2139888f1f78532ccee98 Mon Sep 17 00:00:00 2001
From: chrisjonesBSU <chrisjones4@u.boisestate.edu>
Date: Thu, 11 Jul 2024 22:48:39 -0600
Subject: [PATCH 09/10] always include wall forces when pickling entire sim,
 remake them when loading

---
 flowermd/base/simulation.py            | 20 ++------------------
 flowermd/tests/base/test_simulation.py | 12 ++++++++++++
 2 files changed, 14 insertions(+), 18 deletions(-)

diff --git a/flowermd/base/simulation.py b/flowermd/base/simulation.py
index a713fe81..4e62b832 100644
--- a/flowermd/base/simulation.py
+++ b/flowermd/base/simulation.py
@@ -1142,7 +1142,7 @@ def save_restart_gsd(self, file_path="restart.gsd"):
         """
         hoomd.write.GSD.write(self.state, filename=file_path)
 
-    def save_simulation(self, file_path="simulation.pickle", save_walls=False):
+    def save_simulation(self, file_path="simulation.pickle"):
         """Save a pickle file with everything needed to retart a simulation.
 
         This method is useful for saving the state of a simulation to a file
@@ -1153,8 +1153,6 @@ def save_simulation(self, file_path="simulation.pickle", save_walls=False):
         ----------
         file_path : str, default "simulation.pickle"
             The path to save the pickle file to.
-        save_walls : bool, default False
-            Determines if any wall forces are saved.
 
         Notes
         -----
@@ -1166,21 +1164,7 @@ def save_simulation(self, file_path="simulation.pickle", save_walls=False):
             'forcefield': list of hoomd forces
             'state': gsd.hoomd.Frame
             'sim_kwargs': dict of flowermd.base.Simulation kwargs
-
-        Wall forces are not able to be reused when starting
-        a simulation. If your simulation has wall forces,
-        set `save_walls` to `False` and manually re-add them
-        in the new simulation if needed.
-
         """
-        # Make list of forces
-        if self._wall_forces and save_walls is False:
-            forces = []
-            for force in self._forcefield:
-                if not isinstance(force, hoomd.md.external.wall.LJ):
-                    forces.append(force)
-        else:
-            forces = self._forcefield
         # Make a temp restart gsd file.
         with tempfile.TemporaryDirectory() as tmp_dir:
             temp_file_path = os.path.join(tmp_dir, "temp.gsd")
@@ -1199,7 +1183,7 @@ def save_simulation(self, file_path="simulation.pickle", save_walls=False):
         # Create the final dict that holds everything.
         sim_dict = {
             "reference_values": self.reference_values,
-            "forcefield": forces,
+            "forcefield": self._forcefield,
             "state": snap,
             "sim_kwargs": sim_kwargs,
         }
diff --git a/flowermd/tests/base/test_simulation.py b/flowermd/tests/base/test_simulation.py
index 145cf249..8df1ab27 100644
--- a/flowermd/tests/base/test_simulation.py
+++ b/flowermd/tests/base/test_simulation.py
@@ -70,6 +70,18 @@ def test_initialize_from_simulation_pickle(self, benzene_system):
         )
         new_sim.run_NVT(n_steps=2, kT=1.0, tau_kt=0.001)
 
+    def test_initialize_from_simulation_pickle_with_walls(self, benzene_system):
+        sim = Simulation.from_snapshot_forces(
+            initial_state=benzene_system.hoomd_snapshot,
+            forcefield=benzene_system.hoomd_forcefield,
+            reference_values=benzene_system.reference_values,
+        )
+        sim.add_walls(wall_axis=(1, 0, 0), sigma=1, epsilon=1, r_cut=1)
+        sim.save_simulation("simulation.pickle")
+        new_sim = Simulation.from_simulation_pickle("simulation.pickle")
+        assert len(new_sim.forces) == len(sim.forces)
+        new_sim.run_NVT(n_steps=2, kT=1.0, tau_kt=0.001)
+
     def test_initialize_from_bad_pickle(self, benzene_system):
         sim = Simulation.from_snapshot_forces(
             initial_state=benzene_system.hoomd_snapshot,

From 6dbdfccffc1231baebaeb6f0d6bfc4d1f3f967dd Mon Sep 17 00:00:00 2001
From: chrisjonesBSU <chrisjones4@u.boisestate.edu>
Date: Thu, 11 Jul 2024 23:02:07 -0600
Subject: [PATCH 10/10] update test to check for positions

---
 flowermd/tests/base/test_simulation.py | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/flowermd/tests/base/test_simulation.py b/flowermd/tests/base/test_simulation.py
index 8df1ab27..9ffa365a 100644
--- a/flowermd/tests/base/test_simulation.py
+++ b/flowermd/tests/base/test_simulation.py
@@ -49,8 +49,11 @@ def test_initialize_from_simulation_pickle(self, benzene_system):
             forcefield=benzene_system.hoomd_forcefield,
             reference_values=benzene_system.reference_values,
         )
+        sim.run_NVT(n_steps=1e3, kT=1.0, tau_kt=0.001)
         sim.save_simulation("simulation.pickle")
+        sim.save_restart_gsd("sim.gsd")
         new_sim = Simulation.from_simulation_pickle("simulation.pickle")
+        new_sim.save_restart_gsd("new_sim.gsd")
         assert new_sim.dt == sim.dt
         assert new_sim.gsd_write_freq == sim.gsd_write_freq
         assert new_sim.log_write_freq == sim.log_write_freq
@@ -63,11 +66,12 @@ def test_initialize_from_simulation_pickle(self, benzene_system):
         assert new_sim.reference_mass == sim.reference_mass
         assert new_sim.reference_energy == sim.reference_energy
         assert new_sim.reference_length == sim.reference_length
-        snap = sim.state.get_snapshot()
-        new_snap = new_sim.state.get_snapshot()
-        assert np.array_equal(
-            snap.particles.position, new_snap.particles.position
-        )
+        with gsd.hoomd.open("sim.gsd") as sim_traj:
+            with gsd.hoomd.open("new_sim.gsd") as new_sim_traj:
+                assert np.array_equal(
+                    sim_traj[0].particles.position,
+                    new_sim_traj[0].particles.position,
+                )
         new_sim.run_NVT(n_steps=2, kT=1.0, tau_kt=0.001)
 
     def test_initialize_from_simulation_pickle_with_walls(self, benzene_system):