diff --git a/resmda/__init__.py b/resmda/__init__.py
index a7f2bb0..730e088 100644
--- a/resmda/__init__.py
+++ b/resmda/__init__.py
@@ -16,11 +16,13 @@
 
 from resmda import utils
 from resmda.utils import Report
-from resmda.data_assimilation import esmda
+from resmda.data_assimilation import esmda, build_localization_matrix
 from resmda.reservoir_simulator import Simulator, RandomPermeability
 
 
 __all__ = ['reservoir_simulator', 'data_assimilation', 'utils',
-           'Simulator', 'RandomPermeability', 'esmda', 'Report']
+           'Simulator', 'RandomPermeability',
+           'esmda', 'build_localization_matrix',
+           'Report']
 
 __version__ = utils.__version__
diff --git a/resmda/data_assimilation.py b/resmda/data_assimilation.py
index e50dd80..8950184 100644
--- a/resmda/data_assimilation.py
+++ b/resmda/data_assimilation.py
@@ -26,42 +26,47 @@ def __dir__():
 
 
 def esmda(model_prior, forward, data_obs, sigma, alphas=4, data_prior=None,
-          callback_post=None, return_post_data=False, random=None):
-    """ESMDA algorithm (Emerick and Reynolds, 2013).
+          callback_post=None, return_post_data=False, random=None,
+          enable_localization=False, localization_matrix=None):
+    """ESMDA algorithm (Emerick and Reynolds, 2013) with optional localization.
 
 
     Parameters
     ----------
-
-    model_post : ndarray
+    model_prior : ndarray
         Prior models of dimension ``(ne, ...)``, where ``ne`` is the number of
         ensembles.
-
     forward : callable
         Forward model that takes an ndarray of the shape of the prior models
         ``(ne, ...)``, and returns a ndarray of the shape of the prior data
         ``(ne, nd)``; ``ne`` is the number of ensembles, ``nd`` the number of
         data.
-
     data_obs : ndarray
         Observed data of shape ``(nd)``.
-
-    sigma : float, ndarray
-
+    sigma : {float, ndarray}
+        Standard deviation(s) of the observation noise.
     alphas : {int, ndarray}, default: 4
-        es-mda inflation factor
-
+        Inflation factors for ES-MDA.
     data_prior : ndarray, default: None
-        (ne, nd)
-
-    vmin, vmax : float, default: None
-
-    return_data_post : bool, default: False
+        Prior data ensemble, of shape (ne, nd).
+    callback_post : function, default: None
+        Function to be executed after each ESMDA iteration.
+    return_post_data : bool, default: False
+        If true, returns data
+    random : {None, int,  np.random.Generator}, default: None
+        Seed or random generator for reproducibility.
+    enable_localization : bool, default: False
+        If True, apply localization to the Kalman gain matrix.
+    localization_matrix : ???, default: None
+        Localization matrix.
 
 
     Returns
     -------
-
+    model_post : ndarray
+        Posterior model ensemble.
+    data_post : ndarray, only returned if ``return_post_data=True``
+        Posterior simulated data ensemble.
 
     """
     # Get number of ensembles and time steps
@@ -122,6 +127,14 @@ def esmda(model_prior, forward, data_obs, sigma, alphas=4, data_prior=None,
         # Calculate the Kalman gain
         K = CMD@Cinv
 
+        # TODO: Implement localization
+        if enable_localization and localization_matrix is not None:
+            localization_matrix = localization_matrix
+            K = K.reshape(-1, nd)
+            K = K * localization_matrix
+            # Reshape back to original shape
+            K = K.reshape(model_prior.shape[1], model_prior.shape[2], nd)
+
         # Update the ensemble parameters
         model_post += np.moveaxis(K @ (data_pert - data_prior).T, -1, 0)
 
@@ -131,3 +144,55 @@ def esmda(model_prior, forward, data_obs, sigma, alphas=4, data_prior=None,
 
     # Return posterior model and corresponding data
     return model_post, forward(model_post)
+
+
+def convert_positions_to_indices(positions, grid_dimensions):
+    """Convert 2D grid positions to 1D indices assuming zero-indexed positions.
+
+    Parameters
+    ----------
+    positions : ndarray
+        Array of (x, y) positions in the grid.
+    grid_dimensions : tuple
+        Dimensions of the grid (nx, ny).
+
+    Returns
+    -------
+    indices : ndarray
+        Array of indices corresponding to the positions in a flattened array.
+
+    """
+    nx, ny = grid_dimensions
+    # Ensure the positions are zero-indexed and correctly calculated for
+    # row-major order
+    return positions[:, 1] * nx + positions[:, 0]
+
+
+def build_localization_matrix(cov_matrix, data_positions, grid_dimensions):
+    """Build a localization matrix
+
+    Build a localization matrix from a full covariance matrix based on specific
+    data positions.
+
+    Parameters
+    ----------
+    cov_matrix : ndarray
+        The full (nx*ny) x (nx*ny) covariance matrix.
+    data_positions : ndarray
+        Positions in the grid for each data point (e.g., wells), zero-indexed.
+    grid_dimensions : tuple
+        Dimensions of the grid (nx, ny).
+
+    Returns
+    -------
+    loc_matrix : ndarray
+        Localization matrix of shape (nx*ny, number of data positions).
+
+    """
+    # Convert 2D positions of data points to 1D indices suitable for accessing
+    # the covariance matrix
+    indices = convert_positions_to_indices(
+            data_positions, grid_dimensions).astype(int)
+    # Extract the columns from the covariance matrix corresponding to each data
+    # point's position
+    return cov_matrix[:, indices]
diff --git a/resmda/reservoir_simulator.py b/resmda/reservoir_simulator.py
index 68ca29e..99bf14f 100644
--- a/resmda/reservoir_simulator.py
+++ b/resmda/reservoir_simulator.py
@@ -27,34 +27,44 @@ def __dir__():
 
 
 class Simulator:
-    """A small Reservoir Simulator.
+    """A small 2D Reservoir Simulator.
 
+    Created by following the course
+    **AESM304A - Flow and Simulation of Subsurface processes** at
+    Delft University of Technology (TUD); this particular part was taught by
+    Dr. D.V. Voskov, https://orcid.org/0000-0002-5399-1755.
 
     """
 
