From 12c36273fb6b2cb5ff58d3f243c7bbc77337fb60 Mon Sep 17 00:00:00 2001
From: Alisa Galishnikova <55898700+alisagk@users.noreply.github.com>
Date: Sat, 8 Nov 2025 13:53:55 -0500
Subject: [PATCH] Changed behavior of spatial distribution in NonUniform
 injector. it takes ppc and weight distributions as parameters and fills them

---
 src/archetypes/spatial_dist.h | 20 +++++++++++++++-----
 src/kernels/injectors.hpp     | 30 ++++++++++++++++++------------
 2 files changed, 33 insertions(+), 17 deletions(-)

diff --git a/src/archetypes/spatial_dist.h b/src/archetypes/spatial_dist.h
index 55c84ddf..37494b31 100644
--- a/src/archetypes/spatial_dist.h
+++ b/src/archetypes/spatial_dist.h
@@ -40,8 +40,12 @@ namespace arch {
   struct Uniform : public SpatialDistribution<S, M> {
     Uniform(const M& metric) : SpatialDistribution<S, M> { metric } {}
 
-    Inline auto operator()(const coord_t<M::Dim>&) const -> real_t {
-      return ONE;
+    Inline auto operator()(const coord_t<M::Dim>&,
+                           real_t& nppc_distribution,
+                           real_t& weight_distribution) const {
+      nppc_distribution = ONE;
+      weight_distribution = ONE;
+      // return nppc_distribution;
     }
   };
 
@@ -65,7 +69,9 @@ namespace arch {
       , target_density { target_density }
       , target_max_density { target_max_density } {}
 
-    Inline auto operator()(const coord_t<M::Dim>& x_Ph) const -> real_t {
+    Inline auto operator()(const coord_t<M::Dim>& x_Ph,
+                           real_t& nppc_distribution,
+                           real_t& weight_distribution) const {
       coord_t<M::Dim> x_Cd { ZERO };
       metric.template convert<Crd::Ph, Crd::Cd>(x_Ph, x_Cd);
       real_t dens { ZERO };
@@ -85,9 +91,13 @@ namespace arch {
       }
       const auto target = target_density(x_Ph);
       if (0.9 * target > dens) {
-        return (target - dens) / target_max_density;
+        nppc_distribution = (target - dens) / target_max_density;
+        weight_distribution = ONE;
+        // return nppc_distribution;
       } else {
-        return ZERO;
+        nppc_distribution = ZERO;
+        weight_distribution = ONE;
+        // return nppc_distribution;
       }
     }
   };
diff --git a/src/kernels/injectors.hpp b/src/kernels/injectors.hpp
index c321d18f..a4f9f126 100644
--- a/src/kernels/injectors.hpp
+++ b/src/kernels/injectors.hpp
@@ -675,7 +675,9 @@ namespace kernel {
         coord_t<Dim::_1D> x_Cd { i1_ + HALF };
         coord_t<Dim::_1D> x_Ph { ZERO };
         metric.template convert<Crd::Cd, Crd::Ph>(x_Cd, x_Ph);
-        const auto ppc = static_cast<npart_t>(ppc0 * spatial_dist(x_Ph));
+        auto ppc_dist { ZERO }, weight_dist { ONE };
+        spatial_dist(x_Ph, ppc_dist, weight_dist);
+        const auto ppc = static_cast<npart_t>(ppc0 * ppc_dist);
         if (ppc == 0) {
           return;
         }
@@ -704,10 +706,10 @@ namespace kernel {
           tags_1(index + offset1) = ParticleTag::alive;
           tags_2(index + offset2) = ParticleTag::alive;
           if (M::CoordType == Coord::Cart) {
-            weights_1(index + offset1) = ONE;
-            weights_2(index + offset2) = ONE;
+            weights_1(index + offset1) = ONE * weight_dist;
+            weights_2(index + offset2) = ONE * weight_dist;
           } else {
-            const auto wei = metric.sqrt_det_h({ i1_ + HALF }) * inv_V0;
+            const auto wei = metric.sqrt_det_h({ i1_ + HALF }) * inv_V0 * weight_dist;
             weights_1(index + offset1) = wei;
             weights_2(index + offset2) = wei;
           }
@@ -731,7 +733,9 @@ namespace kernel {
           x_Cd_[2] = ZERO;
         }
         metric.template convert<Crd::Cd, Crd::Ph>(x_Cd, x_Ph);
-        const auto ppc = static_cast<npart_t>(ppc0 * spatial_dist(x_Ph));
+        auto ppc_dist { ZERO }, weight_dist { ONE };
+        spatial_dist(x_Ph, ppc_dist, weight_dist);
+        const auto ppc = static_cast<npart_t>(ppc0 * ppc_dist);
         if (ppc == 0) {
           return;
         }
@@ -774,10 +778,10 @@ namespace kernel {
           tags_1(index + offset1) = ParticleTag::alive;
           tags_2(index + offset2) = ParticleTag::alive;
           if (M::CoordType == Coord::Cart) {
-            weights_1(index + offset1) = ONE;
-            weights_2(index + offset2) = ONE;
+            weights_1(index + offset1) = ONE * weight_dist;
+            weights_2(index + offset2) = ONE * weight_dist;
           } else {
-            const auto wei = metric.sqrt_det_h({ i1_ + HALF, i2_ + HALF }) * inv_V0;
+            const auto wei = metric.sqrt_det_h({ i1_ + HALF, i2_ + HALF }) * inv_V0 * weight_dist;
             weights_1(index + offset1) = wei;
             weights_2(index + offset2) = wei;
           }
@@ -798,7 +802,9 @@ namespace kernel {
         coord_t<Dim::_3D> x_Cd { i1_ + HALF, i2_ + HALF, i3_ + HALF };
         coord_t<Dim::_3D> x_Ph { ZERO };
         metric.template convert<Crd::Cd, Crd::Ph>(x_Cd, x_Ph);
-        const auto ppc = static_cast<npart_t>(ppc0 * spatial_dist(x_Ph));
+        auto ppc_dist { ZERO }, weight_dist { ONE };
+        spatial_dist(x_Ph, ppc_dist, weight_dist);
+        const auto ppc = static_cast<npart_t>(ppc0 * ppc_dist);
         if (ppc == 0) {
           return;
         }
@@ -847,12 +853,12 @@ namespace kernel {
           tags_1(index + offset1) = ParticleTag::alive;
           tags_2(index + offset2) = ParticleTag::alive;
           if (M::CoordType == Coord::Cart) {
-            weights_1(index + offset1) = ONE;
-            weights_2(index + offset2) = ONE;
+            weights_1(index + offset1) = ONE * weight_dist;
+            weights_2(index + offset2) = ONE * weight_dist;
           } else {
             const auto wei = metric.sqrt_det_h(
                                { i1_ + HALF, i2_ + HALF, i3_ + HALF }) *
-                             inv_V0;
+                             inv_V0 * weight_dist;
             weights_1(index + offset1) = wei;
             weights_2(index + offset2) = wei;
           }