diff --git a/main.py b/main.py
index 1d8eca0240720512b84e2821673c595619a300c7..cd132985b44606526004366fd67f16eb5d5d9790 100644
--- a/main.py
+++ b/main.py
@@ -26,11 +26,15 @@ def cahn_hilliard_fdiff(it, c):
         cp[:, -1] = cp[:, 1]
         
         f[1:-1, 1:-1] = cp[1:-1, 1:-1]**3 - cp[1:-1, 1:-1] - (a / dx)**2 * (cp[0:-2, 1:-1] + cp[2:, 1:-1] + cp[1:-1, 0:-2] + cp[1:-1, 2:] - 4 * cp[1:-1, 1:-1])
-        cp[1:-1, 1:-1] = cp[1:-1, 1:-1] + (dt / dx**2) * (f[0:-2, 1:-1] + f[2:, 1:-1] + f[1:-1, 0:-2] + f[1:-1, 2:] - 4 * f[1:-1, 1:-1]) 
 
-    c = cp[1:-1, 1:-1]
+        f[0, :]  = f[-2, :]
+        f[-1, :] = f[1, :] 
+        f[:, 0]  = f[:, -2]
+        f[:, -1] = f[:, 1]
 
-    return c
+        cp[1:-1, 1:-1] = cp[1:-1, 1:-1] + (dt / dx**2) * (f[0:-2, 1:-1] + f[2:, 1:-1] + f[1:-1, 0:-2] + f[1:-1, 2:] - 4 * f[1:-1, 1:-1]) 
+
+    return cp[1:-1, 1:-1]
     
 
 def cahn_hilliard_spectral(it, c):
@@ -56,13 +60,13 @@ if __name__ == "__main__":
     c2 = cahn_hilliard_spectral(it, c)
 
     toc = time.time() - tic
-    
+
     print(f"Elapsed time is {toc} seconds.")
     
     fig, (ax1, ax2) = plt.subplots(1, 2)
     fig.suptitle(f"Cahn-Hilliard integration for {it} iterations ")
-    ax1.imshow(c1, cmap="jet", vmin=-1, vmax=1)
+    ax1.contourf(xy, xy, c1, cmap="jet", vmin=-1, vmax=1)
     ax1.set_title("Finite differences method")
-    ax2.imshow(c2, cmap="jet", vmin=-1, vmax=1)
+    ax2.contourf(xy, xy, c2, cmap="jet", vmin=-1, vmax=1)
     ax2.set_title("Spectral method")
     plt.show()