Commitin

test-data/heat_equation/definition.py

@@ -15,6 +15,11 @@
 u_init = np.zeros((k_end, nx, nx))
 u_init[:, 5:7, 5:7] = 100
 
+# Initialize 3D arrays (with depth 1 to simulate 2D)
+u_3d = np.zeros((k_end, 1, nx, nx))
+u_3d[0, 0] = u_init[0]
+alpha_3d = np.ones((1, nx, nx))
+
 # Resolution in the x-direction (nothing to worry about really)
 dx = 1
 

test-data/heat_equation/heat_equation.py

@@ -1,13 +1,20 @@
 #!/usr/bin/env python3
 """Partial Differential Equations to use as forward models."""
 
-import sys
 from typing import Optional
 
 import geostat
 import numpy as np
 import numpy.typing as npt
-from definition import dx, k_end, k_start, nx, obs_coordinates, obs_times, u_init
+from definition import (
+    dx,
+    k_end,
+    k_start,
+    nx,
+    obs_coordinates,
+    obs_times,
+    u_3d,
+)
 
 
 def heat_equation(
@@ -51,23 +58,116 @@ def heat_equation(
     return _u
 
 
-def sample_prior_conductivity(ensemble_size, nx, rng):
-    mesh = np.meshgrid(np.linspace(0, 1, nx), np.linspace(0, 1, nx))
-    return np.exp(geostat.gaussian_fields(mesh, rng, ensemble_size, r=0.8))
+def heat_equation_3d(
+    u: npt.NDArray[np.float64],
+    alpha: npt.NDArray[np.float64],
+    dx: int,
+    dt: float,
+    k_start: int,
+    k_end: int,
+    rng: np.random.Generator,
+    scale: Optional[float] = None,
+) -> npt.NDArray[np.float64]:
+    """3D heat equation that supports field of heat coefficients."""
+    _u = u.copy()
+    nz, ny, nx = u.shape[1:]  # number of grid cells in each dimension
+    print(alpha.shape)
+    assert alpha.shape == (nx, ny, nz)
+    gamma = (alpha * dt) / (dx**2)
+
+    for k in range(k_start, k_end - 1, 1):
+        if nz == 1:  # 2D case embedded in 3D
+            for j in range(1, ny - 1, dx):
+                for l in range(1, nx - 1, dx):
+                    if scale is not None:
+                        noise = rng.normal(scale=scale)
+                    else:
+                        noise = 0
+                    _u[k + 1, 0, j, l] = (
+                        gamma[l, j, 0]  # Changed from gamma[0, j, l]
+                        * (
+                            _u[k][0][j + 1][l]
+                            + _u[k][0][j - 1][l]
+                            + _u[k][0][j][l + 1]
+                            + _u[k][0][j][l - 1]
+                            - 4 * _u[k][0][j][l]
+                        )
+                        + _u[k][0][j][l]
+                        + noise
+                    )
+        else:  # Full 3D case
+            for i in range(1, nz - 1, dx):
+                for j in range(1, ny - 1, dx):
+                    for l in range(1, nx - 1, dx):
+                        if scale is not None:
+                            noise = rng.normal(scale=scale)
+                        else:
+                            noise = 0
+                        _u[k + 1, i, j, l] = (
+                            gamma[l, j, i]  # Changed from gamma[i, j, l]
+                            * (
+                                _u[k][i + 1][j][l]
+                                + _u[k][i - 1][j][l]
+                                + _u[k][i][j + 1][l]
+                                + _u[k][i][j - 1][l]
+                                + _u[k][i][j][l + 1]
+                                + _u[k][i][j][l - 1]
+                                - 6 * _u[k][i][j][l]
+                            )
+                            + _u[k][i][j][l]
+                            + noise
+                        )
+
+    return _u
+
+
+def sample_prior_conductivity(nx, ny, nz=None, rng=None):
+    """
+    Sample prior conductivity for 2D or 3D cases.
+
+    Parameters:
+    - nx, ny: int, number of grid points in x and y directions
+    - nz: int or None, number of grid points in z direction (None for 2D case)
+    - rng: numpy.random.Generator object or None
+
+    Returns:
+    - numpy array of shape (nx, ny, nz) for 3D or (nx, ny) for 2D
+    """
+    if rng is None:
+        rng = np.random.default_rng()
+
+    if nz is None:
+        # 2D case
+        mesh = np.meshgrid(np.linspace(0, 1, nx), np.linspace(0, 1, ny))
+    else:
+        # 3D case
+        mesh = np.meshgrid(
+            np.linspace(0, 1, nx), np.linspace(0, 1, ny), np.linspace(0, 1, nz)
+        )
+
+    fields = geostat.gaussian_fields(mesh, rng, N=1, r=0.8)
+
+    if nz is None:
+        # Reshape for 2D case
+        return np.exp(fields.reshape(ny, nx)).T
+    else:
+        # Reshape for 3D case
+        return np.exp(fields.reshape(nz, ny, nx)).transpose(2, 1, 0)
 
 
 if __name__ == "__main__":
-    iens = int(sys.argv[1])
+    # iens = int(sys.argv[1])
+    iens = 42
     rng = np.random.default_rng(iens)
-    cond = sample_prior_conductivity(ensemble_size=1, nx=nx, rng=rng).reshape(nx, nx)
+    cond = sample_prior_conductivity(nx=nx, ny=nx, nz=1, rng=rng)
 
     # Write the array to a GRDECL formatted file
     with open("cond.grdecl", "w", encoding="utf-8") as f:
         f.write("COND\n")  # Write the property name
         f.write("-- Conductivity data\n")  # Optional comment line
 
         # Write the data
-        for row in cond:
+        for row in cond[:, :, 0]:
             for value in row:
                 f.write(f"{value:.6f} ")
             f.write("\n")
@@ -79,9 +179,9 @@ def sample_prior_conductivity(ensemble_size, nx, rng):
     # Note that this could be avoided if we used an implicit solver.
     dt = dx**2 / (4 * max(np.max(cond), np.max(cond)))
 
-    response = heat_equation(u_init, cond, dx, dt, k_start, k_end, rng)
+    response = heat_equation_3d(u_3d, cond, dx, dt, k_start, k_end, rng)
 
-    index = sorted((obs.x, obs.y) for obs in obs_coordinates)
+    index = sorted((0, obs.x, obs.y) for obs in obs_coordinates)
     for time_step in obs_times:
         with open(f"gen_data_{time_step}.out", "w", encoding="utf-8") as f:
             for i in index: