Replace uses of __setitem__ with fill for Storages. closes #1022 (#1014)

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

Author: abulenok

PySDM/backends/impl_numba/storage.py

@@ -187,3 +187,9 @@ def to_ndarray(self):
 
     def upload(self, data):
         np.copyto(self.data, data, casting="safe")
+
+    def fill(self, other):
+        if isinstance(other, Storage):
+            self.data[:] = other.data
+        else:
+            self.data[:] = other

PySDM/backends/impl_thrust_rtc/storage.py

@@ -497,4 +497,17 @@ def divide_if_not_zero(self, divisor):
             Impl.divide_if_not_zero(self, divisor)
             return self
 
+        def fill(self, value):
+            if isinstance(value, Storage):
+                trtc.Copy(value.data, self.data)
+            else:
+                if isinstance(value, int):
+                    dvalue = trtc.DVInt64(value)
+                elif isinstance(value, float):
+                    dvalue = BACKEND._get_floating_point(value)
+                else:
+                    raise TypeError("Only Storage, int and float are supported.")
+                trtc.Fill(self.data, dvalue)
+            return self
+
     return Storage

PySDM/dynamics/collisions/breakup_efficiencies/constEb.py

@@ -12,4 +12,4 @@ def register(self, builder):
         self.particulator = builder.particulator
 
     def __call__(self, output, is_first_in_pair):
-        output[:] = self.Eb
+        output.fill(self.Eb)

PySDM/dynamics/collisions/breakup_fragmentations/always_n.py

@@ -11,7 +11,7 @@ def __init__(self, n, vmin=0.0, nfmax=None):
         self.nfmax = nfmax
 
     def __call__(self, nf, frag_size, u01, is_first_in_pair):
-        nf[:] = self.N
+        nf.fill(self.N)
         frag_size.sum(self.particulator.attributes["volume"], is_first_in_pair)
         frag_size /= self.N
 

PySDM/dynamics/collisions/breakup_fragmentations/straub2010.py

@@ -47,7 +47,7 @@ def __call__(self, nf, frag_size, u01, is_first_in_pair):
 
         # compute the dimensionless numbers and CW=CKE * We
         self.arrays["tmp"].sum(self.particulator.attributes["volume"], is_first_in_pair)
-        self.arrays["Sc"][:] = self.arrays["tmp"][:]  # TODO #976
+        self.arrays["Sc"].fill(self.arrays["tmp"])
         self.arrays["Sc"] **= 2 / 3
         self.arrays["Sc"] *= (
             self.const.PI * self.const.sgm_w * (6 / self.const.PI) ** (2 / 3)
@@ -64,12 +64,12 @@ def __call__(self, nf, frag_size, u01, is_first_in_pair):
         self.arrays["CKE"] *= self.arrays["tmp2"]
         self.arrays["CKE"] *= self.const.rho_w
 
-        self.arrays["We"][:] = self.arrays["CKE"][:]  # TODO #976
+        self.arrays["We"].fill(self.arrays["CKE"])
         self.arrays["We"].divide_if_not_zero(self.arrays["Sc"])
 
-        self.arrays["CW"][:] = self.arrays["We"][:]  # TODO #976
+        self.arrays["CW"].fill(self.arrays["We"])
         self.arrays["CW"] *= self.arrays["CKE"]
-        self.arrays["CW"] /= si.joule * 1e-6  # convert to µJ
+        self.arrays["CW"] /= si.uJ
 
         self.arrays["gam"].max(self.particulator.attributes["radius"], is_first_in_pair)
         self.arrays["tmp"].min(self.particulator.attributes["radius"], is_first_in_pair)

PySDM/dynamics/collisions/coalescence_efficiencies/constEc.py

@@ -10,4 +10,4 @@ def register(self, builder):
         self.particulator = builder.particulator
 
     def __call__(self, output, is_first_in_pair):
-        output[:] = self.Ec
+        output.fill(self.Ec)

PySDM/dynamics/collisions/coalescence_efficiencies/straub2010.py

@@ -25,7 +25,7 @@ def register(self, builder):
 
     def __call__(self, output, is_first_in_pair):
         self.arrays["tmp"].sum(self.particulator.attributes["volume"], is_first_in_pair)
-        self.arrays["Sc"][:] = self.arrays["tmp"][:]
+        self.arrays["Sc"].fill(self.arrays["tmp"])
         self.arrays["tmp"] *= 2
 
         self.arrays["tmp2"].distance(
@@ -47,4 +47,4 @@ def __call__(self, output, is_first_in_pair):
         self.arrays["We"].divide_if_not_zero(self.arrays["Sc"])
         self.arrays["We"] *= -1.15
 
-        output[:] = np.exp(self.arrays["We"])
+        output[:] = np.exp(self.arrays["We"])  # TODO #976

PySDM/dynamics/collisions/collision_kernels/constantK.py

@@ -9,8 +9,7 @@ def __init__(self, a):
         self.particulator = None
 
     def __call__(self, output, is_first_in_pair):
-        output *= 0
-        output += self.a
+        output.fill(self.a)
 
     def register(self, builder):
         self.particulator = builder.particulator