diff --git a/examples/ramses/run_kh.py b/examples/ramses/run_kh.py
index 824439a800..1b33087b9a 100644
--- a/examples/ramses/run_kh.py
+++ b/examples/ramses/run_kh.py
@@ -169,15 +169,12 @@ def rho_map(rmin, rmax):
         else:
             return rho_1
 
-    def rhovel_map(rmin, rmax):
-        rho = rho_map(rmin, rmax)
-
+    def vel_map(rmin, rmax):
         x, y, z = rmin
 
         ampl = 0.01
         n = 2
         pert = np.sin(2 * np.pi * n * x / (xs))
-
         sigma = 0.05 / (2**0.5)
         gauss1 = np.exp(-((y - y_interface) ** 2) / (2 * sigma * sigma))
         gauss2 = np.exp(-((y + y_interface) ** 2) / (2 * sigma * sigma))
@@ -186,29 +183,35 @@ def rhovel_map(rmin, rmax):
         # Alternative formula (See T. Tricco paper)
         # interf_sz = ys/32
         # vx = np.arctan(y/interf_sz)/np.pi
+        # return vx
+
+        if y > y_interface:
+            return vslip / 2, ampl * pert, 0
+        else:
+            return -vslip / 2, ampl * pert, 0
+
+    def P_map(rmin, rmax):
+        x, y, z = rmin
 
-        vx = 0
-        if np.abs(y) > y_interface:
-            vx = vslip / 2
+        if y > y_interface:
+            return P_2
         else:
-            vx = -vslip / 2
+            return P_1
 
-        return (vx * rho, ampl * pert * rho, 0)
+    def rhovel_map(rmin, rmax):
+        rho = rho_map(rmin, rmax)
+        vx, vy, vz = vel_map(rmin, rmax)
+
+        return (vx * rho, vy * rho, vz * rho)
 
     def rhoetot_map(rmin, rmax):
         rho = rho_map(rmin, rmax)
         rhovel = rhovel_map(rmin, rmax)
+        P = P_map(rmin, rmax)
 
         rhovel2 = rhovel[0] * rhovel[0] + rhovel[1] * rhovel[1] + rhovel[2] * rhovel[2]
         rhoekin = 0.5 * rhovel2 / rho
 
-        x, y, z = rmin
-
-        if y > y_interface:
-            P = P_2
-        else:
-            P = P_1
-
         return rhoekin + P / (gamma - 1)
 
     model.set_field_value_lambda_f64("rho", rho_map)