Add more functionality for AeroState (sample/add[particles], make_dry…

…, mobility_diameters, ids) (#341)

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

src/aero_state.F90

@@ -194,6 +194,24 @@ subroutine f_aero_state_dry_diameters(ptr_c, aero_data_ptr_c, diameters, n_parts
 
   end subroutine
 
+  subroutine f_aero_state_mobility_diameters(ptr_c, aero_data_ptr_c, &
+       env_state_ptr_c, diameters, n_parts) bind (C)
+    type(c_ptr), intent(in) :: ptr_c, aero_data_ptr_c, env_state_ptr_c
+    integer(c_int), intent(in) :: n_parts
+    real(c_double), intent(out) :: diameters(n_parts)
+    type(aero_state_t), pointer :: ptr_f => null()
+    type(aero_data_t), pointer :: aero_data_ptr_f => null()
+    type(env_state_t) , pointer :: env_state_ptr_f => null()
+
+    call c_f_pointer(ptr_c, ptr_f)
+    call c_f_pointer(aero_data_ptr_c, aero_data_ptr_f)
+    call c_f_pointer(env_state_ptr_c, env_state_ptr_f)
+
+    diameters = aero_state_mobility_diameters(ptr_f, aero_data_ptr_f, &
+         env_state_ptr_f)
+
+  end subroutine
+
   subroutine f_aero_state_diameters(ptr_c, aero_data_ptr_c, diameters, &
        n_parts, include_len, exclude_len, include, exclude) bind(C)
 
@@ -534,4 +552,91 @@ subroutine f_aero_state_zero(ptr_c) bind(C)
 
   end subroutine
 
+  subroutine f_aero_state_ids(ptr_c, ids, n_parts) bind(C)
+    type(c_ptr), intent(in) :: ptr_c
+    integer(c_int), intent(in) :: n_parts
+    integer(c_int), intent(out) :: ids(n_parts)
+    type(aero_state_t), pointer :: ptr_f => null()
+
+    call c_f_pointer(ptr_c, ptr_f)
+
+    ids = aero_state_ids(ptr_f)
+
+  end subroutine
+
+  subroutine f_aero_state_make_dry(ptr_c, aero_data_ptr_c) bind(C)
+    type(c_ptr), intent(in) :: ptr_c, aero_data_ptr_c
+    type(aero_state_t), pointer :: ptr_f => null()
+    type(aero_data_t), pointer :: aero_data_ptr_f => null()
+
+    call c_f_pointer(ptr_c, ptr_f)
+    call c_f_pointer(aero_data_ptr_c, aero_data_ptr_f)
+
+    call aero_state_make_dry(ptr_f, aero_data_ptr_f)
+
+  end subroutine
+
+  subroutine f_aero_state_add(ptr_c, delta_ptr_c, aero_data_ptr_c) bind(C)
+    type(c_ptr), intent(in) :: ptr_c, delta_ptr_c, aero_data_ptr_c
+    type(aero_state_t), pointer :: ptr_f => null()
+    type(aero_state_t), pointer :: delta_ptr_f => null()
+    type(aero_data_t), pointer :: aero_data_ptr_f => null()
+
+    call c_f_pointer(ptr_c, ptr_f)
+    call c_f_pointer(delta_ptr_c, delta_ptr_f)
+    call c_f_pointer(aero_data_ptr_c, aero_data_ptr_f)
+
+    call aero_state_add(ptr_f, delta_ptr_f, aero_data_ptr_f)
+
+  end subroutine
+
+  subroutine f_aero_state_add_particles(ptr_c, delta_ptr_c, aero_data_ptr_c &
+       ) bind(C)
+    type(c_ptr), intent(in) :: ptr_c, delta_ptr_c, aero_data_ptr_c
+    type(aero_state_t), pointer :: ptr_f => null()
+    type(aero_state_t), pointer :: delta_ptr_f => null()
+    type(aero_data_t), pointer :: aero_data_ptr_f => null()
+
+    call c_f_pointer(ptr_c, ptr_f)
+    call c_f_pointer(delta_ptr_c, delta_ptr_f)
+    call c_f_pointer(aero_data_ptr_c, aero_data_ptr_f)
+
+    call aero_state_add_particles(ptr_f, delta_ptr_f, aero_data_ptr_f)
+
+  end subroutine
+
+  subroutine f_aero_state_sample(ptr_c, aero_state_to_ptr_c, aero_data_ptr_c, &
+       sample_prob) bind(C)
+    type(c_ptr), intent(in) :: ptr_c, aero_state_to_ptr_c, aero_data_ptr_c
+    real(c_double), intent(in) :: sample_prob
+    type(aero_state_t), pointer :: ptr_f => null()
+    type(aero_state_t), pointer :: aero_state_to_ptr_f => null()
+    type(aero_data_t), pointer :: aero_data_ptr_f => null()
+
+    call c_f_pointer(ptr_c, ptr_f)
+    call c_f_pointer(aero_state_to_ptr_c, aero_state_to_ptr_f)
+    call c_f_pointer(aero_data_ptr_c, aero_data_ptr_f)
+
+    call aero_state_sample(ptr_f, aero_state_to_ptr_f, aero_data_ptr_f, & 
+         sample_prob, AERO_INFO_NONE)
+
+  end subroutine
+
+  subroutine f_aero_state_sample_particles(ptr_c, aero_state_to_ptr_c, &
+        aero_data_ptr_c, sample_prob) bind(C)
+    type(c_ptr), intent(in) :: ptr_c, aero_state_to_ptr_c, aero_data_ptr_c
+    real(c_double), intent(in) :: sample_prob
+    type(aero_state_t), pointer :: ptr_f => null()
+    type(aero_state_t), pointer :: aero_state_to_ptr_f => null()
+    type(aero_data_t), pointer :: aero_data_ptr_f => null()
+
+    call c_f_pointer(ptr_c, ptr_f)
+    call c_f_pointer(aero_state_to_ptr_c, aero_state_to_ptr_f)
+    call c_f_pointer(aero_data_ptr_c, aero_data_ptr_f)
+
+    call aero_state_sample_particles(ptr_f, aero_state_to_ptr_f, &
+         aero_data_ptr_f, sample_prob, AERO_INFO_NONE)
+
+  end subroutine
+
 end module

src/aero_state.hpp

@@ -71,6 +71,14 @@ extern "C" void f_aero_state_dry_diameters(
     const int *n_parts
 ) noexcept;
 
+extern "C" void f_aero_state_mobility_diameters(
+    const void *ptr,
+    const void *aero_dataptr,
+    const void *env_stateptr,
+    double *mobility_diameters,
+    const int *n_parts
+) noexcept;
+
 extern "C" void f_aero_state_diameters(
     const void *ptr,
     const void *aero_dataptr,
@@ -101,6 +109,11 @@ extern "C" void f_aero_state_crit_rel_humids(
     const int *n_parts
 ) noexcept;
 
+extern "C" void f_aero_state_make_dry(
+    void *ptr,
+    const void *aero_dataptr
+) noexcept;
+
 extern "C" void f_aero_state_mixing_state_metrics(
     const void *aero_state, 
     const void *aero_data,
@@ -163,6 +176,38 @@ extern "C" void f_aero_state_zero(
     void *ptr_c
 ) noexcept;
 
+extern "C" void f_aero_state_ids(
+    const void *ptr_c,
+    int *ids,
+    const int *n_parts
+) noexcept;
+
+extern "C" void f_aero_state_add(
+     void *ptr_c,
+     const void *delta_ptr_c,
+     const void *aero_data_ptr
+) noexcept;
+
+extern "C" void f_aero_state_add_particles(
+     void *ptr_c,
+     const void *delta_ptr_c,
+     const void *aero_data_ptr
+) noexcept;
+
+extern "C" void f_aero_state_sample(
+     void *ptr_c,
+     void *aero_state_to_ptr_c,
+     const void *aero_data_ptr,
+     const double *sample_prob
+) noexcept;
+
+extern "C" void f_aero_state_sample_particles(
+     void *ptr_c,
+     void *aero_state_to_ptr_c,
+     const void *aero_data_ptr,
+     const double *sample_prob
+) noexcept;
+
 template <typename arr_t, typename arg_t>
 auto pointer_vec_magic(arr_t &data_vec, const arg_t &arg) {
     std::vector<char*> pointer_vec(data_vec.size());
@@ -313,6 +358,25 @@ struct AeroState {
         return dry_diameters;
     }
 
+    static auto mobility_diameters(const AeroState &self, const EnvState &env_state) {
+        int len;
+        f_aero_state_len(
+            self.ptr.f_arg(),
+            &len
+        );
+        std::valarray<double> mobility_diameters(len);
+
+        f_aero_state_mobility_diameters(
+            self.ptr.f_arg(),
+            self.aero_data->ptr.f_arg(),
+            env_state.ptr.f_arg(),
+            begin(mobility_diameters),
+            &len
+        );
+
+        return mobility_diameters;
+    }
+
     static auto diameters(
         const AeroState &self,
         const tl::optional<std::valarray<std::string>> &include,
@@ -401,6 +465,32 @@ struct AeroState {
         return crit_rel_humids;
     }
 
+    static void make_dry(
+        AeroState &self
+    ) {
+        f_aero_state_make_dry(
+            self.ptr.f_arg_non_const(),
+            self.aero_data->ptr.f_arg()
+        );
+    }
+
+    static auto ids(const AeroState &self) {
+        int len;
+        f_aero_state_len(
+            self.ptr.f_arg(),
+            &len
+        );
+        std::valarray<int> ids(len);
+
+        f_aero_state_ids(
+            self.ptr.f_arg(),
+            begin(ids),
+            &len
+        );
+
+        return ids;
+    }
+
     static auto mixing_state(
         const AeroState &self,
         const tl::optional<std::valarray<std::string>> &include,
@@ -538,4 +628,48 @@ struct AeroState {
    ) {
       f_aero_state_zero(self.ptr.f_arg_non_const());
    }
+
+   static void add(
+       AeroState &self,
+       const AeroState &delta
+   )  {
+        f_aero_state_add(self.ptr.f_arg_non_const(),
+            delta.ptr.f_arg(),
+            self.aero_data->ptr.f_arg()
+      );
+   }
+
+   static void add_particles(
+       AeroState &self,
+       const AeroState &delta
+   )  {
+        f_aero_state_add_particles(self.ptr.f_arg_non_const(),
+            delta.ptr.f_arg(),
+            self.aero_data->ptr.f_arg()
+      );
+   }
+
+   static void sample(
+       AeroState &self,
+       AeroState &aero_state_sample,
+       const double sample_prob
+   )  {
+        f_aero_state_sample(self.ptr.f_arg_non_const(),
+            aero_state_sample.ptr.f_arg_non_const(),
+            self.aero_data->ptr.f_arg(),
+            &sample_prob
+      );
+   }
+
+   static void sample_particles(
+       AeroState &self,
+       AeroState &aero_state_sample,
+       const double sample_prob
+   )  {
+        f_aero_state_sample_particles(self.ptr.f_arg_non_const(),
+            aero_state_sample.ptr.f_arg_non_const(),
+            self.aero_data->ptr.f_arg(),
+            &sample_prob
+      );
+   }
 };

src/pypartmc.cpp

@@ -223,11 +223,17 @@ PYBIND11_MODULE(_PyPartMC, m) {
              py::arg("include") = py::none(), py::arg("exclude") = py::none())
         .def_property_readonly("dry_diameters", AeroState::dry_diameters,
             "returns the dry diameter of each particle in the population")
+        .def("mobility_diameters", AeroState::mobility_diameters,
+            "returns the mobility diameter of each particle in the population")
         .def("diameters", AeroState::diameters,
             "returns the diameter of each particle in the population",
             py::arg("include") = py::none(), py::arg("exclude") = py::none())
         .def("crit_rel_humids", AeroState::crit_rel_humids,
             "returns the critical relative humidity of each particle in the population")
+        .def("make_dry", AeroState::make_dry,
+            "Make all particles dry (water set to zero).")
+        .def_property_readonly("ids", AeroState::ids,
+            "returns the IDs of all particles.")
         .def("mixing_state", AeroState::mixing_state,
             "returns the mixing state parameters (d_alpha, d_gamma, chi) of the population",
             py::arg("include") = py::none(), py::arg("exclude") = py::none(),
@@ -243,6 +249,24 @@ PYBIND11_MODULE(_PyPartMC, m) {
             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("add", AeroState::add,
+            R"pbdoc(aero_state += aero_state_delta, including combining the
+            weights, so the new concentration is the weighted average of the
+            two concentrations.)pbdoc")
+        .def("add_particles", AeroState::add_particles,
+            R"pbdoc(aero_state += aero_state_delta, with the weight left unchanged
+             so the new concentration is the sum of the two concentrations.)pbdoc")
+        .def("sample", AeroState::sample,
+             R"pbdoc(Generates a random sample by removing particles from
+             aero_state_from and adding them to aero_state_to, transfering
+             weight as well. This is the equivalent of aero_state_add().)pbdoc")
+        .def("sample_particles", AeroState::sample_particles,
+             R"pbdoc(  !> Generates a random sample by removing particles from
+             aero_state_from and adding them to aero_state_to, which must be
+             already allocated (and should have its weight set).
+
+             None of the weights are altered by this sampling, making this the
+             equivalent of aero_state_add_particles().)pbdoc")
         .def("copy_weight", AeroState::copy_weight,
              "copy weighting from another AeroState")
         .def("remove_particle", AeroState::remove_particle,