diff --git a/include/enrico/openmc_heat_driver.h b/include/enrico/openmc_heat_driver.h
index 4cead3e0..00695b34 100644
--- a/include/enrico/openmc_heat_driver.h
+++ b/include/enrico/openmc_heat_driver.h
@@ -45,7 +45,7 @@ class OpenmcHeatDriver : public CoupledDriver {
   std::unordered_map<int, std::vector<int>> cell_inst_to_ring_;
 
   // Mapping of surrogate fluid elements to OpenMC cell instances and vice versa
-  std::unordered_map<int, std::vector<int>> elem_to_cell_inst_;
+  std::vector<std::vector<int>> elem_to_cell_inst_;
   std::unordered_map<int, std::vector<int>> cell_inst_to_elem_;
 
 protected:
diff --git a/src/openmc_heat_driver.cpp b/src/openmc_heat_driver.cpp
index 38291d0e..13b5f579 100644
--- a/src/openmc_heat_driver.cpp
+++ b/src/openmc_heat_driver.cpp
@@ -57,9 +57,6 @@ void OpenmcHeatDriver::init_mappings()
   // index for rings in solid
   int ring_index = 0;
 
-  // index for elements in fluid
-  int elem_index = 0;
-
   // TODO: Don't hardcode number of azimuthal segments
   int n_azimuthal = 4;
 
@@ -136,10 +133,16 @@ void OpenmcHeatDriver::init_mappings()
     n_solid_cells_ = gsl::narrow<int>(openmc_driver_->cells_.size());
   }
 
+  // index for elements in fluid
+  int elem_index = 0;
+
   // Establish mappings between fluid regions and OpenMC cells; it is assumed that there
   // are no azimuthal divisions in the fluid phase in the OpenMC model so that we
   // can take a point on a 45 degree ray from the pin center. TODO: add a check to make
   // sure that the T/H model is finer than the OpenMC model.
+  int n_elems = heat_driver_->n_pins_ * heat_driver_->n_axial_;
+  elem_to_cell_inst_.resize(n_elems);
+
   for (gsl::index i = 0; i < heat_driver_->n_pins_; ++i) {
     double x_center = heat_driver_->pin_centers_(i, 0);
     double y_center = heat_driver_->pin_centers_(i, 1);
@@ -204,21 +207,14 @@ void OpenmcHeatDriver::init_densities()
       // the density IC based on the densities used in the OpenMC input file. Set
       // the initial solid density to 0.0 because it is not used in the simulation
       // at all anyways.
-      int ring_index = 0;
-      for (gsl::index i = 0; i < heat_driver_->n_pins_; ++i) {
-        for (gsl::index j = 0; j < heat_driver_->n_axial_; ++j) {
-          for (gsl::index k = 0; k < heat_driver_->n_rings(); ++k, ++ring_index) {
-            densities_(ring_index) = 0.0;
-          }
-        }
-      }
+      densities_.fill(0.0);
 
       int elem_index = 0;
-      int solid_offset = ring_index;
+      int solid_offset = heat_driver_->n_pins_ * heat_driver_->n_axial_ * heat_driver_->n_rings();
       for (gsl::index i = 0; i < heat_driver_->n_pins_; ++i) {
         for (gsl::index j = 0; j < heat_driver_->n_axial_; ++j) {
 
-          double rho_avg = 0.0;
+          double mass = 0.0;
           double total_vol = 0.0;
 
           for (const auto& idx : elem_to_cell_inst_[elem_index++]) {
@@ -226,17 +222,15 @@ void OpenmcHeatDriver::init_densities()
             double vol = c.volume_;
 
             total_vol += vol;
-            rho_avg += c.get_density() * vol;
+            mass += c.get_density() * vol;
           }
 
-          densities_(elem_index + solid_offset) = rho_avg / total_vol;
+          densities_(elem_index + solid_offset) = mass / total_vol;
         }
       }
 
       std::copy(densities_.begin(), densities_.end(), densities_prev_.begin());
-    }
-
-    if (density_ic_ == Initial::heat) {
+    } else if (density_ic_ == Initial::heat) {
       throw std::runtime_error{
         "Density initial conditions from surrogate heat-fluids "
         "solver not supported."};
@@ -322,9 +316,7 @@ void OpenmcHeatDriver::init_temperatures()
       }
 
       std::copy(temperatures_.begin(), temperatures_.end(), temperatures_prev_.begin());
-    }
-
-    if (temperature_ic_ == Initial::heat) {
+    } else if (temperature_ic_ == Initial::heat) {
       throw std::runtime_error{
         "Temperature initial conditions from surrogate heat-fluids "
         "solver not supported."};
@@ -377,20 +369,19 @@ void OpenmcHeatDriver::update_density()
   densities_ = heat_driver_->density();
   auto solid_offset = heat_driver_->n_pins_ * heat_driver_->n_rings() * heat_driver_->n_axial_;
 
-  for (gsl::index i = 0; i < n_fluid_cells_; ++i) {
-    int index = i + n_solid_cells_;
-    const auto& c = openmc_driver_->cells_[index];
-    const auto& elems = cell_inst_to_elem_[index];
+  for (gsl::index i = n_solid_cells_; i < n_solid_cells_ + n_fluid_cells_; ++i) {
+    const auto& c = openmc_driver_->cells_[i];
+    const auto& elems = cell_inst_to_elem_[i];
 
-    double average_rho = 0.0;
+    double mass = 0.0;
     double total_vol = 0.0;
     for (auto elem_index : elems) {
       double vol = heat_driver_->z_(elem_index + 1) - heat_driver_->z_(elem_index);
-      average_rho += densities_(elem_index + solid_offset);
+      mass += densities_(elem_index + solid_offset) * vol;
       total_vol += vol;
     }
 
-    c.set_density(average_rho / total_vol);
+    c.set_density(mass / total_vol);
   }
 }