-    def __init__(self, nx, ny, phi=0.2, c_f=1e-5, p0=1, rho0=1, mu_w=1,
-                 rw=0.15, pres_ini=150, wells=None, dx=50, dz=10):
+    def __init__(self, nx, ny, phi=0.2, c_f=1e-5, p0=1.0, rho0=1.0, mu_w=1.0,
+                 rw=0.15, pres_ini=150.0, wells=None, dx=50.0, dz=10.0):
         """Initialize a Simulation instance.
 
         Parameters
         ----------
         nx, ny : int
-            Dimension of field
-        phi : float
-            Porosity (-)
-        c_f : float
-            Formation compressibility (1/kPa)
-        p0 : float
-            Initial pressure (bar or kPa?)
-        rho0 : float
-            Fixed density (kg/m3)
-        mu_w : float
-            Viscosity (cP - Pa s)
-        rw : float
-            Well radius (m)
-        pres_ini : initial pressure [?]
-        wells : location and pressure of wells [?]
-        dx, dz : floats
+            Dimension of field (number of cells).
+        phi : float, default: 0.2
+            Porosity (-).
+        c_f : float, default: 1e-5
+            Formation compressibility (1/kPa).
+        p0 : float, default: 1.0
+            Initial pressure (bar or kPa?).
+        rho0 : float, default: 1.0
+            Fixed density (kg/m3).
+        mu_w : float, default: 1.0
+            Viscosity (cP - Pa s).
+        rw : float, default: 0.15
+            Well radius (m).
+        pres_ini : float, default: 150.0
+            Initial pressure [?].
+        wells : {ndarray, None}, default: None
+            Nd array of shape (nwells, 3), indicating well locations (with cell
+            indices) and pressure. If None, the default is used, which is
+                np.array([[0, 0, 180], [self.nx-1, self.ny-1, 120]])
+            corresponding to a well in the first and in the last cell, with a
+            pressure of 180 and 120, respectively.
+        dx, dz : floats, default: 50.0, 10.0
             Cell dimensions in horizontal (dx) and vertical (dz)
             directions (m).
 
@@ -80,6 +90,13 @@ def __init__(self, nx, ny, phi=0.2, c_f=1e-5, p0=1, rho0=1, mu_w=1,
         self._vol_phi_cf = self.volumes * self.phi * self.c_f
 
         if wells is None:
+            # Default wells setup if none provided. Each well is specified by
+            # its grid coordinates followed by its pressure. The first well
+            # ([0, 0, 180]) is placed at the top-left corner with a pressure of
+            # 180 units, representing an injection pressure. The second well
+            # ([self.nx-1, self.ny-1, 120]), is located at the bottom-right
+            # corner, and has a pressure of 120 units, possibly a lower
+            # pressure or production scenario.
             self.wells = np.array([[0, 0, 180], [self.nx-1, self.ny-1, 120]])
         else:
             self.wells = np.array(wells)
@@ -89,7 +106,11 @@ def __init__(self, nx, ny, phi=0.2, c_f=1e-5, p0=1, rho0=1, mu_w=1,
 
     @property
     def _set_well_terms(self):
-        # Get well terms
+        """Set well terms.
+
+        Calculate well terms based on current permeability field, to be used in
+        the simulation. Adjust well impacts using calculated terms.
+        """
         wi = 2 * np.pi * self.perm_field[self.locs] * self.dz
         wi /= self.mu_w * np.log(0.208 * self.dx / self.rw)
 
@@ -98,7 +119,22 @@ def _set_well_terms(self):
         self._add_wells_d = wi
 
     def solve(self, pressure, dt):
-        """Construct K-matrix."""
+        """Construct & solve K-matrix for the simulation of pressure over time.
+
+        Parameters
+        ----------
+        pressure : ndarray
+            Current pressure state of the reservoir of size ``self.size``.
+
+        dt : float
+            Time step for the simulation.
+
+        Returns
+        -------
+        pressure : ndarray
+            Pressure state after applying the time step, of size ``self.size``.
+
+        """
 
         # Mobility ratio without permeability
         phi = self.rho0 * (1 + self.c_f * (pressure - self.p0)) / self.mu_w
@@ -146,10 +182,26 @@ def solve(self, pressure, dt):
     def __call__(self, perm_fields, dt=np.ones(10)*0.0001, data=False):
         """Run simulator.
 
