vectorized perm calc

    return thickness_pvalues, permeability_pvalues, transmissivity_pvalues_sampled

from numba import njit, prange

@njit(parallel=True)

def calculate_transmissivity(

        thickness_mean,

        thickness_sd,

        ln_perm_mean,

        ln_perm_sd,

    p_value,

    n_samples=10_000,

):

    if (

        np.isnan(thickness_mean)

        or np.isnan(thickness_sd)

        or np.isnan(ln_perm_mean)

        or np.isnan(ln_perm_sd)

        q,  # quantile (e.g., 0.05 for p95)

        z,  # precomputed ndtri(q)

        n_samples,

):

        return np.nan, np.nan, np.nan

    q = 1.0 - p_value / 100.0

    # Thickness

    if thickness_sd == 0:

        thickness_p = thickness_mean

        thickness_samples = np.full(n_samples, thickness_mean)

    else:

        t_dist = stats.norm(thickness_mean, thickness_sd)

        thickness_p = t_dist.ppf(q)

        thickness_samples = np.clip(t_dist.rvs(n_samples), 0.01, None)

    shape = thickness_mean.shape

    tm_flat = thickness_mean.ravel()

    ts_flat = thickness_sd.ravel()

    lpm_flat = ln_perm_mean.ravel()

    lps_flat = ln_perm_sd.ravel()

    # Permeability

    k_dist = stats.norm(ln_perm_mean, ln_perm_sd)

    permeability_p = np.exp(k_dist.ppf(q))

    # Transmissivity via sampling

    tr_samples = np.sort(np.exp(k_dist.rvs(n_samples) + np.log(thickness_samples)))

    n_cells = tm_flat.size

    idx = int(q * n_samples)

    transmissivity_p = tr_samples[idx]

    return thickness_p, permeability_p, transmissivity_p

 No newline at end of file

    thickness_p = np.empty(n_cells)

    permeability_p = np.empty(n_cells)

    transmissivity_p = np.empty(n_cells)

    for i in prange(n_cells):

        tm = tm_flat[i]

        ts = ts_flat[i]

        lpm = lpm_flat[i]

        lps = lps_flat[i]

        if np.isnan(tm) or np.isnan(ts) or np.isnan(lpm) or np.isnan(lps):

            thickness_p[i] = np.nan

            permeability_p[i] = np.nan

            transmissivity_p[i] = np.nan

            continue

        # Generate samples for this cell

        thickness_samples = np.empty(n_samples)

        ln_perm_samples = np.empty(n_samples)

        for j in range(n_samples):

            if ts == 0:

                thickness_samples[j] = tm

            else:

                thickness_samples[j] = max(0.01, np.random.normal(tm, ts))

            ln_perm_samples[j] = np.random.normal(lpm, lps)

        # Transmissivity samples

        tr_samples = np.exp(ln_perm_samples) * thickness_samples

        tr_samples.sort()

        transmissivity_p[i] = tr_samples[idx]

        # Analytical percentiles

        thickness_p[i] = tm + z * ts if ts > 0 else tm

        permeability_p[i] = np.exp(lpm + z * lps)

    return (

        thickness_p.reshape(shape),

        permeability_p.reshape(shape),

        transmissivity_p.reshape(shape),

    )

 No newline at end of file

    return aquifer_ds

from scipy.special import ndtri

def compute_hydro_properties_mc(

    aquifer_ds: xr.Dataset,

    p_value: float,

    n_samples: int = 10_000,

) -> xr.Dataset:

    thickness_mean = aquifer_ds["thickness"]

    thickness_sd = aquifer_ds["thickness_sd"]

    ln_perm_mean = np.log(aquifer_ds["permeability"])

    ln_perm_sd = aquifer_ds["permeability_lnsd"]

    start = time.perf_counter()

    thickness_p, permeability_p, transmissivity_p = xr.apply_ufunc(

        calculate_transmissivity,

        thickness_mean,

        thickness_sd,

        ln_perm_mean,

        ln_perm_sd,

        p_value,

    q = 1.0 - p_value / 100.0

    z = ndtri(q)

    thickness_p, permeability_p, transmissivity_p = calculate_transmissivity(

        aquifer_ds["thickness"].values,

        aquifer_ds["thickness_sd"].values,

        np.log(aquifer_ds["permeability"].values),

        aquifer_ds["permeability_lnsd"].values,

        q,

        z,

        n_samples,

        input_core_dims=[[], [], [], [], [], []],

        output_core_dims=[[], [], []],

        vectorize=True,

        dask="parallelized",

        output_dtypes=[np.float64, np.float64, np.float64],

    )

    end = time.perf_counter()

    print(end-start)

    return xr.Dataset(

        {

            "thickness": thickness_p,

            "permeability": permeability_p,

            "transmissivity": transmissivity_p,

            "transmissivity_with_ntg": transmissivity_with_ntg,

        }

            "thickness": (aquifer_ds["thickness"].dims, thickness_p),

            "permeability": (aquifer_ds["thickness"].dims, permeability_p),

            "transmissivity": (aquifer_ds["thickness"].dims, transmissivity_p),

            "transmissivity_with_ntg": transmissivity_with_ntg,  # this one is already a DataArray

        },

        coords=aquifer_ds.coords,

    )

        results_dir="C:/Users/knappersfy/work/thermogis/pythermogis/workdir/ROSL_ROSLU/output",

        aquifers=["ROSL_ROSLU"],

        segment_length=1,

        p_values=[95] # For now only calculating this value, because it is way faster than 50

        p_values=[50] # For now only calculating this value, because it is way faster than 50

    )

    result = run_utc_workflow(config)

    result.ROSL_ROSLU.utc_p95.plot() # should match p value in config

    result.ROSL_ROSLU.utc_p50.plot() # should match p value in config

    plt.show()

    print(result)

 No newline at end of file