diff --git a/app/poet.cpp b/app/poet.cpp
index 1af116370..086948bc7 100644
--- a/app/poet.cpp
+++ b/app/poet.cpp
@@ -51,6 +51,22 @@ poet::ChemistryModule::SingleCMap DFToHashMap(const Rcpp::DataFrame &df) {
   return out_map;
 }
 
+// HACK: this is a step back as the order and also the count of fields is
+// predefined, but it will change in the future
+void writeFieldsToR(RInside &R, const Field &trans, const Field &chem) {
+  R["TMP"] = Rcpp::wrap(trans.AsVector());
+  R["TMP_PROPS"] = Rcpp::wrap(trans.GetProps());
+  R.parseEval(std::string(
+      "mysetup$state_T <- setNames(data.frame(matrix(TMP, nrow=" +
+      std::to_string(trans.GetRequestedVecSize()) + ")), TMP_PROPS)"));
+
+  R["TMP"] = Rcpp::wrap(chem.AsVector());
+  R["TMP_PROPS"] = Rcpp::wrap(chem.GetProps());
+  R.parseEval(std::string(
+      "mysetup$state_C <- setNames(data.frame(matrix(TMP, nrow=" +
+      std::to_string(chem.GetRequestedVecSize()) + ")), TMP_PROPS)"));
+}
+
 void set_chem_parameters(poet::ChemistryModule &chem, uint32_t wp_size,
                          const std::string &database_path) {
   chem.SetErrorHandlerMode(1);
@@ -95,12 +111,12 @@ void set_chem_parameters(poet::ChemistryModule &chem, uint32_t wp_size,
   chem.LoadDatabase(database_path);
 }
 
-inline double RunMasterLoop(SimParams &params, RInside &R, Grid &grid,
+inline double RunMasterLoop(SimParams &params, RInside &R,
                             ChemistryParams &chem_params,
                             const GridParams &g_params, uint32_t nxyz_master) {
 
   DiffusionParams d_params{R};
-  DiffusionModule diffusion(d_params, grid);
+  DiffusionModule diffusion(d_params, g_params);
   /* Iteration Count is dynamic, retrieving value from R (is only needed by
    * master for the following loop) */
   uint32_t maxiter = R.parseEval("mysetup$iterations");
@@ -113,7 +129,7 @@ inline double RunMasterLoop(SimParams &params, RInside &R, Grid &grid,
   chem.RunInitFile(chem_params.input_script);
 
   poet::ChemistryModule::SingleCMap init_df = DFToHashMap(d_params.initial_t);
-  chem.mergeFieldWithModule(init_df, grid.GetTotalCellCount());
+  chem.initializeField(diffusion.getField());
 
   if (params.getNumParams().dht_enabled) {
     chem.SetDHTEnabled(true, params.getNumParams().dht_size_per_process);
@@ -131,8 +147,6 @@ inline double RunMasterLoop(SimParams &params, RInside &R, Grid &grid,
     }
   }
 
-  StateMemory *chem_state = grid.RegisterState("state_C", chem.GetPropNames());
-  chem_state->mem = chem.GetField();
   /* SIMULATION LOOP */
 
   double dStartTime = MPI_Wtime();
@@ -154,19 +168,17 @@ inline double RunMasterLoop(SimParams &params, RInside &R, Grid &grid,
     // TODO: transport to diffusion
     diffusion.simulate(dt);
 
-    grid.PreModuleFieldCopy(tick++);
+    chem.getField().UpdateFromField(diffusion.getField());
 
     cout << "CPP: Chemistry" << endl;
 
     chem.SetTimeStep(dt);
 
-    chem.SetField(chem_state->mem);
     chem.SetTimeStep(dt);
     chem.RunCells();
-    chem_state->mem = chem.GetField();
 
-    grid.WriteFieldsToR(R);
-    grid.PreModuleFieldCopy(tick++);
+    writeFieldsToR(R, diffusion.getField(), chem.GetField());
+    diffusion.getField().UpdateFromField(chem.GetField());
 
     R["req_dt"] = dt;
     R["simtime"] = (sim_time += dt);
@@ -302,14 +314,9 @@ int main(int argc, char *argv[]) {
   std::string master_init_code = "mysetup <- master_init(setup=setup)";
   R.parseEval(master_init_code);
 
-  Grid grid;
   GridParams g_params(R);
 
-  grid.InitModuleFromParams(g_params);
-  grid.PushbackModuleFlow(poet::DIFFUSION_MODULE_NAME, CHEMISTRY_MODULE_NAME);
-  grid.PushbackModuleFlow(CHEMISTRY_MODULE_NAME, poet::DIFFUSION_MODULE_NAME);
-
-  params.initVectorParams(R, grid.GetSpeciesCount());
+  params.initVectorParams(R);
 
   // MDL: store all parameters
   if (world_rank == 0) {
@@ -331,9 +338,9 @@ int main(int argc, char *argv[]) {
   MPI_Bcast(&c_size, 1, MPI_UINT32_T, 0, MPI_COMM_WORLD);
 
   MPI_Bcast(db_path.data(), c_size, MPI_CHAR, 0, MPI_COMM_WORLD);
-  uint32_t nxyz_master = (world_size == 1 ? grid.GetTotalCellCount() : 1);
+  uint32_t nxyz_master = (world_size == 1 ? g_params.total_n : 1);
 
-  dSimTime = RunMasterLoop(params, R, grid, chem_params, g_params, nxyz_master);
+  dSimTime = RunMasterLoop(params, R, chem_params, g_params, nxyz_master);
 
   cout << "CPP: finished simulation loop" << endl;
 
diff --git a/bench/dolo_diffu_inner/dolo_diffu_inner.R b/bench/dolo_diffu_inner/dolo_diffu_inner.R
index 33c5114a2..4821160e8 100644
--- a/bench/dolo_diffu_inner/dolo_diffu_inner.R
+++ b/bench/dolo_diffu_inner/dolo_diffu_inner.R
@@ -1,7 +1,7 @@
-## Time-stamp: "Last modified 2023-04-17 12:25:25 mluebke"
+## Time-stamp: "Last modified 2023-04-24 14:11:01 mluebke"
 
 database <- normalizePath("../share/poet/examples/phreeqc_kin.dat")
-input_script <- normalizePath("../share/poet/bench/dolo_inner.pqi")
+input_script <- normalizePath("../share/poet/bench/dolo/dolo_inner.pqi")
 
 #################################################################
 ##                          Section 1                          ##
diff --git a/include/poet/ChemistryModule.hpp b/include/poet/ChemistryModule.hpp
index fc096c862..07b854b1b 100644
--- a/include/poet/ChemistryModule.hpp
+++ b/include/poet/ChemistryModule.hpp
@@ -1,6 +1,7 @@
 #ifndef CHEMISTRYMODULE_H_
 #define CHEMISTRYMODULE_H_
 
+#include "Field.hpp"
 #include "IrmResult.h"
 #include "PhreeqcRM.h"
 #include "poet/DHT_Wrapper.hpp"
@@ -77,7 +78,7 @@ public:
   /**
    * Returns the chemical field.
    */
-  auto GetField() const { return this->field; }
+  auto GetField() const { return this->chem_field; }
   /**
    * Returns all known species names, including not only aqueous species, but
    * also equilibrium, exchange, surface and kinetic reactants.
@@ -129,14 +130,15 @@ public:
   };
 
   /**
+   * **This function has to be run!**
+   *
    * Merge initial values from existing module with the chemistry module and set
    * according internal variables.
    *
-   * \param input_map Map with name and initial values of another module like
-   * Diffusion.
-   * \param n_cells number of cells used to initialize the field with.
+   * \param other Field to merge chemistry with. Most likely it is something
+   * like the diffusion field.
    */
-  void mergeFieldWithModule(const SingleCMap &input_map, std::uint32_t n_cells);
+  void initializeField(const Field &other);
 
   /**
    * **Only called by workers!** Start the worker listening loop.
@@ -187,16 +189,6 @@ public:
    */
   void ReadDHTFile(const std::string &input_file);
 
-  /**
-   * Overwrite the current field state by another field.
-   *
-   * There are no checks if new vector dimensions matches expected sizes etc.
-   *
-   * \param field New input field. Current field state of the instance will be
-   * overwritten.
-   */
-  void SetField(const std::vector<double> &field) { this->field = field; }
-
   /**
    * **Master only** Return count of grid cells.
    */
@@ -271,6 +263,13 @@ public:
    * \return Vector of all count of DHT evictions.
    */
   std::vector<uint32_t> GetWorkerDHTEvictions() const;
+
+  /**
+   * **Master only** Returns the current state of the chemical field.
+   *
+   * \return Reference to the chemical field.
+   */
+  Field &getField() { return this->chem_field; }
 #endif
 protected:
 #ifdef POET_USE_PRM
@@ -350,6 +349,13 @@ protected:
   std::vector<uint32_t> CalculateWPSizesVector(uint32_t n_cells,
                                                uint32_t wp_size) const;
 
+  std::vector<double> shuffleField(const std::vector<double> &in_field,
+                                   uint32_t size_per_prop, uint32_t prop_count,
+                                   uint32_t wp_count);
+  void unshuffleField(const std::vector<double> &in_buffer,
+                      uint32_t size_per_prop, uint32_t prop_count,
+                      uint32_t wp_count, std::vector<double> &out_field);
+
   int comm_size, comm_rank;
   MPI_Comm group_comm;
 
@@ -384,7 +390,9 @@ protected:
   uint32_t n_cells = 0;
   uint32_t prop_count = 0;
   std::vector<std::string> prop_names;
-  std::vector<double> field;
+
+  Field chem_field{0};
+
   static constexpr int MODULE_COUNT = 5;
 
   std::array<std::uint32_t, MODULE_COUNT> speciesPerModule{};
diff --git a/include/poet/DiffusionModule.hpp b/include/poet/DiffusionModule.hpp
index e2ac32466..e9b9c61f5 100644
--- a/include/poet/DiffusionModule.hpp
+++ b/include/poet/DiffusionModule.hpp
@@ -21,9 +21,12 @@
 #ifndef DIFFUSION_MODULE_H
 #define DIFFUSION_MODULE_H
 
+#include "Field.hpp"
+#include "SimParams.hpp"
 #include "poet/SimParams.hpp"
 #include <array>
 #include <cmath>
+#include <cstdint>
 #include <poet/Grid.hpp>
 #include <string>
 #include <tug/BoundaryCondition.hpp>
@@ -49,7 +52,8 @@ public:
    *
    * @param R RRuntime object
    */
-  DiffusionModule(poet::DiffusionParams diffu_args, Grid &grid_);
+  DiffusionModule(const poet::DiffusionParams &diffu_args,
+                  const poet::GridParams &grid_params);
 
   /**
    * @brief Run simulation for one iteration
@@ -74,6 +78,13 @@ public:
    */
   double getTransportTime();
 
+  /**
+   * \brief Returns the current diffusion field.
+   *
+   * \return Reference to the diffusion field.
+   */
+  Field &getField() { return this->t_field; }
+
 private:
   /**
    * @brief Instance of RRuntime
@@ -83,9 +94,8 @@ private:
 
   enum { DIM_1D = 1, DIM_2D };
 
-  void initialize(poet::DiffusionParams args);
+  void initialize(const poet::DiffusionParams &args);
 
-  Grid &grid;
   uint8_t dim;
 
   uint32_t prop_count;
@@ -96,7 +106,7 @@ private:
   std::vector<std::string> prop_names;
 
   std::vector<tug::bc::BoundaryCondition> bc_vec;
-  poet::StateMemory *state;
+  Field t_field;
 
   uint32_t n_cells_per_prop;
 
diff --git a/include/poet/Field.hpp b/include/poet/Field.hpp
index 4a35c99e2..cb4e46996 100644
--- a/include/poet/Field.hpp
+++ b/include/poet/Field.hpp
@@ -63,6 +63,12 @@ public:
   void InitFromVec(const std::vector<double> &input,
                    const std::vector<std::string> &prop_vec);
 
+  Field &operator=(const Field &rhs) {
+    this->req_vec_size = rhs.req_vec_size;
+    this->props = rhs.props;
+    return *this;
+  }
+
   /**
    * Returns a reference to the column vector with given name. Creates a new
    * vector if prop was not found. The prop name will be added to the end of the
@@ -79,14 +85,16 @@ public:
    *
    * \return Vector containing all species names in output order.
    */
-  auto GetProps() const { return this->props; };
+  auto GetProps() const -> std::vector<std::string> { return this->props; }
 
   /**
    * Return the requested vector size.
    *
    * \return Requested vector size set in the instanciation of the object.
    */
-  auto GetRequestedVecSize() const { return this->req_vec_size; };
+  auto GetRequestedVecSize() const -> std::uint32_t {
+    return this->req_vec_size;
+  };
 
   /**
    * Updates all species with values from another field. If one element of the
@@ -142,7 +150,7 @@ public:
   void SetFromVector(const std::vector<FieldColumn> &cont_field);
 
 private:
-  const uint32_t req_vec_size;
+  std::uint32_t req_vec_size;
 
   std::vector<std::string> props;
 };
diff --git a/include/poet/SimParams.hpp b/include/poet/SimParams.hpp
index d4f2efb76..44ffb1263 100644
--- a/include/poet/SimParams.hpp
+++ b/include/poet/SimParams.hpp
@@ -21,6 +21,7 @@
 #ifndef PARSER_H
 #define PARSER_H
 
+#include <cstdint>
 #include <string>
 #include <string_view>
 #include <vector>
@@ -69,8 +70,10 @@ typedef struct {
 } t_simparams;
 
 using GridParams = struct s_GridParams {
-  std::vector<uint32_t> n_cells;
-  std::vector<double> s_cells;
+  std::array<uint32_t, 2> n_cells;
+  std::array<double, 2> s_cells;
+  std::uint8_t dim;
+  std::uint32_t total_n;
 
   std::string type;
 
@@ -161,10 +164,8 @@ public:
    * 'init_chemistry' must be run beforehand.
    *
    * @param R R runtime
-   * @param col_count Count of variables per grid cell (typically the count of
-   * columns of each grid cell)
    */
-  void initVectorParams(RInside &R, int col_count);
+  void initVectorParams(RInside &R);
 
   /**
    * @brief Get the numerical params struct
diff --git a/src/ChemistryModule/ChemistryModule.cpp b/src/ChemistryModule/ChemistryModule.cpp
index 3357e0397..497fdb517 100644
--- a/src/ChemistryModule/ChemistryModule.cpp
+++ b/src/ChemistryModule/ChemistryModule.cpp
@@ -104,8 +104,7 @@ void poet::ChemistryModule::RunInitFile(const std::string &input_script_path) {
 }
 
 #ifndef POET_USE_PRM
-void poet::ChemistryModule::mergeFieldWithModule(const SingleCMap &input_map,
-                                                 std::uint32_t n_cells) {
+void poet::ChemistryModule::initializeField(const Field &trans_field) {
 
   if (is_master) {
     int f_type = CHEM_INIT_SPECIES;
@@ -119,14 +118,13 @@ void poet::ChemistryModule::mergeFieldWithModule(const SingleCMap &input_map,
       this->prop_names.begin(), this->prop_names.begin() + speciesPerModule[0]};
 
   if (is_master) {
-    for (const auto &init_val : input_map) {
-      std::string name = init_val.first;
+    for (auto &prop : trans_field.GetProps()) {
       if (std::find(new_solution_names.begin(), new_solution_names.end(),
-                    name) == new_solution_names.end()) {
-        int size = name.size();
+                    prop) == new_solution_names.end()) {
+        int size = prop.size();
         ChemBCast(&size, 1, MPI_INT);
-        ChemBCast(name.data(), name.size(), MPI_CHAR);
-        new_solution_names.push_back(name);
+        ChemBCast(prop.data(), prop.size(), MPI_CHAR);
+        new_solution_names.push_back(prop);
       }
     }
     int end = 0;
@@ -162,21 +160,27 @@ void poet::ChemistryModule::mergeFieldWithModule(const SingleCMap &input_map,
   this->SetPOETSolutionNames(new_solution_names);
 
   if (is_master) {
-    this->n_cells = n_cells;
+    this->n_cells = trans_field.GetRequestedVecSize();
+    chem_field = Field(n_cells);
 
-    this->field.clear();
-    this->field.reserve(this->prop_count * n_cells);
+    std::vector<double> phreeqc_init;
+    this->getDumpedField(phreeqc_init);
 
-    std::vector<double> init_values;
-    this->getDumpedField(init_values);
-
-    const std::vector<std::string> ess_names = old_prop_names;
-
-    for (int i = 0; i < prop_names.size(); i++) {
-      double value = init_values[i];
-      std::vector<double> curr_vec(n_cells, value);
-      this->field.insert(field.end(), curr_vec.begin(), curr_vec.end());
+    if (is_sequential) {
+      std::vector<double> init_vec{phreeqc_init};
+      this->unshuffleField(phreeqc_init, n_cells, prop_count, 1, init_vec);
+      chem_field.InitFromVec(init_vec, prop_names);
+      return;
     }
+
+    std::vector<std::vector<double>> initial_values;
+
+    for (int i = 0; i < phreeqc_init.size(); i++) {
+      std::vector<double> init(n_cells, phreeqc_init[i]);
+      initial_values.push_back(std::move(init));
+    }
+
+    chem_field.InitFromVec(initial_values, prop_names);
   }
 }
 
@@ -283,6 +287,42 @@ void poet::ChemistryModule::ReadDHTFile(const std::string &input_file) {
   }
 }
 
+std::vector<double>
+poet::ChemistryModule::shuffleField(const std::vector<double> &in_field,
+                                    uint32_t size_per_prop, uint32_t prop_count,
+                                    uint32_t wp_count) {
+  std::vector<double> out_buffer(in_field.size());
+  uint32_t write_i = 0;
+  for (uint32_t i = 0; i < wp_count; i++) {
+    for (uint32_t j = i; j < size_per_prop; j += wp_count) {
+      for (uint32_t k = 0; k < prop_count; k++) {
+        out_buffer[(write_i * prop_count) + k] =
+            in_field[(k * size_per_prop) + j];
+      }
+      write_i++;
+    }
+  }
+  return out_buffer;
+}
+
+void poet::ChemistryModule::unshuffleField(const std::vector<double> &in_buffer,
+                                           uint32_t size_per_prop,
+                                           uint32_t prop_count,
+                                           uint32_t wp_count,
+                                           std::vector<double> &out_field) {
+  uint32_t read_i = 0;
+
+  for (uint32_t i = 0; i < wp_count; i++) {
+    for (uint32_t j = i; j < size_per_prop; j += wp_count) {
+      for (uint32_t k = 0; k < prop_count; k++) {
+        out_field[(k * size_per_prop) + j] =
+            in_buffer[(read_i * prop_count) + k];
+      }
+      read_i++;
+    }
+  }
+}
+
 #else // POET_USE_PRM
 
 inline void poet::ChemistryModule::PrmToPoetField(std::vector<double> &field) {
diff --git a/src/ChemistryModule/MasterFunctions.cpp b/src/ChemistryModule/MasterFunctions.cpp
index 4cd47d6d7..c97c50824 100644
--- a/src/ChemistryModule/MasterFunctions.cpp
+++ b/src/ChemistryModule/MasterFunctions.cpp
@@ -77,40 +77,6 @@ std::vector<uint32_t> poet::ChemistryModule::GetWorkerDHTEvictions() const {
   return MasterGatherWorkerMetrics(WORKER_DHT_EVICTIONS);
 }
 
-inline std::vector<double> shuffleField(const std::vector<double> &in_field,
-                                        uint32_t size_per_prop,
-                                        uint32_t prop_count,
-                                        uint32_t wp_count) {
-  std::vector<double> out_buffer(in_field.size());
-  uint32_t write_i = 0;
-  for (uint32_t i = 0; i < wp_count; i++) {
-    for (uint32_t j = i; j < size_per_prop; j += wp_count) {
-      for (uint32_t k = 0; k < prop_count; k++) {
-        out_buffer[(write_i * prop_count) + k] =
-            in_field[(k * size_per_prop) + j];
-      }
-      write_i++;
-    }
-  }
-  return out_buffer;
-}
-
-inline void unshuffleField(const std::vector<double> &in_buffer,
-                           uint32_t size_per_prop, uint32_t prop_count,
-                           uint32_t wp_count, std::vector<double> &out_field) {
-  uint32_t read_i = 0;
-
-  for (uint32_t i = 0; i < wp_count; i++) {
-    for (uint32_t j = i; j < size_per_prop; j += wp_count) {
-      for (uint32_t k = 0; k < prop_count; k++) {
-        out_field[(k * size_per_prop) + j] =
-            in_buffer[(read_i * prop_count) + k];
-      }
-      read_i++;
-    }
-  }
-}
-
 inline void printProgressbar(int count_pkgs, int n_wp, int barWidth = 70) {
   /* visual progress */
   double progress = (float)(count_pkgs + 1) / n_wp;
@@ -234,11 +200,13 @@ void poet::ChemistryModule::RunCells() {
 
 void poet::ChemistryModule::MasterRunSequential() {
   std::vector<double> shuffled_field =
-      shuffleField(field, n_cells, prop_count, 1);
+      shuffleField(chem_field.AsVector(), n_cells, prop_count, 1);
   this->setDumpedField(shuffled_field);
   PhreeqcRM::RunCells();
   this->getDumpedField(shuffled_field);
-  unshuffleField(shuffled_field, n_cells, prop_count, 1, this->field);
+  std::vector<double> out_vec{shuffled_field};
+  unshuffleField(shuffled_field, n_cells, prop_count, 1, out_vec);
+  chem_field.SetFromVector(out_vec);
 }
 
 void poet::ChemistryModule::MasterRunParallel() {
@@ -270,8 +238,9 @@ void poet::ChemistryModule::MasterRunParallel() {
 
   /* shuffle grid */
   // grid.shuffleAndExport(mpi_buffer);
-  std::vector<double> mpi_buffer = shuffleField(
-      this->field, this->n_cells, this->prop_count, wp_sizes_vector.size());
+  std::vector<double> mpi_buffer =
+      shuffleField(chem_field.AsVector(), this->n_cells, this->prop_count,
+                         wp_sizes_vector.size());
 
   /* setup local variables */
   pkg_to_send = wp_sizes_vector.size();
@@ -317,8 +286,10 @@ void poet::ChemistryModule::MasterRunParallel() {
 
   /* unshuffle grid */
   // grid.importAndUnshuffle(mpi_buffer);
+  std::vector<double> out_vec{mpi_buffer};
   unshuffleField(mpi_buffer, this->n_cells, this->prop_count,
-                 wp_sizes_vector.size(), this->field);
+                       wp_sizes_vector.size(), out_vec);
+  chem_field.SetFromVector(out_vec);
 
   /* do master stuff */
 
diff --git a/src/ChemistryModule/WorkerFunctions.cpp b/src/ChemistryModule/WorkerFunctions.cpp
index 5022b2f38..c4f940767 100644
--- a/src/ChemistryModule/WorkerFunctions.cpp
+++ b/src/ChemistryModule/WorkerFunctions.cpp
@@ -44,9 +44,8 @@ void poet::ChemistryModule::WorkerLoop() {
       break;
     }
     case CHEM_INIT_SPECIES: {
-      SingleCMap dummy_map;
-      std::uint32_t n_cells_dummy;
-      mergeFieldWithModule(dummy_map, n_cells_dummy);
+      Field dummy{0};
+      initializeField(dummy);
       break;
     }
     case CHEM_DHT_ENABLE: {
diff --git a/src/DiffusionModule.cpp b/src/DiffusionModule.cpp
index c3d37ee75..609d9cf91 100644
--- a/src/DiffusionModule.cpp
+++ b/src/DiffusionModule.cpp
@@ -55,21 +55,19 @@ inline const char *convert_bc_to_config_file(uint8_t in) {
   return "";
 }
 
-DiffusionModule::DiffusionModule(poet::DiffusionParams diffu_args, Grid &grid_)
-    : grid(grid_) {
-  this->diff_input.setGridCellN(grid_.GetGridCellsCount(GRID_X_DIR),
-                                grid_.GetGridCellsCount(GRID_Y_DIR));
-  this->diff_input.setDomainSize(grid_.GetGridSize(GRID_X_DIR),
-                                 grid_.GetGridSize(GRID_Y_DIR));
+DiffusionModule::DiffusionModule(const poet::DiffusionParams &diffu_args,
+                                 const poet::GridParams &grid_params)
+    : t_field{grid_params.total_n}, n_cells_per_prop(grid_params.total_n) {
+  this->diff_input.setGridCellN(grid_params.n_cells[0], grid_params.n_cells[1]);
+  this->diff_input.setDomainSize(grid_params.s_cells[0],
+                                 grid_params.s_cells[1]);
 
-  this->dim = grid_.GetGridDimension();
-
-  this->n_cells_per_prop = grid_.GetTotalCellCount();
+  this->dim = grid_params.dim;
 
   this->initialize(diffu_args);
 }
 
-void DiffusionModule::initialize(poet::DiffusionParams args) {
+void DiffusionModule::initialize(const poet::DiffusionParams &args) {
   // const poet::DiffusionParams args(this->R);
 
   // name of props
@@ -77,35 +75,32 @@ void DiffusionModule::initialize(poet::DiffusionParams args) {
   this->prop_names = Rcpp::as<std::vector<std::string>>(args.initial_t.names());
   this->prop_count = this->prop_names.size();
 
-  this->state =
-      this->grid.RegisterState(DIFFUSION_MODULE_NAME, this->prop_names);
-
-  std::vector<double> &field = this->state->mem;
-
   // initialize field and alphas
-  field.resize(this->n_cells_per_prop * this->prop_count);
   this->alpha.reserve(this->prop_count);
 
+  std::vector<std::vector<double>> initial_values;
+
   for (uint32_t i = 0; i < this->prop_count; i++) {
     // get alphas - we cannot assume alphas are provided in same order as
     // initial input
     this->alpha.push_back(args.alpha[this->prop_names[i]]);
 
-    double val = args.initial_t[prop_names[i]];
+    const double val = args.initial_t[prop_names[i]];
+    std::vector<double> init_val(t_field.GetRequestedVecSize(), val);
+    initial_values.push_back(std::move(init_val));
 
-    std::vector<double> prop_vec(n_cells_per_prop, val);
-    std::copy(prop_vec.begin(), prop_vec.end(),
-              field.begin() + (i * this->n_cells_per_prop));
     if (this->dim == this->DIM_2D) {
-      tug::bc::BoundaryCondition bc(this->grid.GetGridCellsCount(GRID_X_DIR),
-                                    this->grid.GetGridCellsCount(GRID_Y_DIR));
+      tug::bc::BoundaryCondition bc(diff_input.grid.grid_cells[0],
+                                    diff_input.grid.grid_cells[1]);
       this->bc_vec.push_back(bc);
     } else {
-      tug::bc::BoundaryCondition bc(this->grid.GetGridCellsCount(GRID_X_DIR));
+      tug::bc::BoundaryCondition bc(diff_input.grid.grid_cells[0]);
       this->bc_vec.push_back(bc);
     }
   }
 
+  t_field.InitFromVec(initial_values, prop_names);
+
   // apply boundary conditions to each ghost node
   uint8_t bc_size = (this->dim == this->DIM_1D ? 2 : 4);
   for (uint8_t i = 0; i < bc_size; i++) {
@@ -133,14 +128,6 @@ void DiffusionModule::initialize(poet::DiffusionParams args) {
   // apply inner grid constant cells
   // NOTE: opening a scope here for distinguish variable names
   if (args.vecinj_inner.size() != 0) {
-    // get indices of constant grid cells
-    // Rcpp::NumericVector indices_const_cells = args.vecinj_inner(Rcpp::_, 0);
-    // this->index_constant_cells =
-    //     Rcpp::as<std::vector<uint32_t>>(indices_const_cells);
-
-    // // get indices to vecinj for constant cells
-    // Rcpp::NumericVector vecinj_indices = args.vecinj_inner(Rcpp::_, 1);
-
     // apply inner constant cells for every concentration
     for (int i = 0; i < this->prop_count; i++) {
       std::vector<double> bc_vec = args.vecinj[this->prop_names[i]];
@@ -169,23 +156,25 @@ void DiffusionModule::simulate(double dt) {
   std::cout << "DiffusionModule::simulate(): Starting diffusion ..."
             << std::flush;
 
-  std::vector<double> &curr_field = this->state->mem;
+  std::vector<std::vector<double>> field_2d = t_field.As2DVector();
 
   this->diff_input.setTimestep(dt);
 
-  double *field = curr_field.data();
-  for (int i = 0; i < this->prop_count; i++) {
-    double *in_field = &field[i * this->n_cells_per_prop];
+  for (int i = 0; i < field_2d.size(); i++) {
     std::vector<double> in_alpha(this->n_cells_per_prop, this->alpha[i]);
     this->diff_input.setBoundaryCondition(this->bc_vec[i]);
 
     if (this->dim == this->DIM_1D) {
-      tug::diffusion::BTCS_1D(this->diff_input, in_field, in_alpha.data());
+      tug::diffusion::BTCS_1D(this->diff_input, field_2d[i].data(),
+                              in_alpha.data());
     } else {
-      tug::diffusion::ADI_2D(this->diff_input, in_field, in_alpha.data());
+      tug::diffusion::ADI_2D(this->diff_input, field_2d[i].data(),
+                             in_alpha.data());
     }
   }
 
+  t_field.SetFromVector(field_2d);
+
   std::cout << " done!\n";
 
   sim_a_transport = MPI_Wtime();
diff --git a/src/SimParams.cpp b/src/SimParams.cpp
index 847cf31d2..7a536134c 100644
--- a/src/SimParams.cpp
+++ b/src/SimParams.cpp
@@ -19,6 +19,8 @@
 */
 
 #include "poet/DHT_Types.hpp"
+#include <algorithm>
+#include <cassert>
 #include <poet/SimParams.hpp>
 
 #include <RInside.h>
@@ -34,10 +36,24 @@ using namespace std;
 using namespace Rcpp;
 
 poet::GridParams::s_GridParams(RInside &R) {
-  this->n_cells =
+  auto tmp_n_cells =
       Rcpp::as<std::vector<uint32_t>>(R.parseEval("mysetup$grid$n_cells"));
-  this->s_cells =
+  assert(tmp_n_cells.size() < 3);
+
+  this->dim = tmp_n_cells.size();
+
+  std::copy(tmp_n_cells.begin(), tmp_n_cells.end(), this->n_cells.begin());
+
+  auto tmp_s_cells =
       Rcpp::as<std::vector<double>>(R.parseEval("mysetup$grid$s_cells"));
+
+  assert(tmp_s_cells.size() == this->dim);
+
+  std::copy(tmp_s_cells.begin(), tmp_s_cells.end(), this->s_cells.begin());
+
+  this->total_n =
+      (dim == 1 ? this->n_cells[0] : this->n_cells[0] * this->n_cells[1]);
+
   this->type = Rcpp::as<std::string>(R.parseEval("mysetup$grid$type"));
   this->init_df =
       Rcpp::as<Rcpp::DataFrame>(R.parseEval("mysetup$grid$init_cell"));
@@ -191,7 +207,7 @@ int SimParams::parseFromCmdl(char *argv[], RInside &R) {
   return poet::PARSER_OK;
 }
 
-void SimParams::initVectorParams(RInside &R, int col_count) {
+void SimParams::initVectorParams(RInside &R) {
   if (simparams.dht_enabled) {
     /*Load significance vector from R setup file (or set default)*/
     bool signif_vector_exists = R.parseEval("exists('signif_vector')");