Testing the speed of the doublet calculations

import numpy as np

from scipy import stats

import xarray as xr

def generate_thickness_permeability_transmissivity_for_pvalue(thickness_mean: float, thickness_sd: float, ln_permeability_mean: float, ln_permeability_sd: float, Pvalue: float, nSamples: int = 10000) -> float:

def generate_thickness_permeability_transmissivity_for_pvalue(thickness_mean: float, thickness_sd: float, ln_permeability_mean: float, ln_permeability_sd: float, p_values: float, nSamples: int = 10000) -> float:

    """

    Given thickness provided as a normal distribution and ln(permeability) provided as a normal distribution, and a specific p-value, then generate the values for that p-value of the

    -thickness

    if np.isnan(thickness_mean) | np.isnan(thickness_sd) | np.isnan(ln_permeability_mean) | np.isnan(ln_permeability_sd):

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

    p_value_fraction = Pvalue / 100

    p_value_fractions = p_values / 100

    thickness_dist = stats.norm(loc=thickness_mean, scale=thickness_sd)

    thickness_pvalues = thickness_dist.ppf(1 - p_value_fractions)

    ln_permeability_dist = stats.norm(loc=ln_permeability_mean, scale=ln_permeability_sd)

    permeability_pvalues = np.exp(ln_permeability_dist.ppf(1 - p_value_fractions))

    # Sampling method for transmissivity

    thickness_samples = thickness_dist.rvs(nSamples)

    thickness_samples = np.where(thickness_samples < 0.0, 0.1, thickness_samples)   # Ensure no negative thicknesses

    transmissivity_samples = np.sort(np.exp(ln_permeability_dist.rvs(nSamples) + np.log(thickness_samples)))

    transmissivity_pvalues_sampled = transmissivity_samples[int((1 - p_value_fractions) * nSamples)]

    return thickness_pvalues, permeability_pvalues, transmissivity_pvalues_sampled

def generate_thickness_permeability_transmissivity_for_pvalues(thickness_mean: float, thickness_sd: float, ln_permeability_mean: float, ln_permeability_sd: float, p_values: xr.DataArray, nSamples: int = 10000) -> float:

    """

    Given thickness provided as a normal distribution and ln(permeability) provided as a normal distribution, and a specific p-value, then generate the values for that p-value of the

    -thickness

    -permeability

    -transmissivity

    Note: Transmissivity is a function of ln(permeability) * thickness, and so no analytical solution exists to combine these two probability distributions and so the transmissivity distribution has to be sampled.

    :param thickness_mean:

    :param thickness_sd:

    :param ln_permeability_mean:

    :param ln_permeability_sd:

    :param p_values:

    :param nSamples:

    :return:

    thickness, permeability, transmissivity

    """

    if np.isnan(thickness_mean) | np.isnan(thickness_sd) | np.isnan(ln_permeability_mean) | np.isnan(ln_permeability_sd):

        return np.full(len(p_values), np.nan), np.full(len(p_values), np.nan), np.full(len(p_values), np.nan)

    p_value_fractions = 1 - p_values / 100

    thickness_dist = stats.norm(loc=thickness_mean, scale=thickness_sd)

    thickness_pvalue = thickness_dist.ppf(1 - p_value_fraction)

    thickness_pvalues = thickness_dist.ppf(p_value_fractions)

    ln_permeability_dist = stats.norm(loc=ln_permeability_mean, scale=ln_permeability_sd)

    permeability_pvalue = np.exp(ln_permeability_dist.ppf(1 - p_value_fraction))

    permeability_pvalues = np.exp(ln_permeability_dist.ppf(p_value_fractions))

    # Sampling method for transmissivity

    thickness_samples = thickness_dist.rvs(nSamples)

    thickness_samples = np.where(thickness_samples < 0.0, 0.1, thickness_samples)   # Ensure no negative thicknesses

    transmissivity_samples = np.sort(np.exp(ln_permeability_dist.rvs(nSamples) + np.log(thickness_samples)))

    transmissivity_pvalue_sampled = transmissivity_samples[int((1 - p_value_fraction) * nSamples)]

    return thickness_pvalue, permeability_pvalue, transmissivity_pvalue_sampled

    sample_indexes = p_value_fractions * nSamples

    sample_indexes = sample_indexes.astype(int)

    transmissivity_pvalues_sampled = transmissivity_samples[sample_indexes]

    return thickness_pvalues, permeability_pvalues, transmissivity_pvalues_sampled

 No newline at end of file

from jpype import JClass

from pythermogis.physics.temperature_grid_calculation import calculate_temperature_from_gradient

from pythermogis.statistics.calculate_thick_perm_trans import generate_thickness_permeability_transmissivity_for_pvalue

from pythermogis.statistics.calculate_thick_perm_trans import generate_thickness_permeability_transmissivity_for_pvalue, \

    generate_thickness_permeability_transmissivity_for_pvalues

from pythermogis.thermogis_classes.jvm_start import start_jvm

                                  input_params: dict = None,

                                  p_values: List[float] = [50.0]) -> xr.Dataset:

    """

    Perform a ThermoGIS Doublet performance assessment. This will occur across all dimensions of the input_data (ie. input data can have a single value for each required variable, or it can be 1Dimensional or a 2Dimensional grid)

    Perform a ThermoGIS Doublet performance simulation. This will occur across all dimensions of the input_data (ie. input data can have a single value for each required variable, or it can be 1Dimensional or a 2Dimensional grid)

    The default doublet simulation settings are for the ThermoGIS BaseCase; but major input parameters can be set using the input_params variable. This is a dictonary which by default contain the following keys:

    input_params = {"hp_minimum_injection_temperature": 15,

         "return_temperature": 30,

         "surface_temperature": 10,

         "degrees_per_km": 31,

         "max_cooling_temperature_range": 100,

         "stimKhMax": 20,

         "use_stimulation": False,

         "use_heat_pump": False,

         "calculate_cop": True,

         "hp_application_mode": False,

         "hp_direct_heat_input_temp": 70.0,

         "utc_cutoff_shallow": 5.1,

         "utc_cutoff_deep": 6.5,

         "utc_cutoff_depth": 4000.0,

         "rng_seed": np.random.randint(low=0, high=10000)}

    To change only one of these input parameters you can provide a dictionary with a single key to this method, and it will use the default values for the rest of the model parameters.

    e.g. input_params = {"use_stimulation":True}

    :param input_data:

        A xr.Dataset (input_data) with the required variables: "thickness_mean", "thickness_sd", "porosity", "ntg", "depth", "ln_permeability_mean", "ln_permeability_sd", Performance will be calculated across all dimensions.

        "power", "heat_pump_power", "capex", "opex", "utc", "npv", "hprod", "cop", "cophp", "pres", "flow_rate", "welld"

        The output will have the same dimensions as the input_data class, with the additional p_value dimension

    """

    # Check that all essential variables are provided

    output_data["temperature"] = input_data["temperature"].copy()

    # Calculate Thickness, Permeability and Transmissivity for each P-value

    start = timeit.default_timer()

    output_data["thickness"], output_data["permeability"], output_data["transmissivity"] = xr.apply_ufunc(generate_thickness_permeability_transmissivity_for_pvalue,

    output_data["thickness"], output_data["permeability"], output_data["transmissivity"] = xr.apply_ufunc(generate_thickness_permeability_transmissivity_for_pvalues,

                                                                 input_data.thickness_mean,

                                                                 input_data.thickness_sd,

                                                                 input_data.ln_permeability_mean,

                                                                 input_data.ln_permeability_sd,

                                                                 p_values,

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

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

                                                                 input_core_dims=[[], [], [], [], ["p_value"]],

                                                                 output_core_dims=[["p_value"], ["p_value"], ["p_value"]],

                                                                 vectorize=True

                                                             )

    stop = timeit.default_timer()

    print('\nTime elapsed: %.1f seconds' % (stop - start))

    # Calculate transmissivity scaled by ntg and converted to Dm

    output_data[f"transmissivity_with_ntg"] = (output_data[f"transmissivity"] * input_data.ntg) / 1e3

    doublet = instantiate_thermogis_doublet(input_params)

    # Calculate the doublet performance across all dimensions

    start = timeit.default_timer()

    output_data_arrays = xr.apply_ufunc(calculate_performance_of_single_location,

                                        input_data.mask,

                                        input_data.depth,

                                        kwargs={"doublet": doublet, "input_params": input_params},

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

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

                                        vectorize=True,

                                        dask="parallelized"

                                        vectorize=True

                                        )

    stop = timeit.default_timer()

    print('\nTime elapsed: %.1f seconds' % (stop - start))

    # Assign output DataArrays to the output_data object

    output_data["power"] = output_data_arrays[0]

import sys

from importlib.resources import files

import jpype