diff --git a/src/Init/ChemistryInit.cpp b/src/Init/ChemistryInit.cpp
index 962fc3b73..2975eec2c 100644
--- a/src/Init/ChemistryInit.cpp
+++ b/src/Init/ChemistryInit.cpp
@@ -6,6 +6,8 @@ InitialList::ChemistryInit InitialList::getChemistryInit() const {
 
   chem_init.initial_grid = Field(initial_grid);
 
+  chem_init.total_grid_cells = this->n_cols * this->n_rows;
+
   chem_init.database = database;
   chem_init.pqc_scripts = pqc_scripts;
   chem_init.pqc_ids = pqc_ids;
diff --git a/src/Init/DiffusionInit.cpp b/src/Init/DiffusionInit.cpp
index 93046575d..2d83d1091 100644
--- a/src/Init/DiffusionInit.cpp
+++ b/src/Init/DiffusionInit.cpp
@@ -174,4 +174,27 @@ void InitialList::initDiffusion(const Rcpp::List &diffusion_input) {
   // R.parseEval("print(alpha_x)");
   // R.parseEval("print(alpha_y)");
 }
+
+InitialList::DiffusionInit InitialList::getDiffusionInit() const {
+  DiffusionInit diff_init;
+
+  diff_init.dim = this->dim;
+
+  diff_init.n_cols = this->n_cols;
+  diff_init.n_rows = this->n_rows;
+
+  diff_init.s_cols = this->s_cols;
+  diff_init.s_rows = this->s_rows;
+
+  diff_init.constant_cells = this->constant_cells;
+  diff_init.transport_names = this->transport_names;
+
+  diff_init.initial_grid = Field(this->initial_grid);
+  diff_init.boundaries = Field(this->boundaries);
+  diff_init.alpha_x = Field(this->alpha_x);
+  diff_init.alpha_y = Field(this->alpha_y);
+
+  return diff_init;
+}
+
 } // namespace poet
\ No newline at end of file
diff --git a/src/Init/GridInit.cpp b/src/Init/GridInit.cpp
index ba7e1716a..2b9d030ab 100644
--- a/src/Init/GridInit.cpp
+++ b/src/Init/GridInit.cpp
@@ -179,24 +179,4 @@ void InitialList::initGrid(const Rcpp::List &grid_input) {
   }
 }
 
-InitialList::DiffusionInit InitialList::getDiffusionInit() const {
-  DiffusionInit diff_init;
-
-  diff_init.n_cols = this->n_cols;
-  diff_init.n_rows = this->n_rows;
-
-  diff_init.s_cols = this->s_cols;
-  diff_init.s_rows = this->s_rows;
-
-  diff_init.constant_cells = this->constant_cells;
-  diff_init.transport_names = this->transport_names;
-
-  diff_init.initial_grid = Field(this->initial_grid);
-  diff_init.boundaries = Field(this->boundaries);
-  diff_init.alpha_x = Field(this->alpha_x);
-  diff_init.alpha_y = Field(this->alpha_y);
-
-  return diff_init;
-}
-
 } // namespace poet
\ No newline at end of file
diff --git a/src/Init/InitialList.hpp b/src/Init/InitialList.hpp
index 7a97bbead..0f3b631b1 100644
--- a/src/Init/InitialList.hpp
+++ b/src/Init/InitialList.hpp
@@ -93,6 +93,7 @@ private:
 
 public:
   struct DiffusionInit {
+    uint8_t dim;
     std::uint32_t n_cols;
     std::uint32_t n_rows;
 
@@ -152,6 +153,8 @@ public:
   struct ChemistryInit {
     Field initial_grid;
 
+    uint32_t total_grid_cells;
+
     std::string database;
     std::vector<std::string> pqc_scripts;
     std::vector<int> pqc_ids;