diff --git a/examples/multi_stage/analysis_script.py b/examples/multi_stage/analysis_script.py
index fecc96b8..f8d810d4 100644
--- a/examples/multi_stage/analysis_script.py
+++ b/examples/multi_stage/analysis_script.py
@@ -9,10 +9,9 @@
 def get_emittance(ts, t):
     """Calculate the beam emittance at the given time step."""
     w, x, ux = ts.get_particle(["w", "x", "ux"], t=t)
-    x2 = np.average(x**2, weights=w)
-    u2 = np.average(ux**2, weights=w)
+    x2 = np.average( x**2, weights=w)
     xu = np.average(x * ux, weights=w)
-    return np.sqrt(x2 * u2 - xu**2)
+    return np.sqrt( x2 * np.average( (ux - xu/x2 * x)**2, weights=w) )
 
 
 def analyze_simulation(simulation_directory, output_params):