From 0f910c06a3e3255cbb1601c33e79a273baadcbfe Mon Sep 17 00:00:00 2001 From: Sylwester Arabas Date: Mon, 5 Aug 2024 10:04:20 +0200 Subject: [PATCH] fixing solver instantiation logic for kinematic 1D env + some code cleanups --- PySDM/dynamics/eulerian_advection.py | 12 ++--- PySDM/environments/impl/moist.py | 14 +++--- PySDM/environments/kinematic_1d.py | 2 +- .../Shipway_and_Hill_2012/mpdata_1d.py | 49 +++++++++++++------ .../Shipway_and_Hill_2012/simulation.py | 9 ++-- 5 files changed, 50 insertions(+), 36 deletions(-) diff --git a/PySDM/dynamics/eulerian_advection.py b/PySDM/dynamics/eulerian_advection.py index 55cae0155e..970253ad80 100644 --- a/PySDM/dynamics/eulerian_advection.py +++ b/PySDM/dynamics/eulerian_advection.py @@ -15,12 +15,8 @@ def register(self, builder): self.particulator = builder.particulator def __call__(self): - self.particulator.environment.get_predicted( - "water_vapour_mixing_ratio" - ).download( - self.particulator.environment.get_water_vapour_mixing_ratio(), reshape=True - ) - self.particulator.environment.get_predicted("thd").download( - self.particulator.environment.get_thd(), reshape=True - ) + for field in ("water_vapour_mixing_ratio", "thd"): + self.particulator.environment.get_predicted(field).download( + getattr(self.particulator.environment, f"get_{field}")(), reshape=True + ) self.solvers() diff --git a/PySDM/environments/impl/moist.py b/PySDM/environments/impl/moist.py index c42663205b..b56f74e716 100644 --- a/PySDM/environments/impl/moist.py +++ b/PySDM/environments/impl/moist.py @@ -44,21 +44,21 @@ def _allocate(self, variables): result[var] = self.particulator.Storage.empty((self.mesh.n_cell,), float) return result - def __getitem__(self, index): - """returns a Storage representing the variable (field) at a given index or - otherwise a NaN-filled Storage if the index is not found (in order to simplify + def __getitem__(self, key: str): + """returns a Storage representing the variable (field) at a given key or + otherwise a NaN-filled Storage if the key is not found (in order to simplify generic code which uses optional variables, e.g. air viscosity, etc.)""" - if index in self._values["current"]: - return self._values["current"][index] + if key in self._values["current"]: + return self._values["current"][key] return self._nan_field - def get_predicted(self, index): + def get_predicted(self, key: str): if self._values["predicted"] is None: raise AssertionError( "It seems the AmbientThermodynamics dynamic was not added" " when building particulator" ) - return self._values["predicted"][index] + return self._values["predicted"][key] def _recalculate_temperature_pressure_relative_humidity(self, target): self.particulator.backend.temperature_pressure_rh( diff --git a/PySDM/environments/kinematic_1d.py b/PySDM/environments/kinematic_1d.py index 79c3dceaf1..6556e45830 100644 --- a/PySDM/environments/kinematic_1d.py +++ b/PySDM/environments/kinematic_1d.py @@ -26,7 +26,7 @@ def register(self, builder): self._tmp["rhod"] = rhod def get_water_vapour_mixing_ratio(self) -> np.ndarray: - return self.particulator.dynamics["EulerianAdvection"].solvers.advectee.get() + return self.particulator.dynamics["EulerianAdvection"].solvers.advectee def get_thd(self) -> np.ndarray: return self.thd0 diff --git a/examples/PySDM_examples/Shipway_and_Hill_2012/mpdata_1d.py b/examples/PySDM_examples/Shipway_and_Hill_2012/mpdata_1d.py index 3ae26de500..fdbfac640e 100644 --- a/examples/PySDM_examples/Shipway_and_Hill_2012/mpdata_1d.py +++ b/examples/PySDM_examples/Shipway_and_Hill_2012/mpdata_1d.py @@ -1,4 +1,7 @@ +from functools import cached_property + import numpy as np + from PyMPDATA import Options, ScalarField, Solver, Stepper, VectorField from PyMPDATA.boundary_conditions import Extrapolated @@ -19,41 +22,55 @@ def __init__( self.dt = dt self.advector_of_t = advector_of_t - grid = (nz,) - options = Options( + self._grid = (nz,) + self._options = Options( n_iters=mpdata_settings["n_iters"], infinite_gauge=mpdata_settings["iga"], nonoscillatory=mpdata_settings["fct"], third_order_terms=mpdata_settings["tot"], ) - stepper = Stepper(options=options, grid=grid, non_unit_g_factor=True) bcs = (Extrapolated(),) - zZ_scalar = arakawa_c.z_scalar_coord(grid) / nz - g_factor = ScalarField( + zZ_scalar = arakawa_c.z_scalar_coord(self._grid) / nz + self._g_factor = ScalarField( data=g_factor_of_zZ(zZ_scalar), - halo=options.n_halo, + halo=self._options.n_halo, boundary_conditions=bcs, ) - advector = VectorField( + self._advector = VectorField( data=(np.full(nz + 1, advector_of_t(0)),), - halo=options.n_halo, + halo=self._options.n_halo, boundary_conditions=bcs, ) - self.advectee = ScalarField( + self._advectee = ScalarField( data=advectee_of_zZ_at_t0(zZ_scalar), - halo=options.n_halo, + halo=self._options.n_halo, boundary_conditions=bcs, ) - self.solver = Solver( - stepper=stepper, - advectee=self.advectee, - advector=advector, - g_factor=g_factor, + + @cached_property + def solver(self): + return Solver( + stepper=Stepper( + options=self._options, grid=self._grid, non_unit_g_factor=True + ), + advectee=self._advectee, + advector=self._advector, + g_factor=self._g_factor, ) + @property + def solver_cached(self): + return "solver" in self.__dict__ + + @property + def advectee(self): + return (self.solver.advectee if self.solver_cached else self._advectee).get() + @property def advector(self): - return self.solver.advector.get_component(0) + return ( + self.solver.advector if self.solver_cached else self._advector + ).get_component(0) def update_advector_field(self): self.__t += 0.5 * self.dt diff --git a/examples/PySDM_examples/Shipway_and_Hill_2012/simulation.py b/examples/PySDM_examples/Shipway_and_Hill_2012/simulation.py index 85baf9996e..ffa9790d59 100644 --- a/examples/PySDM_examples/Shipway_and_Hill_2012/simulation.py +++ b/examples/PySDM_examples/Shipway_and_Hill_2012/simulation.py @@ -46,7 +46,7 @@ def zZ_to_z_above_reservoir(zZ): z_above_reservoir = zZ * (settings.nz * settings.dz) + self.z0 return z_above_reservoir - self.mpdata = MPDATA_1D( + mpdata = MPDATA_1D( nz=settings.nz, dt=settings.dt, mpdata_settings=settings.mpdata_settings, @@ -79,7 +79,7 @@ def zZ_to_z_above_reservoir(zZ): update_thd=settings.condensation_update_thd, ) ) - self.builder.add_dynamic(EulerianAdvection(self.mpdata)) + self.builder.add_dynamic(EulerianAdvection(mpdata)) self.products = [] if settings.precip: @@ -257,10 +257,11 @@ def run(self): self.save(0) for step in range(self.nt): - self.mpdata.update_advector_field() + mpdata = self.particulator.dynamics["EulerianAdvection"].solvers + mpdata.update_advector_field() if "Displacement" in self.particulator.dynamics: self.particulator.dynamics["Displacement"].upload_courant_field( - (self.mpdata.advector / self.g_factor_vec,) + (mpdata.advector / self.g_factor_vec,) ) self.particulator.run(steps=1) self.save(step + 1)