PySDM.backends.numba

19class Numba(  # pylint: disable=too-many-ancestors,duplicate-code
20    methods.CollisionsMethods,
21    methods.FragmentationMethods,
22    methods.PairMethods,
23    methods.IndexMethods,
24    methods.PhysicsMethods,
25    methods.CondensationMethods,
26    methods.ChemistryMethods,
27    methods.MomentsMethods,
28    methods.FreezingMethods,
29    methods.DisplacementMethods,
30    methods.TerminalVelocityMethods,
31    methods.IsotopeMethods,
32    methods.SeedingMethods,
33    methods.DepositionMethods,
34):
35    Storage = ImportedStorage
36    Random = ImportedRandom
37
38    default_croupier = "local"
39
40    def __init__(self, formulae=None, double_precision=True, override_jit_flags=None):
41        if not double_precision:
42            raise NotImplementedError()
43        self.formulae = formulae or Formulae()
44        self.formulae_flattened = self.formulae.flatten
45
46        parallel_default = True
47        if platform.machine() == "arm64":
48            if "CI" not in os.environ:
49                warnings.warn(
50                    "Disabling Numba threading due to ARM64 CPU (atomics do not work yet)"
51                )
52            parallel_default = False  # TODO #1183 - atomics don't work on ARM64!
53
54        try:
55            numba.parfors.parfor.ensure_parallel_support()
56        except numba.core.errors.UnsupportedParforsError:
57            if "CI" not in os.environ:
58                warnings.warn(
59                    "Numba version used does not support parallel for (32 bits?)"
60                )
61            parallel_default = False
62
63        assert "fastmath" not in (override_jit_flags or {})
64        self.default_jit_flags = {
65            **JIT_FLAGS,  # here parallel=False (for out-of-backend code)
66            **{"fastmath": self.formulae.fastmath, "parallel": parallel_default},
67            **(override_jit_flags or {}),
68        }
69
70        methods.CollisionsMethods.__init__(self)
71        methods.FragmentationMethods.__init__(self)
72        methods.PairMethods.__init__(self)
73        methods.IndexMethods.__init__(self)
74        methods.PhysicsMethods.__init__(self)
75        methods.CondensationMethods.__init__(self)
76        methods.ChemistryMethods.__init__(self)
77        methods.MomentsMethods.__init__(self)
78        methods.FreezingMethods.__init__(self)
79        methods.DisplacementMethods.__init__(self)
80        methods.TerminalVelocityMethods.__init__(self)
81        methods.IsotopeMethods.__init__(self)
82        methods.SeedingMethods.__init__(self)
83        methods.DepositionMethods.__init__(self)

 17class Storage(StorageBase):
 18    FLOAT = np.float64
 19    INT = np.int64
 20    BOOL = np.bool_
 21
 22    def __getitem__(self, item):
 23        dim = len(self.shape)
 24        if isinstance(item, slice):
 25            step = item.step or 1
 26            if step != 1:
 27                raise NotImplementedError("step != 1")
 28            start = item.start or 0
 29            if dim == 1:
 30                stop = item.stop or len(self)
 31                result_data = self.data[item]
 32                result_shape = (stop - start,)
 33            elif dim == 2:
 34                stop = item.stop or self.shape[0]
 35                result_data = self.data[item]
 36                result_shape = (stop - start, self.shape[1])
 37            else:
 38                raise NotImplementedError(
 39                    "Only 2 or less dimensions array is supported."
 40                )
 41            if stop > self.data.shape[0]:
 42                raise IndexError(
 43                    f"requested a slice ({start}:{stop}) of Storage"
 44                    f" with first dim of length {self.data.shape[0]}"
 45                )
 46            result = Storage(StorageSignature(result_data, result_shape, self.dtype))
 47        elif isinstance(item, tuple) and dim == 2 and isinstance(item[1], slice):
 48            result = Storage(
 49                StorageSignature(self.data[item[0]], (*self.shape[1:],), self.dtype)
 50            )
 51        else:
 52            result = self.data[item]
 53        return result
 54
 55    def __setitem__(self, key, value):
 56        if hasattr(value, "data"):
 57            self.data[key] = value.data
 58        else:
 59            self.data[key] = value
 60        return self
 61
 62    def __iadd__(self, other):
 63        if isinstance(other, Storage):
 64            impl.add(self.data, other.data)
 65        elif (
 66            isinstance(other, tuple)
 67            and len(other) == 3
 68            and isinstance(other[0], float)
 69            and other[1] == "*"
 70            and isinstance(other[2], Storage)
 71        ):
 72            impl.add_with_multiplier(self.data, other[2].data, other[0])
 73        else:
 74            impl.add(self.data, other)
 75        return self
 76
 77    def __isub__(self, other):
 78        impl.subtract(self.data, other.data)
 79        return self
 80
 81    def __imul__(self, other):
 82        if hasattr(other, "data"):
 83            impl.multiply(self.data, other.data)
 84        else:
 85            impl.multiply(self.data, other)
 86        return self
 87
 88    def __itruediv__(self, other):
 89        if hasattr(other, "data"):
 90            self.data[:] /= other.data[:]
 91        else:
 92            self.data[:] /= other
 93        return self
 94
 95    def __imod__(self, other):
 96        impl.row_modulo(self.data, other.data)
 97        return self
 98
 99    def __ipow__(self, other):
