introducing @register_dynamic decorator to solve dynamic-reuse issue (#1372)

slayoo · web-flow · commit 118f47182134 · 2024-08-05T23:35:09.000+02:00
diff --git a/PySDM/builder.py b/PySDM/builder.py
@@ -121,8 +121,8 @@ def build(
             self._resolve_attribute(attr_name)
         self.req_attr_names = None
 
-        for dynamic in self.particulator.dynamics.values():
-            dynamic.register(self)
+        for key, dynamic in self.particulator.dynamics.items():
+            self.particulator.dynamics[key] = dynamic.instantiate(builder=self)
 
         single_buffer_for_all_products = np.empty(self.particulator.mesh.grid)
         for product in products:
diff --git a/PySDM/dynamics/ambient_thermodynamics.py b/PySDM/dynamics/ambient_thermodynamics.py
@@ -2,7 +2,10 @@
 environment-sync triggering class
 """
 
+from PySDM.dynamics.impl import register_dynamic
 
+
+@register_dynamic()
 class AmbientThermodynamics:
     def __init__(self):
         self.particulator = None
diff --git a/PySDM/dynamics/aqueous_chemistry.py b/PySDM/dynamics/aqueous_chemistry.py
@@ -13,12 +13,14 @@
     M,
     SpecificGravities,
 )
+from PySDM.dynamics.impl import register_dynamic
 
 DEFAULTS = namedtuple("_", ("pH_min", "pH_max", "pH_rtol", "ionic_strength_threshold"))(
     pH_min=-1.0, pH_max=14.0, pH_rtol=1e-6, ionic_strength_threshold=0.02 * M
 )
 
 
+@register_dynamic()
 class AqueousChemistry:  # pylint: disable=too-many-instance-attributes
     def __init__(
         self,
diff --git a/PySDM/dynamics/collisions/collision.py b/PySDM/dynamics/collisions/collision.py
@@ -23,6 +23,7 @@
     RandomGeneratorOptimizerNoPair,
 )
 from PySDM.physics import si
+from PySDM.dynamics.impl import register_dynamic
 
 # pylint: disable=too-many-lines
 
@@ -36,6 +37,7 @@
 )
 
 
+@register_dynamic()
 class Collision:  # pylint: disable=too-many-instance-attributes
     def __init__(
         self,
@@ -288,6 +290,7 @@ def compute_gamma(self, prob, rand, is_first_in_pair, out):
         )
 
 
+@register_dynamic()
 class Coalescence(Collision):
     def __init__(
         self,
@@ -316,6 +319,7 @@ def __init__(
         )
 
 
+@register_dynamic()
 class Breakup(Collision):
     def __init__(
         self,
diff --git a/PySDM/dynamics/condensation.py b/PySDM/dynamics/condensation.py
@@ -7,7 +7,8 @@
 
 import numpy as np
 
-from ..physics import si
+from PySDM.physics import si
+from PySDM.dynamics.impl import register_dynamic
 
 DEFAULTS = namedtuple("_", ("rtol_x", "rtol_thd", "cond_range", "schedule"))(
     rtol_x=1e-6,
@@ -17,6 +18,7 @@
 )
 
 
+@register_dynamic()
 class Condensation:  # pylint: disable=too-many-instance-attributes
     def __init__(
         self,
diff --git a/PySDM/dynamics/displacement.py b/PySDM/dynamics/displacement.py
@@ -11,9 +11,12 @@
 
 import numpy as np
 
+from PySDM.dynamics.impl import register_dynamic
+
 DEFAULTS = namedtuple("_", ("rtol", "adaptive"))(rtol=1e-2, adaptive=True)
 
 
+@register_dynamic()
 class Displacement:  # pylint: disable=too-many-instance-attributes
     def __init__(
         self,
diff --git a/PySDM/dynamics/eulerian_advection.py b/PySDM/dynamics/eulerian_advection.py
@@ -2,7 +2,10 @@
 wrapper class for triggering integration in the Eulerian advection solver
 """
 
+from PySDM.dynamics.impl import register_dynamic
 
+
+@register_dynamic()
 class EulerianAdvection:
     def __init__(self, solvers):
         self.solvers = solvers
@@ -12,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
-        )
-        self.solvers()
+        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(self.particulator.dynamics["Displacement"])
diff --git a/PySDM/dynamics/freezing.py b/PySDM/dynamics/freezing.py
@@ -7,8 +7,10 @@
     TimeDependentAttributes,
 )
 from PySDM.physics.heterogeneous_ice_nucleation_rate import Null
