diff --git a/lib/test.py b/lib/test.py
index 129c773..cfe09e8 100644
--- a/lib/test.py
+++ b/lib/test.py
@@ -51,7 +51,7 @@
     "simulation_id": 1,
     "model_type": MODEL["No_Shear_Thinning"],
     "reservoir_geometry": GEOMETRY["Rectilinear"],
-    "permeability": PERMEABILITY["Homogeneous"],
+    "permeability": PERMEABILITY["Heterogeneous"],
     "polymer_type": POLYMER["Xanthane"],
     "polymer_concentration": 0.001,
     "surfactant_type": SURFACTANT["Alkyl_Ether_Sulfate"],
diff --git a/lib/water.py b/lib/water.py
index e972933..825426f 100644
--- a/lib/water.py
+++ b/lib/water.py
@@ -580,7 +580,7 @@ def compute_water_saturation(
                                     + g1 * (1 - f[cnt][i])
                                     + (
                                         (D_g[cnt][i] + D_g[cnt][i + 1]) / (dx**2)
-                                        + (D_g[cnt + 1][i] + D_g[cnt][i]) / (dy**2)
+                                        + (D_g[cnt + 1][i] + D_g[cnt][i]) / (dx**2)
                                     )
                                     * surfactant.concentration_matrix[cnt][i]
                                     - (D_g[cnt][i] + D_g[cnt][i + 1])
@@ -595,7 +595,7 @@ def compute_water_saturation(
 
                                 BB[j][i] = (
                                     1 / dt_array[cnt][i]
-                                    - (D_s[cnt][i] + D_s[cnt][i + 1]) / (dx**2)
+                                    - (D_s[cnt+1][i] + D_s[cnt][i + 1]) / (dx**2)
                                     - (D_s[cnt + 1][i] + D_s[cnt][i]) / (dy**2)
                                 )
 
@@ -730,7 +730,7 @@ def compute_water_saturation(
                                     * surfactant.concentration_matrix[cnt - 1][i]
                                 )
 
-                                BB[j][i - 1] = (D_s[cnt][i] + D_s[cnt][i - 1]) / (dx**2)
+                                BB[j][i - 1] = (D_s[cnt][i] + D_s[cnt][i - 1]) / (dy**2)
 
                                 BB[j][i] = (
                                     1 / dt_array[cnt][i]
@@ -1037,7 +1037,7 @@ def compute_water_saturation(
             warnings.warn(f"BiCGSTAB: convergence issue (info={info}) in water saturation solver")
         Qnew = Qnew_flat = Qnew_flat.reshape(m, n)
 
-        Qnew[Qnew < 0] = 0
+        Qnew[Qnew > 1] = 1
 
         self.water_saturation = Qnew