PySDM.builder

  1"""
  2The Builder class handling creation of  `PySDM.particulator.Particulator` instances
  3"""
  4
  5import inspect
  6import warnings
  7
  8import numpy as np
  9
 10from PySDM.attributes.impl.attribute_registry import get_attribute_class
 11from PySDM.impl.particle_attributes_factory import ParticleAttributesFactory
 12from PySDM.impl.wall_timer import WallTimer
 13from PySDM.initialisation.discretise_multiplicities import (  # TODO #324
 14    discretise_multiplicities,
 15)
 16from PySDM.particulator import Particulator
 17from PySDM.physics.particle_shape_and_density import LiquidSpheres, MixedPhaseSpheres
 18
 19
 20def _warn_env_as_ctor_arg():
 21    warnings.warn(
 22        "PySDM > v2.31 Builder expects environment instance as argument",
 23        DeprecationWarning,
 24    )
 25
 26
 27class Builder:
 28    def __init__(self, n_sd, backend, environment=None):
 29        assert not inspect.isclass(backend)
 30        self.formulae = backend.formulae
 31        self.particulator = Particulator(n_sd, backend)
 32        self.req_attr_names = ["multiplicity", "water mass", "cell id"]
 33        self.req_attr = None
 34        self.aerosol_radius_threshold = 0
 35        self.condensation_params = None
 36
 37        if environment is None:
 38            _warn_env_as_ctor_arg()
 39        else:
 40            self._set_environment(environment)
 41
 42    def _set_condensation_parameters(self, **kwargs):
 43        self.condensation_params = kwargs
 44
 45    def set_environment(self, environment):
 46        _warn_env_as_ctor_arg()
 47        self._set_environment(environment)
 48
 49    def _set_environment(self, environment):
 50        if self.particulator.environment is not None:
 51            raise AssertionError("environment has already been set")
 52        self.particulator.environment = environment.instantiate(builder=self)
 53
 54    def add_dynamic(self, dynamic):
 55        assert self.particulator.environment is not None
 56        key = inspect.getmro(type(dynamic))[-2].__name__
 57        assert key not in self.particulator.dynamics
 58        self.particulator.dynamics[key] = dynamic
 59
 60    def _register_product(self, product, buffer):
 61        if product.name in self.particulator.products:
 62            raise ValueError(f'product name "{product.name}" already registered')
 63        self.particulator.products[product.name] = product.instantiate(
 64            builder=self, buffer=buffer
 65        )
 66
 67    def _resolve_attribute(self, attr_name):
 68        if attr_name not in self.req_attr:
 69            self.req_attr[attr_name] = get_attribute_class(
 70                attr_name,
 71                self.particulator.dynamics.keys(),
 72                self.formulae,
 73            )(self)
 74            assert self.req_attr is not None
 75
 76    def get_attribute(self, attribute_name):
 77        """intended for obtaining attribute instances during build() logic,
 78        from within register() methods"""
 79        self._resolve_attribute(attribute_name)
 80        return self.req_attr[attribute_name]
 81
 82    def request_attribute(self, attribute_name):
 83        """can be called either before or during build()"""
 84        if self.req_attr_names is not None:
 85            self.req_attr_names.append(attribute_name)
 86        else:
 87            self._resolve_attribute(attribute_name)
 88
 89    def build(
 90        self,
 91        attributes: dict,
 92        products: tuple = (),
 93        int_caster=discretise_multiplicities,
 94    ):
 95        assert self.particulator.environment is not None
 96
 97        if "n" in attributes and "multiplicity" not in attributes:
 98            attributes["multiplicity"] = attributes["n"]
 99            del attributes["n"]
100            warnings.warn(
101                'renaming attributes["n"] to attributes["multiplicity"]',
102                DeprecationWarning,
103            )
104
105        if "volume" in attributes and "water mass" not in attributes:
106            assert self.particulator.formulae.particle_shape_and_density.__name__ in (
107                LiquidSpheres.__name__,
108                MixedPhaseSpheres.__name__,
109            ), "implied volume-to-mass conversion is only supported for spherical particles"
110            attributes["water mass"] = (
111                self.particulator.formulae.particle_shape_and_density.volume_to_mass(
112                    attributes.pop("volume")
113                )
114            )
115            self.request_attribute("volume")
116
117        if (
118            "water mass" in attributes
119            and "signed water mass" not in attributes
120            and not self.particulator.formulae.particle_shape_and_density.supports_mixed_phase()
121        ):
122            attributes["signed water mass"] = attributes.pop("water mass")
123            self.request_attribute("water mass")
124
125        self.req_attr = {}
126        for attr_name in self.req_attr_names:
127            self._resolve_attribute(attr_name)
128        self.req_attr_names = None
129
130        for key, dynamic in self.particulator.dynamics.items():
131            self.particulator.dynamics[key] = dynamic.instantiate(builder=self)
132
133        single_buffer_for_all_products = np.empty(self.particulator.mesh.grid)
134        for product in products:
135            self._register_product(product, single_buffer_for_all_products)
136
137        for attribute in attributes:
138            self.request_attribute(attribute)
139        if "Condensation" in self.particulator.dynamics:
140            self.particulator.condensation_solver = (
141                self.particulator.backend.make_condensation_solver(
142                    self.particulator.dt,
143                    self.particulator.mesh.n_cell,
144                    **self.condensation_params,
145                )
146            )
147        attributes["multiplicity"] = int_caster(attributes["multiplicity"])
148        if self.particulator.mesh.dimension == 0:
149            attributes["cell id"] = np.zeros_like(
150                attributes["multiplicity"], dtype=np.int64
151            )
152        self.particulator.attributes = ParticleAttributesFactory.attributes(
153            self.particulator, self.req_attr, attributes
154        )
155        self.particulator.recalculate_cell_id()
156
157        for key in self.particulator.dynamics:
158            self.particulator.timers[key] = WallTimer()
159
160        if (attributes["multiplicity"] == 0).any():
161            self.particulator.attributes.healthy = False
162            self.particulator.attributes.sanitize()
163
164        return self.particulator