PySDM.builder

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