introducing @register_environment decorator to solve environment-reus…

…e issue (#1373)
open-atmos · Aug 10, 2024 · 723e62c · 723e62c
1 parent 55e0cce
commit 723e62c
Show file tree

Hide file tree

Showing 46 changed files with 4,828 additions and 192 deletions.
diff --git a/PySDM/builder.py b/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

diff --git a/PySDM/environments/box.py b/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

diff --git a/PySDM/environments/impl/__init__.py b/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
diff --git a/PySDM/environments/impl/register_environment.py b/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
diff --git a/PySDM/environments/kinematic_1d.py b/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, [])

diff --git a/PySDM/environments/kinematic_2d.py b/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)

diff --git a/PySDM/environments/parcel.py b/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,

diff --git a/README.md b/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),

diff --git a/examples/PySDM_examples/Abade_and_Albuquerque_2024/simulation.py b/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 = (

diff --git a/examples/PySDM_examples/Abdul_Razzak_Ghan_2000/run_ARG_parcel.py b/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)

diff --git a/examples/PySDM_examples/Alpert_and_Knopf_2016/simulation.py b/examples/PySDM_examples/Alpert_and_Knopf_2016/simulation.py
@@ -215,8 +215,11 @@ def simulation(
         constants=constants,
         particle_shape_and_density="MixedPhaseSpheres",
     )
-    env = Box(dt=time_step, dv=volume)
-    builder = Builder(n_sd=n_sd, backend=CPU(formulae=formulae), environment=env)
+    builder = Builder(
+        n_sd=n_sd,
+        backend=CPU(formulae=formulae),
+        environment=Box(dt=time_step, dv=volume),
+    )
     builder.add_dynamic(Freezing(singular=False))
 
     if hasattr(spectrum, "s_geom") and spectrum.s_geom == 1:
@@ -236,6 +239,7 @@ def simulation(
         TotalUnfrozenImmersedSurfaceArea(name="A_tot"),
     )
     particulator = builder.build(attributes=attributes, products=products)
+    env = particulator.environment
 
     env["T"] = initial_temperature
     env["a_w_ice"] = np.nan

diff --git a/examples/PySDM_examples/Arabas_et_al_2023/aida.ipynb b/examples/PySDM_examples/Arabas_et_al_2023/aida.ipynb
diff --git a/examples/PySDM_examples/Arabas_et_al_2023/fig_A2.ipynb b/examples/PySDM_examples/Arabas_et_al_2023/fig_A2.ipynb
@@ -104,18 +104,17 @@
     "        output[backend_key][key] = {\"unfrozen_fraction\": [], \"dt\": case[\"dt\"], \"N\": case[\"N\"]}\n",
     "\n",
     "        builder = Builder(n_sd=n_sd, backend=backend_class(formulae=formulae, double_precision=double_precision))\n",
-    "        env = Box(dt=case[\"dt\"], dv=d_v)\n",
-    "        builder.set_environment(env)\n",
+    "        builder.set_environment(Box(dt=case[\"dt\"], dv=d_v))\n",
     "        builder.add_dynamic(Freezing(singular=False))\n",
     "        attributes = {\n",
     "            \"multiplicity\": np.full(n_sd, int(case[\"N\"])),\n",
     "            \"immersed surface area\": np.full(n_sd, immersed_surface_area),\n",
     "            \"volume\": np.full(n_sd, vol),\n",
     "        }\n",
     "        particulator = builder.build(attributes=attributes, products=products)\n",
-    "        env[\"RH\"] = 1.0001\n",
-    "        env[\"a_w_ice\"] = np.nan\n",
-    "        env[\"T\"] = np.nan\n",
+    "        particulator.environment[\"RH\"] = 1.0001\n",
+    "        particulator.environment[\"a_w_ice\"] = np.nan\n",
+    "        particulator.environment[\"T\"] = np.nan\n",
     "\n",
     "        cell_id = 0\n",
     "        for i in range(int(total_time / case[\"dt\"]) + 1):\n",

diff --git a/examples/PySDM_examples/Arabas_et_al_2023/make_particulator.py b/examples/PySDM_examples/Arabas_et_al_2023/make_particulator.py
@@ -57,9 +57,9 @@ def make_particulator(
         ) = sampling.sample(backend=backend, n_sd=n_sd)
     attributes["multiplicity"] *= total_particle_number
 
-    env = Box(dt, volume)
-    builder = Builder(n_sd, backend, env)
+    builder = Builder(n_sd=n_sd, backend=backend, environment=Box(dt, volume))
 
+    env = builder.particulator.environment
     env["T"] = initial_temperature
     env["RH"] = A_VALUE_LARGER_THAN_ONE
     env["rhod"] = 1.0

diff --git a/examples/PySDM_examples/Bieli_et_al_2022/simulation.py b/examples/PySDM_examples/Bieli_et_al_2022/simulation.py
@@ -9,11 +9,12 @@
 def make_core(settings, coal_eff):
     backend = CPU
 
-    env = Box(dv=settings.dv, dt=settings.dt)
     builder = Builder(
-        n_sd=settings.n_sd, backend=backend(settings.formulae), environment=env
+        n_sd=settings.n_sd,
+        backend=backend(settings.formulae),
+        environment=Box(dv=settings.dv, dt=settings.dt),
     )
