geoloop 0.0.1__py3-none-any.whl → 1.0.0b1__py3-none-any.whl

geoloop/geoloopcore/soilproperties.py

@@ -0,0 +1,152 @@
+import numpy as np
+
+from geoloop.configuration import SingleRunConfig
+from geoloop.geoloopcore.strat_interpolator import StratInterpolator, TgInterpolator
+from geoloop.lithology.process_lithology import ProcessLithologyToThermalConductivity
+
+
+class SoilProperties:
+    """
+    Store soil thermal properties, including geotherm (temperature vs. depth)
+    and thermal conductivity profiles (depth-dependent or lithology-based).
+
+    Supports:
+    - Depth-based interpolation of ground temperature Tg(z)
+    - Depth-based or lithology-based interpolation of k_s(z)
+    - Scaling of conductivity profiles
+    """
+
+    def __init__(
+        self,
+        Tg: list | float,
+        k_s: list | float,
+        z_Tg: list = None,
+        Tgrad: float = 0.0,
+        z_k_s: list = None,
+        k_s_scale: float = 1.0,
+        lithology_to_k: ProcessLithologyToThermalConductivity = None,
+        alfa: float = 1e-6,
+    ) -> None:
+        """
+        Initialize SoilProperties, containing ground temperature and thermal
+        conductivity as functions of depth.
+
+        Parameters
+        ----------
+        Tg : float or list
+            Ground temperature(s) used for interpolation.
+            - If array-list: temperature at depths specified by `z_Tg`.
+            - If float: constant surface temperature for linear gradient computation.
+        k_s : float or list
+            Thermal conductivity values (W/m·K), possibly depth-dependent.
+            If list-like, entries correspond to depths in `z_k_s`.
+        z_Tg : list or None, optional
+            Depth values (m) corresponding to `Tg` samples.
+            If None, temperature is assumed constant or determined by `Tgrad`.
+        Tgrad : float, optional
+            Vertical geothermal gradient (K/m). Default is 0 (no gradient).
+        z_k_s : list or None, optional
+            Depth values (m) corresponding to `k_s` samples.
+            If None, conductivity is assumed constant at all depths.
+        k_s_scale : float, optional
+            Scaling factor applied to thermal conductivity profiles (default 1.0).
+        lithology_to_k : lithology_to_k or None, optional
+            If provided, conductivity is derived from lithology-based profiles
+            using the lithology_to_k lookup structure instead of `k_s`/`z_k_s`.
+        alfa : float, optional
+            Thermal diffusivity of the ground (m²/s). Default is 1e-6.
+
+        Notes
+        -----
+        - If `lithology_to_k` is provided, the lithology-based thermal conductivity
+          profile overrides `k_s` and `z_k_s`.
+        - This constructor initializes two interpolators:
+            * `interpolatorTg` – for temperature vs. depth
+            * `interpolatorKs` – for conductivity vs. depth
+        """
+        # subsurface temperature and/or geothermal gradient
+        self.Tg = Tg
+        self.z_Tg = z_Tg
+        self.Tgrad = Tgrad
+        self.interpolatorTg = TgInterpolator(self.z_Tg, self.Tg, self.Tgrad)
+
+        # subsurface bulk thermal conductivity
+        self.k_s = np.asarray(k_s)
+        if z_k_s == None:
+            self.z_k_s = np.ones_like(k_s)
+        else:
+            self.z_k_s = np.asarray(z_k_s)
+        self.k_s_scale = k_s_scale
+        self.lithology_to_k = lithology_to_k
+
+        self.alfa = alfa  # thermal diffusivity
+
+        # interpolate thermal conductivity
+        if isinstance(self.lithology_to_k, ProcessLithologyToThermalConductivity):
+            # lithology-based thermal conductivity profile
+            zstart, zend = self.lithology_to_k.get_start_end_depths()
+            zval = self.lithology_to_k.get_thermcon_sample_profile(0).kh_bulk
+            self.interpolatorKs = StratInterpolator(zend, zval)
+        else:
+            self.interpolatorKs = StratInterpolator(
+                self.z_k_s, self.k_s, stepvalue=False
+            )
+
+    def getTg(self, z: float | np.ndarray) -> float | np.ndarray:
+        """Return ground temperature at depth `z`."""
+        return self.interpolatorTg.getTg(z)
+
+    def get_k_s(
+        self, zstart: np.ndarray, zend: np.ndarray, isample: int = 0
+    ) -> np.ndarray:
+        """
+        Retrieves and interpolates thermal conductivity-depth profile for depth-segments in the simulation, for the basecase
+        in a single simulation or for the specified sample in a stochastic simulation.
+
+        Parameters
+        ----------
+        zstart : np.ndarray
+            Array of starting depths of each segment.
+        zend : np.ndarray
+            Array of ending depths of each segment.
+        isample : int, optional
+            Index selecting conductivity sample from lithology_to_k (default 0).
+
+        Returns
+        -------
+        np.ndarray
+            Thermal conductivity-depth profile with depth-resolution for the depth-segments in the simulation.
+        """
+        if isinstance(self.lithology_to_k, ProcessLithologyToThermalConductivity):
+            zval = self.lithology_to_k.get_thermcon_sample_profile(isample).kh_bulk
+            self.interpolatorKs.zval = zval
+
+        ks = self.interpolatorKs.interp(zstart, zend)
+        return ks * self.k_s_scale
+
+    @classmethod
+    def from_config(cls, config: SingleRunConfig) -> "SoilProperties":
+        """
+        Build a SoilProperties object from a configuration object.
+
+        Parameters
+        ----------
+        config : SingleRunConfig
+            Object containing entries including Tg, k_s, stratigraphy, scaling,
+            lithology parameters, and diffusivity.
+
+        Returns
+        -------
+        SoilProperties
+        """
+
+        return SoilProperties(
+            Tg=config.Tg,
+            z_Tg=config.z_Tg,
+            Tgrad=config.Tgrad,
+            k_s=config.k_s,
+            z_k_s=config.z_k_s,
+            k_s_scale=config.k_s_scale,
+            alfa=config.alfa,
+            lithology_to_k=config.lithology_to_k,  # <-- already built upstream
+        )

