use hens transmissivity calc

    transmissivity_pvalues_sampled = transmissivity_samples[sample_indexes.astype(int)]

    return thickness_pvalues, permeability_pvalues, transmissivity_pvalues_sampled

def _tpkt_kernel(

    thickness_mean,

    thickness_sd,

    ln_perm_mean,

    ln_perm_sd,

    p_values,

    n_samples,

):

    if (

        np.isnan(thickness_mean)

        or np.isnan(thickness_sd)

        or np.isnan(ln_perm_mean)

        or np.isnan(ln_perm_sd)

    ):

        n = len(p_values)

        return (

            np.full(n, np.nan),

            np.full(n, np.nan),

            np.full(n, np.nan),

        )

    p = np.asarray(p_values)

    q = 1.0 - p / 100.0

    # P50 analytic shortcut

    if p.size == 1 and p[0] == 50:

        t = thickness_mean

        k = np.exp(ln_perm_mean)

        return np.array([t]), np.array([k]), np.array([t * k])

    # Thickness

    if thickness_sd == 0:

        thickness_p = np.full(p.size, 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)

    # 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)))

    idx = (q * n_samples).astype(int)

    transmissivity_p = tr_samples[idx]

    return thickness_p, permeability_p, transmissivity_p

 No newline at end of file

from pythermogis.workflow.utc.configuration import UTCConfiguration

from pythermogis.workflow.utc.doublet import DoubletInput, calculate_doublet_performance

from pythermogis.workflow.utc.stats import lognormal_quantile, normal_quantile

from pythermogis.transmissivity.calculate_thick_perm_trans import _tpkt_kernel

def run_utc_workflow(config: UTCConfiguration, coarsen_dataset_by: int | None = None ) -> xr.DataTree:

    root = xr.DataTree()

    # TODO: Maybe save the results as datasets iso dataarrays, and have p value be a dim?

    # calc transmissivity(_with_ntg)

    pvalue = 50

    hydro = compute_hydro_properties(aquifer_ds, pvalue)

    p_values = xr.DataArray([50], dims="p")

    hydro = compute_hydro_properties_mc(aquifer_ds, p_values)

    transmissivity = hydro["transmissivity"]

    transmissivity_with_ntg = hydro["transmissivity_with_ntg"]

    aquifer_ds["transmissivity"] = transmissivity

    return aquifer_ds

def compute_hydro_properties(aquifer_ds: xr.Dataset, pvalue: float) -> xr.Dataset:

    logk = aquifer_ds["permeability"]

    logk_sd = aquifer_ds["permeability_lnsd"]

    h_mean = aquifer_ds["thickness"]

    h_sd = aquifer_ds["thickness_sd"]

    permeability = lognormal_quantile(logk, logk_sd, pvalue)

    permeability = xr.where(permeability < 0.01, 0.01, permeability)

    thickness = normal_quantile(h_mean, h_sd, pvalue)

    thickness = xr.where(thickness < 0.01, 0.01, thickness)

    transmissivity = permeability * thickness

def compute_hydro_properties_mc(

    aquifer_ds: xr.Dataset,

    p_values: xr.DataArray,

    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"]

    thickness_p, permeability_p, transmissivity_p = xr.apply_ufunc(

        _tpkt_kernel,

        thickness_mean,

        thickness_sd,

        ln_perm_mean,

        ln_perm_sd,

        p_values,

        n_samples,

        input_core_dims=[[], [], [], [], ["p"], []],

        output_core_dims=[["p"], ["p"], ["p"]],

        vectorize=True,

        dask="parallelized",

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

    )

    transmissivity_with_ntg = transmissivity * aquifer_ds["ntg"] / 1000.0

    transmissivity_with_ntg = xr.where(transmissivity_with_ntg < 0, 0, transmissivity_with_ntg)

    transmissivity_with_ntg = transmissivity_p * aquifer_ds["ntg"] / 1000.0

    transmissivity_with_ntg = transmissivity_with_ntg.where(transmissivity_with_ntg >= 0)

    return xr.Dataset(

        {

            "permeability": permeability,

            "thickness": thickness,

            "transmissivity": transmissivity,

            "thickness": thickness_p,

            "permeability": permeability_p,

            "transmissivity": transmissivity_p,

            "transmissivity_with_ntg": transmissivity_with_ntg,

        }

    )