add ability to add and remove particles from an AeroState (#322)

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

Author: jcurtis2

src/aero_state.F90

@@ -482,4 +482,56 @@ subroutine f_aero_state_add_aero_dist_sample(ptr_c, ptr_aero_data_c, &
 
   end subroutine
 
+  subroutine f_aero_state_add_particle(ptr_c, ptr_aero_data_c, &
+       ptr_aero_particle_c) bind(C)
+
+    type(c_ptr) :: ptr_c, ptr_aero_data_c, ptr_aero_particle_c
+    type(aero_state_t), pointer :: ptr_f => null()
+    type(aero_data_t), pointer :: ptr_aero_data_f => null()
+    type(aero_particle_t), pointer :: ptr_aero_particle_f => null()
+
+    call c_f_pointer(ptr_c, ptr_f)
+    call c_f_pointer(ptr_aero_data_c, ptr_aero_data_f)
+    call c_f_pointer(ptr_aero_particle_c, ptr_aero_particle_f)
+
+    call aero_state_add_particle(ptr_f, ptr_aero_particle_f, ptr_aero_data_f, &
+        .true.)
+
+  end subroutine
+
+  subroutine f_aero_state_copy_weight(ptr_c, ptr_aero_state_to_c) bind(C)
+
+    type(c_ptr) :: ptr_c, ptr_aero_state_to_c
+    type(aero_state_t), pointer :: ptr_f => null()
+    type(aero_state_t), pointer :: ptr_aero_state_to_f => null()
+
+    call c_f_pointer(ptr_c, ptr_f)
+    call c_f_pointer(ptr_aero_state_to_c, ptr_aero_state_to_f)
+
+    call aero_state_copy_weight(ptr_f, ptr_aero_state_to_f)
+
+  end subroutine
+
+  subroutine f_aero_state_remove_particle(ptr_c, i_part) bind(C)
+
+    type(c_ptr) :: ptr_c
+    type(aero_state_t), pointer :: ptr_f => null()
+    integer(c_int) :: i_part
+
+    call c_f_pointer(ptr_c, ptr_f)
+
+    call aero_state_remove_particle_no_info(ptr_f, i_part + 1)
+
+  end subroutine
+
+  subroutine f_aero_state_zero(ptr_c) bind(C)
+    type(c_ptr) :: ptr_c
+    type(aero_state_t), pointer :: ptr_f => null()
+
+    call c_f_pointer(ptr_c, ptr_f)
+
+    call aero_state_zero(ptr_f)
+
+  end subroutine
+
 end module

src/aero_state.hpp

@@ -143,6 +143,25 @@ extern "C" void f_aero_state_add_aero_dist_sample(
     int *n_part_add
 ) noexcept;
 
+extern "C" void f_aero_state_add_particle(
+    void *ptr_c,
+    const void *ptr_aero_data_c,
+    const void *ptr_aero_particle_c
+) noexcept;
+
+extern "C" void f_aero_state_copy_weight(
+    const void *ptr_c,
+    void *ptr_aero_state_to_c
+) noexcept;
+
+extern "C" void f_aero_state_remove_particle(
+    void *ptr_c,
+    const int *i_part
+) noexcept;
+
+extern "C" void f_aero_state_zero(
+    void *ptr_c
+) noexcept;
 
 template <typename arr_t, typename arg_t>
 auto pointer_vec_magic(arr_t &data_vec, const arg_t &arg) {
@@ -485,4 +504,38 @@ struct AeroState {
        return n_part_add;
    }
 
+   static void add_particle(
+       AeroState &self,
+       const AeroParticle &particle
+   ) {
+
+       f_aero_state_add_particle(self.ptr.f_arg_non_const(),
+            self.aero_data->ptr.f_arg(),
+            particle.ptr.f_arg()
+       );
+
+   } 
+
+   static void copy_weight(
+      AeroState &self,
+      const AeroState &aero_state_from
+   ) {
+
+      f_aero_state_copy_weight(aero_state_from.ptr.f_arg(),
+           self.ptr.f_arg_non_const()
+      );
+   }
+
+   static void remove_particle(
+      AeroState &self,
+      const int &i_part
+   ) {
+     f_aero_state_remove_particle(self.ptr.f_arg_non_const(), &i_part);
+   }
+
+   static void zero(
+      AeroState &self
+   ) {
+      f_aero_state_zero(self.ptr.f_arg_non_const());
+   }
 };

src/pypartmc.cpp

@@ -235,6 +235,12 @@ PYBIND11_MODULE(_PyPartMC, m) {
             "sample particles for AeroState from an AeroDist",
             py::arg("AeroDist"), py::arg("sample_prop") = 1.0, py::arg("create_time") = 0.0,
             py::arg("allow_doubling") = true, py::arg("allow_halving") = true)
+        .def("add_particle", AeroState::add_particle, "add a particle to an AeroState")
+        .def("copy_weight", AeroState::copy_weight,
+             "copy weighting from another AeroState")
+        .def("remove_particle", AeroState::remove_particle,
+            "remove particle of a given index")
+        .def("zero", AeroState::zero, "remove all particles from an AeroState")
     ;
 
     py::class_<GasData, std::shared_ptr<GasData>>(m, "GasData",

tests/test_aero_state.py

@@ -322,6 +322,43 @@ def test_get_particle_out_of_range(
         # assert
         assert False
 
+    @staticmethod
+    def test_remove_particle(sut_minimal):  # pylint: disable=redefined-outer-name
+        diameters = sut_minimal.diameters()
+        sut_minimal.remove_particle(len(sut_minimal) - 1)
+
+        assert diameters[0:-1] == sut_minimal.diameters()
+
+    @staticmethod
+    def test_zero(sut_minimal):  # pylint: disable=redefined-outer-name
+        # act
+        sut_minimal.zero()
+
+        # assert
+        assert len(sut_minimal) == 0
+
+    @staticmethod
+    def test_add_particle(sut_minimal):  # pylint: disable=redefined-outer-name
+        particle = sut_minimal.particle(1)
+        aero_data = ppmc.AeroData(AERO_DATA_CTOR_ARG_MINIMAL)
+        sut = ppmc.AeroState(aero_data, *AERO_STATE_CTOR_ARG_MINIMAL)
+
+        sut.add_particle(particle)
+
+        assert len(sut) == 1
+        assert sut.particle(0).diameter == sut_minimal.particle(1).diameter
+
+    @staticmethod
+    def test_copy_weight(sut_minimal):  # pylint: disable=redefined-outer-name
+        aero_data = ppmc.AeroData(AERO_DATA_CTOR_ARG_MINIMAL)
+        sut = ppmc.AeroState(aero_data, *AERO_STATE_CTOR_ARG_MINIMAL)
+
+        sut.copy_weight(sut_minimal)
+        sut.add_particle(sut_minimal.particle(0))
+
+        # pylint: disable=unsubscriptable-object
+        assert sut.num_concs[0] == sut_minimal.num_concs[0]
+
     @staticmethod
     @pytest.mark.parametrize(
         "args",