Merge pull request #1303 from drewejohnson/capi-energy-function-filter

Expose/improve EnergyFunctionFilter through C-API

include/openmc/capi.h

@@ -21,6 +21,10 @@ extern "C" {
   int openmc_cell_set_temperature(int32_t index, double T, const int32_t* instance);
   int openmc_energy_filter_get_bins(int32_t index, const double** energies, size_t* n);
   int openmc_energy_filter_set_bins(int32_t index, size_t n, const double* energies);
+  int openmc_energyfunc_filter_get_energy(int32_t index, size_t* n, const double** energy);
+  int openmc_energyfunc_filter_get_y(int32_t index, size_t* n, const double** y);
+  int openmc_energyfunc_filter_set_data(int32_t index, size_t n,
+                                        const double* energies, const double* y);
   int openmc_extend_cells(int32_t n, int32_t* index_start, int32_t* index_end);
   int openmc_extend_filters(int32_t n, int32_t* index_start, int32_t* index_end);
   int openmc_extend_materials(int32_t n, int32_t* index_start, int32_t* index_end);

include/openmc/tallies/filter_energyfunc.h

@@ -40,6 +40,13 @@ public:
 
   std::string text_label(int bin) const override;
 
+  //----------------------------------------------------------------------------
+  // Accessors
+
+  const std::vector<double>& energy() const { return energy_; }
+  const std::vector<double>& y() const { return y_; }
+  void set_data(gsl::span<const double> energy, gsl::span<const double> y);
+
 private:
   //----------------------------------------------------------------------------
   // Data members

openmc/capi/filter.py

@@ -34,6 +34,18 @@ _dll.openmc_energy_filter_get_bins.errcheck = _error_handler
 _dll.openmc_energy_filter_set_bins.argtypes = [c_int32, c_size_t, POINTER(c_double)]
 _dll.openmc_energy_filter_set_bins.restype = c_int
 _dll.openmc_energy_filter_set_bins.errcheck = _error_handler
+_dll.openmc_energyfunc_filter_set_data.restype = c_int
+_dll.openmc_energyfunc_filter_set_data.errcheck = _error_handler
+_dll.openmc_energyfunc_filter_set_data.argtypes = [
+    c_int32, c_size_t, POINTER(c_double), POINTER(c_double)]
+_dll.openmc_energyfunc_filter_get_energy.resttpe = c_int
+_dll.openmc_energyfunc_filter_get_energy.errcheck = _error_handler
+_dll.openmc_energyfunc_filter_get_energy.argtypes = [
+    c_int32, POINTER(c_size_t), POINTER(POINTER(c_double))]
+_dll.openmc_energyfunc_filter_get_y.resttpe = c_int
+_dll.openmc_energyfunc_filter_get_y.errcheck = _error_handler
+_dll.openmc_energyfunc_filter_get_y.argtypes = [
+    c_int32, POINTER(c_size_t), POINTER(POINTER(c_double))]
 _dll.openmc_filter_get_id.argtypes = [c_int32, POINTER(c_int32)]
 _dll.openmc_filter_get_id.restype = c_int
 _dll.openmc_filter_get_id.errcheck = _error_handler
@@ -201,6 +213,48 @@ class DistribcellFilter(Filter):
 class EnergyFunctionFilter(Filter):
     filter_type = 'energyfunction'
 
+    def __new__(cls, energy=None, y=None, uid=None, new=True, index=None):
+        return super().__new__(cls, uid=uid, new=new, index=index)
+
+    def __init__(self, energy=None, y=None, uid=None, new=True, index=None):
+        if (energy is None) != (y is None):
+            raise AttributeError("Need both energy and y or neither")
+        super().__init__(uid, new, index)
+        if energy is not None:
+            self.set_data(energy, y)
+
+    def set_data(self, energy, y):
+        """Set the interpolation information for the filter
+
+        Parameters
+        ----------
+        energy : numpy.ndarray
+            Independent variable for the interpolation
+        y : numpy.ndarray
+            Dependent variable for the interpolation
+        """
+        energy_array = np.asarray(energy)
+        y_array = np.asarray(y)
+        energy_p = energy_array.ctypes.data_as(POINTER(c_double))
+        y_p = y_array.ctypes.data_as(POINTER(c_double))
+
+        _dll.openmc_energyfunc_filter_set_data(
+            self._index, len(energy_array), energy_p, y_p)
+
+    @property
+    def energy(self):
+        return self._get_attr(_dll.openmc_energyfunc_filter_get_energy)
+
+    @property
+    def y(self):
+        return self._get_attr(_dll.openmc_energyfunc_filter_get_y)
+
+    def _get_attr(self, cfunc):
+        array_p = POINTER(c_double)()
+        n = c_size_t()
+        cfunc(self._index, n, array_p)
+        return as_array(array_p, (n.value, ))
+
 
 class LegendreFilter(Filter):
     filter_type = 'legendre'

src/tallies/filter_energyfunc.cpp

@@ -21,12 +21,37 @@ EnergyFunctionFilter::from_xml(pugi::xml_node node)
   if (!check_for_node(node, "energy"))
     fatal_error("Energy grid not specified for EnergyFunction filter.");
 
-  energy_ = get_node_array<double>(node, "energy");
+  auto energy = get_node_array<double>(node, "energy");
 
   if (!check_for_node(node, "y"))
     fatal_error("y values not specified for EnergyFunction filter.");
 
-  y_ = get_node_array<double>(node, "y");
+  auto y = get_node_array<double>(node, "y");
+
+  this->set_data(energy, y);
+}
+
+void
+EnergyFunctionFilter::set_data(gsl::span<const double> energy,
+                               gsl::span<const double> y)
+{
+  // Check for consistent sizes with new data
+  if (energy.size() != y.size()) {
+    fatal_error("Energy grid and y values are not consistent");
+  }
+  energy_.clear();
+  energy_.reserve(energy.size());
+  y_.clear();
+  y_.reserve(y.size());
+
+  // Copy over energy values, ensuring they are valid
+  for (gsl::index i = 0; i < energy.size(); ++i) {
+    if (i > 0 && energy[i] <= energy[i - 1]) {
+      throw std::runtime_error{"Energy bins must be monotonically increasing."};
+    }
+    energy_.push_back(energy[i]);
+    y_.push_back(y[i]);
+  }
 }
 
 void