+        Run the simulation across multiple time steps and possibly multiple
+        permeability scenarios.
+
         Parameters
         ----------
-        dt : array
-            Time steps.
+        perm_fields : ndarray
+            Permeability fields to simulate, either of dimension
+            (ne, nx, ny), or of dimension (nx, ny).
+
+        dt : ndarray, default: np.ones(10)*0.0001
+            Time steps to use for simulation.
+
+        data : {False, [ndarray, ndarray]}, default: False
+            Specific indices [nx, ny] to output data for; if False, return all
+            data
+
+        Returns
+        -------
+        simulation : ndarray
+            Simulation results over time for given permeability fields.
 
         """
         if perm_fields.ndim == 2:
@@ -169,6 +221,7 @@ def __call__(self, perm_fields, dt=np.ones(10)*0.0001, data=False):
             for i, d in enumerate(dt):
                 pressure[i+1, :] = self.solve(pressure[i, :], d)
             out[n, ...] = pressure.reshape((dt.size+1, *self.shape), order='F')
+
         if ne == 1:
             out = out[0, ...]
 
@@ -179,22 +232,46 @@ def __call__(self, perm_fields, dt=np.ones(10)*0.0001, data=False):
 
 
 class RandomPermeability:
+    """Return random permeability fields with certain statistical props."""
 
-    def __init__(self, nx, ny, perm_mean, perm_min, perm_max, length=(10, 10),
-                 theta=45, sigma_pr2=1.0, dtype='float32'):
-        self.nx = nx
-        self.ny = ny
-        self.nc = nx*ny
-        self.perm_mean = perm_mean
-        self.perm_min = perm_min
-        self.perm_max = perm_max
-        self.length = length
-        self.theta = theta
-        self.sigma_pr2 = sigma_pr2
-        self.dtype = dtype
+    def __init__(self, nx, ny, perm_mean, perm_min, perm_max,
+                 length=(10.0, 10.0), theta=45.0, sigma_pr2=1.0,
+                 dtype='float32'):
+        """Initialize parameters for generating random permeability fields.
+
+        Parameters
+        ----------
+        nx, ny : int
+            Dimensions of the grid.
+        perm_mean : float
+            Mean permeability.
+        perm_min, perm_max : float
+            Minimum and maximum values for permeability.
+        length : tuple of two floats, default: (10.0, 10.0)
+            Length scales for the correlation of permeability.
+        theta : float, default: 45.0
+            Rotation angle for the anisotropy in the permeability field.
+        sigma_pr2 : float, default: 1.0
+            Variance scale for the permeability.
+        dtype : str, default: 'float32'
+            Data type for computations, for precision and performance tuning.
+
+        """
+        self.nx, self.ny = nx, ny  # Grid dimensions
+        self.nc = nx * ny  # Total number of cells
+        self.perm_mean = perm_mean  # Permeability statistics
+        self.perm_min, self.perm_max = perm_min, perm_max
+        self.length, self.theta = length, theta  # Anisotropy parameters
+        self.sigma_pr2 = sigma_pr2  # Variance
+        self.dtype = dtype          # Data type
 
     @property
     def cov(self):
+        """Covariance matrix
+
+        Lazy-loaded covariance matrix, calculated based on anisotropy and
+        statistical parameters.
+        """
         if not hasattr(self, '_cov'):
             self._cov = covariance(
                 nx=self.nx, ny=self.ny, length=self.length,
@@ -204,58 +281,110 @@ def cov(self):
 
     @property
     def lcho(self):
+        """Lower Cholesky decomposition
+
+        Lower Cholesky decomposition of the covariance matrix, used for
+        generating random fields.
+        """
         if not hasattr(self, '_lcho'):
             self._lcho = sp.linalg.cholesky(self.cov, lower=True)
         return self._lcho
 
     def __call__(self, n, perm_mean=None, perm_min=None, perm_max=None,
                  random=None):
+        """Gerenate n random permeadility fields
+
+        Generate n random permeability fields using the specified statistical
+        parameters and random seed.
+
+        Parameters
+        ----------
+        n : int
+            Number of fields to generate
+        perm_mean : {float, None}, default: None
+            Mean permeability to override the initialized value.
+        perm_min, perm_max : {float, None}, default: None
+            Min and max permeability values to clip the fields.
+        random : {None, int,  np.random.Generator}, default: None
+            Seed or random generator for reproducibility.
+
+        Returns
+        -------
+        perm_fields : ndarray
+            An array of generated permeability fields.
+
+        """
+
         if perm_mean is None:
             perm_mean = self.perm_mean
         if perm_min is None:
             perm_min = self.perm_min
         if perm_max is None:
             perm_max = self.perm_max
+
+        # Initialize fields with mean permeability
         out = np.full((n, self.nx, self.ny), perm_mean, order='F')
         for i in range(n):
-            z = rng(random).normal(size=self.nc)
+            z = rng(random).normal(size=self.nc)  # Generate random numbers
+            # Apply the Cholesky transform
             out[i, ...] += (self.lcho @ z).reshape(
                     (self.nx, self.ny), order='F')
 
+        # Clip the results to stay within specified bounds
         return out.clip(perm_min, perm_max)
 
 
 def covariance(nx, ny, length, theta, sigma_pr2, dtype='float32'):
-    """Build co-variance matrix for permeability.
+    """Return covariance matrix
+
+    Generate covariance matrix based on grid size, anisotropy, and statistical
+    parameters.
+
 
     Parameters
     ----------
     nx, ny : int
-        Number of cells in x and y
+        Dimensions of the grid.
+    length : float
+        Length scales for the correlation of permeability.
     theta : float
-        Angle (in degrees)
+        Rotation angle for the anisotropy in the permeability field.
+    sigma_pr2 : float
+        Variance scale for the permeability.
+    dtype : str, default: 'float32'
+        Data type for computations.
+
+
+    Returns
+    -------
+    cov : ndarray
+        Covariance matrix for the permeability field.
+
     """
