Add support for tracing multiple energy groups in Rover

src/libs/rover/domain.cpp

@@ -86,9 +86,15 @@ 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();
+}
+
+bool
+Domain::get_field_mismatch_error()
+{
+  return m_engine->get_field_mismatch_error();
 }
 
 void

src/libs/rover/domain.hpp

@@ -39,7 +39,8 @@ 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();
+  bool get_field_mismatch_error();
 protected:
   std::shared_ptr<Engine> m_engine;
   vtkmDataSet             m_dataset;

src/libs/rover/engine.cpp

@@ -9,13 +9,15 @@
 #include <rover_exceptions.hpp>
 #include <utils/rover_logging.hpp>
 #include <vtkm/cont/DefaultTypes.h>
+#include <ascent_logging.hpp>
 
 namespace rover
 {
 
 Engine::Engine()
 {
   m_tracer = nullptr;
+  m_field_mismatch_error = false;
 }
 
 Engine::~Engine()
@@ -36,7 +38,7 @@ Engine::validate_tracer()
   // this in 1 or 2 spots and won't need a helper.
   if (!m_tracer)
   {
-    ROVER_ERROR("Error - Engine::validate_tracer: data was not set before tracing");
+    ASCENT_LOG_ERROR("Error - Engine::validate_tracer: data was not set before tracing");
   }
 }
 
@@ -57,13 +59,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 +85,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,40 +119,33 @@ 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");
-  }
+  const vtkm::cont::Field &absorption_field = m_dataset.GetField(absorption);
+  vtkm::Id num_absorption_bins = absorption_field.GetData().GetNumberOfComponentsFlat();
 
-  vtkm::Id modulo = absorption_size % num_cells;
-  if (modulo != 0)
+  // If the emission field is set, verify that it has the same number of energy groups
+  // as the absorption field
+  if (rover::settings.has_child("emission"))
   {
-    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");
+    const std::string emission = rover::settings["emission"].as_string();
+    const vtkm::cont::Field &emission_field = m_dataset.GetField(emission);
+    vtkm::Id num_emission_bins = emission_field.GetData().GetNumberOfComponentsFlat();
+
+    if (num_absorption_bins != num_emission_bins)
+    {
+      m_field_mismatch_error = true;
+    }
   }
-  vtkm::Id num_bins = absorption_size / num_cells;
-  ROVER_INFO("Engine::get_num_channels: Detected " << num_bins << " bins");
-  return static_cast<int>(num_bins);
+
+  return static_cast<int>(num_absorption_bins);
+}
+
+bool
+Engine::get_field_mismatch_error()
+{
+  return m_field_mismatch_error;
 }
 
 vtkmRange

src/libs/rover/engine.hpp

@@ -43,14 +43,16 @@ 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();
+  bool get_field_mismatch_error();
   vtkmRange get_primary_range();
   void set_primary_range(const vtkmRange &range);
   void set_composite_background(bool on);
 
 protected:
   vtkmDataSet m_dataset;
   vtkh::rendering::ConnectivityProxy *m_tracer;
+  bool m_field_mismatch_error;
 
   template<typename Precision>
   void init_emission(vtkmRayTracing::Ray<Precision> &rays,

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);

src/libs/rover/rover.cpp

@@ -4,7 +4,9 @@
 // other details. No copyright assignment is required to contribute to Ascent.
 //~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~//
 
-#include <logging/ascent_annotations.hpp>
+#include <ascent_annotations.hpp>
+#include <ascent_logging.hpp>
+
 #include <typed_scheduler.hpp>
 #include <rover.hpp>
 #include <rover_exceptions.hpp>
@@ -13,6 +15,7 @@
 #include <utils/rover_logging.hpp>
 #include <settings.hpp>
 
+
 #ifdef ROVER_PARALLEL
 #include <mpi.h>
 #endif
@@ -141,7 +144,7 @@ Rover::create_scheduler()
   // Check to see if we have been initialized
   if(-1 == m_rank)
   {
-    ROVER_ERROR("Error - Rover::create_scheduler: MPI was not initialized");
+    ASCENT_LOG_ERROR("Error - Rover::create_scheduler: MPI was not initialized");
   }
   m_scheduler->set_comm_handle(m_comm_handle);
 #endif
@@ -285,7 +288,7 @@ Rover::execute()
   // but a developer would prefer to be made aware.
   if (!m_scheduler)
   {
-    ROVER_ERROR("Error - Rover::execute: Execute called before adding a dataset");
+    ASCENT_LOG_ERROR("Error - Rover::execute: Execute called before adding a dataset");
   }
 
   // Applies the user-supplied parameters and then begins the ray trace