TNO Intern

Commit dd6ff6de authored by Hen Brett's avatar Hen Brett 🐔
Browse files

Huh it is actually working now I think

parent d607b6cc

src/pythermogis/dask_utils/__init__.py

0 → 100644
+0 −0

Empty file added.

src/pythermogis/dask_utils/dask_utils.py

0 → 100644
+43 −0
Original line number Diff line number Diff line 
import numpy as np

import xarray as xr



def auto_chunk_xarray(xobj, target_chunk_mb=50):

    """

    Automatically chunks a Dataset or DataArray to aim for ~target_chunk_mb per chunk.



    Parameters:

        xobj: xr.Dataset or xr.DataArray

        target_chunk_mb: desired chunk size in megabytes (default: 50 MB)



    Returns:

        Dask-chunked version of the input xarray object.

    """

    def compute_chunk_shape(shape, dtype_size, ndim):

        total_elements = np.prod(shape)

        total_size = total_elements * dtype_size / 1e6  # in MB

        scale = max(total_size / target_chunk_mb, 1)

        return tuple(max(1, int(s / scale**(1/ndim))) for s in shape)



    if isinstance(xobj, xr.DataArray):

        shape = xobj.shape

        dims = xobj.dims

        dtype_size = xobj.dtype.itemsize

        chunk_shape = compute_chunk_shape(shape, dtype_size, len(shape))

        chunk_dict = dict(zip(dims, chunk_shape))

        return chunk_dict



    elif isinstance(xobj, xr.Dataset):

        # For datasets, compute union of shapes/dtypes across variables

        largest_var = max(

            xobj.data_vars.values(),

            key=lambda v: np.prod(v.shape) * v.dtype.itemsize

        )

        shape = largest_var.shape

        dims = largest_var.dims

        dtype_size = largest_var.dtype.itemsize

        chunk_shape = compute_chunk_shape(shape, dtype_size, len(shape))

        chunk_dict = dict(zip(dims, chunk_shape))

        return chunk_dict



    else:

        raise TypeError("Expected xarray.DataArray or xarray.Dataset")
 
 No newline at end of file

src/pythermogis/thermogis_classes/doublet.py

+4 −5
Original line number Diff line number Diff line
@@ -6,9 +6,6 @@ def simulate_doublet(output_data: xr.Dataset, reservoir_properties: xr.Dataset, 
    # Calculate transmissivity scaled by ntg and converted to Dm

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



    # Instantiate ThermoGIS doublet

    doublet = instantiate_thermogis_doublet(utc_properties, rng_seed)



    # Calculate the doublet performance across all dimensions

    output_data_arrays = xr.apply_ufunc(calculate_performance_of_single_location,

                                        reservoir_properties.mask,
@@ -19,7 +16,8 @@ def simulate_doublet(output_data: xr.Dataset, reservoir_properties: xr.Dataset, 
                                        output_data.temperature,

                                        output_data.transmissivity,

                                        output_data.transmissivity_with_ntg,

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

                                        rng_seed,

                                        kwargs={"utc_properties": utc_properties},

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

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

                                        vectorize=True,
@@ -45,7 +43,7 @@ def simulate_doublet(output_data: xr.Dataset, reservoir_properties: xr.Dataset, 
    return output_data





def calculate_performance_of_single_location(mask: float, depth: float, thickness: float, porosity: float, ntg: float, temperature: float, transmissivity: float, transmissivity_with_ntg: float, doublet: JClass = None ,

def calculate_performance_of_single_location(mask: float, depth: float, thickness: float, porosity: float, ntg: float, temperature: float, transmissivity: float, transmissivity_with_ntg: float, rng_seed: int,

                                             utc_properties: JClass = None) -> float:

    """

    Calculate the performance of a doublet at a single location.
@@ -97,6 +95,7 @@ def calculate_performance_of_single_location(mask: float, depth: float, thicknes 
    if not np.isnan(mask):

        return 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0



    doublet = instantiate_thermogis_doublet(utc_properties, rng_seed)

    set_doublet_parameters(doublet, transmissivity_with_ntg, depth, porosity, ntg, temperature, utc_properties)



    # The Java routine which calculates DoubletPerformance, for more detail on the simulation inspect the Java source code

tests/resources/test_output/exceedance_example/exceedance_probabilities.png

+1.09 KiB (82.1 KiB)
Loading image diff...

tests/resources/test_output/exceedance_example/reservoir_distributions.png

−1.33 KiB (45.8 KiB)
Loading image diff...