Add operator[] for TwodSubgridCalculator

celeritas-project · Nov 29, 2024 · bc67ab3 · bc67ab3
1 parent 6e61bc5
commit bc67ab3
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 6 deletions.
diff --git a/src/celeritas/grid/InverseCdfFinder.hh b/src/celeritas/grid/InverseCdfFinder.hh
@@ -52,7 +52,7 @@ CELER_FUNCTION InverseCdfFinder<G, C>::InverseCdfFinder(G&& grid, C&& calc_cdf)
     , calc_cdf_(celeritas::forward<C>(calc_cdf))
 {
     CELER_EXPECT(grid_.size() >= 2);
-    CELER_EXPECT(calc_cdf_(0) == 0);
+    CELER_EXPECT(calc_cdf_(0) == 0 && calc_cdf(grid_.size() - 1) == 1);
 }
 
 //---------------------------------------------------------------------------//

diff --git a/src/celeritas/neutron/interactor/detail/CascadeCollider.hh b/src/celeritas/neutron/interactor/detail/CascadeCollider.hh
@@ -147,13 +147,11 @@ CELER_FUNCTION auto CascadeCollider::operator()(Engine& rng) -> FinalState
     if (kin_energy_ < energy_grid.back())
     {
         // Find cos\theta from tabulated angular data for a given CDF
-        Grid costheta_grid(cdf_grid.y, shared_.reals);
         auto calc_cdf
             = TwodGridCalculator(cdf_grid, shared_.reals)(kin_energy_);
-        cos_theta = InverseCdfFinder(Grid(cdf_grid.y, shared_.reals),
-                                     [&calc_cdf, &costheta_grid](size_type i) {
-                                         return calc_cdf(costheta_grid[i]);
-                                     })(cdf);
+        cos_theta = InverseCdfFinder(
+            Grid(cdf_grid.y, shared_.reals),
+            [&calc_cdf](size_type i) { return calc_cdf[i]; })(cdf);
     }
     else
     {

diff --git a/src/corecel/grid/TwodSubgridCalculator.hh b/src/corecel/grid/TwodSubgridCalculator.hh
@@ -41,6 +41,9 @@ class TwodSubgridCalculator
     // Calculate the value at the given y coordinate
     inline CELER_FUNCTION real_type operator()(real_type y) const;
 
+    // Calculate the value at the given y grid point
+    inline CELER_FUNCTION real_type operator[](size_type i) const;
+
     //! Index of the preselected lower x value
     CELER_FORCEINLINE_FUNCTION size_type x_index() const
     {
@@ -108,6 +111,17 @@ CELER_FUNCTION real_type TwodSubgridCalculator::operator()(real_type y) const
                     + (y_loc.fraction) * at_corner(1, 1));
 }
 
+//---------------------------------------------------------------------------//
+/*!
+ * Calculate the value at the given y grid point for preselected x.
+ */
+CELER_FUNCTION real_type TwodSubgridCalculator::operator[](size_type i) const
+{
+    CELER_EXPECT(i < grids_.y.size());
+    return (1 - x_loc_.fraction) * this->at(x_loc_.index, i)
+           + x_loc_.fraction * this->at(x_loc_.index + 1, i);
+}
+
 //---------------------------------------------------------------------------//
 /*!
  * Get the value at the specified x/y coordinate.

diff --git a/test/corecel/grid/TwodGridCalculator.test.cc b/test/corecel/grid/TwodGridCalculator.test.cc
@@ -140,6 +140,10 @@ TEST_F(TwodGridCalculatorTest, subgrid)
         {
             EXPECT_SOFT_EQ(calc_expected(x, y), interpolate(y));
         }
+        for (size_type i : range(ygrid_.size()))
+        {
+            EXPECT_SOFT_EQ(calc_expected(x, ygrid_[i]), interpolate[i]);
+        }
     }
 
     unsigned int const expected_lower_idx[] = {0u, 1u, 2u};