Add support for tracing multiple energy groups in Rover

src/libs/rover/domain.cpp

@@ -86,9 +86,9 @@ Domain::init()
 }
 
 const int
-Domain::get_num_channels()
+Domain::get_num_energy_groups()
 {
-  return m_engine->get_num_channels();
+  return m_engine->get_num_energy_groups();
 }
 
 void

src/libs/rover/domain.hpp

@@ -39,7 +39,7 @@ class Domain
   vtkm::Bounds& get_domain_bounds();
   vtkmRange get_primary_range();
   void set_global_bounds(vtkm::Bounds bounds);
-  const int get_num_channels();
+  const int get_num_energy_groups();
 protected:
   std::shared_ptr<Engine> m_engine;
   vtkmDataSet             m_dataset;

src/libs/rover/engine.cpp

@@ -57,13 +57,13 @@ Engine::set_dataset(vtkm::cont::DataSet &dataset)
 template<typename Precision>
 void
 Engine::init_emission(vtkmRayTracing::Ray<Precision> &rays,
-                      const int num_bins)
+                      const int num_energy_groups)
 {
   if (rover::settings.has_child("emission"))
   {
     const std::string emission = rover::settings["emission"].as_string();
     m_tracer->SetEmissionField(emission);
-    rays.AddBuffer(num_bins, "emission");
+    rays.AddBuffer(num_energy_groups, "emission");
     rays.GetBuffer("emission").InitConst(0.0f);
   }
 }
@@ -83,25 +83,25 @@ void
 Engine::init_rays(Ray32 &rays)
 {
   validate_tracer();
-  const int num_bins = get_num_channels();
-  rays.Buffers.at(0).SetNumChannels(num_bins);
+  const int num_energy_groups = get_num_energy_groups();
+  rays.Buffers.at(0).SetNumChannels(num_energy_groups);
   // TODO: I think this should be init with background intensities
   rays.Buffers.at(0).InitConst(1.0f);
-  init_emission(rays, num_bins);
-  rays.AddBuffer(num_bins, "optical_depths");
+  init_emission(rays, num_energy_groups);
+  rays.AddBuffer(num_energy_groups, "optical_depths");
   rays.GetBuffer("optical_depths").InitConst(0.0f);
 }
 
 void
 Engine::init_rays(Ray64 &rays)
 {
   validate_tracer();
-  const int num_bins = get_num_channels();
-  rays.Buffers.at(0).SetNumChannels(num_bins);
+  const int num_energy_groups = get_num_energy_groups();
+  rays.Buffers.at(0).SetNumChannels(num_energy_groups);
   // TODO: I think this should be init with background intensities
   rays.Buffers.at(0).InitConst(1.0f);
-  init_emission(rays, num_bins);
-  rays.AddBuffer(num_bins, "optical_depths");
+  init_emission(rays, num_energy_groups);
+  rays.AddBuffer(num_energy_groups, "optical_depths");
   rays.GetBuffer("optical_depths").InitConst(0.0f);
 }
 
@@ -117,39 +117,11 @@ Engine::partial_trace(Ray64 &rays, PartialVector64 &partials)
 }
 
 int
-Engine::get_num_channels()
+Engine::get_num_energy_groups()
 {
-  vtkm::Id absorption_size = 0;
-  ArraySizeFunctor functor(&absorption_size);
   const std::string absorption = rover::settings["absorption"].as_string();
-  m_dataset.GetField(absorption).
-                     GetData().
-                     CastAndCallForTypes<vtkm::TypeListAll, VTKM_DEFAULT_STORAGE_LIST>(functor);
-  vtkm::Id num_cells = m_dataset.GetCellSet().GetNumberOfCells();
-
-  // TODO: Seemingly redundant assert followed by a check that num_cells == 0
-  assert(num_cells > 0);
-  assert(absorption_size > 0);
-  if (num_cells == 0)
-  {
-    ROVER_ERROR("Error - Engine::get_num_channels: num cells is 0"
-                << "\n        num cells " << num_cells
-                << "\n        field size " <<a bsorption_size);
-    m_dataset.PrintSummary(std::cerr);
-    throw RoverException("Failed to detect bins. Num cells cannot be 0\n");
-  }
-
-  vtkm::Id modulo = absorption_size % num_cells;
-  if (modulo != 0)
-  {
-    ROVER_ERROR("Error - Engine::get_num_channels: absorption field size is not evenly divided by num_cells"
-                << "\n       modulo " << modulo
-                << "\n       num cells " << num_cells
-                << "\n       field size " << absorption_size);
-    throw RoverException("absorption field size is not evenly divided by num_cells\n");
-  }
-  vtkm::Id num_bins = absorption_size / num_cells;
-  ROVER_INFO("Engine::get_num_channels: Detected " << num_bins << " bins");
+  const auto &field = m_dataset.GetField(absorption);
+  vtkm::Id num_bins = field.GetData().GetNumberOfComponentsFlat();
   return static_cast<int>(num_bins);
 }
 

src/libs/rover/engine.hpp

@@ -43,7 +43,7 @@ class Engine
   void init_rays(Ray64 &rays);
   void partial_trace(Ray32 &rays, PartialVector32 &partials);
   void partial_trace(Ray64 &rays, PartialVector64 &partials);
-  int  get_num_channels();
+  int  get_num_energy_groups();
   vtkmRange get_primary_range();
   void set_primary_range(const vtkmRange &range);
   void set_composite_background(bool on);

src/libs/rover/image.cpp

@@ -139,8 +139,8 @@ void cast_array_handle(vtkm::cont::ArrayHandle<T> &cast_to,
 //
 template<typename T, typename O> void init_from_image(Image<T> &left, Image<O> &right)
 {
-  const size_t channels = right.m_intensity_values.size();
-  for(size_t i = 0; i < channels; ++i)
+  const size_t num_energy_groups = right.m_intensity_values.size();
+  for(size_t i = 0; i < num_energy_groups; ++i)
   {
     cast_array_handle(left.m_intensity_values[i], right.m_intensity_values[i]);
     cast_array_handle(left.m_optical_depth_values[i], right.m_optical_depth_values[i]);
@@ -171,7 +171,7 @@ Image<FloatType>::operator=(Image<O> &other)
 
 template<typename FloatType>
 int
-Image<FloatType>::get_num_channels() const
+Image<FloatType>::get_num_energy_groups() const
 {
   return static_cast<int>(m_intensity_values.size());
 }
@@ -200,7 +200,7 @@ Image<FloatType>::init_from_partial(PartialImage<FloatType> &partial)
   const int64 width = rover::settings["width"].to_int64();
   const int64 height = rover::settings["height"].to_int64();
   const int64 channel_size = width * height;
-  const int num_channels = partial.m_transmission.GetNumChannels();
+  const int num_energy_groups = partial.m_transmission.GetNumChannels();
 
   // Helper lambda to expand a channel and push its buffer to the output vector
   auto expand_and_push = [&](int channel_index,
@@ -213,7 +213,7 @@ Image<FloatType>::init_from_partial(PartialImage<FloatType> &partial)
     output_vector.push_back(expanded.Buffer);
   };
 
-  for (int i = 0; i < num_channels; i++)
+  for (int i = 0; i < num_energy_groups; i++)
   {
     // Intensities
     expand_and_push(i,
@@ -255,15 +255,15 @@ template<typename FloatType>
 vtkm::cont::ArrayHandle<FloatType>
 Image<FloatType>::flatten_intensity_values()
 {
-  const int num_channels = this->get_num_channels();
+  const int num_energy_groups = get_num_energy_groups();
 
   HandleType res;
   const int64 width = rover::settings["width"].to_int64();
   const int64 height = rover::settings["height"].to_int64();
   const int64 size = width * height;
-  res.Allocate(num_channels * size);
+  res.Allocate(num_energy_groups * size);
   auto output = res.WritePortal();
-  for(int c = 0; c < num_channels; ++c)
+  for(int c = 0; c < num_energy_groups; ++c)
   {
     auto channel = m_intensity_values[c].ReadPortal();
 
@@ -272,7 +272,8 @@ Image<FloatType>::flatten_intensity_values()
 #endif
     for(int i = 0; i < size; ++i)
     {
-      output.Set( i * num_channels + c, channel.Get(i));
+      // Deinterleave the output: all pixels for group 0, then all pixels for group 1, etc.
+      output.Set(c * size + i, channel.Get(i));
     }
   }
   return res;
@@ -282,23 +283,24 @@ template<typename FloatType>
 vtkm::cont::ArrayHandle<FloatType>
 Image<FloatType>::flatten_optical_depth_values()
 {
-  const int num_channels = this->get_num_channels();
+  const int num_energy_groups = get_num_energy_groups();
 
   HandleType res;
   const int64 width = rover::settings["width"].to_int64();
   const int64 height = rover::settings["height"].to_int64();
   const int64 size = width * height;
-  res.Allocate(num_channels * size);
+  res.Allocate(num_energy_groups * size);
   auto output = res.WritePortal();
-  for(int c = 0; c < num_channels; ++c)
+  for(int c = 0; c < num_energy_groups; ++c)
   {
     auto channel = m_optical_depth_values[c].ReadPortal();
 #ifdef ROVER_OPENMP_ENABLED
     #pragma omp parallel for
 #endif
     for(int i = 0; i < size; ++i)
     {
-      output.Set( i * num_channels + c, channel.Get(i));
+      // Deinterleave the output: all pixels for group 0, then all pixels for group 1, etc.
+      output.Set( c * size + i, channel.Get(i));
     }
   }
   return res;

src/libs/rover/image.hpp

@@ -27,7 +27,7 @@ class Image
 
   HandleType get_intensity(const int &channel_num);
   HandleType get_optical_depth(const int &channel_num);
-  int get_num_channels() const;
+  int get_num_energy_groups() const;
   bool has_intensity(const int &channel_num) const;
   bool has_optical_depth(const int &channel_num) const;
   void normalize_intensity(const int &channel_num);