diff --git a/src/one-body-sample.cc b/src/one-body-sample.cc
index 621639a9641a6180fd2b07fb6106af72ab75361b..998d5c1c363f5506a5cbf9b0c660ba9b116c9f4f 100644
--- a/src/one-body-sample.cc
+++ b/src/one-body-sample.cc
@@ -109,9 +109,8 @@ int main(int argc, char *argv[]) {
       Python::run("import sys");
       Python::run("sys.path.append('" srcdir "')");
 
-      Python::import("one-body-sample_neumann")
-          .get("Functions")
-          .toC<FunctionMap::Base>(functions);
+      Python::import("one-body-sample").get("Functions").toC<FunctionMap::Base>(
+          functions);
     }
 
     Dune::ParameterTree parset;
@@ -232,8 +231,14 @@ int main(int argc, char *argv[]) {
       if (parset.get<bool>("solver.tnnmg.use")) {
         assemble_neumann<GridType, GridView, SmallVector, P1Basis>(
             leafView, p1Basis, neumannNodes, b4,
-            functions.get("sampleFunction"), h * run);
+            functions.get("neumannCondition"), h * run);
         stiffnessMatrix.mmv(u4, b4);
+        // Apply Dirichlet condition
+        for (int i = 0; i < finestSize; ++i)
+          if (ignoreNodes[i].count() == dim)
+            functions.get("dirichletCondition")
+                .evaluate(h * run, u4_diff[i][0]);
+
         for (int state_fpi = 0;
              state_fpi < parset.get<int>("solver.tnnmg.fixed_point_iterations");
              ++state_fpi) {
@@ -267,7 +272,7 @@ int main(int argc, char *argv[]) {
 
         if (parset.get<bool>("printEvolution"))
           print_evolution<SingletonVectorType, VectorType>(
-              frictionalNodes, *s4_new, u4, functions.get("sampleFunction"),
+              frictionalNodes, *s4_new, u4, functions.get("neumannCondition"),
               run, h * run, octave_writer);
       }
 
@@ -277,7 +282,7 @@ int main(int argc, char *argv[]) {
       if (parset.get<bool>("benchmarks.fpi.enable")) {
         assemble_neumann<GridType, GridView, SmallVector, P1Basis>(
             leafView, p1Basis, neumannNodes, b5,
-            functions.get("sampleFunction"), h * run);
+            functions.get("neumannCondition"), h * run);
         stiffnessMatrix.mmv(u5, b5);
         for (int state_fpi = 0;
              state_fpi < parset.get<int>("benchmarks.fpi.iterations");
diff --git a/src/one-body-sample.py b/src/one-body-sample.py
new file mode 100644
index 0000000000000000000000000000000000000000..0640a256c16dd15aa630efc4cdcbe40dbfca7caf
--- /dev/null
+++ b/src/one-body-sample.py
@@ -0,0 +1,30 @@
+class neumannCondition:
+    def __call__(self, x):
+        # return 0
+        fst = 0.3
+        snd = 0.1
+        trd = 0.3
+        if x < 1.0/3:
+            return fst * x
+        elif x < 2.0/3:
+            return snd * (x - 1.0/3) + fst * 1.0/3
+        else:
+            return trd * (x - 2.0/3) + (fst + snd) * 1.0/3
+
+class dirichletCondition:
+    def __call__(self, x):
+        return 0
+        # fst = 0.3e-3
+        # snd = 0.1e-3
+        # trd = 0.3e-3
+        # if x < 1.0/3:
+        #     return fst * x
+        # elif x < 2.0/3:
+        #     return snd * (x - 1.0/3) + fst * 1.0/3
+        # else:
+        #     return trd * (x - 2.0/3) + (fst + snd) * 1.0/3
+
+Functions = {
+    'neumannCondition' : neumannCondition(),
+    'dirichletCondition' : dirichletCondition()
+}
diff --git a/src/one-body-sample_neumann.py b/src/one-body-sample_neumann.py
deleted file mode 100644
index a553be95e84495560c9f34c111ab85c3150c65da..0000000000000000000000000000000000000000
--- a/src/one-body-sample_neumann.py
+++ /dev/null
@@ -1,13 +0,0 @@
-class sampleFunction:
-    def __call__(self, x):
-        fst = 0.3
-        snd = 0.1
-        trd = 0.3
-        if x < 1.0/3:
-            return fst * x
-        elif x < 2.0/3:
-            return snd * (x - 1.0/3) + fst * 1.0/3
-        else:
-            return trd * (x - 2.0/3) + (fst + snd) * 1.0/3
-
-Functions = {'sampleFunction' : sampleFunction()}