diff --git a/docs/src/api.md b/docs/src/api.md
index c2564f6..bf88b2b 100644
--- a/docs/src/api.md
+++ b/docs/src/api.md
@@ -5,6 +5,7 @@
 Modules = [HyperFEM.ComputationalModels]
 Order   = [:module, :type, :function, :macro]
 Public  = true
+Private = false
 ```
 
 ## DiscreteModeling
@@ -13,6 +14,7 @@ Public  = true
 Modules = [HyperFEM.DiscreteModeling, HyperFEM.DiscreteModeling.CartesianTags, HyperFEM.DiscreteModeling.EvolutionFunctions]
 Order   = [:module, :type, :function, :macro]
 Public  = true
+Private = false
 ```
 
 ## PhysicalModels
@@ -21,6 +23,7 @@ Public  = true
 Modules = [HyperFEM.PhysicalModels]
 Order   = [:module, :type, :function, :macro]
 Public  = true
+Private = false
 ```
 
 ## Solvers
@@ -29,6 +32,7 @@ Public  = true
 Modules = [HyperFEM.Solvers]
 Order   = [:module, :type, :function, :macro]
 Public  = true
+Private = false
 ```
 
 ## TensorAlgebra
@@ -37,6 +41,7 @@ Public  = true
 Modules = [HyperFEM.TensorAlgebra]
 Order   = [:module, :type, :function, :macro]
 Public  = true
+Private = false
 ```
 
 ## WeakForms
@@ -45,4 +50,5 @@ Public  = true
 Modules = [HyperFEM.WeakForms]
 Order   = [:module, :type, :function, :macro]
 Public  = true
+Private = false
 ```
diff --git a/src/ComputationalModels/Drivers.jl b/src/ComputationalModels/Drivers.jl
index 1622623..50b1c88 100644
--- a/src/ComputationalModels/Drivers.jl
+++ b/src/ComputationalModels/Drivers.jl
@@ -284,6 +284,8 @@ get_assemblers(m::DynamicNonlinearModel) = (m.caches[4])
 
 
 """
+    update_velocity!
+
 Update the velocity field `vh` based on the current displacement field `xh⁺`
 and the previous displacement `xh⁻` using a midpoint time-stepping scheme.
 The velocity is updated in place.
@@ -298,6 +300,8 @@ end
 
 
 """
+    update_displacements!
+
 Update the old displacement field `xh⁻` with the new displacement field `xh⁺`.
 The update is performed in place, modifying `xh⁻` to match both the free dof values
 and the dirichlet dof values of `xh⁺`.
diff --git a/src/Exports.jl b/src/Exports.jl
index 0417462..d0f4dbf 100644
--- a/src/Exports.jl
+++ b/src/Exports.jl
@@ -96,7 +96,6 @@ end
 @publish PhysicalModels FlexoElectro
 @publish PhysicalModels ThermoElectroMechano
 @publish PhysicalModels EnergyInterpolationScheme
-@publish PhysicalModels update_state!
 @publish PhysicalModels Kinematics
 @publish PhysicalModels Solid
 @publish PhysicalModels KinematicModel
@@ -114,6 +113,9 @@ end
 
 @publish PhysicalModels SecondPiola
 @publish PhysicalModels Dissipation
+
+@publish PhysicalModels CellState      # reexport form Gridap
+@publish PhysicalModels update_state!  # reexport from Gridap
 @publish PhysicalModels initialize_state
 @publish PhysicalModels return_mapping
 @publish PhysicalModels update_time_step!
diff --git a/src/PhysicalModels/ElectroMechanicalModels.jl b/src/PhysicalModels/ElectroMechanicalModels.jl
index dfe36ee..f1f3b61 100644
--- a/src/PhysicalModels/ElectroMechanicalModels.jl
+++ b/src/PhysicalModels/ElectroMechanicalModels.jl
@@ -81,7 +81,7 @@ function initialize_state(obj::ElectroMechModel)
   initialize_state(obj.mechano)
 end
 