geoloop/geoloopcore/strat_interpolator.py

@@ -0,0 +1,194 @@
+from typing import Union
+
+import numpy as np
+
+ArrayLike = Union[np.ndarray, list, float, int]
+
+
+class TgInterpolator:
+    """
+    class to manage interpolation of temperature over depth.
+    """
+
+    def __init__(self, z_Tg: ArrayLike | None, Tg: ArrayLike, Tgrad: float) -> None:
+        """
+        Parameters
+        ----------
+        z_Tg : array_like or None
+            Depth values (m) at which temperature samples in `Tg` are defined.
+            If None, `Tg` is treated as a scalar.
+        Tg : float or array_like
+            Ground temperature value(s).
+            - If float: temperature at surface.
+            - If array_like: temperature profile at depths `z_Tg`.
+        Tgrad : float
+            Geothermal gradient (C/m). Used only when `Tg` is scalar.
+        """
+        self.z_Tg = z_Tg
+        self.Tg = Tg
+        self.Tgrad = Tgrad
+
+    def interp_Tg(self, z: ArrayLike) -> np.ndarray:
+        """
+        Interpolate temperature at depth.
+
+        Parameters
+        ----------
+        z : array_like
+            Depth(s) in meters.
+
+        Returns
+        -------
+        np.ndarray
+            Interpolated temperature values.
+        """
+        return np.interp(z, self.z_Tg, self.Tg)
+
+    def getTg(self, z: ArrayLike) -> ArrayLike:
+        """
+        Get ground temperature at depth.
+
+        Parameters
+        ----------
+        z : float or array_like
+            Depth(s) in meters.
+
+        Returns
+        -------
+        float or np.ndarray
+            Temperature at depth.
+
+        Notes
+        -----
+        - If `Tg` is scalar: Tg + Tgrad * 0.01 * z
+        - If `Tg` is array: interpolated from depth-temperature profile
+        """
+        if np.isscalar(self.Tg):
+            return self.Tg + self.Tgrad * 0.01 * z
+        else:
+            return self.interp_Tg(z)
+
+
+class StratInterpolator:
+    """
+    class to manage interpolation of soil properties to averaged segment properties.
+    """
+
+    def __init__(
+        self,
+        zcoord: ArrayLike,
+        zval: ArrayLike,
+        nz: int = 10000,
+        stepvalue: bool = True,
+    ) -> None:
+        """
+        Parameters
+        ----------
+        zcoord : array_like
+            Depth coordinates (m) defining the property profile.
+        zval : array_like
+            Property values corresponding to `zcoord`.
+        nz : int, optional
+            Number of points for the internal fine-resolution discretization.
+        stepvalue : bool, optional
+            If True:
+                Interpret `zval` as piecewise-constant over intervals
+                (stepwise profile).
+            If False:
+                Use linear interpolation between points.
+        """
+        self.zcoord = zcoord
+        self.zval = zval
+        self.nz = nz
+
+        # high-resolution depth array
+        self.zp = np.linspace(0, self.zcoord[-1], num=nz)
+        self.dz = self.zp[1] - self.zp[0]
+
+        self.stepvalue = stepvalue
+        if self.stepvalue:
+            self.init_indices()
+
+    @property
+    def zval(self) -> np.ndarray:
+        return self._zval
+
+    @zval.setter
+    def zval(self, value: ArrayLike) -> None:
+        self._zval = value
+
+    def init_indices(self) -> None:
+        """
+        For the high resolution array self.zp get indices on the courser intervalled zcoord array.
+        The indices are determined by searchsorted of np, giving the lower index when on interval zlith[0] to zlith[1].
+        This is only used if stepvalue=True.
+        """
+        self.indexl = self.zcoord.searchsorted(self.zp)
+
+    def interp(self, zstart: np.ndarray, zend: np.ndarray) -> np.ndarray:
+        """
+        Interpolate for zstart, zend arrays of the segments, the average value of zval.
+        Compute average property value for each depth segment.
+
+        Parameters
+        ----------
+        zstart : np.ndarray
+            Segment start depths (m).
+        zend : np.ndarray
+            Segment end depths (m).
+
+        Returns
+        -------
+        np.ndarray
+            Average property value per segment.
+        """
+        val = zstart * 0.0
+        ilstart = max(0, int(zstart[0] / self.dz))
+        for i in range(len(zend)):
+            valsum = 0
+            ilend = min(self.nz, int(round(zend[i] / self.dz)))
+            ilrange = max(1, ilend - ilstart)
+            if self.stepvalue:
+                for il in range(ilrange):
+                    valsum += self.zval[self.indexl[il + ilstart]]
+                val[i] = valsum / ilrange
+            else:
+                vals = np.interp(self.zp[ilstart:ilend], self.zcoord, self.zval)
+                val[i] = np.sum(vals) / ilrange
+            ilstart = min(self.nz - 1, ilend)
+        return val
+
+    def interp_plot(self, zstart: np.ndarray, zend: np.ndarray) -> np.ndarray:
+        """
+        Generate a high-resolution vector of interpolated property values.
+
+        Parameters
+        ----------
+        zstart : np.ndarray
+            Segment start depths (m).
+        zend : np.ndarray
+            Segment end depths (m).
+
+        Returns
+        -------
+        np.ndarray
+            Interpolated property values along internal grid `zp`.
+        """
+        val = np.zeros_like(self.zp)
+        ilstart = max(0, int(zstart[0] / self.dz))
+
+        for i in range(len(zend)):
+            ilend = min(self.nz, int(round(zend[i] / self.dz)))
+            if ilend > ilstart:
+                if self.stepvalue:
+                    val[ilstart:ilend] = self.zval[self.indexl[ilstart]]
+                else:
+                    vals = np.interp(self.zp[ilstart:ilend], self.zcoord, self.zval)
+                    val[ilstart:ilend] = vals
+            ilstart = ilend
+
+        # Fill the remaining 0 values with the last used value
+        last_nonzero_index = (val != 0).nonzero()[0][-1]
+        val[last_nonzero_index + 1 :] = val[last_nonzero_index]
+
+        return val