-    env["rhod"] = 1.0
+    builder.particulator.environment["rhod"] = 1.0
     attributes = {}
     attributes["volume"], attributes["multiplicity"] = ConstantMultiplicity(
         settings.spectrum

diff --git a/examples/PySDM_examples/Grabowski_and_Pawlowska_2023/simulation.py b/examples/PySDM_examples/Grabowski_and_Pawlowska_2023/simulation.py
@@ -19,20 +19,19 @@ def __init__(
         scipy_solver=False,
         sampling_class=ConstantMultiplicity,
     ):
-        env = Parcel(
-            dt=settings.timestep,
-            p0=settings.initial_pressure,
-            initial_water_vapour_mixing_ratio=settings.initial_vapour_mixing_ratio,
-            T0=settings.initial_temperature,
-            w=settings.vertical_velocity,
-            mass_of_dry_air=44 * si.kg,
-        )
         builder = Builder(
             n_sd=settings.n_sd,
             backend=CPU(
                 formulae=settings.formulae, override_jit_flags={"parallel": False}
             ),
-            environment=env,
+            environment=Parcel(
+                dt=settings.timestep,
+                p0=settings.initial_pressure,
+                initial_water_vapour_mixing_ratio=settings.initial_vapour_mixing_ratio,
+                T0=settings.initial_temperature,
+                w=settings.vertical_velocity,
+                mass_of_dry_air=44 * si.kg,
+            ),
         )
         builder.add_dynamic(AmbientThermodynamics())
         builder.add_dynamic(
@@ -45,6 +44,7 @@ def __init__(
         ):
             builder.request_attribute(attribute)
 
+        env = builder.particulator.environment
         volume = env.mass_of_dry_air / settings.initial_air_density
         attributes = {
             k: np.empty(0)

diff --git a/examples/PySDM_examples/Graf_et_al_2019/figure_4.ipynb b/examples/PySDM_examples/Graf_et_al_2019/figure_4.ipynb
@@ -179,7 +179,7 @@
    "cell_type": "code",
    "outputs": [],
    "source": [
-    "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)"
    ],
    "metadata": {
     "collapsed": false,
@@ -196,7 +196,7 @@
    "outputs": [],
    "source": [
     "attributes = {\n",
-    "    'multiplicity': discretise_multiplicities(specific_concentration * env.mass_of_dry_air),\n",
+    "    'multiplicity': discretise_multiplicities(specific_concentration * builder.particulator.environment.mass_of_dry_air),\n",
     "    'dry volume': v_dry,\n",
     "    'kappa times dry volume': kappa * v_dry,\n",
     "    'volume': formulae.trivia.volume(radius=r_wet)\n",

diff --git a/examples/PySDM_examples/Jensen_and_Nugent_2017/simulation.py b/examples/PySDM_examples/Jensen_and_Nugent_2017/simulation.py
@@ -34,24 +34,22 @@ def __init__(
             settings.p0 / settings.RH0 / pvs_Celsius(settings.T0 - const.T0) - 1
         )
 
-        env = Parcel(
-            dt=settings.dt,
-            mass_of_dry_air=666 * si.kg,
-            p0=settings.p0,
-            initial_water_vapour_mixing_ratio=initial_water_vapour_mixing_ratio,
-            T0=settings.T0,
-            w=settings.vertical_velocity,
-            z0=settings.z0,
-        )
-
         n_gccn = np.count_nonzero(table_3.NA) if gccn else 0
 
         builder = Builder(
             n_sd=N_SD_NON_GCCN + n_gccn,
             backend=CPU(
                 formulae=settings.formulae, override_jit_flags={"parallel": False}
             ),
-            environment=env,
+            environment=Parcel(
+                dt=settings.dt,
+                mass_of_dry_air=666 * si.kg,
+                p0=settings.p0,
+                initial_water_vapour_mixing_ratio=initial_water_vapour_mixing_ratio,
+                T0=settings.T0,
+                w=settings.vertical_velocity,
+                z0=settings.z0,
+            ),
         )
 
         additional_derived_attributes = ("radius", "equilibrium supersaturation")
@@ -83,8 +81,10 @@ def __init__(
         )
         rhod0 = settings.formulae.state_variable_triplet.rhod_of_pd_T(pd0, settings.T0)
 
-        attributes = env.init_attributes(
-            n_in_dv=n_in_unit_volume * env.mass_of_dry_air / rhod0,
+        attributes = builder.particulator.environment.init_attributes(
+            n_in_dv=n_in_unit_volume
+            * builder.particulator.environment.mass_of_dry_air
+            / rhod0,
             kappa=settings.kappa,
             r_dry=self.r_dry,
         )

diff --git a/examples/PySDM_examples/Kreidenweis_et_al_2003/simulation.py b/examples/PySDM_examples/Kreidenweis_et_al_2003/simulation.py
@@ -12,24 +12,22 @@
 
 class Simulation(BasicSimulation):
     def __init__(self, settings, products=None):
-        env = Parcel(
-            dt=settings.dt,
-            mass_of_dry_air=settings.mass_of_dry_air,
-            p0=settings.p0,
-            initial_water_vapour_mixing_ratio=settings.initial_water_vapour_mixing_ratio,
-            T0=settings.T0,
-            w=settings.w,
-        )
-
         builder = Builder(
             n_sd=settings.n_sd,
             backend=CPU(
                 formulae=settings.formulae, override_jit_flags={"parallel": False}
             ),
-            environment=env,
+            environment=Parcel(
+                dt=settings.dt,
+                mass_of_dry_air=settings.mass_of_dry_air,
+                p0=settings.p0,
+                initial_water_vapour_mixing_ratio=settings.initial_water_vapour_mixing_ratio,
+                T0=settings.T0,
+                w=settings.w,
+            ),
         )
 
-        attributes = env.init_attributes(
+        attributes = builder.particulator.environment.init_attributes(
             n_in_dv=settings.n_in_dv,
             kappa=settings.kappa,
             r_dry=settings.r_dry,
Original file line number	Diff line number	Diff line change
		@@ -1 +1,3 @@
		""" common internals not intended to be imported from user code """

		from PySDM.environments.impl.register_environment import register_environment