Pass environment to Builder constructor (#1178)

Co-authored-by: Sylwester Arabas <[email protected]>
Co-authored-by: Sylwester Arabas <[email protected]>

PySDM/builder.py

@@ -19,8 +19,15 @@
 from PySDM.physics.particle_shape_and_density import LiquidSpheres, MixedPhaseSpheres
 
 
+def _warn_env_as_ctor_arg():
+    warnings.warn(
+        "PySDM > v2.31 Builder expects environment instance as argument",
+        DeprecationWarning,
+    )
+
+
 class Builder:
-    def __init__(self, n_sd, backend):
+    def __init__(self, n_sd, backend, environment=None):
         assert not inspect.isclass(backend)
         self.formulae = backend.formulae
         self.particulator = Particulator(n_sd, backend)
@@ -32,10 +39,19 @@ def __init__(self, n_sd, backend):
         self.aerosol_radius_threshold = 0
         self.condensation_params = None
 
+        if environment is None:
+            _warn_env_as_ctor_arg()
+        else:
+            self._set_environment(environment)
+
     def _set_condensation_parameters(self, **kwargs):
         self.condensation_params = kwargs
 
     def set_environment(self, environment):
+        _warn_env_as_ctor_arg()
+        self._set_environment(environment)
+
+    def _set_environment(self, environment):
         if self.particulator.environment is not None:
             raise AssertionError("environment has already been set")
         self.particulator.environment = environment

README.md

@@ -240,8 +240,8 @@ ParticleVolumeVersusRadiusLogarithmSpectrum = pyimport("PySDM.products").Particl
 
 radius_bins_edges = 10 .^ range(log10(10*si.um), log10(5e3*si.um), length=32) 
 
-builder = Builder(n_sd=n_sd, backend=CPU())
-builder.set_environment(Box(dt=1 * si.s, dv=1e6 * si.m^3))
+env = Box(dt=1 * si.s, dv=1e6 * si.m^3)
+builder = Builder(n_sd=n_sd, backend=CPU(), environment=env)
 builder.add_dynamic(Coalescence(collision_kernel=Golovin(b=1.5e3 / si.s)))
 products = [ParticleVolumeVersusRadiusLogarithmSpectrum(radius_bins_edges=radius_bins_edges, name="dv/dlnr")] 
 particulator = builder.build(attributes, products)
@@ -260,8 +260,8 @@ ParticleVolumeVersusRadiusLogarithmSpectrum = py.importlib.import_module('PySDM.
 
 radius_bins_edges = logspace(log10(10 * si.um), log10(5e3 * si.um), 32);
 
-builder = Builder(pyargs('n_sd', int32(n_sd), 'backend', CPU()));
-builder.set_environment(Box(pyargs('dt', 1 * si.s, 'dv', 1e6 * si.m ^ 3)));
+env = Box(pyargs('dt', 1 * si.s, 'dv', 1e6 * si.m ^ 3));
+builder = Builder(pyargs('n_sd', int32(n_sd), 'backend', CPU(), 'environment', env));
 builder.add_dynamic(Coalescence(pyargs('collision_kernel', Golovin(1.5e3 / si.s))));
 products = py.list({ ParticleVolumeVersusRadiusLogarithmSpectrum(pyargs( ...
   'radius_bins_edges', py.numpy.array(radius_bins_edges), ...
@@ -284,8 +284,8 @@ from PySDM.products import ParticleVolumeVersusRadiusLogarithmSpectrum
 
 radius_bins_edges = np.logspace(np.log10(10 * si.um), np.log10(5e3 * si.um), num=32)
 
-builder = Builder(n_sd=n_sd, backend=CPU())
-builder.set_environment(Box(dt=1 * si.s, dv=1e6 * si.m ** 3))
+env = Box(dt=1 * si.s, dv=1e6 * si.m ** 3)
+builder = Builder(n_sd=n_sd, backend=CPU(), environment=env)
 builder.add_dynamic(Coalescence(collision_kernel=Golovin(b=1.5e3 / si.s)))
 products = [ParticleVolumeVersusRadiusLogarithmSpectrum(radius_bins_edges=radius_bins_edges, name='dv/dlnr')]
 particulator = builder.build(attributes, products)
@@ -387,14 +387,13 @@ The component submodules used to create this simulation are visualized below:
     COAL[":Coalescence"] --->|passed as arg to| BUILDER_ADD_DYN(["Builder.add_dynamic()"])
     BUILDER_INSTANCE["builder :Builder"] -...-|has a method| BUILDER_BUILD(["Builder.build()"])
     ATTRIBUTES[attributes: dict] -->|passed as arg to| BUILDER_BUILD
-    N_SD["n_sd :int"] -->|passed as arg to| BUILDER_INIT
-    BUILDER_INSTANCE -..-|has a method| BUILDER_SET_ENV(["Builder.set_environment()"])
-    BUILDER_INIT(["Builder.__init__()"]) ----->|instantiates| BUILDER_INSTANCE
+    N_SD["n_sd :int"] ---->|passed as arg to| BUILDER_INIT
+    BUILDER_INIT(["Builder.__init__()"]) --->|instantiates| BUILDER_INSTANCE
     BUILDER_INSTANCE -..-|has a method| BUILDER_ADD_DYN(["Builder.add_dynamic()"])
-    ENV_INIT(["Box.__init__()"]) --->|instantiates| ENV
-    DT[dt :float] ----->|passed as arg to| ENV_INIT
-    DV[dv :float] ----->|passed as arg to| ENV_INIT
-    ENV[":Box"] -->|passed as arg to| BUILDER_SET_ENV
+    ENV_INIT(["Box.__init__()"]) -->|instantiates| ENV
+    DT[dt :float] -->|passed as arg to| ENV_INIT
+    DV[dv :float] -->|passed as arg to| ENV_INIT
+    ENV[":Box"] -->|passed as arg to| BUILDER_INIT
     B["b: float"] --->|passed as arg to| KERNEL_INIT(["Golovin.__init__()"])
     KERNEL_INIT -->|instantiates| KERNEL
     KERNEL[collision_kernel: Golovin] -->|passed as arg to| COAL_INIT(["Coalesncence.__init__()"])
@@ -413,7 +412,7 @@ The component submodules used to create this simulation are visualized below:
     volume -->|added as element of| ATTRIBUTES
     BUILDER_BUILD -->|returns| PARTICULATOR_INSTANCE["particulator :Particulator"]
     PARTICULATOR_INSTANCE -.-|has a field| PARTICULATOR_PROD(["Particulator.products:dict"])
-    BACKEND_INSTANCE["backend :CPU"] -->|passed as arg to| BUILDER_INIT
+    BACKEND_INSTANCE["backend :CPU"] ---->|passed as arg to| BUILDER_INIT
     PRODUCTS -.-|accessible via| PARTICULATOR_PROD
     NP_LOGSPACE(["np.logspace()"]) -->|returns| EDGES 
     EDGES[radius_bins_edges: np.ndarray] -->|passed as arg to| SPECTRUM_INIT
@@ -424,7 +423,6 @@ The component submodules used to create this simulation are visualized below:
     click BUILDER_INSTANCE "https://open-atmos.github.io/PySDM/PySDM/builder.html"
     click BUILDER_INIT "https://open-atmos.github.io/PySDM/PySDM/builder.html"
     click BUILDER_ADD_DYN "https://open-atmos.github.io/PySDM/PySDM/builder.html"
-    click BUILDER_SET_ENV "https://open-atmos.github.io/PySDM/PySDM/builder.html"
     click ENV_INIT "https://open-atmos.github.io/PySDM/PySDM/environments/index.html"
     click ENV "https://open-atmos.github.io/PySDM/PySDM/environments/index.html"
     click KERNEL_INIT "https://open-atmos.github.io/PySDM/PySDM/dynamics/collisions/collision_kernels/index.html"
@@ -495,8 +493,7 @@ output_points = 40
 n_sd = 256
 
 formulae = Formulae()
-builder = Builder(backend=CPU(formulae), n_sd=n_sd)
-builder.set_environment(env)
+builder = Builder(backend=CPU(formulae), n_sd=n_sd, environment=env)
 builder.add_dynamic(AmbientThermodynamics())
 builder.add_dynamic(Condensation())
 
@@ -577,8 +574,7 @@ output_points = 40;
 n_sd = 256;
 
 formulae = Formulae();
-builder = Builder(pyargs('backend', CPU(formulae), 'n_sd', int32(n_sd)));
-builder.set_environment(env);
+builder = Builder(pyargs('backend', CPU(formulae), 'n_sd', int32(n_sd), 'environment', env));
 builder.add_dynamic(AmbientThermodynamics());
 builder.add_dynamic(Condensation());
 
@@ -678,8 +674,7 @@ output_points = 40
 n_sd = 256
 
 formulae = Formulae()
-builder = Builder(backend=CPU(formulae), n_sd=n_sd)
-builder.set_environment(env)
+builder = Builder(backend=CPU(formulae), n_sd=n_sd, environment=env)
 builder.add_dynamic(AmbientThermodynamics())
 builder.add_dynamic(Condensation())
 

examples/PySDM_examples/Abdul_Razzak_Ghan_2000/run_ARG_parcel.py

@@ -50,8 +50,7 @@ def run_parcel(
     aerosol = AerosolARG(M2_sol=sol2, M2_N=N2, M2_rad=rad2)
     n_sd = n_sd_per_mode * len(aerosol.modes)
 
-    builder = Builder(backend=CPU(), n_sd=n_sd)
-    builder.set_environment(env)
+    builder = Builder(backend=CPU(), n_sd=n_sd, environment=env)
     builder.add_dynamic(AmbientThermodynamics())
     builder.add_dynamic(Condensation())
     builder.request_attribute("critical supersaturation")

examples/PySDM_examples/Alpert_and_Knopf_2016/simulation.py

@@ -215,9 +215,8 @@ def simulation(
         constants=constants,
         particle_shape_and_density="MixedPhaseSpheres",
     )
-    builder = Builder(n_sd=n_sd, backend=CPU(formulae=formulae))
     env = Box(dt=time_step, dv=volume)
-    builder.set_environment(env)
+    builder = Builder(n_sd=n_sd, backend=CPU(formulae=formulae), environment=env)
     builder.add_dynamic(Freezing(singular=False))
 
     if hasattr(spectrum, "s_geom") and spectrum.s_geom == 1:

examples/PySDM_examples/Arabas_and_Shima_2017/simulation.py

@@ -18,16 +18,17 @@ def __init__(self, settings, backend=CPU):
         while dt_output / self.n_substeps >= settings.dt_max:  # TODO #334 dt_max
             self.n_substeps += 1
 
-        builder = Builder(backend=backend(formulae=settings.formulae), n_sd=1)
-        builder.set_environment(
-            Parcel(
+        builder = Builder(
+            backend=backend(formulae=settings.formulae),
+            n_sd=1,
+            environment=Parcel(
                 dt=dt_output / self.n_substeps,
                 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.add_dynamic(AmbientThermodynamics())

examples/PySDM_examples/Arabas_et_al_2023/aida.ipynb

@@ -139,7 +139,6 @@
     "    \"\"\" orchestrates AIDA-like simulation setup and run \"\"\"\n",
     "    def __init__(self, settings: Settings):\n",
     "        self.settings = settings\n",
-    "        builder = Builder(n_sd=settings.n_sd, backend=CPU(settings.formulae))\n",
     "        env = AIDA(\n",
     "            dt = settings.dt,\n",
     "            mass_of_dry_air = settings.mass_of_dry_air,\n",
@@ -149,7 +148,7 @@
     "            w = vertical_velocity,\n",
     "            mixed_phase = True\n",
     "        )\n",
-    "        builder.set_environment(env)\n",
+    "        builder = Builder(n_sd=settings.n_sd, backend=CPU(settings.formulae), environment=env)\n",
     "        builder.add_dynamic(AmbientThermodynamics())\n",
     "        builder.add_dynamic(Condensation())\n",
     "        builder.add_dynamic(Freezing(singular=settings.singular))\n",

examples/PySDM_examples/Arabas_et_al_2023/fig_convergence.ipynb

@@ -107,9 +107,8 @@
     "        key = f\"{case['dt']}:{case['N']}\"\n",
     "        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))\n",
     "        env = Box(dt=case[\"dt\"], dv=d_v)\n",
-    "        builder.set_environment(env)\n",
+    "        builder = Builder(n_sd=n_sd, backend=backend_class(formulae=formulae), environment=env)\n",
     "        builder.add_dynamic(Freezing(singular=False))\n",
     "        attributes = {\n",
     "            \"multiplicity\": np.full(n_sd, int(case[\"N\"])),\n",

examples/PySDM_examples/Arabas_et_al_2023/make_particulator.py

@@ -57,10 +57,9 @@ def make_particulator(
         ) = sampling.sample(backend=backend, n_sd=n_sd)
     attributes["multiplicity"] *= total_particle_number
 
-    builder = Builder(n_sd, backend)
-
     env = Box(dt, volume)
-    builder.set_environment(env)
+    builder = Builder(n_sd, backend, env)
+
     env["T"] = initial_temperature
     env["RH"] = A_VALUE_LARGER_THAN_ONE
     env["rhod"] = 1.0

examples/PySDM_examples/Bieli_et_al_2022/simulation.py

@@ -9,9 +9,10 @@
 def make_core(settings, coal_eff):
     backend = CPU
 
-    builder = Builder(n_sd=settings.n_sd, backend=backend(settings.formulae))
     env = Box(dv=settings.dv, dt=settings.dt)
-    builder.set_environment(env)
+    builder = Builder(
+        n_sd=settings.n_sd, backend=backend(settings.formulae), environment=env
+    )
     env["rhod"] = 1.0
     attributes = {}
     attributes["volume"], attributes["multiplicity"] = ConstantMultiplicity(

examples/PySDM_examples/Grabowski_and_Pawlowska_2023/simulation.py

@@ -27,8 +27,9 @@ def __init__(
             w=settings.vertical_velocity,
             mass_of_dry_air=44 * si.kg,
         )
-        builder = Builder(n_sd=settings.n_sd, backend=CPU(formulae=settings.formulae))
-        builder.set_environment(env)
+        builder = Builder(
+            n_sd=settings.n_sd, backend=CPU(formulae=settings.formulae), environment=env
+        )
         builder.add_dynamic(AmbientThermodynamics())
         builder.add_dynamic(
             Condensation(rtol_thd=settings.rtol_thd, rtol_x=settings.rtol_x)

examples/PySDM_examples/Kreidenweis_et_al_2003/simulation.py

@@ -21,8 +21,9 @@ def __init__(self, settings, products=None):
             w=settings.w,
         )
 
-        builder = Builder(n_sd=settings.n_sd, backend=CPU(formulae=settings.formulae))
-        builder.set_environment(env)
+        builder = Builder(
+            n_sd=settings.n_sd, backend=CPU(formulae=settings.formulae), environment=env
+        )
 
         attributes = env.init_attributes(
             n_in_dv=settings.n_in_dv,

examples/PySDM_examples/Lowe_et_al_2019/simulation.py

@@ -22,8 +22,9 @@ def __init__(self, settings, products=None):
             w=settings.w,
         )
         n_sd = settings.n_sd_per_mode * len(settings.aerosol.modes)
-        builder = Builder(n_sd=n_sd, backend=CPU(formulae=settings.formulae))
-        builder.set_environment(env)
+        builder = Builder(
+            n_sd=n_sd, backend=CPU(formulae=settings.formulae), environment=env
+        )
 
         attributes = {
             "dry volume": np.empty(0),

examples/PySDM_examples/Niedermeier_et_al_2014/simulation.py

@@ -16,18 +16,19 @@
 class Simulation(BasicSimulation):
     def __init__(self, settings: Settings):
         n_particles = settings.ccn_sampling_n - 1 + settings.in_sampling_n
-        builder = Builder(n_sd=n_particles, backend=CPU(settings.formulae))
-        builder.set_environment(
-            Parcel(
-                dt=settings.timestep,
-                p0=settings.p0,
-                T0=settings.T0,
-                initial_water_vapour_mixing_ratio=settings.initial_water_vapour_mixing_ratio,
-                mass_of_dry_air=settings.mass_of_dry_air,
-                w=settings.vertical_velocity,
-                mixed_phase=True,
-            )
+        env = Parcel(
+            dt=settings.timestep,
+            p0=settings.p0,
+            T0=settings.T0,
+            initial_water_vapour_mixing_ratio=settings.initial_water_vapour_mixing_ratio,
+            mass_of_dry_air=settings.mass_of_dry_air,
+            w=settings.vertical_velocity,
+            mixed_phase=True,
+        )
+        builder = Builder(
+            n_sd=n_particles, backend=CPU(settings.formulae), environment=env
         )
+
         builder.add_dynamic(AmbientThermodynamics())
         builder.add_dynamic(Condensation())
         builder.add_dynamic(Freezing(singular=False))

examples/PySDM_examples/Pyrcel/simulation.py

@@ -24,8 +24,9 @@ def __init__(
             mass_of_dry_air=44 * si.kg,
         )
         n_sd = sum(settings.n_sd_per_mode)
-        builder = Builder(n_sd=n_sd, backend=CPU(formulae=settings.formulae))
-        builder.set_environment(env)
+        builder = Builder(
+            n_sd=n_sd, backend=CPU(formulae=settings.formulae), environment=env
+        )
         builder.add_dynamic(AmbientThermodynamics())
         builder.add_dynamic(Condensation(rtol_thd=rtol_thd, rtol_x=rtol_x))
 

examples/PySDM_examples/Shima_et_al_2009/example.py

@@ -14,8 +14,10 @@
 
 
 def run(settings, backend=CPU, observers=()):
-    builder = Builder(n_sd=settings.n_sd, backend=backend(formulae=settings.formulae))
-    builder.set_environment(Box(dv=settings.dv, dt=settings.dt))
+    env = Box(dv=settings.dv, dt=settings.dt)
+    builder = Builder(
+        n_sd=settings.n_sd, backend=backend(formulae=settings.formulae), environment=env
+    )
     attributes = {}
     sampling = ConstantMultiplicity(settings.spectrum)
     attributes["volume"], attributes["multiplicity"] = sampling.sample(settings.n_sd)

examples/PySDM_examples/Shima_et_al_2009/example_timing.py

@@ -12,8 +12,8 @@
 
 
 def run(settings, backend):
-    builder = Builder(n_sd=settings.n_sd, backend=backend)
-    builder.set_environment(Box(dv=settings.dv, dt=settings.dt))
+    env = Box(dv=settings.dv, dt=settings.dt)
+    builder = Builder(n_sd=settings.n_sd, backend=backend, environment=env)
     attributes = {}
     sampling = ConstantMultiplicity(settings.spectrum)
     attributes["volume"], attributes["multiplicity"] = sampling.sample(settings.n_sd)

examples/PySDM_examples/Shipway_and_Hill_2012/simulation.py

@@ -30,9 +30,6 @@ def __init__(self, settings, backend=CPU):
         self.output_attributes = None
         self.output_products = None
 
-        self.builder = Builder(
-            n_sd=settings.n_sd, backend=backend(formulae=settings.formulae)
-        )
         self.mesh = Mesh(
             grid=(settings.nz,),
             size=(settings.z_max + settings.particle_reservoir_depth,),
@@ -66,7 +63,11 @@ def zZ_to_z_above_reservoir(zZ):
         )
         self.g_factor_vec = settings.rhod(_z_vec)
 
-        self.builder.set_environment(self.env)
+        self.builder = Builder(
+            n_sd=settings.n_sd,
+            backend=backend(formulae=settings.formulae),
+            environment=self.env,
+        )
         self.builder.add_dynamic(AmbientThermodynamics())
         self.builder.add_dynamic(
             Condensation(

examples/PySDM_examples/Srivastava_1982/simulation.py

@@ -22,6 +22,7 @@ def __init__(
         }
 
     def build(self, n_sd, seed, products):
+        env = Box(dt=self.settings.dt, dv=self.settings.dv)
         builder = Builder(
             backend=self.settings.backend_class(
                 formulae=Formulae(
@@ -32,8 +33,8 @@ def build(self, n_sd, seed, products):
                 double_precision=self.double_precision,
             ),
             n_sd=n_sd,
+            environment=env,
         )
-        builder.set_environment(Box(dt=self.settings.dt, dv=self.settings.dv))
         builder.add_dynamic(self.collision_dynamic)
         particulator = builder.build(
             products=products,

examples/PySDM_examples/Szumowski_et_al_1998/simulation.py

@@ -35,15 +35,16 @@ def products(self):
     def reinit(self, products=None):
         formulae = self.settings.formulae
         backend = self.backend_class(formulae=formulae)
-        builder = Builder(n_sd=self.settings.n_sd, backend=backend)
         environment = Kinematic2D(
             dt=self.settings.dt,
             grid=self.settings.grid,
             size=self.settings.size,
             rhod_of=self.settings.rhod_of_zZ,
             mixed_phase=self.settings.processes["freezing"],
         )
-        builder.set_environment(environment)
+        builder = Builder(
+            n_sd=self.settings.n_sd, backend=backend, environment=environment
+        )
 
         if products is not None:
             products = list(products)