-function update_state!(obj::ElectroMechModel, state, F, E, args...)
+function Gridap.CellData.update_state!(obj::ElectroMechModel, state, F, E, args...)
   update_state!(obj.mechano, state, F, args...)
 end
 
diff --git a/src/PhysicalModels/PhysicalModels.jl b/src/PhysicalModels/PhysicalModels.jl
index 9690547..185a98e 100644
--- a/src/PhysicalModels/PhysicalModels.jl
+++ b/src/PhysicalModels/PhysicalModels.jl
@@ -12,7 +12,7 @@ using ..TensorAlgebra: _∂H∂F_2D
 using ..TensorAlgebra: trAA
 
 import Base: +
-import Gridap: update_state!
+import Gridap.CellData: CellState, update_state!
 
 export Yeoh3D
 export PlaneStressIncompressible_I1PD 
@@ -85,6 +85,8 @@ export EnergyInterpolationScheme
 export SecondPiola
 export Dissipation
 
+export CellState      # reexport from Gridap
+export update_state!  # reexport form Gridap
 export initialize_state
 export return_mapping
 export update_time_step!
@@ -158,7 +160,11 @@ include("PINNs.jl")
 Base.broadcastable(m::PhysicalModel) = Ref(m) # Allows to use the @. syntax for passing a single constitutive model into a vectorized function
 
 """
+    CellState(model, dΩ)
+
 Initialize the state variables for the given constitutive model and discretization.
+The constitutive model passed to the function will determine the type of the state variables,
+e.g., a vector, tensor, tuple of state variables...
 """
 function Gridap.CellData.CellState(::PhysicalModel, args...)
   return nothing
@@ -166,7 +172,10 @@ end
 
 
 """
-Define the state variable at a Gauss point. Unlike the function 'CellState', the state variable is represented by a number or a tensor.
+    initialize_state(model)
+
+Define the state variable at a Gauss point. Unlike the function [`CellState`](@ref), the returned
+state variable is represented by a number or a tensor.
 """
 function initialize_state(::PhysicalModel)
   return nothing
@@ -174,20 +183,30 @@ end
 
 
 """
-Update the state variables. The state variables must be initialized using the function 'CellState' with the constitutive model.
+    update_state!(model, A, F, Fn)
+
+Update the state variables. The state variables must be initialized using the function
+[`CellState`](@ref) with the constitutive model.
+
+NOTE: The Gridap function expects the following order of arguments: `update_state!(updater, cell_states, cell_fields)`,
+hence, the order of the arguments differ from the standar energy function, like `Ψ(F, Fn, A...)`
 """
-function update_state!(::PhysicalModel, vars...)
+function Gridap.CellData.update_state!(::PhysicalModel, vars...)
 end
 
 
 """
-Update the state variables at a Gauss point. Unlike the function 'update_state!', the state variables are represented by a number or a tensor.
+    return_mapping(model, F, Fn, A...)
+
+Update the state variables at a Gauss point. Unlike the function [`update_state!`](@ref), the state variables are represented by a number or a tensor.
 """
 function return_mapping(::PhysicalModel, vars...)
 end
 
 
 """
+    Dissipation(model)
+
 Return the dissipation and its derivatives if any.
 """
 function Dissipation(::PhysicalModel, args...)
@@ -196,6 +215,8 @@ end
 
 
 """
+    SecondPiola(model)
+
 Return the energy density and its derivatives as functions of C instead of F.
 """
 function SecondPiola(::T, args...) where {T<:PhysicalModel}
@@ -204,7 +225,11 @@ end
 
 
 """
+    update_time_step!(model, Δt)
+
 Set the time step to be used internally by the constitutive model.
+The time step is a reference, hence, the weak forms derived from the 
+constitutive model will be automatically updated with the new time step.
 """
 function update_time_step!(::PhysicalModel, Δt::Float64)
   Δt