+from PySDM.dynamics.impl import register_dynamic
 
 
+@register_dynamic()
 class Freezing:
     def __init__(self, *, singular=True, record_freezing_temperature=False, thaw=False):
         assert not (record_freezing_temperature and singular)
diff --git a/PySDM/dynamics/impl/__init__.py b/PySDM/dynamics/impl/__init__.py
@@ -1 +1,3 @@
 """ stuff not intended to be imported from user code """
+
+from .register_dynamic import register_dynamic
diff --git a/PySDM/dynamics/impl/register_dynamic.py b/PySDM/dynamics/impl/register_dynamic.py
@@ -0,0 +1,21 @@
+""" decorator for dynamics classes
+ensuring that their instances can be re-used with multiple builders """
+
+from copy import deepcopy
+
+
+def _instantiate(self, *, builder):
+    copy = deepcopy(self)
+    copy.register(builder=builder)
+    return copy
+
+
+def register_dynamic():
+    def decorator(cls):
+        if hasattr(cls, "instantiate"):
+            assert cls.instantiate is _instantiate
+        else:
+            setattr(cls, "instantiate", _instantiate)
+        return cls
+
+    return decorator
diff --git a/PySDM/dynamics/isotopic_fractionation.py b/PySDM/dynamics/isotopic_fractionation.py
@@ -4,11 +4,13 @@
 """
 
 from PySDM.dynamics.condensation import Condensation
+from PySDM.dynamics.impl import register_dynamic
 
 LIGHT_ISOTOPES = ("1H", "16O")
 HEAVY_ISOTOPES = ("2H", "3H", "17O", "18O")
 
 
+@register_dynamic()
 class IsotopicFractionation:
     def __init__(self, isotopes: tuple = HEAVY_ISOTOPES):
         self.isotopes = isotopes
diff --git a/PySDM/dynamics/relaxed_velocity.py b/PySDM/dynamics/relaxed_velocity.py
@@ -6,8 +6,10 @@
 
 from PySDM.attributes.impl.attribute import Attribute
 from PySDM.particulator import Particulator
+from PySDM.dynamics.impl import register_dynamic
 
 
+@register_dynamic()
 class RelaxedVelocity:  # pylint: disable=too-many-instance-attributes
     """
     A dynamic which updates the fall momentum according to a relaxation timescale
diff --git 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
@@ -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/PySDM/environments/kinematic_2d.py b/PySDM/environments/kinematic_2d.py
@@ -85,16 +85,12 @@ def init_attributes(
         return attributes
 
     def get_thd(self):
-        return (
-            self.particulator.dynamics["EulerianAdvection"].solvers["th"].advectee.get()
-        )
+        return self.particulator.dynamics["EulerianAdvection"].solvers["th"]
 
     def get_water_vapour_mixing_ratio(self):
-        return (
-            self.particulator.dynamics["EulerianAdvection"]
-            .solvers["water_vapour_mixing_ratio"]
-            .advectee.get()
-        )
+        return self.particulator.dynamics["EulerianAdvection"].solvers[
+            "water_vapour_mixing_ratio"
+        ]
 
     def sync(self):
         self.particulator.dynamics["EulerianAdvection"].solvers.wait()
diff --git 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,47 +22,61 @@ 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
         self.advector[:] = self.advector_of_t(self.__t)
         np.testing.assert_array_less(np.abs(self.advector), 1)
         self.__t += 0.5 * self.dt
 
-    def __call__(self):
+    def __call__(self, _):
         self.solver.advance(1)
diff --git 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)
diff --git a/examples/PySDM_examples/Szumowski_et_al_1998/gui_settings.py b/examples/PySDM_examples/Szumowski_et_al_1998/gui_settings.py
@@ -348,10 +348,15 @@ def mpdata_iters(self):
                 return widget.value
         raise NotImplementedError()
 
-    def stream_function(self, xX, zZ, _):
+    @property
+    def stream_function(self):
         assert hasattr(self.__settings, "rhod_w_max")
         self.__settings.rhod_w_max = self.ui_rhod_w_max.value
