WIP: 3D heat-equation

test-data/ert/heat_equation/definition.py

@@ -5,15 +5,20 @@
 # Number of grid-cells in x and y direction
 nx = 10
 
+# Number of grid-cells in z direction (depth layers)
+nz = 3
+
 # time steps
 k_start = 0
 k_end = 500
 
 # Define initial condition, i.e., the initial temperature distribution.
 # How you define initial conditions will effect the spread of results,
 # i.e., how similar different realisations are.
-u_init = np.zeros((k_end, nx, nx))
-u_init[:, 5:7, 5:7] = 100
+# Heat source is placed only in the middle z-layer(s) so that
+# z-diffusion creates genuine layer differentiation.
+u_init = np.zeros((k_end, nx, nx, nz))
+u_init[:, 5:7, 5:7, nz // 2] = 100
 
 # Resolution in the x-direction (nothing to worry about really)
 dx = 1
@@ -33,6 +38,8 @@ class Coordinate(NamedTuple):
     Coordinate(7, 2),
 ]
 
-summary_names = [f"HEAT_{coord.x}_{coord.y}" for coord in obs_coordinates]
+summary_names = [
+    f"HEAT_{coord.x}_{coord.y}_{z}" for coord in obs_coordinates for z in range(nz)
+]
 
 obs_times = np.linspace(10, k_end, 8, endpoint=False, dtype=int)

test-data/ert/heat_equation/generate_files.py

@@ -11,7 +11,7 @@
 import pandas as pd
 import resfo
 import xtgeo
-from definition import Coordinate, obs_coordinates, obs_times
+from definition import Coordinate, nz, obs_coordinates, obs_times, u_init
 from heat_equation import heat_equation, sample_prior_conductivity
 
 # Some seeds produce priors that yield poor results.
@@ -20,7 +20,7 @@
 
 NCOL = 10
 NROW = 10
-NLAY = 1
+NLAY = nz
 
 
 def create_egrid_file():
@@ -47,6 +47,7 @@ def make_observations(
             "k": pd.Series(dtype=int),
             "x": pd.Series(dtype=int),
             "y": pd.Series(dtype=int),
+            "z": pd.Series(dtype=int),
             "value": pd.Series(dtype=float),
             "sd": pd.Series(dtype=float),
         }
