simplifying how xarrays calculates across the P-value dimension

def calculate_doublet_performance_across_dimensions(input_data: xr.Dataset,

                                                    input_params: dict = None,

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

                                                    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)

    # Check that all essential variables are provided

    validate_input_data(input_data)

    # If no p_values provided, calculate only the P50

    if p_values is None:

        p_values = [50.0]

    # convert p_values list to a xarray DataArray; needed to ensure the dimensionality of the calculations

    p_values = xr.DataArray(

        data=p_values,

        dims=["p_value"],

        coords=dict(

            p_value=(["p_value"], p_values),

        ))

    # Apply provided input_params; or return the base case

    input_params = apply_input_params(input_params)

    if "mask" not in input_data:

        input_data["mask"] = np.nan

    # Instantiate ThermoGIS doublet

    doublet = instantiate_thermogis_doublet(input_params)

    # Setup output_data dataset

    output_data = input_data.thickness_mean.copy().to_dataset(name="thickness")

    output_data = output_data.expand_dims({"p_value": p_values})

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

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

    # TODO: Find a less ugly way to handle all these outputs than instantiating different arrays; probably by setting the output_core_dims in the apply_ufunc

    thickness_data = []

    permeability_data = []

    transmissivity_data = []

    for i, p in enumerate(p_values):

        thickness, permeability, 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_pvalue,

                                                             input_data.thickness_mean,

                                                             input_data.thickness_sd,

                                                             input_data.ln_permeability_mean,

                                                             input_data.ln_permeability_sd,

                                                                 p,

                                                             p_values,

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

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

                                                                 vectorize=True,

                                                             vectorize=True

                                                             )

        thickness_data.append(thickness.data)

        permeability_data.append(permeability.data)

        transmissivity_data.append(transmissivity.data)

    output_data["thickness"] = (output_data.coords, thickness_data)

    output_data["permeability"] = (output_data.coords, permeability_data)

    output_data["transmissivity"] = (output_data.coords, transmissivity_data)

    # Calculate transmissivity scaled by ntg and converted to Dm

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

    # Calculate performance for each P-value

    power_data = []

    heat_pump_power_data = []

    capex_data = []

    opex_data = []

    utc_data = []

    npv_data = []

    hprod_data = []

    cop_data = []

    cophp_data = []

    pres_data = []

    flow_rate_data = []

    welld_data = []

    # TODO: Find a less ugly way to handle all these outputs than instantiating different arrays; probably by setting the output_core_dims in the apply_ufunc

    for i, p in enumerate(p_values):

    # Instantiate ThermoGIS doublet

    doublet = instantiate_thermogis_doublet(input_params)

    # Calculate the doublet performance across all dimensions

    output_data_arrays = xr.apply_ufunc(calculate_performance_of_single_location,

                                        input_data.mask,

                                        input_data.depth,

                                            output_data.thickness.isel(p_value=i),

                                        output_data.thickness,

                                        input_data.porosity,

                                        input_data.ntg,

                                        input_data.temperature,

                                            output_data.transmissivity.isel(p_value=i),

                                            output_data.transmissivity_with_ntg.isel(p_value=i),

                                        output_data.transmissivity,

                                        output_data.transmissivity_with_ntg,

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

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

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

                                        vectorize=True,

                                        )

        # Assign values from calculate performance to their grids for each p-value

        power_data.append(output_data_arrays[0])

        heat_pump_power_data.append(output_data_arrays[1])

        capex_data.append(output_data_arrays[2])

        opex_data.append(output_data_arrays[3])

        utc_data.append(output_data_arrays[4])

        npv_data.append(output_data_arrays[5])

        hprod_data.append(output_data_arrays[6])

        cop_data.append(output_data_arrays[7])

        cophp_data.append(output_data_arrays[8])

        pres_data.append(output_data_arrays[9])

        flow_rate_data.append(output_data_arrays[10])

        welld_data.append(output_data_arrays[11])

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

    output_data["power"] = (output_data.coords, power_data)

    output_data["heat_pump_power"] = (output_data.coords, heat_pump_power_data)

    output_data["capex"] = (output_data.coords, capex_data)

    output_data["opex"] = (output_data.coords, opex_data)

    output_data["utc"] = (output_data.coords, utc_data)

    output_data["npv"] = (output_data.coords, npv_data)

    output_data["hprod"] = (output_data.coords, hprod_data)

    output_data["cop"] = (output_data.coords, cop_data)

    output_data["cophp"] = (output_data.coords, cophp_data)

    output_data["pres"] = (output_data.coords, pres_data)

    output_data["flow_rate"] = (output_data.coords, flow_rate_data)

    output_data["welld"] = (output_data.coords, welld_data)

    # Assign output DataArrays to the output_data object

    output_data["power"] = output_data_arrays[0]

    output_data["heat_pump_power"] = output_data_arrays[1]

    output_data["capex"] = output_data_arrays[2]

    output_data["opex"] = output_data_arrays[3]

    output_data["utc"] = output_data_arrays[4]

    output_data["npv"] = output_data_arrays[5]

    output_data["hprod"] = output_data_arrays[6]

    output_data["cop"] = output_data_arrays[7]

    output_data["cophp"] = output_data_arrays[8]

    output_data["pres"] = output_data_arrays[9]

    output_data["flow_rate"] = output_data_arrays[10]

    output_data["welld"] = output_data_arrays[11]

    return output_data

        return {key: input_params.get(key, input_params_basecase[key]) for key in input_params_basecase}

def instantiate_thermogis_doublet(input_params):

def instantiate_thermogis_doublet(input_params: dict):

    """

    Instantiate a ThermoGIS Doublet class, with a set random seed if provided

    :param rng_seed:

        "cophp"], output_values["pres"], output_values["flow_rate"], output_values["welld"]

def set_doublet_parameters(doublet, transmissivity_with_ntg, depth, porosity, ntg, temperature, input_params):

def set_doublet_parameters(doublet, transmissivity_with_ntg: float, depth: float, porosity: float, ntg: float, temperature: float, input_params: dict):

    """

    For a single location sets the necessary data on the doublet class, to then run a doublet simulation

    :param doublet: