Index: /inundation/ga/storm_surge/pyvolution/config.py
===================================================================
--- /inundation/ga/storm_surge/pyvolution/config.py	(revision 220)
+++ /inundation/ga/storm_surge/pyvolution/config.py	(revision 221)
@@ -59,6 +59,6 @@
 #use_extensions = False   #Do not use C-extensions
 
-use_psyco = True  #Use psyco optimisations
-#use_psyco = False  #Do not use psyco optimisations
+#use_psyco = True  #Use psyco optimisations
+use_psyco = False  #Do not use psyco optimisations
 
 
Index: /inundation/ga/storm_surge/pyvolution/quantity.py
===================================================================
--- /inundation/ga/storm_surge/pyvolution/quantity.py	(revision 220)
+++ /inundation/ga/storm_surge/pyvolution/quantity.py	(revision 221)
@@ -78,5 +78,5 @@
             v2 = self.vertex_values[i, 2]
             
-            self.centroid_values[i] = (v0 + v1 + v2)/3
+            self.centroid_values[i] = (v0 + v1 + v2)/3.0
             
             self.edge_values[i, 0] = 0.5*(v1 + v2)
@@ -139,11 +139,13 @@
                 #Special case where we have only one neighbouring volume.
 
-                k0 = k  #Self
                 #Find index of the one neighbour
-                for k1 in self.domain.neighbours[k,:]:
-                    if k1 >= 0:
+                for k0 in self.domain.neighbours[k,:]:
+                    if k0 >= 0:
                         break
-                assert k1 != k0
-                assert k1 >= 0
+
+                assert k0 != k
+                assert k0 >= 0
+
+                k1 = k  #Self
 
                 #Get data
@@ -180,4 +182,10 @@
                 x2, y2 = self.domain.centroids[k2] #V2 centroid
 
+                #if k < 4:
+                #    print 'k, k0, k1, k2 = %d, %d, %d, %d' %(k, k0, k1, k2)
+                #    #print 'k = %d: (%f, %f), (%f, %f), (%f, %f)'\
+                #    #      %(k, x0, y0, x1, y1, x2, y2)
+                #    print 'k = %d: (%f, %f, %f)' %(k, q0, q1, q2)   
+                
                 #Gradient
                 a[k], b[k] = gradient(x0, y0, x1, y1, x2, y2, q0, q1, q2)
@@ -266,4 +274,6 @@
         
         a, b = self.compute_gradients()
+        ##for k in range(4):
+        ##    print 'Gradients %d: %16.12f, %16.12f' %(k, a[k], b[k])
 
         V = self.domain.get_vertex_coordinates()
Index: /inundation/ga/storm_surge/pyvolution/shallow_water.py
===================================================================
--- /inundation/ga/storm_surge/pyvolution/shallow_water.py	(revision 220)
+++ /inundation/ga/storm_surge/pyvolution/shallow_water.py	(revision 221)
@@ -272,9 +272,11 @@
 
 def distribute_to_vertices_and_edges(domain):
-    
-    protect_against_infinitesimal_heights_centroid(domain)
+
+    #print 'Calling distrib within SW'
     if domain.order == 1:
+        protect_against_infinitesimal_heights_centroid(domain)        
         extrapolate_first_order(domain)
     elif domain.order == 2:
+        protect_against_negative_heights_centroid(domain)        
         extrapolate_second_order(domain)        
     else:
@@ -291,4 +293,5 @@
     minimum allowed height
     """
+    #FIXME: Used only in first order
     
     #Water levels at centroids
@@ -316,4 +319,31 @@
                 
 
+def protect_against_negative_heights_centroid(domain):
+    """Adjust height and momentum at centroid if height is less than zero
+    """
+
+    #FIXME: USed only in second order
+    #Water levels at centroids
+    wc = domain.quantities['level'].centroid_values
+
+    #Bed elevations at centroids
+    zc = domain.quantities['elevation'].centroid_values
+
+    #Water depths at centroids    
+    hc = wc - zc
+
+    #Momentums at centroids
+    xmomc = domain.quantities['xmomentum'].centroid_values
+    ymomc = domain.quantities['ymomentum'].centroid_values        
+
+    for k in range(domain.number_of_elements):
+        #Protect against infinitesimal heights and high velocities
+        if hc[k] < 0.0:
+            #Control level and height
+            wc[k] = zc[k]; hc[k] = 0.0
+
+            #Control momentum
+            xmomc[k] = ymomc[k] = 0.0
+                
 
 
Index: /inundation/ga/storm_surge/pyvolution/test_shallow_water.py
===================================================================
--- /inundation/ga/storm_surge/pyvolution/test_shallow_water.py	(revision 220)
+++ /inundation/ga/storm_surge/pyvolution/test_shallow_water.py	(revision 221)
@@ -960,4 +960,222 @@
         #assert allclose(domain.max_timestep, 0.0396825396825) 
         #print domain.quantities['level'].centroid_values
+
+
+    def test_flatbed_first_order(self):
+        from mesh_factory import rectangular
+        from shallow_water import Domain,\
+             Reflective_boundary, Dirichlet_boundary
+
+        from Numeric import array
+
+        #Create basic mesh
+        N = 8
+        points, vertices, boundary = rectangular(N, N)
+
+        #Create shallow water domain
+        domain = Domain(points, vertices, boundary)
+        domain.smooth = False
+        domain.visualise = False
+        domain.default_order=1
+
+        # Boundary conditions
+        Br = Reflective_boundary(domain)
+        Bd = Dirichlet_boundary([0.2,0.,0.])
+        
+        domain.set_boundary({'left': Bd, 'right': Br, 'top': Br, 'bottom': Br})
+        domain.check_integrity()
+
+  
+        #Evolution
+        for t in domain.evolve(yieldstep = 0.02, finaltime = 0.5):
+            pass
+            #domain.write_time()
+
+        assert allclose(domain.min_timestep, 0.0140413643926)
+        assert allclose(domain.max_timestep, 0.0140947355753)
+
+        for i in range(3):
+            assert allclose(domain.quantities['level'].edge_values[:4,i],
+                            [0.10730244,0.12337617,0.11200126,0.12605666])
+
+            assert allclose(domain.quantities['xmomentum'].edge_values[:4,i],
+                            [0.07610894,0.06901572,0.07284461,0.06819712])
+
+            assert allclose(domain.quantities['ymomentum'].edge_values[:4,i],
+                            [-0.0060238, -0.00157404, -0.00309633, -0.0001637])                         
+        
+    def test_flatbed_second_order(self):
+        from mesh_factory import rectangular
+        from shallow_water import Domain,\
+             Reflective_boundary, Dirichlet_boundary
+
+        from Numeric import array
+
+        #Create basic mesh
+        N = 8
+        points, vertices, boundary = rectangular(N, N)
+
+        #Create shallow water domain
+        domain = Domain(points, vertices, boundary)
+        domain.smooth = False
+        domain.visualise = False
+        domain.default_order=2
+
+        # Boundary conditions
+        Br = Reflective_boundary(domain)
+        Bd = Dirichlet_boundary([0.2,0.,0.])
+
+        domain.set_boundary({'left': Bd, 'right': Br, 'top': Br, 'bottom': Br})
+        domain.check_integrity()
+  
+        #Evolution
+        for t in domain.evolve(yieldstep = 0.01, finaltime = 0.03):
+            pass
+
+
+        assert allclose(domain.min_timestep, 0.0210448446782)
+        assert allclose(domain.max_timestep, 0.0210448446782) 
+
+        #print domain.quantities['level'].vertex_values[:4,0]
+        #print domain.quantities['xmomentum'].vertex_values[:4,0]
+        #print domain.quantities['ymomentum'].vertex_values[:4,0]
+        
+        assert allclose(domain.quantities['level'].vertex_values[:4,0],
+                        [0.00101913,0.05352143,0.00104852,0.05354394]) 
+        assert allclose(domain.quantities['xmomentum'].vertex_values[:4,0],
+                        [0.00090581,0.03685719,0.00088303,0.03687313])
+        assert allclose(domain.quantities['ymomentum'].vertex_values[:4,0],
+                        [-0.00139463,0.0006156,-0.00060364,0.00061827])
+
+
+        
+    def test_flatbed_second_order_distribute(self):
+        #Use real data from pyvolution 2
+        #painfully setup and extracted.
+        from mesh_factory import rectangular
+        from shallow_water import Domain,\
+             Reflective_boundary, Dirichlet_boundary
+
+        from Numeric import array
+
+        #Create basic mesh
+        N = 8
+        points, vertices, boundary = rectangular(N, N)
+
+        #Create shallow water domain
+        domain = Domain(points, vertices, boundary)
+        domain.smooth = False
+        domain.visualise = False
+        domain.default_order=domain.order=2
+
+        # Boundary conditions
+        Br = Reflective_boundary(domain)
+        Bd = Dirichlet_boundary([0.2,0.,0.])
+
+        domain.set_boundary({'left': Bd, 'right': Br, 'top': Br, 'bottom': Br})
+        domain.check_integrity()
+
+
+
+        for V in [False, True]:
+            if V:
+                #Set centroids as if system had been evolved            
+                L = zeros(2*N*N, Float)
+                L[:32] = [7.21205592e-003, 5.35214298e-002, 1.00910824e-002,
+                          5.35439433e-002, 1.00910824e-002, 5.35439433e-002,
+                          1.00910824e-002, 5.35439433e-002, 1.00910824e-002,
+                          5.35439433e-002, 1.00910824e-002, 5.35439433e-002,
+                          1.00910824e-002, 5.35393928e-002, 1.02344264e-002,
+                          5.59605058e-002, 0.00000000e+000, 3.31027800e-004,
+                          0.00000000e+000, 4.37962142e-005, 0.00000000e+000,
+                          4.37962142e-005, 0.00000000e+000, 4.37962142e-005,
+                          0.00000000e+000, 4.37962142e-005, 0.00000000e+000,
+                          4.37962142e-005, 0.00000000e+000, 4.37962142e-005,
+                          0.00000000e+000, 5.57305948e-005]
+        
+                X = zeros(2*N*N, Float)
+                X[:32] = [6.48351607e-003, 3.68571894e-002, 8.50733285e-003,
+                          3.68731327e-002, 8.50733285e-003, 3.68731327e-002,
+                          8.50733285e-003, 3.68731327e-002, 8.50733285e-003,
+                          3.68731327e-002, 8.50733285e-003, 3.68731327e-002,
+                          8.50733285e-003, 3.68693861e-002, 8.65220973e-003,
+                          3.85055387e-002, 0.00000000e+000, 2.86060840e-004,
+                          0.00000000e+000, 3.58905503e-005, 0.00000000e+000,
+                          3.58905503e-005, 0.00000000e+000, 3.58905503e-005,
+                          0.00000000e+000, 3.58905503e-005, 0.00000000e+000,
+                          3.58905503e-005, 0.00000000e+000, 3.58905503e-005,
+                          0.00000000e+000, 4.57662812e-005]
+
+                Y = zeros(2*N*N, Float)
+                Y[:32]=[-1.39463104e-003, 6.15600298e-004, -6.03637382e-004,
+                        6.18272251e-004, -6.03637382e-004, 6.18272251e-004,
+                        -6.03637382e-004, 6.18272251e-004, -6.03637382e-004,
+                        6.18272251e-004, -6.03637382e-004, 6.18272251e-004,
+                        -6.03637382e-004, 6.18599320e-004, -6.74622797e-004,
+                        -1.48934756e-004, 0.00000000e+000, -5.35079969e-005,
+                        0.00000000e+000, -2.57264987e-005, 0.00000000e+000,
+                        -2.57264987e-005, 0.00000000e+000, -2.57264987e-005,
+                        0.00000000e+000, -2.57264987e-005, 0.00000000e+000,
+                        -2.57264987e-005, 0.00000000e+000, -2.57264987e-005,
+                        0.00000000e+000, -2.57635178e-005]
+
+        
+                domain.set_quantity('level', L, 'centroids')
+                domain.set_quantity('xmomentum', X, 'centroids')
+                domain.set_quantity('ymomentum', Y, 'centroids')
+        
+                domain.check_integrity()
+            else:    
+                #Evolution
+                for t in domain.evolve(yieldstep = 0.01, finaltime = 0.03):
+                    pass
+                assert allclose(domain.min_timestep, 0.0210448446782)
+                assert allclose(domain.max_timestep, 0.0210448446782)
+        
+
+            #Centroids were correct but not vertices.
+            #Hence the check of distribute below.
+            assert allclose(domain.quantities['level'].centroid_values[:4],
+                            [0.00721206,0.05352143,0.01009108,0.05354394])
+
+            assert allclose(domain.quantities['xmomentum'].centroid_values[:4],
+                            [0.00648352,0.03685719,0.00850733,0.03687313])
+
+            assert allclose(domain.quantities['ymomentum'].centroid_values[:4],
+                            [-0.00139463,0.0006156,-0.00060364,0.00061827])
+
+            #print 'C17=', domain.quantities['xmomentum'].centroid_values[17]
+            #print 'C19=', domain.quantities['xmomentum'].centroid_values[19]
+
+            assert allclose(domain.quantities['xmomentum'].centroid_values[17],
+                            0.00028606084)
+
+            import copy
+            XX = copy.copy(domain.quantities['xmomentum'].centroid_values)
+            assert allclose(domain.quantities['xmomentum'].centroid_values, XX)
+        
+            domain.distribute_to_vertices_and_edges()
+
+            assert allclose(domain.quantities['xmomentum'].centroid_values, XX)
+
+            assert allclose(domain.quantities['xmomentum'].centroid_values[17],
+                            0.00028606084)
+
+
+            assert allclose(domain.quantities['level'].vertex_values[:4,0],
+                            [0.00101913,0.05352143,0.00104852,0.05354394])
+
+            assert allclose(domain.quantities['ymomentum'].vertex_values[:4,0],
+                            [-0.00139463,0.0006156,-0.00060364,0.00061827])
+
+            #This was the culprit. First triangles vertex 0 had an
+            #x-momentum of 0.0064835 instead of 0.00090581 and
+            #third triangle had 0.00850733 instead of 0.00088303
+            #print domain.quantities['xmomentum'].vertex_values[:4,0]
+        
+            assert allclose(domain.quantities['xmomentum'].vertex_values[:4,0],
+                            [0.00090581,0.03685719,0.00088303,0.03687313])        
+
+
 
         