diff --git a/src/PhysicalModels/ThermoElectroMechanicalModels.jl b/src/PhysicalModels/ThermoElectroMechanicalModels.jl
index 1940343..05d1d29 100644
--- a/src/PhysicalModels/ThermoElectroMechanicalModels.jl
+++ b/src/PhysicalModels/ThermoElectroMechanicalModels.jl
@@ -7,7 +7,7 @@ function initialize_state(obj::ThermoElectroMechano)
   initialize_state(obj.mechano)
 end
 
-function update_state!(obj::ThermoElectroMechano, state, F, E, θ, args...)
+function Gridap.CellData.update_state!(obj::ThermoElectroMechano, state, F, E, θ, args...)
   update_state!(obj.mechano, state, F, args...)
 end
 
diff --git a/src/PhysicalModels/ThermoMechanicalModels.jl b/src/PhysicalModels/ThermoMechanicalModels.jl
index 509e43a..96fe9f9 100644
--- a/src/PhysicalModels/ThermoMechanicalModels.jl
+++ b/src/PhysicalModels/ThermoMechanicalModels.jl
@@ -11,7 +11,7 @@ function initialize_state(obj::TM) where {TM<:ThermoMechano}
   initialize_state(obj.mechano)
 end
 
-function update_state!(obj::TM, state, F, θ, args...) where {TM<:ThermoMechano}
+function Gridap.CellData.update_state!(obj::TM, state, F, θ, args...) where {TM<:ThermoMechano}
   update_state!(obj.mechano, state, F, args...)
 end
 
diff --git a/src/PhysicalModels/ViscousModels.jl b/src/PhysicalModels/ViscousModels.jl
index a53d5e6..e37cc1d 100644
--- a/src/PhysicalModels/ViscousModels.jl
+++ b/src/PhysicalModels/ViscousModels.jl
@@ -40,7 +40,7 @@ function initialize_state(::ViscousIncompressible)
   VectorValue(I3..., 0.0)
 end
 
-function update_state!(obj::ViscousIncompressible, A, F, Fn)
+function Gridap.CellData.update_state!(obj::ViscousIncompressible, A, F, Fn)
   state_updater(Aᵅ, Fᵅ, Fnᵅ) = (true, return_mapping(obj, Fᵅ, Fnᵅ, Aᵅ))
   update_state!(state_updater, A, F, Fn)
 end
@@ -103,7 +103,7 @@ function initialize_state(obj::NVisco)
   map(initialize_state, obj)
 end
 
-function update_state!(obj::NVisco, states, F, Fn)
+function Gridap.CellData.update_state!(obj::NVisco, states, F, Fn)
   @assert length(obj) == length(states)
   map((b, s) -> update_state!(b, s, F, Fn), obj, states)
 end
@@ -157,11 +157,11 @@ function initialize_state(obj::GeneralizedMaxwell)
   initialize_state(obj.branches)
 end
 
-function update_state!(obj::GeneralizedMaxwell{<:IsoElastic}, states, F, Fn)
+function Gridap.CellData.update_state!(obj::GeneralizedMaxwell{<:IsoElastic}, states, F, Fn)
   update_state!(obj.branches, states, F, Fn)
 end
 
-function update_state!(obj::GeneralizedMaxwell{<:AnisoElastic}, states, F, n, Fn)
+function Gridap.CellData.update_state!(obj::GeneralizedMaxwell{<:AnisoElastic}, states, F, n, Fn)
   update_state!(obj.branches, states, F, Fn)
 end
 
diff --git a/src/PhysicalModels/ViscousPolyconvex.jl b/src/PhysicalModels/ViscousPolyconvex.jl
index e9202e4..48bb126 100644
--- a/src/PhysicalModels/ViscousPolyconvex.jl
+++ b/src/PhysicalModels/ViscousPolyconvex.jl
@@ -12,7 +12,7 @@ where `F` is the deformation gradient, `J` is the jacobian and `Cᵥ` is the vis
 - `μ::Float64`: Shear modulus.
 - `τ::Float64`: Relaxation time.
 - `Δt::Base.RefValue{Float64}`: `Reference` to the time step.
- """
+"""
 struct ViscousPolyconvex <: Visco
   μ::Float64
   τ::Float64