Index: wflow-py/wflow/wflow_sbm.py
===================================================================
diff -u -r1a45d207f12a3a68e619b5f40a0c6c4c42841501 -ra26946a23ad13222449074a7ec217e0c1320b4b4
--- wflow-py/wflow/wflow_sbm.py	(.../wflow_sbm.py)	(revision 1a45d207f12a3a68e619b5f40a0c6c4c42841501)
+++ wflow-py/wflow/wflow_sbm.py	(.../wflow_sbm.py)	(revision a26946a23ad13222449074a7ec217e0c1320b4b4)
@@ -1296,7 +1296,10 @@
         self.CellStorage = self.UStoreDepth + self.FirstZoneDepth
         self.DeltaStorage = self.CellStorage - self.OrgStorage
         OutFlow = self.FirstZoneFlux
-        CellInFlow = upstream(self.TopoLdd, scalar(self.FirstZoneFlux))
+        if self.waterdem:
+            CellInFlow = upstream(self.waterLdd, scalar(self.FirstZoneFlux))
+        else:
+            CellInFlow = upstream(self.TopoLdd, scalar(self.FirstZoneFlux))
 
         self.CumOutFlow = self.CumOutFlow + OutFlow
         self.CumActInfilt = self.CumActInfilt + self.ActInfilt
@@ -1314,10 +1317,9 @@
         self.CumExfiltWater = self.CumExfiltWater + self.ExfiltWater
 
 
-        self.SoilWatbal = self.ActInfilt - self.Transpiration - self.soilevap  - self.ExfiltWater  -\
-                      self.SubCellGWRunoff + self.reinfiltwater - \
-                      self.DeltaStorage - \
-                      self.FirstZoneFlux + CellInFlow
+        self.SoilWatbal = self.ActInfilt + self.reinfiltwater  + CellInFlow - self.Transpiration - self.soilevap  -\
+                          self.ExfiltWater  -  self.SubCellGWRunoff  - self.DeltaStorage -\
+                          self.FirstZoneFlux
 
         self.InterceptionWatBal = self.PrecipitationPlusMelt - self.Interception -self.StemFlow - self.ThroughFall -\
                              (self.OldCanopyStorage - self.CanopyStorage)