geoloop/lithology/__init__.py

@@ -0,0 +1,3 @@
+"""
+This package contains the modules with the main classes and functions for creating subsurface thermal conductivity models for Geoloop simulations.
+"""

geoloop/lithology/plot_lithology.py

@@ -0,0 +1,277 @@
+import sys
+from pathlib import Path
+
+import numpy as np
+import xarray as xr
+from matplotlib import pyplot as plt
+from matplotlib.colors import to_hex, to_rgb
+
+from geoloop.constants import lithology_colors, lithology_names_english
+from geoloop.geoloopcore.strat_interpolator import StratInterpolator
+
+
+def mix_hex_colors(colors: list[str], fractions: list[float]) -> str:
+    """
+    Combine multiple hex colors using given fractional weights.
+
+    Parameters
+    ----------
+    colors : list of str
+        Hex color strings (e.g. '#aabbcc').
+    fractions : list of float
+        Fractions for each color. Typically sums to 1. Values may be 0..1.
+
+    Returns
+    -------
+    str
+        Resulting mixed color as a hex string.
+    """
+    mixed_color = np.zeros(3)
+    for color, fraction in zip(colors, fractions):
+        mixed_color += np.array(to_rgb(color)) * fraction
+    return to_hex(mixed_color)
+
+
+def plot_lithology_and_thermcon(
+    litho_ds: xr.Dataset,
+    kx_plotting_base: np.ndarray,
+    interp_obj: StratInterpolator,
+    kx_plotting_max: np.ndarray,
+    kx_plotting_min: np.ndarray,
+    litho_h5path: str | Path,
+    out_dir: Path,
+) -> None:
+    """
+    Create and save a combined lithology "layer-cake" plot with thermal
+    conductivity profiles plotted alongside.
+
+    The main panel shows lithology as vertically stacked colored bands
+    (mixed color when two lithologies are present). A narrow side panel
+    shows horizontally stacked bars for lithology fractions. Thermal
+    conductivity curves are plotted on a twin x-axis (using the internal
+    grid from the StratInterpolator).
+
+    Parameters
+    ----------
+    litho_ds : xarray.Dataset
+        Dataset containing fields `depth`, lithology type a `lithology_a`, lithology type b `lithology_b`,
+        and the fraction of type a `lithology_a_fraction`.
+    kx_plotting_base : np.ndarray
+        Fine-grid, interpolated base-case thermal conductivity values (aligned with interp_obj.zp).
+    interp_obj : StratInterpolator
+        Interpolator object; used for the internal z-grid (zp) and interpolation.
+    kx_plotting_max : np.ndarray
+        Fine-grid, interpolated maximum thermal conductivity values.
+    kx_plotting_min : np.ndarray
+        Fine-grid, interpolated minimum thermal conductivity values.
+    litho_h5path : str | Path
+        Path to the source lithology .h5 file (used to build output filename).
+    out_dir : Path
+        Directory where the figure will be written.
+
+    Returns
+    -------
+    None
+
+    """
+    plt.rcParams.update({"font.size": 14})
+
+    # Create main + side axes: wide main column + narrow side bar
+    fig, (ax_main, ax_side) = plt.subplots(
+        1, 2, figsize=(8, 8), gridspec_kw={"width_ratios": [4, 0.8]}
+    )
+
+    # Prepare lists for plotting
+    current_depth = 0
+    depths = []
+    fractions_a = []
+    fractions_b = []
+    colors_main = []
+
+    # Build mixed colors and fraction lists
+    for i in range(len(litho_ds["depth"])):
+        # depth value for this layer (end depth)
+        depth = litho_ds["depth"].values[i]
+
+        lithology_a = litho_ds["lithology_a"].isel(depth=i).item()
+        lithology_b = litho_ds["lithology_b"].isel(depth=i).item()
+        fraction_a = litho_ds["lithology_a_fraction"].isel(depth=i).item()
+        fraction_b = 1 - fraction_a
+
+        # Prepare data for the side plot
+        depths.append((current_depth + depth) / 2)  # Center of the bar
+        fractions_a.append(fraction_a)
+        fractions_b.append(fraction_b)
+
+        # Prepare data for the main plot
+        mixed_color = mix_hex_colors(
+            [lithology_colors[lithology_a], lithology_colors[lithology_b]],
+            [fraction_a, fraction_b],
+        )
+        colors_main.append(mixed_color)
+
+        current_depth = depth
+
+    # Plot the main lithology column with mixed colors
+    for i in range(len(litho_ds["depth"])):
+        depth_start = litho_ds["depth"].values[i - 1] if i > 0 else 0
+        depth_end = litho_ds["depth"].values[i]
+        ax_main.fill_betweenx(
+            [depth_start, depth_end], 0, 1, color=colors_main[i], edgecolor=None
+        )
+
+    # Adjust axis limits to cover all the plotted data
+    ax_main.set_xlim(0, 1)  # Set x-axis to range from 0 to 1 for consistency
+    ax_main.set_ylim(
+        min(litho_ds["depth"].values), max(litho_ds["depth"].values)
+    )  # Make sure the depth range covers all data
+
+    # Calculate depth ranges
+    depth_ranges = [
+        (litho_ds["depth"].values[i - 1] if i > 0 else 0, litho_ds["depth"].values[i])
+        for i in range(len(litho_ds["depth"]))
+    ]
+
+    # Create the side plot (horizontally stacked bars for fractions)
+    for i, (depth_start, depth_end) in enumerate(depth_ranges):
+        bar_height = depth_end - depth_start
+        # Left segment (fraction A)
+        ax_side.barh(
+            y=depth_start,  # Start of the depth range
+            width=fractions_a[i],
+            height=bar_height,
+            color=lithology_colors[litho_ds["lithology_a"].isel(depth=i).item()],
+            align="edge",
+        )
+        # Right segment (fraction B), stacked by using left=fraction_a
+        ax_side.barh(
+            y=depth_start,  # Start of the depth range
+            width=fractions_b[i],
+            height=bar_height,
+            left=fractions_a[i],
+            color=lithology_colors[litho_ds["lithology_b"].isel(depth=i).item()],
+            align="edge",
+        )
+
+    # Style the side axis
+    ax_side.set_xlabel("Lithology fractions")
+    ax_side.set_xlim(0, 1)
+    ax_side.set_xticks([0, 0.5, 1])
+    ax_side.set_xticklabels(["0", "0.5", "1.0"])
+    ax_side.set_ylim(ax_main.get_ylim())  # Align y-axes
+    ax_side.invert_yaxis()
+
+    # Remove y-axis ticks and labels
+    ax_side.set_yticks([])  # Remove ticks
+    ax_side.set_yticklabels([])  # Remove labels
+
+    # Spines visibility
+    ax_side.spines["top"].set_visible(True)
+    ax_side.spines["right"].set_visible(True)
+    ax_side.spines["bottom"].set_visible(True)
+    ax_side.spines["left"].set_visible(True)
+
+    # Plot thermal conductivity on a secondary x-axis for the main plot
+    ax_main.set_xticks([])
+    ax_main.set_ylabel("Depth [m]")
+    ax_kx = ax_main.twiny()
+    ax_kx.plot(
+        kx_plotting_base, interp_obj.zp, color="black", linestyle="-", label="Basecase"
+    )
+    ax_kx.plot(kx_plotting_max, interp_obj.zp, color="blue", linestyle=":", label="Max")
+    ax_kx.plot(kx_plotting_min, interp_obj.zp, color="blue", linestyle=":", label="Min")
+    ax_kx.set_xlabel("Thermal Conductivity [W/mk]")
+    ax_kx.invert_yaxis()
+    ax_kx.legend(bbox_to_anchor=(0.5, -0.01), title="Thermal Conductivity")
+
+    # Create legend for lithologies
+    handles = []
+    labels = []
+    used_lithologies = set(litho_ds["lithology_a"].values) | set(
+        litho_ds["lithology_b"].values
+    )
+    for lithology_type in used_lithologies:
+        english_name = lithology_names_english.get(lithology_type, lithology_type)
+        handles.append(
+            plt.Rectangle((0, 0), 1, 1, color=lithology_colors[lithology_type])
+        )
+        labels.append(lithology_type)
+    ax_main.legend(
+        handles,
+        labels,
+        loc="upper right",
+        bbox_to_anchor=(0.9, -0.01),
+        title="Lithology",
+    )
+
+    plt.tight_layout()
+
+    # Save the figure
+    h5_file_name = Path(litho_h5path).name
+    fig_name = Path(h5_file_name).stem
+    fig_path = out_dir / fig_name
+
+    plt.savefig(fig_path)
+
+
+def main_Plot_lithology(litho_h5_path: str) -> None:
+    """
+    Main function to load data, process thermal conductivity profiles, and generate lithology plots.
+
+    Parameters
+    ----------
+    litho_h5_path : str
+        Path to the .h5 file containing the dataset of subsurface properties.
+
+    Returns
+    --------
+    None
+
+    """
+    base_dir = Path(litho_h5_path).parent
+
+    # Open dataset (group 'litho_k' expected)
+    litho_k_ds = xr.open_dataset(litho_h5_path, group="litho_k", engine="netcdf4")
+
+    # select only basecase, sample 0
+    litho_k_ds_base = litho_k_ds.sel(n_samples=0)
+
+    z = litho_k_ds_base.depth.values
+    zval = litho_k_ds_base.kx.values
+
+    zend = z
+    zstart = z[:-1] * 1
+
+    # Append 0 at the beginning of zstart (so first segment starts at 0)
+    zstart = np.insert(zstart, 0, 0)
+
+    # Build interpolator & compute fine-grid, interpolated conductivity profiles
+    interp_obj = StratInterpolator(zend, zval)
+
+    # Interpolate basecase thermal conductivity values
+    kx_plotting_base = interp_obj.interp_plot(zstart, zend)
+
+    # Calculate max and min thermal conductivity profiles and interpolate
+    kx_max = litho_k_ds.kx.max(dim="n_samples").values
+    interp_obj.zval = kx_max
+    kx_plotting_max = interp_obj.interp_plot(zstart, zend)
+
+    kx_min = litho_k_ds.kx.min(dim="n_samples").values
+    interp_obj.zval = kx_min
+    kx_plotting_min = interp_obj.interp_plot(zstart, zend)
+
+    # Create the layer cake plot
+    plot_lithology_and_thermcon(
+        litho_k_ds_base,
+        kx_plotting_base,
+        interp_obj,
+        kx_plotting_max,
+        kx_plotting_min,
+        litho_h5_path,
+        base_dir,
+    )
+
+
+if __name__ == "__main__":
+    litho_h5_path = sys.argv[1]
+    main_Plot_lithology(litho_h5_path)