@@ -65,4 +90,72 @@ EnergyFunctionFilter::text_label(int bin) const
   return out.str();
 }
 
+//==============================================================================
+// C-API functions
+//==============================================================================
+
+extern "C" int
+openmc_energyfunc_filter_set_data(int32_t index, size_t n, const double* energy,
+                                  const double* y)
+{
+  // Ensure this is a valid index to allocated filter
+  if (int err = verify_filter(index)) return err;
+
+  // Get a pointer to the filter
+  const auto& filt_base = model::tally_filters[index].get();
+  // Downcast to EnergyFunctionFilter
+  auto* filt = dynamic_cast<EnergyFunctionFilter*>(filt_base);
+
+  // Check if a valid filter was produced
+  if (!filt) {
+    set_errmsg("Tried to set interpolation data for non-energy function filter.");
+    return OPENMC_E_INVALID_TYPE;
+  }
+
+  filt->set_data({energy, n}, {y, n});
+  return 0;
+}
+
+extern "C" int
+openmc_energyfunc_filter_get_energy(int32_t index, size_t *n, const double** energy)
+{
+  // ensure this is a valid index to allocated filter
+  if (int err = verify_filter(index)) return err;
+
+  // get a pointer to the filter
+  const auto& filt_base = model::tally_filters[index].get();
+  // downcast to EnergyFunctionFilter
+  auto* filt = dynamic_cast<EnergyFunctionFilter*>(filt_base);
+
+  // check if a valid filter was produced
+  if (!filt) {
+    set_errmsg("Tried to set interpolation data for non-energy function filter.");
+    return OPENMC_E_INVALID_TYPE;
+  }
+  *energy = filt->energy().data();
+  *n = filt->energy().size();
+  return 0;
+}
+
+extern "C" int
+openmc_energyfunc_filter_get_y(int32_t index, size_t *n, const double** y)
+{
+  // ensure this is a valid index to allocated filter
+  if (int err = verify_filter(index)) return err;
+
+  // get a pointer to the filter
+  const auto& filt_base = model::tally_filters[index].get();
+  // downcast to EnergyFunctionFilter
+  auto* filt = dynamic_cast<EnergyFunctionFilter*>(filt_base);
+
+  // check if a valid filter was produced
+  if (!filt) {
+    set_errmsg("Tried to set interpolation data for non-energy function filter.");
+    return OPENMC_E_INVALID_TYPE;
+  }
+  *y = filt->y().data();
+  *n = filt->y().size();
+  return 0;
+}
+
 } // namespace openmc

