13 02 part 2

Author: Bodzio-2

PySDM/backends/impl_jax/methods/moments_methods.py

@@ -120,18 +120,14 @@ def body(
             attr_data,
             cell_id,
             idx,
-            length,
             rank,
-            x_bins,
-            x_attr,
             weighting_attribute,
             weighting_rank,
             bin_to_count,
             idx_i
         ):
 
             i = idx[idx_i]
-
             moment_0 = moment_0.at[bin_to_count, cell_id[i]].add(multiplicity[i] * weighting_attribute[i] ** weighting_rank)
             moments = moments.at[bin_to_count, cell_id[i]].add(multiplicity[i] * weighting_attribute[i] ** weighting_rank * attr_data[i] ** rank)
 
@@ -170,34 +166,35 @@ def cond_fun(k):
             return bin_to_calculate
         # TODO: what happens if k == x_bins.shape[0] - 1
 
-
-        moment_0.data = moment_0.data.at[:, :].set(0)
-        moments.data = moments.data.at[:,:].set(0)
+        new_moment_0 = jax.numpy.zeros((moment_0.shape[0]+1, moment_0.shape[1]))
+        new_moments = jax.numpy.zeros((moment_0.shape[0]+1, moment_0.shape[1]))
         idx_idxs = jax.numpy.arange(length)
 
         count_bins_func = jax.vmap(spectrum_moments_helper, (None, None, None, 0))
         bins_to_count = count_bins_func(x_bins.data, x_attr.data, idx.data, idx_idxs)
-        mapped_spectrum = jax.vmap(self._spectrum_moments_body, (None, None, None, None, None, None, None, None, None, None, None, None, 0, 0))
-
-        moment_0.data, moments.data = mapped_spectrum(
-            moment_0.data,
-            moments.data,            
+        # TODO: bins_to_count > len()-1 is not handled
+        print(bins_to_count)
+        assert all(bins_to_count < new_moments.shape[0])
+        mapped_spectrum = jax.vmap(self._spectrum_moments_body, (None, None, None, None, None, None, None, None, None, 0, 0))
+
+        new_moment_0, new_moments = mapped_spectrum(
+            new_moment_0,
+            new_moments,            
             multiplicity.data,
             attr_data.data,
             cell_id.data,
             idx.data,
-            length,
             rank,
-            x_bins.data,
-            x_attr.data,
             weighting_attribute.data,
             weighting_rank,
             bins_to_count,
             idx_idxs
         )
 
-        moments.data = moments.data.sum(0)
-        moment_0.data = moment_0.data.sum(0)
+        # moments.data = new_moments.sum(0)
+        # moment_0.data = new_moment_0.sum(0)
+        moments.data =  new_moments.at[:-1, :].sum(0)
+        moment_0.data =  new_moment_0.at[:-1, :].sum(0)
 
         if not skip_division_by_m0:
             moments.data = jax.numpy.where(moment_0.data != 0, moments.data / moment_0.data, 0.0)

tests/unit_tests/impl/test_moments.py

@@ -135,3 +135,84 @@ def test_spectrum_moment_0d(backend_class_with_jax, skip_division_by_m0):
             actual=sum(actual if skip_division_by_m0 else actual * spectrum_moment_0.to_ndarray()),
             significant=4,
         )
+
+
+    @staticmethod
+    @pytest.mark.parametrize('skip_division_by_m0', (True, False))
+    def test_spectrum_moment_0d_one_drop_outside_of_bins(backend_class_with_jax, skip_division_by_m0):
+        # Arrange
+        particulator = DummyParticulator(backend_class_with_jax, n_sd=1)
+        attribute = {
+            "multiplicity": np.ones(1),
+            "volume": np.ones(1),
+        }
+        particulator.build(attribute)
+
+        v_bins = np.linspace(0, .5, num=2, endpoint=True)
+
+        assert len(v_bins) == 2
+        spectrum_moment_0 = particulator.backend.Storage.empty(
+            (len(v_bins) - 1, 1), dtype=float
+        )
+        spectrum_moments = particulator.backend.Storage.empty(
+            (len(v_bins) - 1, 1), dtype=float
+        )
+        v_bins_edges = particulator.backend.Storage.from_ndarray(v_bins)
+
+        # Act
+        particulator.spectrum_moments(
+            moment_0=spectrum_moment_0,
+            moments=spectrum_moments,
+            attr="volume",
+            attr_name='volume',
+            rank=1,
+            attr_bins=v_bins_edges,
+            skip_division_by_m0=skip_division_by_m0,
+        )
+
+        # Assert
+        assert spectrum_moment_0.to_ndarray() == [0]
+        assert spectrum_moments.to_ndarray() == [0]
+
+    @staticmethod
+    def test_moment_1d(backend_class_with_jax):
+        # Arrange
+        grid=(2,)
+        particulator = DummyParticulator(
+            backend_class_with_jax,
+            n_sd=1,
+            grid=grid
+        )
+        
+        attribute = {
+            "multiplicity": np.ones(1),
+            "volume": np.ones(1),
+            "cell id": np.ones(1, dtype=int),
+        }
+        particulator.build(attribute)
+
+        v_bins = np.linspace(0, 2, num=2, endpoint=True)
+
+        assert len(v_bins) == 2
+        spectrum_moment_0 = particulator.backend.Storage.empty(
+            (len(v_bins) - 1, np.prod(grid)), dtype=float
+        )
+        spectrum_moments = particulator.backend.Storage.empty(
+            (len(v_bins) - 1, np.prod(grid)), dtype=float
+        )
+        v_bins_edges = particulator.backend.Storage.from_ndarray(v_bins)
+
+        # Act
+        particulator.spectrum_moments(
+            moment_0=spectrum_moment_0,
+            moments=spectrum_moments,
+            attr="volume",
+            attr_name='volume',
+            rank=1,
+            attr_bins=v_bins_edges,
+            skip_division_by_m0=True,
+        )
+
+        # Assert
+        assert np.prod(spectrum_moment_0.to_ndarray() == [0, 1])
+        assert np.prod(spectrum_moments.to_ndarray() == [0, 1])