introducing @register_environment decorator to solve environment-reuse issue (#1373)

Author: slayoo

PySDM/builder.py

@@ -49,8 +49,7 @@ def set_environment(self, environment):
     def _set_environment(self, environment):
         if self.particulator.environment is not None:
             raise AssertionError("environment has already been set")
-        self.particulator.environment = environment
-        self.particulator.environment.register(self)
+        self.particulator.environment = environment.instantiate(builder=self)
 
     def add_dynamic(self, dynamic):
         assert self.particulator.environment is not None

PySDM/environments/box.py

@@ -5,8 +5,10 @@
 import numpy as np
 
 from PySDM.impl.mesh import Mesh
+from PySDM.environments.impl import register_environment
 
 
+@register_environment()
 class Box:
     def __init__(self, dt, dv):
         self.dt = dt

PySDM/environments/impl/__init__.py

@@ -1 +1,3 @@
 """ common internals not intended to be imported from user code """
+
+from PySDM.environments.impl.register_environment import register_environment

PySDM/environments/impl/register_environment.py

@@ -0,0 +1,24 @@
+""" decorator for environment 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_environment():
+    def decorator(cls):
+        if hasattr(cls, "instantiate"):
+            if cls.instantiate is not _instantiate:
+                raise AttributeError(
+                    "decorated class has a different instantiate method"
+                )
+        else:
+            setattr(cls, "instantiate", _instantiate)
+        return cls
+
+    return decorator

PySDM/environments/kinematic_1d.py

@@ -7,10 +7,12 @@
 
 from PySDM.environments.impl.moist import Moist
 
-from ..impl import arakawa_c
-from ..initialisation.equilibrate_wet_radii import equilibrate_wet_radii
+from PySDM.impl import arakawa_c
+from PySDM.initialisation.equilibrate_wet_radii import equilibrate_wet_radii
+from PySDM.environments.impl import register_environment
 
 
+@register_environment()
 class Kinematic1D(Moist):
     def __init__(self, *, dt, mesh, thd_of_z, rhod_of_z, z0=0):
         super().__init__(dt, mesh, [])

PySDM/environments/kinematic_2d.py

@@ -12,10 +12,11 @@
     equilibrate_wet_radii,
 )
 from PySDM.initialisation.sampling.spectral_sampling import ConstantMultiplicity
-
-from ..impl import arakawa_c
+from PySDM.impl import arakawa_c
+from PySDM.environments.impl import register_environment
 
 
+@register_environment()
 class Kinematic2D(Moist):
     def __init__(self, *, dt, grid, size, rhod_of, mixed_phase=False):
         super().__init__(dt, Mesh(grid, size), [], mixed_phase=mixed_phase)

PySDM/environments/parcel.py

@@ -12,8 +12,10 @@
     default_rtol,
     equilibrate_wet_radii,
 )
+from PySDM.environments.impl import register_environment
 
 
+@register_environment()
 class Parcel(Moist):  # pylint: disable=too-many-instance-attributes
     def __init__(
         self,

README.md

@@ -510,7 +510,7 @@ builder.add_dynamic(Condensation())
 
 r_dry, specific_concentration = spectral_sampling.Logarithmic(spectrum).sample(n_sd)
 v_dry = formulae.trivia.volume(radius=r_dry)
-r_wet = equilibrate_wet_radii(r_dry=r_dry, environment=env, kappa_times_dry_volume=kappa * v_dry)
+r_wet = equilibrate_wet_radii(r_dry=r_dry, environment=builder.particulator.environment, kappa_times_dry_volume=kappa * v_dry)
 
 attributes = Dict()
 attributes["multiplicity"] = discretise_multiplicities(specific_concentration * env.mass_of_dry_air)
@@ -595,7 +595,7 @@ v_dry = formulae.trivia.volume(pyargs('radius', r_dry));
 specific_concentration = tmp{2};
 r_wet = equilibrate_wet_radii(pyargs(...
     'r_dry', r_dry, ...
-    'environment', env, ...
+    'environment', builder.particulator.environment, ...
     'kappa_times_dry_volume', kappa * v_dry...
 ));
 
@@ -691,7 +691,7 @@ builder.add_dynamic(Condensation())
 
 r_dry, specific_concentration = spectral_sampling.Logarithmic(spectrum).sample(n_sd)
 v_dry = formulae.trivia.volume(radius=r_dry)
-r_wet = equilibrate_wet_radii(r_dry=r_dry, environment=env, kappa_times_dry_volume=kappa * v_dry)
+r_wet = equilibrate_wet_radii(r_dry=r_dry, environment=builder.particulator.environment, kappa_times_dry_volume=kappa * v_dry)
 
 attributes = {
   'multiplicity': discretise_multiplicities(specific_concentration * env.mass_of_dry_air),

examples/PySDM_examples/Abade_and_Albuquerque_2024/simulation.py

@@ -12,26 +12,25 @@
 
 class Simulation(BasicSimulation):
     def __init__(self, settings):
-        parcel = Parcel(
-            dt=settings.timestep,
-            mass_of_dry_air=settings.mass_of_dry_air,
-            p0=settings.initial_total_pressure,
-            initial_water_vapour_mixing_ratio=settings.initial_water_vapour_mixing_ratio,
-            T0=settings.initial_temperature,
-            w=settings.updraft,
-        )
         builder = Builder(
             backend=CPU(settings.formulae, override_jit_flags={"parallel": False}),
             n_sd=settings.n_sd,
-            environment=parcel,
+            environment=Parcel(
+                dt=settings.timestep,
+                mass_of_dry_air=settings.mass_of_dry_air,
+                p0=settings.initial_total_pressure,
+                initial_water_vapour_mixing_ratio=settings.initial_water_vapour_mixing_ratio,
+                T0=settings.initial_temperature,
+                w=settings.updraft,
+            ),
         )
         builder.add_dynamic(AmbientThermodynamics())
         builder.add_dynamic(Condensation())
 
         r_dry, n_in_dv = ConstantMultiplicity(settings.soluble_aerosol).sample(
             n_sd=settings.n_sd
         )
-        attributes = parcel.init_attributes(
+        attributes = builder.particulator.environment.init_attributes(
             n_in_dv=n_in_dv, kappa=settings.kappa, r_dry=r_dry
         )
         self.products = (

examples/PySDM_examples/Abdul_Razzak_Ghan_2000/run_ARG_parcel.py

@@ -65,7 +65,8 @@ def run_parcel(
         v_dry = builder.formulae.trivia.volume(radius=r_dry)
         specific_concentration = concentration / builder.formulae.constants.rho_STP
         attributes["multiplicity"] = np.append(
-            attributes["multiplicity"], specific_concentration * env.mass_of_dry_air
+            attributes["multiplicity"],
+            specific_concentration * builder.particulator.environment.mass_of_dry_air,
         )
         attributes["dry volume"] = np.append(attributes["dry volume"], v_dry)
         attributes["kappa times dry volume"] = np.append(
@@ -74,7 +75,7 @@ def run_parcel(
 
     r_wet = equilibrate_wet_radii(
         r_dry=builder.formulae.trivia.radius(volume=attributes["dry volume"]),
-        environment=env,
+        environment=builder.particulator.environment,
         kappa_times_dry_volume=attributes["kappa times dry volume"],
     )
     attributes["volume"] = builder.formulae.trivia.volume(radius=r_wet)