tests/unit_tests/test_capi.py

@@ -35,6 +35,14 @@ def pincell_model():
     zernike_tally.scores = ['fission']
     pincell.tallies.append(zernike_tally)
 
+    # Add an energy function tally
+    energyfunc_tally = openmc.Tally()
+    energyfunc_filter = openmc.EnergyFunctionFilter(
+        [0.0, 20e6], [0.0, 20e6])
+    energyfunc_tally.scores = ['fission']
+    energyfunc_tally.filters = [energyfunc_filter]
+    pincell.tallies.append(energyfunc_tally)
+
     # Write XML files in tmpdir
     with cdtemp():
         pincell.export_to_xml()
@@ -170,12 +178,21 @@ def test_settings(capi_init):
 def test_tally_mapping(capi_init):
     tallies = openmc.capi.tallies
     assert isinstance(tallies, Mapping)
-    assert len(tallies) == 2
+    assert len(tallies) == 3
     for tally_id, tally in tallies.items():
         assert isinstance(tally, openmc.capi.Tally)
         assert tally_id == tally.id
 
 
+def test_energy_function_filter(capi_init):
+    """Test special __new__ and __init__ for EnergyFunctionFilter"""
+    efunc = openmc.capi.EnergyFunctionFilter([0.0, 1.0], [0.0, 2.0])
+    assert len(efunc.energy) == 2
+    assert (efunc.energy == [0.0, 1.0]).all()
+    assert len(efunc.y) == 2
+    assert (efunc.y == [0.0, 2.0]).all()
+
+
 def test_tally(capi_init):
     t = openmc.capi.tallies[1]
     assert t.type == 'volume'
@@ -211,6 +228,16 @@ def test_tally(capi_init):
     assert len(t2.filters[1].bins) == 3
     assert t2.filters[0].order == 5
 
+    t3 = openmc.capi.tallies[3]
+    assert len(t3.filters) == 1
+    t3_f = t3.filters[0]
+    assert isinstance(t3_f, openmc.capi.EnergyFunctionFilter)
+    assert len(t3_f.energy) == 2
+    assert len(t3_f.y) == 2
+    t3_f.set_data([0.0, 1.0, 2.0], [0.0, 1.0, 4.0])
+    assert len(t3_f.energy) == 3
+    assert len(t3_f.y) == 3
+
 
 def test_new_tally(capi_init):
     with pytest.raises(exc.AllocationError):
@@ -219,7 +246,7 @@ def test_new_tally(capi_init):
     new_tally.scores = ['flux']
     new_tally_with_id = openmc.capi.Tally(10)
     new_tally_with_id.scores = ['flux']
-    assert len(openmc.capi.tallies) == 4
+    assert len(openmc.capi.tallies) == 5
 
 
 def test_tally_activate(capi_simulation_init):