diff --git a/examples/KelvinHelmholtz_2D/makeMovie.py b/examples/KelvinHelmholtz_2D/makeMovie.py
index 72eafbfeb8dcb3aab1e9d5949c67e3e9fed14ef4..84fe99106bf607830e89b6aa663135b48b6c0744 100644
--- a/examples/KelvinHelmholtz_2D/makeMovie.py
+++ b/examples/KelvinHelmholtz_2D/makeMovie.py
@@ -9,15 +9,10 @@ Edit this file near the bottom with the number of snaps you have.
 Written by Josh Borrow (joshua.borrow@durham.ac.uk)
 """
 
-import matplotlib
-matplotlib.use("Agg")
-
 import os
 import h5py as h5
 import numpy as np
-import scipy as sp
 import scipy.interpolate as si
-import matplotlib.pyplot as plt
 
 
 def load_and_extract(filename):
@@ -67,14 +62,32 @@ def frame(n, *args):
 
 
 if __name__ == "__main__":
+    import matplotlib
+    matplotlib.use("Agg")
+
     from tqdm import tqdm
     from matplotlib.animation import FuncAnimation
     from scipy.stats import gaussian_kde
 
+    import matplotlib.pyplot as plt
+
     filename = "kelvinhelmholtz"
     dpi = 512
 
-    frames = tqdm(np.arange(0, 448))
+
+    # Look for the number of files in the directory.
+    i = 0
+    while True:
+        if os.path.isfile("{}_{:04d}.hdf5".format(filename, i)):
+            i += 1
+        else:
+            break
+
+        if i > 10000:
+            raise FileNotFoundError(
+                "Could not find the snapshots in the directory")
+
+    frames = tqdm(np.arange(0, i))
 
     # Creation of first frame
     fig, ax = plt.subplots(1, 1, figsize=(1, 1), frameon=False)
diff --git a/examples/KelvinHelmholtz_2D/run.sh b/examples/KelvinHelmholtz_2D/run.sh
index 3d83bcc6dd7c226885ed83c3188f3177c4807154..dbb39caf383279dbc71c2baa125499d115538654 100755
--- a/examples/KelvinHelmholtz_2D/run.sh
+++ b/examples/KelvinHelmholtz_2D/run.sh
@@ -8,7 +8,8 @@ then
 fi
 
 # Run SWIFT
-../swift -s -t 1 kelvinHelmholtz.yml 2>&1 | tee output.log
+../swift -s -t 4 kelvinHelmholtz.yml 2>&1 | tee output.log
 
 # Plot the solution
 python plotSolution.py 6
+python makeMovie.py