diff --git a/Poisson_equation_FDM.py b/Poisson_equation_FDM.py
index b9f5470c25aea917e1d3d99768b7d86774d55c58..b436939d8566c414daa71568ac0531ce0388ba01 100644
--- a/Poisson_equation_FDM.py
+++ b/Poisson_equation_FDM.py
@@ -10,34 +10,31 @@ u(x; y) = (x^4)*(y^5) + 17*sin(xy)
 import numpy as np
 import matplotlib.pyplot as plt
 
-class FivePointStencil_2D:
+class FivePointStencil:
     '''
     This class solve the Poisson equation numerically in 2D
     using the five-point stencil finite difference method.
     '''
 
-    def __init__ (grid_number_x:'int',grid_number_y:'int', boundary_x:'tuple', boundary_y:'tuple'):
+    def __init__ (self,nx,ny):
         
         '''
         Parameters:
-                    grid_number_x (int): the number of grid in the x direction
-                    grid_number_y (int): the number of grid in the y direction
-                    boundary_x (tuple): the boundary condition in x direction like [lower_bound,upper_bound]
-                    boundary_y (tuple): the boundary condition in y direction like [lower_bound,upper_bound]
+                    nx (int): the number of grid in the x direction
+                    ny (int): the number of grid in the y direction
+
         '''
         
-        self.grid_number_x = grid_number_x
-        self.grid_number_y = grid_number_y
-        self.boundary_x = boundary_x
-        self.boundary_y = boundary_y
+        self.nx = nx
+        self.ny = ny
 
     def boundary_initialization (self):
         '''
         Initialization of boundary condition based on the right hand side function
         right hand side function: u(x; y) = (x^4)*(y^5) + 17*sin(xy)
         '''
-        x_grid = np.linspace(self.boundary_x[0], self.boundary_x[1], self.grid_number_x+1)
-        y_grid = np.linspace(self.boundary_y[0], self.boundary_y[1], self.grid_number_y+1)
+        x_grid = np.linspace(0, 1, self.nx+1)
+        y_grid = np.linspace(0, 1, self.ny+1)
 
         xx, yy = np.meshgrid(x_grid, y_grid, sparse=True)
         
@@ -49,17 +46,21 @@ class FivePointStencil_2D:
         values = values.flatten()
 
         # compute hx and hy
-        hx = (self.boundary_x[1] - self.boundary_x[0])/self.grid_number_x
-        hy = (self.boundary_y[1] - self.boundary_y[0])/self.grid_number_y
+        hx =1/self.nx
+        hy = 1/self.ny
 
         # Au - Bg = f
         # compute "A" Matrix   
                 
         diag_coeff = 2*((1/hx**2) + (1/hy**2))
-        eins = (1/hx**2)*np.ones((self.grid_number_x,))
-        block_matrix = diag_coeff * np.eye((self.grid_number_x-1)*(self.grid_number_x-1)) + np.diag(-1*eins, k=1)[0:-1, 0:-1] + np.diag(eins, k=-1)[0:-1, 0:-1]
-
-        corner_matrix = (1/hy**2)*np.eye((self.grid_number_x-1)*(self.grid_number_x-1))
+        eins = (1/hx**2)*np.ones((self.nx+1,))
+        block_matrix = diag_coeff * np.eye((self.nx-1)) 
+        off_diag_upper = np.diag(-1*eins, k=1)
+        off_diag_lower = np.diag(eins, k=-1)
+        
+        block_matrix = block_matrix + off_diag_upper + off_diag_lower
+        
+        corner_matrix = (1/hy**2)*np.eye((self.nx-1)*(self.ny-1))
 
         half_matrix_upper = np.hstack((block_matrix,corner_matrix))
         half_matrix_lower = np.hstack((corner_matrix,block_matrix))
@@ -71,7 +72,9 @@ class FivePointStencil_2D:
 
 if __name__=="__main__":
 
-    fdm_obj = FivePointStencil_2D(grid_number_x=10,grid_number_y=10, boundary_x=[0,1], boundary_y=[0,1])
+    nx=4
+    ny=3
+    fdm_obj = FivePointStencil(nx,ny)
     solution = fdm_obj.boundary_initialization()