diff --git a/test_impestpfa_ad.cpp b/test_impestpfa_ad.cpp
index 41326d194..278da4378 100644
--- a/test_impestpfa_ad.cpp
+++ b/test_impestpfa_ad.cpp
@@ -46,9 +46,11 @@ namespace {
         typedef Vector V;
 
         DerivedGeology(const UnstructuredGrid& grid,
-                       const Geology&          geo)
+                       const Geology&          geo ,
+                       const double*           grav = 0)
             : pvol_ (grid.number_of_cells)
             , trans_(grid.number_of_faces)
+            , gpot_ (grid.cell_facepos[ grid.number_of_cells ])
         {
             // Pore volume
             const typename Vector::Index nc = grid.number_of_cells;
@@ -61,14 +63,35 @@ namespace {
             UnstructuredGrid* ug = const_cast<UnstructuredGrid*>(& grid);
             tpfa_htrans_compute(ug, geo.permeability(), htrans.data());
             tpfa_trans_compute (ug, htrans.data()     , trans_.data());
+
+            if (grav != 0) {
+                const typename Vector::Index nd = grid.dimensions;
+
+                for (typename Vector::Index c = 0; c < nc; ++c) {
+                    const double* const cc = & grid.cell_centroids[c*nd + 0];
+
+                    const int* const p = grid.cell_facepos;
+                    for (int i = p[c]; i < p[c + 1]; ++i) {
+                        const int f = grid.cell_faces[i];
+
+                        const double* const fc = & grid.face_centroids[f*nd + 0];
+
+                        for (typename Vector::Index d = 0; d < nd; ++d) {
+                            gpot_[i] += grav[d] * (fc[d] - cc[d]);
+                        }
+                    }
+                }
+            }
         }
 
         const Vector& poreVolume()       const { return pvol_ ; }
         const Vector& transmissibility() const { return trans_; }
+        const Vector& gravityPotential() const { return gpot_ ; }
 
     private:
         Vector pvol_ ;
         Vector trans_;
+        Vector gpot_ ;
     };
 }