100        impl.power(self.data, other)
101        return self
102
103    def __bool__(self):
104        if len(self) == 1:
105            result = bool(self.data[0] != 0)
106        else:
107            raise NotImplementedError("Logic value of array is ambiguous.")
108        return result
109
110    def detach(self):
111        if self.data.base is not None:
112            self.data = np.array(self.data)
113
114    def download(self, target, reshape=False):
115        if reshape:
116            data = self.data.reshape(target.shape)
117        else:
118            data = self.data
119        np.copyto(target, data, casting="safe")
120
121    @staticmethod
122    def _get_empty_data(shape, dtype):
123        if dtype in (float, Storage.FLOAT):
124            data = np.full(shape, np.nan, dtype=Storage.FLOAT)
125            dtype = Storage.FLOAT
126        elif dtype in (int, Storage.INT):
127            data = np.full(shape, -1, dtype=Storage.INT)
128            dtype = Storage.INT
129        elif dtype in (bool, Storage.BOOL):
130            data = np.full(shape, -1, dtype=Storage.BOOL)
131            dtype = Storage.BOOL
132        else:
133            raise NotImplementedError()
134
135        return StorageSignature(data, shape, dtype)
136
137    @staticmethod
138    def empty(shape, dtype):
139        return empty(shape, dtype, Storage)
140
141    @staticmethod
142    def _get_data_from_ndarray(array):
143        return get_data_from_ndarray(
144            array=array,
145            storage_class=Storage,
146            copy_fun=lambda array_astype: array_astype.copy(),
147        )
148
149    def amin(self):
150        return impl.amin(self.data)
151
152    def amax(self):
153        return impl.amax(self.data)
154
155    def all(self):
156        return self.data.all()
157
158    @staticmethod
159    def from_ndarray(array):
160        result = Storage(Storage._get_data_from_ndarray(array))
161        return result
162
163    def floor(self, other=None):
164        if other is None:
165            impl.floor(self.data)
166        else:
167            impl.floor_out_of_place(self.data, other.data)
168        return self
169
170    def product(self, multiplicand, multiplier):
171        if hasattr(multiplier, "data"):
172            impl.multiply_out_of_place(self.data, multiplicand.data, multiplier.data)
173        else:
174            impl.multiply_out_of_place(self.data, multiplicand.data, multiplier)
175        return self
176
177    def ratio(self, dividend, divisor):
178        impl.divide_out_of_place(self.data, dividend.data, divisor.data)
179        return self
180
181    def divide_if_not_zero(self, divisor):
182        impl.divide_if_not_zero(self.data, divisor.data)
183        return self
184
185    def sum(self, arg_a, arg_b):
186        impl.sum_out_of_place(self.data, arg_a.data, arg_b.data)
187        return self
188
189    def ravel(self, other):
190        if isinstance(other, Storage):
191            self.data[:] = other.data.ravel()
192        else:
193            self.data[:] = other.ravel()
194
195    def urand(self, generator):
196        generator(self)
197
198    def to_ndarray(self):
199        return self.data.copy()
200
201    def upload(self, data):
202        np.copyto(self.data, data, casting="safe")
203
204    def fill(self, other):
205        if isinstance(other, Storage):
206            self.data[:] = other.data
207        else:
208            self.data[:] = other
209
210    def exp(self):
211        self.data[:] = np.exp(self.data)
212
213    def abs(self):
214        self.data[:] = np.abs(self.data)

14class Random(RandomCommon):  # pylint: disable=too-few-public-methods
15    def __init__(self, size, seed):
16        super().__init__(size, seed)
17        self.generator = np.random.default_rng(seed)
18
19    def __call__(self, storage):
20        storage.data[:] = self.generator.uniform(0, 1, storage.shape)