-        return self.__settings.stream_function(xX, zZ, _)
+
+        def fun(xX, zZ, _):
+            return self.__settings.stream_function(xX, zZ, _)
+
+        return fun
 
     @property
     def freezing_inp_spec(self):
diff --git a/examples/PySDM_examples/Szumowski_et_al_1998/mpdata_2d.py b/examples/PySDM_examples/Szumowski_et_al_1998/mpdata_2d.py
diff --git a/examples/PySDM_examples/Szumowski_et_al_1998/simulation.py b/examples/PySDM_examples/Szumowski_et_al_1998/simulation.py
diff --git a/tests/smoke_tests/kinematic_2d/arabas_et_al_2015/test_adaptive_displacement.py b/tests/smoke_tests/kinematic_2d/arabas_et_al_2015/test_adaptive_displacement.py
diff --git a/tests/unit_tests/dynamics/test_eulerian_advection.py b/tests/unit_tests/dynamics/test_eulerian_advection.py
diff --git a/tests/unit_tests/dynamics/test_impl_register_dynamic.py b/tests/unit_tests/dynamics/test_impl_register_dynamic.py

Original file line number	Diff line number	Diff line change
`@@ -13,12 +13,14 @@`
`13`	`13`	`M,`
`14`	`14`	`SpecificGravities,`
`15`	`15`	`)`
	`16`	`+from PySDM.dynamics.impl import register_dynamic`
`16`	`17`
`17`	`18`	`DEFAULTS = namedtuple("_", ("pH_min", "pH_max", "pH_rtol", "ionic_strength_threshold"))(`
`18`	`19`	`pH_min=-1.0, pH_max=14.0, pH_rtol=1e-6, ionic_strength_threshold=0.02 * M`
`19`	`20`	`)`
`20`	`21`
`21`	`22`
	`23`	`+@register_dynamic()`
`22`	`24`	`class AqueousChemistry: # pylint: disable=too-many-instance-attributes`
`23`	`25`	`def __init__(`
`24`	`26`	`self,`
Original file line number	Diff line number	Diff line change
`@@ -23,6 +23,7 @@`
`23`	`23`	`RandomGeneratorOptimizerNoPair,`
`24`	`24`	`)`
`25`	`25`	`from PySDM.physics import si`
	`26`	`+from PySDM.dynamics.impl import register_dynamic`
`26`	`27`
`27`	`28`	`# pylint: disable=too-many-lines`
`28`	`29`
`@@ -36,6 +37,7 @@`
`36`	`37`	`)`
`37`	`38`
`38`	`39`
	`40`	`+@register_dynamic()`
`39`	`41`	`class Collision: # pylint: disable=too-many-instance-attributes`
`40`	`42`	`def __init__(`
`41`	`43`	`self,`
`@@ -288,6 +290,7 @@ def compute_gamma(self, prob, rand, is_first_in_pair, out):`
`288`	`290`	`)`
`289`	`291`
`290`	`292`
	`293`	`+@register_dynamic()`
`291`	`294`	`class Coalescence(Collision):`
`292`	`295`	`def __init__(`
`293`	`296`	`self,`
`@@ -316,6 +319,7 @@ def __init__(`
`316`	`319`	`)`
`317`	`320`
`318`	`321`
	`322`	`+@register_dynamic()`
`319`	`323`	`class Breakup(Collision):`
`320`	`324`	`def __init__(`
`321`	`325`	`self,`
Original file line number	Diff line number	Diff line change
`@@ -7,8 +7,10 @@`
`7`	`7`	`TimeDependentAttributes,`
`8`	`8`	`)`
`9`	`9`	`from PySDM.physics.heterogeneous_ice_nucleation_rate import Null`
	`10`	`+from PySDM.dynamics.impl import register_dynamic`
`10`	`11`
`11`	`12`
	`13`	`+@register_dynamic()`
`12`	`14`	`class Freezing:`
`13`	`15`	`def __init__(self, *, singular=True, record_freezing_temperature=False, thaw=False):`
`14`	`16`	`assert not (record_freezing_temperature and singular)`
Original file line number	Diff line number	Diff line change
`@@ -1 +1,3 @@`
`1`	`1`	`""" stuff not intended to be imported from user code """`
	`2`	`+`
	`3`	`+from .register_dynamic import register_dynamic`