-    nc = nx*ny
-    cost = np.cos(theta)
-    sint = np.sin(theta)
+    nc = nx * ny  # Total number of cells
+    # Precompute cosine and sine of the rotation angle
+    cost, sint = np.cos(theta), np.sin(theta)
 
-    # 1. Fill the first row nx * nc, but vertically
+    # 1. Fill the first row of the covariance matrix
     tmp1 = np.zeros([nx, nc], dtype=dtype)
     for i in range(nx):
-        tmp1[i, 0] = 1.0  # diagonal
+        tmp1[i, 0] = 1.0  # Set diagonal
         for j in range(i+1, nc):
+            # Distance in the x and y directions
             d0 = (j % nx) - i
             d1 = (j // nx)
+            # Rotate coordinates
             rot0 = cost*d0 - sint*d1
             rot1 = sint*d0 + cost*d1
+            # Calculate the scaled distance
             hl = np.sqrt((rot0/length[0])**2 + (rot1/length[1])**2)
 
-            # Calculate value.
-            if sigma_pr2:  # Sphere formula
+            # Sphere formula for covariance, modified for anisotropy
+            if sigma_pr2:  # Non-zero variance scale
                 if hl <= 1:
                     tmp1[i, j-i] = sigma_pr2 * (1 - 1.5*hl + hl**3/2)
 
-            else:  # Gaspari Cohn
+            else:  # Gaspari-Cohn function for smoothness
                 if hl < 1:
                     tmp1[i, j-i] = (-(hl**5)/4 + (hl**4)/2 + (hl**3)*5/8 -
                                     (hl**2)*5/3 + 1)
diff --git a/resmda/utils.py b/resmda/utils.py
index 19bdae5..270a8a6 100644
--- a/resmda/utils.py
+++ b/resmda/utils.py
@@ -55,6 +55,7 @@ def rng(random=None):
 
 
     Returns
+    -------
     rng : random number generator
         A :class:`numpy.random.Generator` instance.