diff --git a/src/reactmicp/systems/unsaturated/transport_stagger.cpp b/src/reactmicp/systems/unsaturated/transport_stagger.cpp
index d4b9724..79be167 100644
--- a/src/reactmicp/systems/unsaturated/transport_stagger.cpp
+++ b/src/reactmicp/systems/unsaturated/transport_stagger.cpp
@@ -1,1181 +1,1157 @@
 /* =============================================================================
 
  Copyright (c) 2014 - 2016
  F. Georget <fabieng@princeton.edu> Princeton University
  All rights reserved.
 
 
 Redistribution and use in source and binary forms, with or without modification,
 are permitted provided that the following conditions are met:
 
     1. Redistributions of source code must retain the above copyright notice,
     this list of conditions and the following disclaimer.
 
     2. Redistributions in binary form must reproduce the above copyright notice,
     this list of conditions and the following disclaimer in the documentation
     and/or other materials provided with the distribution.
 
     3. Neither the name of the copyright holder nor the names of its
     contributors may be used to endorse or promote products derived from this
     software without specific prior written permission.
 
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR
 ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
 (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
 ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. *
 
 ============================================================================= */
 
 #include "transport_stagger.hpp"
 #include "reactmicp/solver/staggers_base/stagger_structs.hpp"
 #include "variables.hpp"
 #include "variables_box.hpp"
 
 #include "boundary_conditions.hpp"
 
 #include "saturation_equation.hpp"
 #include "saturation_pressure_equation.hpp"
 #include "pressure_equation.hpp"
 #include "aqueous_equation.hpp"
 #include "aqueous_pressure_equation.hpp"
 
 #include "specmicp_database/database_holder.hpp"
 
 #include "dfpm/solver/parabolic_driver.hpp"
 #include "dfpm/io/print.hpp"
 
 #include "specmicp_common/log.hpp"
 #include "specmicp_common/config.h"
 
 #include <memory>
 #include <iostream>
 
 namespace specmicp {
 
 namespace dfpmsolver {
     extern template class ParabolicDriver<specmicp::reactmicp::systems::unsaturated::SaturationEquation>;
     extern template class ParabolicDriver<specmicp::reactmicp::systems::unsaturated::PressureEquation>;
     extern template class ParabolicDriver<specmicp::reactmicp::systems::unsaturated::AqueousTransportEquation>;
     extern template class ParabolicDriver<specmicp::reactmicp::systems::unsaturated::AqueousGasTransportEquation>;
     extern template class ParabolicDriver<specmicp::reactmicp::systems::unsaturated::SaturationPressureEquation>;
 } //end namespace dfpmsolver
 
 namespace reactmicp {
 namespace systems {
 namespace unsaturated {
 
 
 //! \brief Type of the equation
 //! \internal
 enum class EquationType {
     Saturation,
     LiquidAqueous,
     Pressure
 };
 
 //! \brief Format type to a string (for message error)
 //! \internal
 static std::string to_string(EquationType eq_type);
 
 // forward declaration of wrapper classes
 // They are defined at the bootom of this file
 class TaskBase;
 class SaturationTask;
 
 using VariablesBase = solver::VariablesBase;
 
 
 // =================================== //
 //                                     //
 //  Declaration of the implementation  //
 //                                     //
 // =================================== //
 
 //! \brief Implementation class for the Unsaturated transport stagger
 //! \internal
 //!
 //! This class does all the work
 //! It was designed to be OpenMP compatible
 class UnsaturatedTransportStagger::UnsaturatedTransportStaggerImpl
 {
 public:
     // Constructor
     // -----------
     UnsaturatedTransportStaggerImpl(
             UnsaturatedVariablesPtr variables,
             std::shared_ptr<BoundaryConditions> boundary_conditions,
             bool merge_saturation_pressure,
             bool merge_aqueous_pressure
             );
 
     // Main functions
     // --------------
 
     //! \brief Return the residual
     scalar_t get_residual(UnsaturatedVariables * const vars);
     //! \brief Return the first residual (start of timestep)
     scalar_t get_residual_0(UnsaturatedVariables * const vars);
     //! \brief Return the update (norm of velocity vector)
     scalar_t get_update(UnsaturatedVariables * const vars);
 
     //! \brief Initialize timestep
     void initialize_timestep(
             scalar_t dt,
             UnsaturatedVariables * const vars);
     //! \brief Restart the timestep
     solver::StaggerReturnCode restart_timestep(UnsaturatedVariables* vars);
 
     // Timestep management
     // -------------------
 
     //! \brief Save the timestep
     void save_dt(scalar_t dt) {m_dt = dt;}
     //! \brief Return the timestep
     scalar_t get_dt() {return m_dt;}
 
     // Task index
     // ----------
     index_t& id_aqueous_task(index_t component) {
         return m_aq_equation_task[static_cast<std::size_t>(component)];
     }
     index_t id_aqueous_task(index_t component) const {
         return m_aq_equation_task[static_cast<std::size_t>(component)];
     }
     index_t& id_gas_task(index_t component) {
         return m_gas_equation_task[static_cast<std::size_t>(component)];
     }
     index_t id_gas_task(index_t component) const {
         return m_gas_equation_task[static_cast<std::size_t>(component)];
     }
 
     //! \brief Return the options of the saturation equation
     dfpmsolver::ParabolicDriverOptions* get_saturation_options();
     //! \brief Return the options of aqueous equation of component
     dfpmsolver::ParabolicDriverOptions* get_aqueous_options(index_t component);
     //! \brief Return the options of gas equation of component
     dfpmsolver::ParabolicDriverOptions* get_gas_options(index_t component);
 
 
     void print_debug_information(
             UnsaturatedVariables * const var
             );
 private:
 
     scalar_t m_norm_0 {1.0};
     // timestep management
     scalar_t m_dt {-1.0};
 
     //! \brief The saturation equations
     std::unique_ptr<TaskBase> m_saturation_equation;
     // Equation and solver
     std::vector<std::unique_ptr<TaskBase>> m_equation_list;
 
+    // Holds the position of the equations in the list
+    // so they can be retrieved by using the component and their type
     std::vector<index_t> m_aq_equation_task;
     std::vector<index_t> m_gas_equation_task;
 
     std::vector<bool> m_merged_gas;
-
-
 };
 
 // Main functions
 // ===============
 
 // call of the implementation
 
 UnsaturatedTransportStagger::UnsaturatedTransportStagger(
         std::shared_ptr<UnsaturatedVariables> variables,
         std::shared_ptr<BoundaryConditions> boundary_conditions,
         bool merge_saturation_pressure,
         bool merge_aqueous_pressure
        ):
     m_impl(utils::make_pimpl<UnsaturatedTransportStaggerImpl>(
                variables, boundary_conditions,
                merge_saturation_pressure, merge_aqueous_pressure
                ))
 {
 }
 
 UnsaturatedTransportStagger::~UnsaturatedTransportStagger() = default;
 
 
-static UnsaturatedVariables * const get_var(VariablesBase * const var) {
+static UnsaturatedVariables * const get_var(VariablesBase * const var)
+{
     return static_cast<UnsaturatedVariables * const >(var);
 }
 
 void
 UnsaturatedTransportStagger::initialize_timestep(
         scalar_t dt,
         VariablesBase* var
         )
 {
     m_impl->initialize_timestep(dt, get_var(var));
 }
 
 solver::StaggerReturnCode
 UnsaturatedTransportStagger::restart_timestep(VariablesBase * const var)
 {
     return m_impl->restart_timestep(get_var(var));
 }
 
 scalar_t
 UnsaturatedTransportStagger::get_residual(VariablesBase * const var)
 {
     return m_impl->get_residual(get_var(var));
 }
 
 scalar_t
 UnsaturatedTransportStagger::get_residual_0(VariablesBase * const var)
 {
     return m_impl->get_residual_0(get_var(var));
 }
 
 scalar_t
 UnsaturatedTransportStagger::get_update(VariablesBase * const var)
 {
     return m_impl->get_update(get_var(var));
 }
 
 
 dfpmsolver::ParabolicDriverOptions*
 UnsaturatedTransportStagger::get_saturation_options()
 {
     return m_impl->get_saturation_options();
 }
 
 dfpmsolver::ParabolicDriverOptions*
 UnsaturatedTransportStagger::get_aqueous_options(index_t component)
 {
     return m_impl->get_aqueous_options(component);
 }
 
 dfpmsolver::ParabolicDriverOptions*
 UnsaturatedTransportStagger::get_gas_options(index_t component)
 {
     return m_impl->get_gas_options(component);
 }
 
 void UnsaturatedTransportStagger::print_debug_information(
-        VariablesBase * const var
+        VariablesBase* const var
         )
 {
     m_impl->print_debug_information(get_var(var));
 }
 
 // =============================== //
 // =============================== //
 //                                 //
 //     Implementation details      //
 //     ----------------------      //
 //                                 //
 // =============================== //
 // =============================== //
 
 // 2 main sections
 //    - Solver wrappers : wrapper around 1 couple equation/solver
 //    - UnsaturatedTransportStaggerImpl : call the wrappers
 
 // =============================== //
 //                                 //
 //        Solver wrappers          //
 //                                 //
 // =============================== //
 
 // This wrappers are used to abstract the residual computation
 // and the methods to solve every equations
 
 
 //! \brief Base class for a task
 //!
 //! A task is how we solve governing equations,
 //! and obtain residuals and update
 //!
 //! \internal
 class SPECMICP_DLL_LOCAL TaskBase
 {
 public:
     TaskBase(index_t component, EquationType eq_type):
         m_type(eq_type),
         m_component(component)
     {}
 
     virtual ~TaskBase() {}
 
     //! \brief Compute the squared residuals
     virtual scalar_t compute_squared_residual(
-            UnsaturatedVariables * const vars
-            ) = 0;
+            UnsaturatedVariables* const vars
+            ) =0;
     //! \brief Compute the squared residuals at the beginning of the timestep
     virtual scalar_t compute_squared_residual_0(
-            UnsaturatedVariables * const vars
-            ) = 0;
+            UnsaturatedVariables* const vars
+            ) =0;
     //! \brief Compute the squared update of the variables
     virtual scalar_t compute_squared_update(
-            UnsaturatedVariables * const vars
-            ) = 0;
+            UnsaturatedVariables* const vars
+            ) =0;
 
     //! \brief Initialize the timestep
     virtual void initialize_timestep(
-            scalar_t dt, UnsaturatedVariables * const vars
-            ) = 0;
+            scalar_t dt, UnsaturatedVariables* const vars
+            ) =0;
 
     //! \brief Solve the governing equation
     virtual dfpmsolver::ParabolicDriverReturnCode restart_timestep(
-            UnsaturatedVariables * const vars
-            ) = 0;
+            UnsaturatedVariables* const vars
+            ) =0;
 
     //! \brief Return the component index (in the db)
     index_t component() {return m_component;}
 
     //! \brief The equation type
     EquationType equation_type() {return m_type;}
 
 private:
     EquationType m_type;
     index_t m_component;
 };
 
 
 //! \brief Base class for a equation solver
 //!
 //! \internal
 template <typename Derived, typename traits>
 class SPECMICP_DLL_LOCAL EquationTask: public TaskBase
 {
 public:
     using EqT = typename traits::EqT;
     using SolverT = typename traits::SolverT;
 
     EquationTask(
             index_t component,
             mesh::Mesh1DPtr the_mesh,
             typename traits::VariableBoxT& variables,
             std::shared_ptr<BoundaryConditions> bcs
             ):
         TaskBase(component, traits::equation_type),
         m_equation(component, the_mesh, variables, bcs),
         m_solver(m_equation)
     {}
 
     EquationTask(
             index_t component,
             mesh::Mesh1DPtr the_mesh,
             typename traits::VariableBoxT& variables,
             std::shared_ptr<BoundaryConditions> bcs,
             scalar_t scaling
             ):
         TaskBase(component, traits::equation_type),
         m_equation(component, the_mesh, variables, bcs),
         m_solver(m_equation)
     {
         m_equation.set_scaling(scaling);
     }
 
     Derived* derived() {return static_cast<Derived*>(this);}
 
     // This function must be implemented by subclass
     MainVariable& get_var(UnsaturatedVariables * const vars) {
         return derived()->get_var(vars);
     }
 
     scalar_t compute_squared_residual(
             UnsaturatedVariables * const vars
             ) override {
         Vector residuals;
         const MainVariable& main = get_var(vars);
         m_equation.compute_residuals(main.variable, main.velocity, residuals, true);
         return residuals.squaredNorm()/m_norm_square_0;
     }
 
     scalar_t compute_squared_residual_0(
             UnsaturatedVariables * const vars
             ) override {
         // This method compute ||R||_0, for the convergence check
         // both for the reactmicp and the dfpm solver.
         //
         // It assumes that the system has been initialized and all bcs are set.
 
         Vector residuals;
         const MainVariable& main = get_var(vars);
         // Initialize the dfpm solver
         m_solver.initialize_timestep(m_dt, get_var(vars).variable);
         // Compute ReactMiCP ||R||_0
         m_equation.compute_residuals(main.variable, main.velocity, residuals, false);
         m_norm_square_0 = residuals.squaredNorm();
         if (m_norm_square_0 < 1e-12) m_norm_square_0 = 1.0;
         return 1.0;
     }
     scalar_t compute_squared_update(
             UnsaturatedVariables * const vars
             ) override {
         const MainVariable& main = get_var(vars);
         return main.velocity.squaredNorm();
     }
 
     void initialize_timestep(
             scalar_t dt,
             UnsaturatedVariables * const vars
             ) override {
         // Register timestep
         m_dt = dt;
         // Initialize main variables
         MainVariable& main = get_var(vars);
         main.predictor = main.variable;
         main.transport_fluxes.setZero();
         main.velocity.setZero();
         // Solver initialisation is done in compute_squared_residual_0
         // so change in bcs/vars in chemistry and upscaling stagger
         // initializtion are taken into account.
     }
 
     dfpmsolver::ParabolicDriverReturnCode restart_timestep(
             UnsaturatedVariables * const vars
             ) override {
         // Restart the timestep for the solver
         // Solves the equations given the current state
         MainVariable& main = get_var(vars);
         m_solver.velocity() = main.velocity; // copy the current state
         auto retcode = m_solver.restart_timestep(main.variable);
         if (retcode > dfpmsolver::ParabolicDriverReturnCode::NotConvergedYet)
         {
             main.velocity = m_solver.velocity(); // copy solution only if its good
             // main.variable is already good
             // if solution is not good, it can be compute using the predictors
         }
         return retcode;
     }
 
     dfpmsolver::ParabolicDriverOptions& get_options() {
         return m_solver.get_options();
     }
 
     const dfpmsolver::ParabolicDriverOptions& get_options() const {
         return m_solver.get_options();
     }
 
     scalar_t get_dt() {return m_dt;}
 
 protected:
     scalar_t m_norm_square_0 {-1};
     scalar_t m_dt {-1};
     EqT m_equation;
     SolverT m_solver;
 };
 
 //! \brief Traits struct for the SaturationTask class
 //!
 //! \internal
 struct SPECMICP_DLL_LOCAL SaturationTaskTraits {
-    using EqT = SaturationEquation;
-    using SolverT = dfpmsolver::ParabolicDriver<EqT>;
+    using EqT          = SaturationEquation;
+    using SolverT      = dfpmsolver::ParabolicDriver<EqT>;
     using VariableBoxT = SaturationVariableBox;
     static constexpr EquationType equation_type {EquationType::Saturation};
 };
 
 //! \brief Wrapper for the saturation equation solver
 //!
 //! \internal
 class SPECMICP_DLL_LOCAL SaturationTask:
         public EquationTask<SaturationTask, SaturationTaskTraits>
 {
     using base = EquationTask<SaturationTask, SaturationTaskTraits>;
 public:
 
     SaturationTask(
             index_t component,
             mesh::Mesh1DPtr the_mesh,
             SaturationTaskTraits::VariableBoxT&& variables,
             std::shared_ptr<BoundaryConditions> bcs
             ):
         EquationTask<SaturationTask, SaturationTaskTraits>(
             component, the_mesh, variables, bcs)
     {}
     SaturationTask(
             index_t component,
             mesh::Mesh1DPtr the_mesh,
             SaturationTaskTraits::VariableBoxT&& variables,
             std::shared_ptr<BoundaryConditions> bcs,
             scalar_t scaling
             ):
         EquationTask<SaturationTask, SaturationTaskTraits>(
             component, the_mesh, variables, bcs, scaling)
     {}
 
-    MainVariable& get_var(UnsaturatedVariables * const vars) {
+    MainVariable& get_var(UnsaturatedVariables* const vars) {
         return vars->get_liquid_saturation();
     }
 
     // Also take into account the solid total concentration
     scalar_t compute_squared_update(
-            UnsaturatedVariables * const vars
+            UnsaturatedVariables* const vars
             ) override {
         scalar_t solid_update = vars->get_solid_concentration(component()).velocity.squaredNorm();
         return solid_update + base::compute_squared_update(vars);
     }
 };
 
 //! \brief Traits struct for the SaturationPressureTask class
 //!
 //! \internal
 struct SPECMICP_DLL_LOCAL SaturationPressureTaskTraits {
-    using EqT = SaturationPressureEquation;
-    using SolverT = dfpmsolver::ParabolicDriver<EqT>;
+    using EqT          = SaturationPressureEquation;
+    using SolverT      = dfpmsolver::ParabolicDriver<EqT>;
     using VariableBoxT = SaturationPressureVariableBox;
     static constexpr EquationType equation_type {EquationType::Saturation};
 };
 
 //! \brief Wrapper for the saturation equation solver
 //!
 //! \internal
 class SPECMICP_DLL_LOCAL SaturationPressureTask:
         public EquationTask<SaturationPressureTask, SaturationPressureTaskTraits>
 {
     using base = EquationTask<SaturationPressureTask, SaturationPressureTaskTraits>;
 public:
 
     SaturationPressureTask(
             index_t component,
             mesh::Mesh1DPtr the_mesh,
             SaturationPressureTaskTraits::VariableBoxT&& variables,
             std::shared_ptr<BoundaryConditions> bcs
             ):
         EquationTask<SaturationPressureTask, SaturationPressureTaskTraits>(
             component, the_mesh, variables, bcs)
     {}
     SaturationPressureTask(
             index_t component,
             mesh::Mesh1DPtr the_mesh,
             SaturationPressureTaskTraits::VariableBoxT&& variables,
             std::shared_ptr<BoundaryConditions> bcs,
             scalar_t scaling
             ):
         EquationTask<SaturationPressureTask, SaturationPressureTaskTraits>(
             component, the_mesh, variables, bcs, scaling)
     {}
 
     MainVariable& get_var(UnsaturatedVariables * const vars) {
         return vars->get_liquid_saturation();
     }
 
     // Also take into account the solid total concentration
     scalar_t compute_squared_update(
             UnsaturatedVariables * const vars
             ) override {
         scalar_t solid_update = vars->get_solid_concentration(component()).velocity.squaredNorm();
         if (vars->component_has_gas(0)) {
             solid_update += vars->get_pressure_main_variables(0).velocity.squaredNorm()/vars->get_rt();
         }
         return solid_update + base::compute_squared_update(vars);
     }
 };
 
 //! \brief Traits struct for the LiquidAqueousTask class
 //!
 //! \internal
 struct SPECMICP_DLL_LOCAL LiquidAqueousTaskTraits {
-    using EqT = AqueousTransportEquation;
-    using SolverT = dfpmsolver::ParabolicDriver<EqT>;
+    using EqT          = AqueousTransportEquation;
+    using SolverT      = dfpmsolver::ParabolicDriver<EqT>;
     using VariableBoxT = LiquidAqueousComponentVariableBox;
     static constexpr EquationType equation_type {EquationType::LiquidAqueous};
 };
 
 //! \brief Wrapper for the liquid transport of aqueous component equation
 //!
 //! \internal
 class SPECMICP_DLL_LOCAL LiquidAqueousTask:
         public EquationTask<LiquidAqueousTask, LiquidAqueousTaskTraits>
 {
     using base = EquationTask<LiquidAqueousTask, LiquidAqueousTaskTraits>;
 
 public:
     LiquidAqueousTask(
             index_t component,
             mesh::Mesh1DPtr the_mesh,
             LiquidAqueousTaskTraits::VariableBoxT&& variables,
             std::shared_ptr<BoundaryConditions> bcs
             ):
         EquationTask<LiquidAqueousTask, LiquidAqueousTaskTraits>(
             component, the_mesh, variables, bcs)
     {}
     LiquidAqueousTask(
             index_t component,
             mesh::Mesh1DPtr the_mesh,
             LiquidAqueousTaskTraits::VariableBoxT&& variables,
             std::shared_ptr<BoundaryConditions> bcs,
             scalar_t scaling
             ):
         EquationTask<LiquidAqueousTask, LiquidAqueousTaskTraits>(
             component, the_mesh, variables, bcs, scaling)
     {}
 
     MainVariable& get_var(UnsaturatedVariables * const vars) {
         return vars->get_aqueous_concentration(component());
     }
 
     // Also take into account the solid total concentration
     scalar_t compute_squared_update(
             UnsaturatedVariables * const vars
             ) override {
         scalar_t solid_update = vars->get_solid_concentration(component()).velocity.squaredNorm();
         return solid_update + base::compute_squared_update(vars);
     }
 };
 
 
 //! \brief Traits struct for the LiquidGasAqueousTask class
 //!
 //! \internal
 struct SPECMICP_DLL_LOCAL LiquidGasAqueousTaskTraits {
-    using EqT = AqueousGasTransportEquation;
-    using SolverT = dfpmsolver::ParabolicDriver<EqT>;
+    using EqT          = AqueousGasTransportEquation;
+    using SolverT      = dfpmsolver::ParabolicDriver<EqT>;
     using VariableBoxT = LiquidGasAqueousVariableBox;
     static constexpr EquationType equation_type {EquationType::LiquidAqueous};
 };
 
 //! \brief Wrapper for the liquid transport of aqueous component equation
 //!
 //! \internal
 class SPECMICP_DLL_LOCAL LiquidGasAqueousTask:
         public EquationTask<LiquidGasAqueousTask, LiquidGasAqueousTaskTraits>
 {
     using base = EquationTask<LiquidGasAqueousTask, LiquidGasAqueousTaskTraits>;
 
 public:
     LiquidGasAqueousTask(
             index_t component,
             mesh::Mesh1DPtr the_mesh,
             LiquidGasAqueousTaskTraits::VariableBoxT&& variables,
             std::shared_ptr<BoundaryConditions> bcs
             ):
         base(component, the_mesh, variables, bcs)
     {}
     LiquidGasAqueousTask(
             index_t component,
             mesh::Mesh1DPtr the_mesh,
             LiquidGasAqueousTaskTraits::VariableBoxT&& variables,
             std::shared_ptr<BoundaryConditions> bcs,
             scalar_t scaling
             ):
         base(component, the_mesh, variables, bcs, scaling)
     {}
 
     MainVariable& get_var(UnsaturatedVariables * const vars) {
         return vars->get_aqueous_concentration(component());
     }
     scalar_t compute_squared_residual(
             UnsaturatedVariables * const vars
             ) override {
         Vector residuals;
         const MainVariable& main = get_var(vars);
         m_equation.compute_residuals(main.variable, main.velocity, residuals, true);
         scalar_t res = residuals.squaredNorm()/m_norm_square_0;
         return res;
         //if (res > 10) return 1e-8;
         //else return res;
     }
     // Also take into account the solid total concentration
     scalar_t compute_squared_update(
             UnsaturatedVariables * const vars
             ) override {
         scalar_t solid_update = vars->get_solid_concentration(component()).velocity.squaredNorm();
         scalar_t pressure_update = vars->get_pressure_main_variables(component()).velocity.squaredNorm()/vars->get_rt();
 
         return pressure_update + solid_update + base::compute_squared_update(vars);
     }
 
     void initialize_timestep(
             scalar_t dt,
             UnsaturatedVariables * const vars
             ) override {
         m_dt = dt;
 
         auto& pres = vars->get_pressure_main_variables(component());
         pres.predictor = pres.variable;
         //pres.velocity.setZero();
 
         base::initialize_timestep(dt, vars);
     }
 
     dfpmsolver::ParabolicDriverReturnCode restart_timestep(
             UnsaturatedVariables * const vars
             ) override {
         MainVariable& main = get_var(vars);
         m_solver.velocity() = main.velocity;
         auto retcode = m_solver.restart_timestep(main.variable);
         if (retcode > dfpmsolver::ParabolicDriverReturnCode::NotConvergedYet)
         {
             main.velocity = m_solver.velocity();
         }
         return retcode;
     }
 };
 
 
 //! \brief Traits struct for the Pressure Task traits
 //!
 //! \internal
 struct SPECMICP_DLL_LOCAL PressureTaskTraits {
-    using EqT = PressureEquation;
-    using SolverT = dfpmsolver::ParabolicDriver<EqT>;
+    using EqT          = PressureEquation;
+    using SolverT      = dfpmsolver::ParabolicDriver<EqT>;
     using VariableBoxT = PressureVariableBox;
     static constexpr EquationType equation_type {EquationType::Pressure};
 };
 
 //! \brief Wrapper for the pressure equation solver
 //!
 //! \internal
 class SPECMICP_DLL_LOCAL PressureTask:
         public EquationTask<PressureTask, PressureTaskTraits>
 {
     using base = EquationTask<PressureTask, PressureTaskTraits>;
 
 public:
     PressureTask(
             index_t component,
             mesh::Mesh1DPtr the_mesh,
             PressureTaskTraits::VariableBoxT&& variables,
             std::shared_ptr<BoundaryConditions> bcs
             ):
         EquationTask<PressureTask, PressureTaskTraits>(
             component, the_mesh, variables, bcs)
     {}
     PressureTask(
             index_t component,
             mesh::Mesh1DPtr the_mesh,
             PressureTaskTraits::VariableBoxT&& variables,
             std::shared_ptr<BoundaryConditions> bcs,
             scalar_t scaling
             ):
         EquationTask<PressureTask, PressureTaskTraits>(
             component, the_mesh, variables, bcs, scaling)
     {}
 
     MainVariable& get_var(UnsaturatedVariables * const vars) {
         return vars->get_pressure_main_variables(component());
     }
 };
 
 // ================================== //
 //                                    //
 //  UnsaturatedTransportStaggerImpl   //
 //                                    //
 // ================================== //
 
 // constructor
 // ===========
 
 UnsaturatedTransportStagger::UnsaturatedTransportStaggerImpl::UnsaturatedTransportStaggerImpl(
         UnsaturatedVariablesPtr variables,
         std::shared_ptr<BoundaryConditions> bcs,
         bool merge_saturation_pressure,
         bool merge_aqueous_pressure
         ):
     m_merged_gas(variables->get_database()->nb_component(), false)
 {
     database::RawDatabasePtr raw_db = variables->get_database();
     mesh::Mesh1DPtr the_mesh = variables->get_mesh();
 
     m_aq_equation_task = std::vector<index_t>(
                 raw_db->nb_component(), no_equation);
     m_gas_equation_task = std::vector<index_t>(
                 raw_db->nb_component(), no_equation);
 
     // Saturation equations
     // ====================
     //
     // There is 2 choices for the saturation equation
     // Either SaturationPressureEquation or SaturationPressure
     if (merge_saturation_pressure and variables->component_has_gas(0))
     {
         m_saturation_equation = make_unique<SaturationPressureTask>(
                     0, the_mesh,
                     variables->get_saturation_pressure_variables(),
                     bcs,
                     variables->get_aqueous_scaling(0)
                     );
         m_merged_gas[0] = true;
 
     }
     else
     {
         m_saturation_equation = make_unique<SaturationTask>(
                     0, the_mesh,
                     variables->get_saturation_variables(),
                     bcs,
                     variables->get_aqueous_scaling(0)
                     );
     }
 
     const index_t size = raw_db->nb_aqueous_components() + variables->nb_gas();
     m_equation_list.reserve(size);
 
     // Liquid aqueous diffusion-advection
     // ==================================
     for (index_t id: raw_db->range_aqueous_component())
     {
         if (merge_aqueous_pressure and variables->component_has_gas(id))
         {
             m_equation_list.emplace_back(
                         make_unique<LiquidGasAqueousTask>(
                             id, the_mesh,
                             variables->get_liquid_gas_aqueous_variables(id),
                             bcs,
                             variables->get_aqueous_scaling(id))
                         );
             m_merged_gas[id] = true;
         } else {
             m_equation_list.emplace_back(
                         make_unique<LiquidAqueousTask>(
                             id, the_mesh,
                             variables->get_liquid_aqueous_component_variables(id),
                             bcs,
                             variables->get_aqueous_scaling(id))
                         );
         }
         id_aqueous_task(id) = m_equation_list.size()-1;
     }
 
     // Water partial pressure
     // ======================
     //
     // Depending on the saturation equation chosen we may or may not include
     // this equations
     if (variables->component_has_gas(0) and (not m_merged_gas[0]))
     {
         m_equation_list.emplace_back(
                     make_unique<PressureTask>(0, the_mesh,
                                               variables->get_pressure_variables(0),
                                               bcs,
                                               variables->get_gaseous_scaling(0)
                     ));
         id_gas_task(0) = m_equation_list.size()-1;
     }
     // Gaseous diffusion
     for (index_t id: raw_db->range_aqueous_component())
     {
         if (variables->component_has_gas(id) and (not m_merged_gas[id])) {
             m_equation_list.emplace_back(
                         make_unique<PressureTask>(id, the_mesh,
                                                   variables->get_pressure_variables(id),
                                                   bcs,
                                                   variables->get_gaseous_scaling(id)
                         ));
             id_gas_task(id) = m_equation_list.size()-1;
         }
     }
 }
 
 
 dfpmsolver::ParabolicDriverOptions*
 UnsaturatedTransportStagger::UnsaturatedTransportStaggerImpl::get_saturation_options()
 {
     dfpmsolver::ParabolicDriverOptions* opts = nullptr;
     if (m_merged_gas[0]) {
         opts = &static_cast<SaturationPressureTask*>(
                     m_saturation_equation.get())->get_options();
     } else {
         opts = &static_cast<SaturationTask*>(
                     m_saturation_equation.get())->get_options();
     }
 
     return opts;
 }
 
 dfpmsolver::ParabolicDriverOptions*
 UnsaturatedTransportStagger::UnsaturatedTransportStaggerImpl::get_aqueous_options(
         index_t component)
 {
     dfpmsolver::ParabolicDriverOptions* opts = nullptr;
     auto id = id_aqueous_task(component);
 
     if (id != no_equation) {
         if (m_merged_gas[component]) {
             opts = &static_cast<LiquidGasAqueousTask*>(
                 m_equation_list[id].get())->get_options();
         } else {
             opts = &static_cast<LiquidAqueousTask*>(
                 m_equation_list[id].get())->get_options();
         }
     }
     return opts;
 }
 
 dfpmsolver::ParabolicDriverOptions*
 UnsaturatedTransportStagger::UnsaturatedTransportStaggerImpl::get_gas_options(
         index_t component)
 {
 
     dfpmsolver::ParabolicDriverOptions* opts = nullptr;
     auto id = id_gas_task(component);
 
     if (id != no_equation) {
         opts = &static_cast<PressureTask*>(
            m_equation_list[id].get())->get_options();
     }
     return opts;
 }
 
 // Residuals
 // =========
 
 scalar_t
 UnsaturatedTransportStagger::UnsaturatedTransportStaggerImpl::get_residual(
         UnsaturatedVariables * const vars
         )
 {
     const scalar_t res_w = m_saturation_equation->compute_squared_residual(vars);
     scalar_t sum = 0.0;
 #ifdef SPECMICP_HAVE_OPENMP
 #pragma omp parallel for \
     reduction(+: sum)
 #endif
     for (std::size_t ideq=0; ideq<m_equation_list.size(); ++ideq) {
         auto& task = m_equation_list[ideq];
         sum += task->compute_squared_residual(vars);
     }
     sum += res_w;
     auto norm = std::sqrt(sum);
     //std::cout << "; " << norm << std::endl;
     return norm;
 }
 
 scalar_t
 UnsaturatedTransportStagger::UnsaturatedTransportStaggerImpl::get_update(
         UnsaturatedVariables * const vars
         )
 {
     scalar_t sum = 0.0;
 #ifdef SPECMICP_HAVE_OPENMP
 #pragma omp parallel for \
     reduction(+: sum)
 #endif
     for (std::size_t ideq=0; ideq<m_equation_list.size(); ++ideq) {
         auto& task = m_equation_list[ideq];
         sum += task->compute_squared_update(vars);
     }
     sum += m_saturation_equation->compute_squared_update(vars);
     return std::sqrt(sum);
 }
 
 // Solving the equations
 // =====================
 
 void
 UnsaturatedTransportStagger::UnsaturatedTransportStaggerImpl::initialize_timestep(
         scalar_t dt,
-        UnsaturatedVariables * const vars
+        UnsaturatedVariables* const vars
         )
 {
-
-//    auto& aqwtilde = vars->get_water_aqueous_concentration();
-//    auto& saturation = vars->get_liquid_saturation();
-//    vars->set_aqueous_scaling(0, saturation.variable.maxCoeff());
-
-//    for (auto aqc: vars->get_database()->range_aqueous_component()) {
-//        auto& aqtilde = vars->get_aqueous_concentration(aqc);
-//        vars->set_aqueous_scaling(0, aqtilde.variable.maxCoeff());
-//    }
-
     save_dt(dt);
 
     // secondary variables initialization
     {
         MainVariable& cbar_w = vars->get_solid_concentration(0);
         cbar_w.predictor = cbar_w.variable;
-        //cbar_w.velocity.setZero();
-        //cbar_w.chemistry_rate.setZero();
 
         SecondaryTransientVariable& ctilde_w = vars->get_water_aqueous_concentration();
         ctilde_w.predictor = ctilde_w.variable;
-        //ctilde_w.velocity.setZero();
 
         if (vars->component_has_gas(0))
         {
             MainVariable& pres_vars = vars->get_pressure_main_variables(0);
             pres_vars.predictor = pres_vars.variable;
-            //pres_vars.velocity.setZero();
-            //pres_vars.chemistry_rate.setZero();
         }
     }
     // aqueous component
     for (index_t component: vars->get_database()->range_aqueous_component())
     {
         MainVariable& cbar_i = vars->get_solid_concentration(component);
         cbar_i.predictor = cbar_i.variable;
-        //cbar_i.velocity.setZero();
-        //cbar_i.chemistry_rate.setZero();
 
         if (vars->component_has_gas(component))
         {
             MainVariable& pres_vars = vars->get_pressure_main_variables(component);
             pres_vars.predictor = pres_vars.variable;
-            //pres_vars.velocity.setZero();
-            //pres_vars.chemistry_rate.setZero();
         }
     }
     // porosity
     {
         SecondaryTransientVariable& porosity = vars->get_porosity();
         porosity.predictor = porosity.variable;
     }
     vars->get_advection_flux().set_zero();
     vars->set_relative_variables();
 
 
-
+    // Initialize each equation
+    // Note : first residuals will be computed in get_residual_0
     m_saturation_equation->initialize_timestep(dt, vars);
 #ifdef SPECMICP_HAVE_OPENMP
 #pragma omp parallel for
 #endif
     for (std::size_t ideq=0; ideq<m_equation_list.size(); ++ideq) {
         auto& task = m_equation_list[ideq];
         task->initialize_timestep(dt, vars);
     }
-    m_norm_0 = 1.0;//std::sqrt(sum);
 }
 
 
 scalar_t
 UnsaturatedTransportStagger::UnsaturatedTransportStaggerImpl::get_residual_0(
-        UnsaturatedVariables * const vars
+        UnsaturatedVariables* const vars
         )
 {
     scalar_t sum = 0.0;
 
 #ifdef SPECMICP_HAVE_OPENMP
 #pragma omp parallel for
 #endif
     for (std::size_t ideq=0; ideq<m_equation_list.size(); ++ideq) {
         auto& task = m_equation_list[ideq];
         task->compute_squared_residual_0(vars);
     }
     m_saturation_equation->compute_squared_residual_0(vars);
-    m_norm_0 = 1.0;//std::sqrt(sum);
+    m_norm_0 = 1.0; // the scaling is taken care by each equations
     return m_norm_0;
 }
 
 
 solver::StaggerReturnCode
 UnsaturatedTransportStagger::UnsaturatedTransportStaggerImpl::restart_timestep(
-        UnsaturatedVariables * const vars
+        UnsaturatedVariables* const vars
         )
 {
     dfpmsolver::ParabolicDriverReturnCode retcode;
 
+    bool flag_fail = false;            // true if one of the equation has failed
+    bool flag_error_minimized = false; // true if all equations are minimized
 
-    //auto& satu = vars->get_liquid_saturation();
-    //for (index_t id=0; id<satu.variable.rows(); ++id) {
-    //    if (satu.variable(id) > 0.95) {
-    //        satu.variable(id) = 0.95;
-    //        WARNING << "Reset saturation at node : " << id;
-    //    }
-    //}
-
-    bool flag_fail =  false;
-    bool flag_error_minimized = false;
-
-    // Saturation equation
-    retcode = m_saturation_equation->restart_timestep(vars);
-    if (dfpmsolver::has_failed(retcode)) {
-        WARNING << "Failed to solve saturation equation, return code :"
-                << to_string(retcode);
-        flag_fail = true;
+    {
+        // Saturation equation
+        retcode = m_saturation_equation->restart_timestep(vars);
+        if (dfpmsolver::has_failed(retcode)) {
+            WARNING << "Failed to solve saturation equation, return code :"
+                    << io::to_string(retcode);
+            flag_fail = true;
+        }
+        if (retcode == dfpmsolver::ParabolicDriverReturnCode::ErrorMinimized)
+            flag_error_minimized = true;
     }
-    if (retcode == dfpmsolver::ParabolicDriverReturnCode::ErrorMinimized)
-        flag_error_minimized = true;
 
-    // Other equation
+    // Other equations
     if (not flag_fail)
     {
         vars->set_relative_variables();
 #ifdef SPECMICP_HAVE_OPENMP
 #pragma omp parallel for \
     reduction(or: flag_fail) \
     reduction(and: flag_error_minimized) \
     schedule(dynamic, 1)
 #endif
         for (std::size_t ideq=0; ideq<m_equation_list.size(); ++ideq) {
             auto& task = m_equation_list[ideq];
             retcode = task->restart_timestep(vars);
             if (dfpmsolver::has_failed(retcode))
             {
                 WARNING << "Equation of type '"
                         << to_string(task->equation_type())
                         << "' for component "
                         << task->component()
                         << " has failed with return code : "
                         << io::to_string(retcode);
                 flag_fail = true;
             }
             flag_error_minimized = flag_error_minimized  and
                     (retcode == dfpmsolver::ParabolicDriverReturnCode::ErrorMinimized);
         }
     }
 
     // Return code
     solver::StaggerReturnCode return_code = solver::StaggerReturnCode::ResidualMinimized;
     if (flag_error_minimized)
     {
         return_code = solver::StaggerReturnCode::ErrorMinimized;
     }
     if (flag_fail)
+    {
         return_code = solver::StaggerReturnCode::UnknownError;
-
+    }
     return return_code;
 
 }
 
 
 void UnsaturatedTransportStagger::UnsaturatedTransportStaggerImpl::print_debug_information(
         UnsaturatedVariables * const vars
         )
 {
     ERROR << "Current residuals per equation : \n"
           << "Saturation equation : " << std::sqrt(m_saturation_equation->compute_squared_residual(vars));
 
     for (std::size_t ideq=0; ideq<m_equation_list.size(); ++ideq) {
 
         TaskBase* task = m_equation_list[ideq].get();
         ERROR << to_string(task->equation_type()) << " - " << task->component()
               << " : "
               << std::sqrt(task->compute_squared_residual(vars));
     }
 }
 
 // ================================ //
 //                                  //
 //          Helper functions        //
 //                                  //
 // ================================ //
 
 static std::string to_string(EquationType eq_type)
 {
     std::string str;
     switch (eq_type) {
     case EquationType::LiquidAqueous:
         str = "Liquid advection-diffusion";
         break;
     case EquationType::Pressure:
         str = "Gaseous diffusion";
         break;
     case EquationType::Saturation:
         str = "Saturation equation";
         break;
     default:
         str = "Unknown equation";
         break;
     }
     return str;
 }
 
 
 } //end namespace unsaturated
 } //end namespace systems
 } //end namespace reactmicp
 } //end namespace specmicp