getting to first running prototype (228488ca) · Commits · AGS_public / pythermogis

pythermogis/pytg3/doublet.py

+36 −6

Original line number	Diff line number	Diff line
		from __future__ import annotations

		import logging
		import timeit
		from typing import TYPE_CHECKING

		import numpy as np
		@@ -7,11 +9,15 @@ import xarray as xr
		from jpype import JClass
		from pydantic import BaseModel, ConfigDict

		from pythermogis.dask import auto_chunk_dataset
		from pythermogis.mock import create_logger_mock
		from pythermogis.pytg3.settings import Settings


		logger = logging.getLogger(__name__)

		if TYPE_CHECKING:
		from pythermogis.pytg3.aquifer import Aquifer
		from pythermogis.pytg3.settings import Settings
		from pythermogis.pytg3.typing import FloatOrArray


		@@ -34,6 +40,9 @@ class ThermoGISDoubletResults(BaseModel):
		prd_temp: FloatOrArray
		transmissivity_with_ntg: FloatOrArray

		def to_dataset(self) -> xr.Dataset:
		return xr.Dataset({field: getattr(self, field) for field in self.model_fields})


		class ThermoGISDoublet:

		@@ -51,13 +60,11 @@ class ThermoGISDoublet:
		)


		# TODO: how to handle chunksize? In settings?

		def simulate(self, mask_value: float = np.nan) -> ThermoGISDoubletResults:
		def simulate(self, mask_value: float = np.nan, chunk_size: int \| None = None) -> ThermoGISDoubletResults:
		start = timeit.default_timer()
		transmissivity_with_ntg = (self.aquifer.transmissivity * self.aquifer.ntg) / 1e3

		output_arrays = xr.apply_ufunc(
		_run_doublet_at_location,
		fields = [
		self.aquifer.mask,
		self.aquifer.depth,
		self.aquifer.thickness,
		@@ -66,6 +73,17 @@ class ThermoGISDoublet:
		self.aquifer.temperature,
		self.aquifer.transmissivity,
		transmissivity_with_ntg,
		]

		if chunk_size is not None:
		fields = [
		auto_chunk_dataset(f, chunk_size) if isinstance(f, xr.DataArray) else f
		for f in fields
		]

		output_arrays = xr.apply_ufunc(
		_run_doublet_at_location,
		*fields,
		kwargs={"utc_input": self.settings.create_utc_input(), "mask_value": mask_value},
		dask="parallelized",
		input_core_dims=[[], [], [], [], [], [], [], []],
		@@ -73,6 +91,18 @@ class ThermoGISDoublet:
		vectorize=True,
		)

		if chunk_size is not None:
		output_arrays = [
		a.load() if isinstance(a, xr.DataArray) else a for a in output_arrays
		]
		transmissivity_with_ntg = (
		transmissivity_with_ntg.load()
		if isinstance(transmissivity_with_ntg, xr.DataArray)
		else transmissivity_with_ntg
		)

		logger.info("Doublet simulation took %.1f seconds", timeit.default_timer() - start)

		return ThermoGISDoubletResults(
		power=output_arrays[0],
		heat_pump_power=output_arrays[1],

pythermogis/pytg3/settings.py

+2 −2

Original line number	Diff line number	Diff line
		@@ -35,8 +35,8 @@ class Settings:

		@property
		def temperature_gradient(self) -> float:
		return self.settings.tempGradient()
		return self.settings.getTemperatureGradient()

		@property
		def surface_temperature(self) -> float:
		return self.settings.surfaceTemperature()
		No newline at end of file
		return self.settings.getSurfaceTemperature()
		No newline at end of file

tests/pytg3/test_doublet.py

0 → 100644

+21 −0

Original line number	Diff line number	Diff line
		from pytg3.aquifer import Aquifer
		from pytg3.doublet import ThermoGISDoublet
		import xarray as xr


		def test_deterministic_doublet():
		aquifer = Aquifer(
		thickness=10,
		porosity=0.2,
		ntg=0.5,
		depth=2000,
		permeability=200,
		)

		doublet = ThermoGISDoublet(aquifer=aquifer)

		results1 = doublet.simulate()
		results2 = doublet.simulate()


		xr.testing.assert_equal(results1.to_dataset(), results2.to_dataset())