@@ -55,22 +56,25 @@ def make_observations(
     # Create dataframe with observations and necessary meta data.
     for coordinate in coordinates:
         for k in times:
-            # The reason for u[k, y, x] instead of the perhaps more natural u[k, x, y],
-            # is due to a convention followed by matplotlib's `pcolormesh`
-            # See documentation for details.
-            value = field[k, coordinate.x, coordinate.y]
-            sd = error(value)
-            df_ = pd.DataFrame(
-                {
-                    "k": [k],
-                    "x": [coordinate.x],
-                    "y": [coordinate.y],
-                    "value": [value + rng.normal(loc=0.0, scale=sd)],
-                    "sd": [sd],
-                }
-            )
-            d = pd.concat([d, df_])
-    d = d.set_index(["k", "x", "y"], verify_integrity=True)
+            for z in range(nz):
+                # The reason for u[k, y, x] instead of the perhaps more
+                # natural u[k, x, y],
+                # is due to a convention followed by matplotlib's `pcolormesh`
+                # See documentation for details.
+                value = field[k, coordinate.x, coordinate.y, z]
+                sd = error(value)
+                df_ = pd.DataFrame(
+                    {
+                        "k": [k],
+                        "x": [coordinate.x],
+                        "y": [coordinate.y],
+                        "z": [z],
+                        "value": [value + rng.normal(loc=0.0, scale=sd)],
+                        "sd": [sd],
+                    }
+                )
+                d = pd.concat([d, df_])
+    d = d.set_index(["k", "x", "y", "z"], verify_integrity=True)
 
     return d
 
@@ -86,7 +90,7 @@ def generate_priors():
     corr_lengths = rng.normal(loc=0.8, scale=0.1, size=10)
     for i in range(10):
         cond = sample_prior_conductivity(
-            ensemble_size=1, nx=nx, rng=rng, corr_length=corr_lengths[i]
+            ensemble_size=1, nx=nx, nz=nz, rng=rng, corr_length=corr_lengths[i]
         )
         resfo.write(
             f"cond_{i}.bgrdecl",
@@ -110,18 +114,19 @@ def create_summary_observations(df_obs: pd.DataFrame):
             obs_date = START_DATE + datetime.timedelta(days=int(t))
 
             for coord in obs_coordinates:
-                idx = (int(t), int(coord.x), int(coord.y))
-                row = df_obs.loc[idx]
-                value = float(row["value"])
-                error = float(row["sd"])
-
-                f.write(
-                    f"""SUMMARY_OBSERVATION HEAT_{coord.x}_{coord.y}_{t}
+                for z in range(nz):
+                    idx = (int(t), int(coord.x), int(coord.y), int(z))
+                    row = df_obs.loc[idx]
+                    value = float(row["value"])
+                    error = float(row["sd"])
+
+                    f.write(
+                        f"""SUMMARY_OBSERVATION HEAT_{coord.x}_{coord.y}_{z}_{t}
 {{
     VALUE   = {value:.16e};
     ERROR   = {error:.16e};
     DATE    = {obs_date:%Y-%m-%d};
-    KEY     = HEAT_{coord.x}_{coord.y};
+    KEY     = HEAT_{coord.x}_{coord.y}_{z};
     LOCALIZATION {{
         EAST = {coord.x + 0.5};
         NORTH = {coord.y + 0.5};
@@ -130,7 +135,7 @@ def create_summary_observations(df_obs: pd.DataFrame):
 }};
 
 """
-                )
+                    )
 
 
 if __name__ == "__main__":
@@ -143,23 +148,18 @@ def create_summary_observations(df_obs: pd.DataFrame):
     k_start = 0
     k_end = 500
 
-    # Define initial condition, i.e., the initial temperature distribution.
-    # How you define initial conditions will effect the spread of results,
-    # i.e., how similar different realisations are.
-    u_init = np.zeros((k_end, nx, nx))
-    u_init[:, 5:7, 5:7] = 100
-
     cond_truth = sample_prior_conductivity(
-        ensemble_size=1, nx=nx, rng=rng, corr_length=0.8
-    ).reshape(nx, nx)
+        ensemble_size=1, nx=nx, nz=nz, rng=rng, corr_length=0.8
+    ).reshape(nx, nx, nz)
 
     # Resolution in the x-direction (nothing to worry about really)
     dx = 1
 
     # Calculate maximum `dt`.
     # If higher values are used, the numerical solution will become unstable.
     # Note that this could be avoided if we used an implicit solver.
-    dt = dx**2 / (4 * np.max(cond_truth))
+    neighbors = 6 if nz > 1 else 4
+    dt = dx**2 / (neighbors * np.max(cond_truth))
 
     u_t = heat_equation(u_init, cond_truth, dx, dt, k_start, k_end, rng=rng)
 

test-data/ert/heat_equation/heat_equation.py

@@ -15,6 +15,7 @@
     k_end,
     k_start,
     nx,
+    nz,
     obs_coordinates,
     obs_times,
     summary_names,
@@ -60,7 +61,7 @@ def create_summary_smspec_unsmry(
     smspec = Smspec(
         nx=nx,
         ny=nx,
-        nz=1,
+        nz=nz,
         restarted_from_step=0,
         num_keywords=1 + len(summary_keys),
         restart="        ",
@@ -88,45 +89,81 @@ def heat_equation(
     rng: np.random.Generator,
     scale: float | None = None,
 ) -> npt.NDArray[np.float64]:
-    """2D heat equation that suppoheat_erts field of heat coefficients.
+    """3D heat equation that supports a field of heat coefficients.
+
+    Solves the heat equation on a (nx, nx, nz) grid using explicit
+    finite differences.  When nz == 1 the z-stencil terms vanish and
+    the solver reduces to the classic 2D formulation.
 
     Based on:
     https://levelup.gitconnected.com/solving-2d-heat-equation-numerically-using-python-3334004aa01a
     """
     u_ = u.copy()
-    nx = u.shape[1]  # number of grid cells
-    assert cond.shape == (nx, nx)
+    nx = u.shape[1]  # number of grid cells in x and y
+    n_z = u.shape[3]  # number of grid cells in z
+    assert cond.shape == (nx, nx, n_z)
     gamma = (cond * dt) / (dx**2)
 
     # Pre-sample all noise at once
     if scale is not None:
         num_steps = k_end - k_start - 1
-        noise_all = rng.normal(0, scale, size=(num_steps, nx - 2, nx - 2))
+        noise_all = rng.normal(0, scale, size=(num_steps, nx - 2, nx - 2, n_z))
 
     for k in range(k_start, k_end - 1):
-        # Vectorized finite difference
-        u_[k + 1, 1:-1, 1:-1] = (
-            gamma[1:-1, 1:-1]
+        # Vectorized finite difference in x and y (interior cells)
+        xy_stencil = (
+            gamma[1:-1, 1:-1, :]
             * (
-                u_[k, 2:, 1:-1]  # i+1
-                + u_[k, :-2, 1:-1]  # i-1
-                + u_[k, 1:-1, 2:]  # j+1
-                + u_[k, 1:-1, :-2]  # j-1
-                - 4 * u_[k, 1:-1, 1:-1]
+                u_[k, 2:, 1:-1, :]  # i+1
+                + u_[k, :-2, 1:-1, :]  # i-1
+                + u_[k, 1:-1, 2:, :]  # j+1
+                + u_[k, 1:-1, :-2, :]  # j-1
+                - 4 * u_[k, 1:-1, 1:-1, :]
             )
-            + u_[k, 1:-1, 1:-1]
+            + u_[k, 1:-1, 1:-1, :]
         )
+
+        if n_z > 1:
+            # z-direction stencil for interior z-layers
+            z_contrib = np.zeros_like(xy_stencil)
+            z_contrib[:, :, 1:-1] = (
+                gamma[1:-1, 1:-1, 1:-1]
+                * (
+                    u_[k, 1:-1, 1:-1, 2:]  # l+1
+                    + u_[k, 1:-1, 1:-1, :-2]  # l-1
+                    - 2 * u_[k, 1:-1, 1:-1, 1:-1]
+                )
+            )
+            # z-boundary layers: zero Dirichlet (u=0 at ghost cell)
+            z_contrib[:, :, 0] = gamma[1:-1, 1:-1, 0] * (
+                u_[k, 1:-1, 1:-1, 1] - 2 * u_[k, 1:-1, 1:-1, 0]
+            )
+            z_contrib[:, :, -1] = gamma[1:-1, 1:-1, -1] * (
+                u_[k, 1:-1, 1:-1, -2] - 2 * u_[k, 1:-1, 1:-1, -1]
+            )
+            xy_stencil += z_contrib
+
+        u_[k + 1, 1:-1, 1:-1, :] = xy_stencil
+
         # Add pre-sampled noise
         if scale is not None:
             noise_idx = k - k_start
-            u_[k + 1, 1:-1, 1:-1] += noise_all[noise_idx]
+            u_[k + 1, 1:-1, 1:-1, :] += noise_all[noise_idx]
 
     return u_
 
 
-def sample_prior_conductivity(ensemble_size, nx, rng, corr_length):
+def sample_prior_conductivity(ensemble_size, nx, nz, rng, corr_length):
     mesh = np.meshgrid(np.linspace(0, 1, nx), np.linspace(0, 1, nx))
-    return np.exp(geostat.gaussian_fields(mesh, rng, ensemble_size, r=corr_length))
+    # Generate an independent 2D conductivity field for each z-layer so
+    # that the prior has genuine z-variation. This is necessary for the
+    # EnKF to produce layer-differentiated posteriors and for the
+    # localization z-ordering to matter.
+    layers = [
+        np.exp(geostat.gaussian_fields(mesh, rng, ensemble_size, r=corr_length))
+        for _ in range(nz)
+    ]
+    return np.stack(layers, axis=-1)
 
 
 def load_parameters(filename):
@@ -143,14 +180,18 @@ def load_parameters(filename):
 
     if iteration == 0:
         cond = sample_prior_conductivity(
-            ensemble_size=1, nx=nx, rng=rng, corr_length=float(parameters["x"]["value"])
-        ).reshape(nx, nx)
+            ensemble_size=1,
+            nx=nx,
+            nz=nz,
+            rng=rng,
+            corr_length=float(parameters["x"]["value"]),
+        ).reshape(nx, nx, nz)
 
         resfo.write(
             "cond.bgrdecl", [("COND    ", cond.flatten(order="F").astype(np.float32))]
         )
     else:
-        cond = resfo.read("cond.bgrdecl")[0][1].reshape(nx, nx)
+        cond = resfo.read("cond.bgrdecl")[0][1].reshape(nx, nx, nz, order="F")
 
     # The update may give non-physical parameter values, which here means
     # negative heat conductivity. Setting negative values to a small positive
@@ -160,30 +201,35 @@ def load_parameters(filename):
     # Calculate maximum `dt`.
     # If higher values are used, the numerical solution will become unstable.
     # Note that this could be avoided if we used an implicit solver.
-    dt = dx**2 / (4 * np.max(cond))
+    # The stability bound for a 3D explicit scheme is dx^2 / (6 * max(cond)),
+    # but when nz == 1 the z-stencil vanishes so dx^2 / (4 * max(cond)) suffices.
+    neighbors = 6 if nz > 1 else 4
+    dt = dx**2 / (neighbors * np.max(cond))
 
     scaled_u_init = u_init * float(parameters["t"]["value"])
 
     response = heat_equation(scaled_u_init, cond, dx, dt, k_start, k_end, rng)
 
+    # Observations are extracted per z-layer
     index = sorted((obs.x, obs.y) for obs in obs_coordinates)
     for time_step in obs_times:
         with Path(f"gen_data_{time_step}.out").open("w", encoding="utf-8") as f:
-            f.writelines(f"{response[time_step][i]}\n" for i in index)
+            for i in index:
+                f.writelines(f"{response[time_step][i][z]}\n" for z in range(nz))
 
     time_map = []
     start_date = datetime.date(2010, 1, 1)
 
     summary_values = {name: [] for name in summary_names}
-    # for time_step in obs_times:
     for time_step in range(k_start, k_end):
         time_map.append(
             (start_date + datetime.timedelta(days=float(time_step))).isoformat()
         )
         for obs in obs_coordinates:
-            summary_values[f"HEAT_{obs.x}_{obs.y}"].append(
-                response[time_step][obs.x, obs.y]
-            )
+            for z in range(nz):
+                summary_values[f"HEAT_{obs.x}_{obs.y}_{z}"].append(
+                    response[time_step][obs.x, obs.y, z]
+                )
 
     smspec, unsmry = create_summary_smspec_unsmry(
         summary_vectors=summary_values,