TNO Intern

Commit 0d2101b9 authored by Hen Brett's avatar Hen Brett 🐔
Browse files

experimenting with chunking of data

parent 6689aa9c

src/pythermogis/doublet_simulation/stochastic_doublet.py

+12 −5
Original line number Diff line number Diff line
@@ -3,7 +3,7 @@ import xarray as xr 
import numpy as np

import timeit



from pythermogis import simulate_doublet

from pythermogis import simulate_doublet, auto_chunk_xarray

from pythermogis.physics.temperature_grid_calculation import calculate_temperature_from_gradient

from pythermogis.transmissivity.calculate_thick_perm_trans import generate_thickness_permeability_transmissivity_for_pvalues

from pythermogis.thermogis_classes.utc_properties import instantiate_utc_properties_builder
@@ -105,9 +105,13 @@ def calculate_doublet_performance_stochastic(reservoir_properties: xr.Dataset, 
    if "mask" not in reservoir_properties:

        reservoir_properties["mask"] = np.nan



    # chunk reservoir properties to enable dask parralelization

    chunk_dict = auto_chunk_xarray(reservoir_properties, target_chunk_mb=0.01)

    reservoir_properties_chunk = reservoir_properties.chunk(chunk_dict)



    # Setup output_data dataset

    output_data = reservoir_properties["temperature"].copy().to_dataset()

    output_data["temperature"] = reservoir_properties["temperature"] * 1.0

    output_data = reservoir_properties_chunk["temperature"].copy().to_dataset()

    output_data["temperature"] = reservoir_properties_chunk["temperature"] * 1.0

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



    # Calculate Thickness, Permeability and Transmissivity for each P-value
@@ -119,11 +123,14 @@ def calculate_doublet_performance_stochastic(reservoir_properties: xr.Dataset, 
                                                                                                          p_values,

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

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

                                                                                                          vectorize=True

                                                                                                          vectorize=True,

                                                                                                          dask="parallelized",

                                                                                                          dask_gufunc_kwargs={"allow_rechunk": True}

                                                                                                          )



    output_data = simulate_doublet(output_data, reservoir_properties_chunk, rng_seed, utc_properties).load()

    if print_execution_duration: print(f"Doublet simulation took {timeit.default_timer() - start:.1f} seconds")

    return simulate_doublet(output_data, reservoir_properties, rng_seed, utc_properties)

    return output_data





def validate_input_data(input_data: xr.Dataset):

tests/test_doublet_speed.py

+1 −1
Original line number Diff line number Diff line
@@ -34,7 +34,7 @@ class PyThermoGIS(TestCase): 




    def read_input_grids(self):

        new_cellsize=1000 # in m

        new_cellsize=5000 # in m

        input_grids = resample_xarray_grid(read_grid(self.test_files_out_path / "ROSL_ROSLU__thick.zmap"), new_cellsize=new_cellsize).to_dataset(name="thickness_mean")

        input_grids["thickness_sd"] = resample_xarray_grid(read_grid(self.test_files_out_path / "ROSL_ROSLU__thick_sd.zmap"), new_cellsize=new_cellsize)

        input_grids["ntg"] = resample_xarray_grid(read_grid(self.test_files_out_path / "ROSL_ROSLU__ntg.zmap"), new_cellsize=new_cellsize)