Slope limiting for the St.Venant equations in DOLFINx

[Testosteronnie]

Hi, I am currently using a RK solver to solve the ST. Venant equations for a square channel given by

\frac{\partial A}{\partial t}+\frac{\partial Q}{\partial s}=0

and
\frac{\partial Q}{\partial t} + \frac{\partial}{\partial s}(\frac{Q^2}{A}+g A^2/(2 w)) + gA(S_f -S)=0
Here Q is the discharge, A the area of the flow, w the width of the channel, S the surface gradient and S_f the surface friction given by the Manning friction law.
To do this I use the discontinious Galerkin method (code given below). The code works fine, until shocks arise. Does anyone know of a good way to implement a slope limiter in DOLFINx?

from mpi4py import MPI
from petsc4py import PETSc
import numpy as np
import ufl
from dolfinx import default_real_type, fem, has_adios2, io, mesh
from dolfinx.plot import vtk_mesh
from dolfinx.fem.petsc import LinearProblem, NonlinearProblem
import time
import matplotlib.pyplot as plt
import pyvista

def butcher(order):
    """Butcher table data"""
    if order==2:
        # Explicit trapezium method
        b_runge=[1/2,1/2]
        c_runge=[0,1]
        a_runge=np.zeros((order,order),dtype=np.float64)
        a_runge[1,0]=1
    elif order==3:
        # Third-order RK3
        b_runge=[1/6,4/6,1/6]
        c_runge=[0,1/2,1]
        a_runge=np.zeros((order,order),dtype=np.float64)
        a_runge[1,0]=1/2
        a_runge[2,0]=-1
        a_runge[2,1]=2
    elif order==4:
        # "Classical" 4th-order Runge-Kutta method
        b_runge=[1/6,1/3,1/3,1/6]
        c_runge=np.array([0,1/2,1/2,1])
        a_runge=np.zeros((order,order),dtype=np.float64)
        a_runge[1,0]=1/2
        a_runge[2,1]=1/2
        a_runge[3,2]=1
    elif order==5:
        # Fehlberg 5th order
        b_runge=[16/135,0,6656/12825,28561/56430,-9/50,2/55]
        c_runge=[0,1/4,3/8,12/13,1,1/2]
        a_runge=np.zeros((order,order),dtype=np.float64)
        a_runge[1,0]=1/4
        a_runge[2,0:2]=[3/32,9/32]
        a_runge[3,0:3]=[1932/2197,-7200/2197,7296/2197]
        a_runge[4,0:4]=[439/216,-8,3680/513,-845/4104]
        a_runge[5,0:5]=[-8/27,2,-3544/2565,1859/4104,-11/40]
    else:
        print("INVALID ORDER")
        exit()

    return a_runge,b_runge,c_runge


# ============================================================
# PARAMETERS
# ============================================================

x_l=0.0
x_r=10.0

n_grid=100

t0=0.0
t_end=1.0
dt=1e-3

n_t=int((t_end-t0)/dt)

g=9.81
width=1.0

#The DG order
p=3

#inflow discharge
q_left = 1.0

#initial areas
A_left = 2.0
A_right = 1.0

#the RK order, which we take as the DG order plus one
n_RK=p+1
# ============================================================
# MESH
# ============================================================

domain=mesh.create_interval(
    MPI.COMM_WORLD,
    n_grid,
    [x_l,x_r])

#generate a vector valued function space that contains (A,q)
V=fem.functionspace(domain,("Discontinuous Lagrange",p,(2,)))

# ============================================================
# FACET TAGGING
# ============================================================

#the dimension of the boundaries
fdim=domain.topology.dim-1

#define the boundaries itself
def left_boundary(x):
    return np.isclose(x[0], x_l)

def right_boundary(x):
    return np.isclose(x[0], x_r)

#locate the boundaries
left_facets=mesh.locate_entities_boundary(domain,fdim,left_boundary)

right_facets=mesh.locate_entities_boundary(domain,fdim,right_boundary)

#store the indices in an array
facet_indices=np.hstack([left_facets,right_facets])

facet_markers=np.hstack([
    np.full(len(left_facets), 1, dtype=np.int32),
    np.full(len(right_facets), 2, dtype=np.int32)])

#sort the indices
sorted_facets=np.argsort(facet_indices)

#place tags on them
facet_tag = mesh.meshtags(
    domain,
    fdim,
    facet_indices[sorted_facets],
    facet_markers[sorted_facets]
)

#define ds, the integral over the boundary
ds=ufl.Measure("ds",domain=domain,subdomain_data=facet_tag)



# ============================================================
# FUNCTIONS
# ============================================================

u=fem.Function(V)
u_tmp=fem.Function(V)

v=ufl.TestFunction(V)
du=ufl.TrialFunction(V)

# ============================================================
# INITIAL CONDITION
# ============================================================

L=x_r-x_l

def initial_condition(x):
    #this generates a linear decrese/increase in A and sets q initially to one
    values = np.ones((2, x.shape[1]))
    # area
    values[0]=(A_left*(1.0-x[0]/L)+A_right*x[0]/L)
    # discharge initially one
    values[1]=1.0
    return values

#interpolate the initial conditions
u.interpolate(initial_condition)

# ============================================================
# INFLOW BC
# ============================================================

q_in=fem.Constant(domain,PETSc.ScalarType(q_left))


#DEFINE THE BED SLOPE

#the manning constant for this river
n_manning=0.01


def dz_ds(x):
    #this returns -e^-x, so the flow should be pushed to the right, as there is a negative height gradient
    return -1*ufl.exp(-x)


# ============================================================
# SHALLOW WATER EQUATIONS
# ============================================================

def velocity(U):
    #computes the velocity q/A, but prevents A getting negative
    A = ufl.max_value(U[0], 1e-6)
    return U[1]/A

def wave_speed(U):
    #calucaltes the wave speed, but once again prevents negative values
    A=ufl.max_value(U[0],1e-6)
    return ufl.sqrt(g*A/width)

def max_eigenvalue(U):
    #calculates the maximum eigenvalue for U
    return abs(velocity(U))+wave_speed(U)

def Fc(U):
    #The st venant equations, currently wihtout surface friction and gradients and for a square channel
    # so g*I_1= g*A*A/width
    A=ufl.max_value(U[0],1e-6)
    q=U[1]
    #this is the d/ds term
    momentum=(q*q/A+0.5*g*A*A/width)
    #now we add pressure correction because of surface height gradient
    #IS THIS EVEN NEEDED? NO
    #x=ufl.SpatialCoordinate(domain)
    #p_cor=g*ufl.exp(-x[0])
    #momentum+=p_cor
    #return the results
    return ufl.as_vector([q,momentum])

def source(u):
    #THIS ADDS THE SOURCE THERMS, S_SLOPE, S_F AND THE INFLOW
    A=ufl.max_value(u[0],1e-6)
    q=u[1]
    x=ufl.SpatialCoordinate(domain)
    #the resulting force caused by a terrain given by exp(-x)
    S_slope=ufl.as_vector([0,-g*A*(dz_ds(x[0]))])
    
    #gives the frictions from manningś law
    S_f=ufl.as_vector([0,g*n_manning**2*q*ufl.sqrt(q*q)/(A**(7/3))])


    #ADD RIVER INFLOW FROM THE SIDE
    return S_slope+S_f


# ============================================================
# INFLOW STATE
# ============================================================

def inflow_state(U):
    #this is the inflow on the left side, so only q_in is prescribed
    return ufl.as_vector([U[0],q_in])

# ============================================================
# RUSANOV FLUX
# ============================================================

n=ufl.FacetNormal(domain)

def rusanov_flux(Up, Um, normal):
    #compute the function at both sides
    Fp=Fc(Up)
    Fm=Fc(Um)
    #caluclate alpha
    alpha=ufl.max_value(max_eigenvalue(Up),max_eigenvalue(Um))
    #return the Rusanov flux
    return 0.5*((Fp+Fm)*normal[0]+alpha*(Up-Um))



# ============================================================
# DG FORMULATION
# ============================================================

# wehich is Fc*d/dx v integrated
volume_term=(ufl.inner(Fc(u_tmp),v.dx(0))*ufl.dx)

#the source term because of the friciton and slope, later also add inflow from the side!!!!!!
source_term=ufl.inner(source(u_tmp),v)*ufl.dx

#the interial numerical flux
interior_flux=(ufl.inner(rusanov_flux(u_tmp("+"),u_tmp("-"),n("+")),ufl.jump(v))*ufl.dS)

#caluclte the flux on the left boundary, where we have an inflow q, defined by the inflow_state
left_flux=(ufl.inner(rusanov_flux(u_tmp,inflow_state(u_tmp),n),v)*ds(1))

#on the right boundary the flow can flow out freely IS THIS CORRECT?????
right_flux=(ufl.inner(Fc(u_tmp)*n[0],v)*ds(2))

# Complete semidiscrete form
L_form=fem.form(volume_term+source_term-interior_flux-left_flux-right_flux)

# ============================================================
# MASS MATRIX
# ============================================================

#define the matrix a
a=ufl.inner(du,v)*ufl.dx

A_matrix=fem.petsc.assemble_matrix(fem.form(a))

A_matrix.assemble()

#define the solver
solver=PETSc.KSP().create(domain.comm)

solver.setOperators(A_matrix)

solver.setType("preonly")
solver.getPC().setType("lu")

# ============================================================
# RHS VECTOR
# ============================================================

b=fem.petsc.assemble_vector(L_form)

# ============================================================
# RK STORAGE
# ============================================================

#this just makes 5 functions just in case, this is kinda wacky, can just make less if not needed I guess
k1=fem.Function(V)
k2=fem.Function(V)
k3=fem.Function(V)
k4=fem.Function(V)
k5=fem.Function(V)

k_arr=[k1,k2,k3,k4,k5]
#this contains k for a certain stage
u_stage=fem.Function(V)

# ============================================================
# RHS ASSEMBLY
# ============================================================

def compute_rhs(Uin, kout):
    #set the values for the temprotray u that are used to compute this step
    u_tmp.x.array[:]=Uin.x.array[:]
    #update the vlaues
    u_tmp.x.scatter_forward()
    with b.localForm() as loc:
        loc.set(0)

    #assemble the vecot
    fem.petsc.assemble_vector(b,L_form)

    #update
    b.ghostUpdate(addv=PETSc.InsertMode.ADD,mode=PETSc.ScatterMode.REVERSE)
    
    #osolver
    solver.solve(b,kout.x.petsc_vec)

    kout.x.scatter_forward()
    #return kout so we return it again and put it back into the array
    return kout

# ============================================================
# INITIAL PLOT
# ============================================================


vals=u.x.array.reshape((-1,2))

A_vals=vals[:,0]
q_vals=vals[:,1]


plt.figure()
plt.title("Initial Area")
plt.plot(A_vals[::n_RK])

plt.figure()
plt.title("Initial Discharge")
plt.plot(q_vals[::n_RK])

plt.show()

# ============================================================
# RK4 TIME STEPPING
# ============================================================
#get the coefficients
a_runge,b_runge,c_runge=butcher(n_RK)


#set the initial time
t=t0


for step in range(n_t):
    #now we compute all the runge kutta steps
    #we first caluctelate ht einital stuff
    k_arr[0]=compute_rhs(u,k_arr[0])
    for i in range(1,n_RK):
        #this computes u_stage for the stage
        u_stage.x.array[:]=u.x.array
        for j in range(i):
            #now loop over the runge kutta coefficients
            u_stage.x.array[:]+=a_runge[i,j]*dt*k_arr[j].x.array[:]
        #nowe we hvae computed u stage, we can solve the equations
        k_arr[i]=compute_rhs(u_stage,k_arr[i])
    
    #now we got all the Kś and can thus update the solotion to the next tiem step
    #now update u_n
    for i in range(n_RK):
        u.x.array[:]+=dt*b_runge[i]*k_arr[i].x.array


    #TODO add protection against negative values

    u.x.scatter_forward()

    t+=dt

        


# ============================================================
# FINAL PLOT
# ============================================================

vals=u.x.array.reshape((-1,2))

A_vals=vals[:,0]
q_vals=vals[:,1]

plt.figure()
plt.title("Final Area")
plt.plot(A_vals[::n_RK])

plt.figure()
plt.title("Final Discharge")
plt.plot(q_vals[::n_RK])

plt.show()

[nate]

One of the few projects I know which has implemented slope limiters for DG methods with FEniCS is Ocellaris. See, e.g.,

Slope limiting for DG

Hi, has anyone attempted to use slope limiters for DG schemes in fenics? Right now I’m solving a 2D euler-like system. Is there any good references on slope limiters for DG methods or projects that use them? There’s a description of a slope limiting procedure in this reference but right now I wouldn’t know how to start coding it. Thanks in advance.

[Testosteronnie]

thanks! I will have a look

[Testosteronnie]

I got it working now, but massively decreased the efficiency of the code. I will try to optimize it and once it works decently upload it here so people can use it as a reference in case they ever need it.

[Testosteronnie]

This is a TVD limited code that is relatively quick, but only works for p=1:

from mpi4py import MPI
from petsc4py import PETSc
import numpy as np
import ufl
from dolfinx import default_real_type, fem, has_adios2, io, mesh
from dolfinx.plot import vtk_mesh
from dolfinx.fem.petsc import LinearProblem, NonlinearProblem
import time
import matplotlib.pyplot as plt
import pyvista

#TODO SSPRK
def butcher(order):
    """Butcher table data"""
    if order==2:
        # Explicit trapezium method
        b_runge=[1/2,1/2]
        c_runge=[0,1]
        a_runge=np.zeros((order,order),dtype=np.float64)
        a_runge[1,0]=1
    elif order==3:
        # Third-order RK3
        b_runge=[1/6,4/6,1/6]
        c_runge=[0,1/2,1]
        a_runge=np.zeros((order,order),dtype=np.float64)
        a_runge[1,0]=1/2
        a_runge[2,0]=-1
        a_runge[2,1]=2
    elif order==4:
        # "Classical" 4th-order Runge-Kutta method
        b_runge=[1/6,1/3,1/3,1/6]
        c_runge=np.array([0,1/2,1/2,1])
        a_runge=np.zeros((order,order),dtype=np.float64)
        a_runge[1,0]=1/2
        a_runge[2,1]=1/2
        a_runge[3,2]=1
    elif order==5:
        # Fehlberg 5th order
        b_runge=[16/135,0,6656/12825,28561/56430,-9/50,2/55]
        c_runge=[0,1/4,3/8,12/13,1,1/2]
        a_runge=np.zeros((order,order),dtype=np.float64)
        a_runge[1,0]=1/4
        a_runge[2,0:2]=[3/32,9/32]
        a_runge[3,0:3]=[1932/2197,-7200/2197,7296/2197]
        a_runge[4,0:4]=[439/216,-8,3680/513,-845/4104]
        a_runge[5,0:5]=[-8/27,2,-3544/2565,1859/4104,-11/40]
    else:
        print("INVALID ORDER")
        exit()

    return a_runge,b_runge,c_runge

# ============================================================
# PARAMETERS
# ============================================================
x_l=0.0
x_r=10.0

n_grid=100

dx=(x_r-x_l)/n_grid

t0=0.0
t_end=3.0
dt=1e-4

n_t=int((t_end-t0)/dt)

g=9.81
width=1.0

#The DG order
p=1

#TODO CURRENT CODE ONLY WORKS CORRECTLY FOR P=1
if p!=1:
    print("THE TVD SLOPE LIMITER DOESNT WORK CORRECTLY FOR THIS YET")
    exit()

#inflow discharge
q_left=1.0

#initial areas
A_left=1.
A_right=1.

#the RK order, which we take as the DG order plus one
n_RK=p+1

#ARTIFICIAL VISCOSITY

artificial_viscosity=False

nu_visc=1e-3


# ============================================================
# MESH
# ============================================================

domain=mesh.create_interval(
    MPI.COMM_WORLD,
    n_grid,
    [x_l,x_r])

#generate a vector valued function space that contains (A,q)
V=fem.functionspace(domain,("Discontinuous Lagrange",p,(2,)))

#the dimension of the boundaries
fdim=domain.topology.dim-1

#define the boundaries itself
def left_boundary(x):
    return np.isclose(x[0], x_l)

def right_boundary(x):
    return np.isclose(x[0], x_r)

#locate the boundaries
left_facets=mesh.locate_entities_boundary(domain,fdim,left_boundary)

right_facets=mesh.locate_entities_boundary(domain,fdim,right_boundary)

#store the indices in an array
facet_indices=np.hstack([left_facets,right_facets])

facet_markers=np.hstack([
    np.full(len(left_facets), 1, dtype=np.int32),
    np.full(len(right_facets), 2, dtype=np.int32)])

#sort the indices
sorted_facets=np.argsort(facet_indices)

#place tags on them
facet_tag = mesh.meshtags(
    domain,
    fdim,
    facet_indices[sorted_facets],
    facet_markers[sorted_facets]
)

#define ds, the integral over the boundary
ds=ufl.Measure("ds",domain=domain,subdomain_data=facet_tag)

# ============================================================
# FUNCTIONS
# ============================================================

u=fem.Function(V)
u_tmp=fem.Function(V)

v=ufl.TestFunction(V)
du=ufl.TrialFunction(V)

# ============================================================
# INITIAL CONDITION
# ============================================================

L=x_r-x_l

def initial_condition(x):
    #this generates a linear decrese/increase in A and sets q initially to one
    #now it just sets everything to one
    values = np.ones((2, x.shape[1]))
    # area
    values[0]=(A_left*(1.0-x[0]/L)+A_right*x[0]/L)
    # discharge initially one
    #values[1]=1.0
    return values

#interpolate the initial conditions
u.interpolate(initial_condition)

# ============================================================
# INFLOW BC
# ============================================================

q_in=fem.Constant(domain,PETSc.ScalarType(q_left))

A_in=fem.Constant(domain,PETSc.ScalarType(A_left))

#DEFINE THE BED SLOPE

#the manning constant for this river
n_manning=0.01

# def z(x):
#     #the height itself, needed for the pressure correction
#     return ufl.exp(-x)

def dz_ds(x):
    #this returns -e^-x, so the flow should be pushed to the right, as there is a negative height gradient
    return -0.3*ufl.exp(-x)


# ============================================================
# SHALLOW WATER EQUATIONS
# ============================================================

def velocity(U):
    #computes the velocity q/A, but prevents A getting too small
    A = ufl.max_value(U[0], 1e-6)
    return U[1]/A

def wave_speed(U):
    #calucaltes the wave speed, but once again prevents too small values, needs to be made cleaner
    A=ufl.max_value(U[0],1e-6)
    return ufl.sqrt(g*A/width)

def max_eigenvalue(U):
    #calculates the maximum eigenvalue for U
    return abs(velocity(U))+wave_speed(U)

def Fc(U):
    #The st venant equations
    #square channel so g*I_1= g*A*A/width
    A=ufl.max_value(U[0],1e-6)
    q=U[1]
    #this is the d/ds term
    momentum=(q*q/A+0.5*g*A*A/width)
    #return the results
    return ufl.as_vector([q,momentum])

def source(u):
    #THIS ADDS THE SOURCE THERMS, S_SLOPE, S_F AND THE INFLOW

    #TODO ADD INFLOW FROM THE SIDE

    A=ufl.max_value(u[0],1e-6)
    q=u[1]
    x=ufl.SpatialCoordinate(domain)
    #the resulting force caused by a terrain given by exp(-x)
    S_slope=ufl.as_vector([0,-g*A*(dz_ds(x[0]))])
    
    #gives the frictions from manningś law
    S_f=ufl.as_vector([0,-g*n_manning**2*q*ufl.sqrt(q*q)/(A**(7/3))])

    return S_slope+S_f

# ============================================================
# INFLOW STATE
# ============================================================

def inflow_state(U):
    #this is the inflow on the left side, so only q_in is prescribed
    return ufl.as_vector([U[0],q_in])

# ============================================================
# RUSANOV FLUX
# ============================================================

n=ufl.FacetNormal(domain)

def rusanov_flux(Up, Um, normal):
    #compute the function at both sides
    Fp=Fc(Up)
    Fm=Fc(Um)
    #caluclate alpha
    alpha=ufl.max_value(max_eigenvalue(Up),max_eigenvalue(Um))
    #return the Rusanov flux
    return 0.5*((Fp+Fm)*normal[0]+alpha*(Up-Um))


# ============================================================
# DG FORMULATION
# ============================================================

# wehich is Fc*d/dx v integrated
volume_term=(ufl.inner(Fc(u_tmp),v.dx(0))*ufl.dx)

#the source term because of the friciton and slope, later also add inflow from the side!!!!!!
source_term=ufl.inner(source(u_tmp),v)*ufl.dx

#the interial numerical flux
interior_flux=(ufl.inner(rusanov_flux(u_tmp("+"),u_tmp("-"),n("+")),ufl.jump(v))*ufl.dS)

#caluclte the flux on the left boundary, where we have an inflow q, defined by the inflow_state
left_flux=(ufl.inner(rusanov_flux(u_tmp,inflow_state(u_tmp),n),v)*ds(1))

#on the right boundary the flow can flow out freely IS THIS CORRECT?????
right_flux=(ufl.inner(Fc(u_tmp)*n[0],v)*ds(2))

#ADD ARTIVICIAL VISCOSITY
if artificial_viscosity:
    #the size of the cell
    h=ufl.CellDiameter(domain)
    #penalty parameter
    alpha=10.0
    #we only apply the viscosity to the q, so in the momentum equation
    q_trial=u_tmp[1]
    q_test=v[1]
    #compute the integral nu* d/dx q d/dx v
    visc_volume=(nu_visc*q_trial.dx(0)*q_test.dx(0)* ufl.dx)
    #TODO THIS IS STILL NOT FULLY RIGHT
    #compute the consistensy term, which is the sum over the integrals
    visc_consistency=(-nu_visc*ufl.avg(q_trial.dx(0))*ufl.jump(q_test)*ufl.dS
        -nu_visc*ufl.avg(q_test.dx(0))*ufl.jump(q_trial)*ufl.dS)
    #and introduce the penalty for the jumps
    jump_penalty=(nu_visc*alpha/ufl.avg(h)*ufl.jump(q_trial)*ufl.jump(q_test)*ufl.dS)
    #add them to the function
    volume_term-=visc_volume
    volume_term+=visc_consistency
    volume_term+=jump_penalty


# Complete semidiscrete form
L_form=fem.form(volume_term+source_term-interior_flux-left_flux-right_flux)

# ============================================================
# MASS MATRIX
# ============================================================

#define the matrix a
a=ufl.inner(du,v)*ufl.dx

A_matrix=fem.petsc.assemble_matrix(fem.form(a))

A_matrix.assemble()

#define the solver
solver=PETSc.KSP().create(domain.comm)

solver.setOperators(A_matrix)

solver.setType("preonly")
solver.getPC().setType("lu")


# ============================================================
# RHS VECTOR
# ============================================================

b=fem.petsc.assemble_vector(L_form)

# ============================================================
# RK STORAGE
# ============================================================

#this just makes 5 functions just in case, this is kinda wacky, can just make less if not needed I guess
k1=fem.Function(V)
k2=fem.Function(V)
k3=fem.Function(V)
k4=fem.Function(V)
k5=fem.Function(V)

k_arr=[k1,k2,k3,k4,k5]
#this contains k for a certain stage
u_stage=fem.Function(V)

#for plotting
V_plot=fem.functionspace(domain,("CG",1,(2,)))

u_plot=fem.Function(V_plot)

# ============================================================
# RHS ASSEMBLY
# ============================================================

def compute_rhs(Uin, kout):
    #set the values for the temprotray u that are used to compute this step
    u_tmp.x.array[:]=Uin.x.array[:]
    #update the vlaues
    u_tmp.x.scatter_forward()
    with b.localForm() as loc:
        loc.set(0)

    #assemble the vecot
    fem.petsc.assemble_vector(b,L_form)

    #update
    b.ghostUpdate(addv=PETSc.InsertMode.ADD,mode=PETSc.ScatterMode.REVERSE)
    
    #osolver
    solver.solve(b,kout.x.petsc_vec)

    kout.x.scatter_forward()
    #return kout so we return it again and put it back into the array
    return kout


#DEFINE THE SLOPE LIMITER
#the amount of degrees of freedom per cell
nodfs_cell=(p+1)
#the number of cells
ncells=domain.topology.index_map(domain.topology.dim).size_local

def sigma_minmod(slope_left,slope_right):
    #eq 4.26
    sigma=(np.sign(slope_left)+np.sign(slope_right))/2*np.minimum(np.abs(slope_left),np.abs(slope_right))
    return sigma


#this is the array with x values, only needs to be created once, needed for equation 4.20
x_factor=np.linspace(-0.5,0.5,nodfs_cell)*dx

def slope_limiter(u):
    #applies slope limiter according to equation 4.26 to the problem
    #we generate two arrays, one with the area and one with the discharge
    #CURRENTLY ONLY WORKS FOR P=1
    arr=u.x.array
    #print(arr)
    A_arr=arr[::2].copy()
    q_arr=arr[1::2].copy()
    #print(A_arr)
    #SET TO ZERO WHENEVER SOMETHING IS TOO SMAL, we first order them per cell
    A_cell=A_arr.reshape(-1,nodfs_cell)
    q_cell=q_arr.reshape(-1,nodfs_cell)
    #apply a mask, this sets the discharge to zero whenever there is a area below the threshold of 1e-6
    mask=np.any(A_cell<1e-6,axis=1)
    q_cell[mask]=0.0
    #now we also set the water height to very low for these cases
    A_cell[mask]=1e-6
    #we now have two arrays, both for A and for q, we now compute the average for each cell
    #TODO THIS IS NOT GENRALLY TRUE
    A_avg=A_cell.mean(axis=1)
    q_avg=q_cell.reshape(-1,nodfs_cell).mean(axis=1)
    #now get the differences of the mean between the cells
    #TODO FLEXIBLE DX, but then it just needs to be an array so not that difficult
    dA=np.diff(A_avg)/dx
    dq=np.diff(q_avg)/dx
    #now get the values for sigma for both
    dAL=dA[:-1]
    dAR=dA[1:]
    dqL=dq[:-1]
    dqR=dq[1:]
    sigma_A=sigma_minmod(dAL,dAR)
    sigma_q=sigma_minmod(dqL,dqR)
    #calculate the TVD limited values
    A_lim=A_avg[1:-1,None]+sigma_A[:,None]*x_factor[None,:]
    q_lim=q_avg[1:-1,None]+sigma_q[:,None]*x_factor[None,:]
    #updated the values per cell
    A_cell[1:-1,:]=A_lim
    q_cell[1:-1,:]=q_lim
    #ENFORCE THE MASK AGAIN AFTER LIMITING
    q_cell[mask]=0.0
    #now we also set the water height to very low for these cases
    A_cell[mask]=1e-6
    #now flatten them again
    arr[::2]=A_cell.flatten()
    arr[1::2]=q_cell.flatten()

    #now scatter forward
    u.x.scatter_forward()


# ============================================================
# INITIAL PLOT
# ============================================================

vals=u.x.array.reshape((-1,2)).copy()

A_vals_init=vals[:,0]
q_vals_init=vals[:,1]


# ============================================================
# RK4 TIME STEPPING
# ============================================================
#get the coefficients
a_runge,b_runge,c_runge=butcher(n_RK)


#set the initial time
t=t0


for step in range(n_t):
    #now we compute all the runge kutta steps
    #we first caluctelate ht einital stuff
    k_arr[0]=compute_rhs(u,k_arr[0])
    for i in range(1,n_RK):
        #this computes u_stage for the stage
        u_stage.x.array[:]=u.x.array
        for j in range(i):
            #now loop over the runge kutta coefficients
            u_stage.x.array[:]+=a_runge[i,j]*dt*k_arr[j].x.array[:]
        #nowe we hvae computed u stage, we can solve the equations

        #apply slope limiting
        slope_limiter(u_stage)

        k_arr[i]=compute_rhs(u_stage,k_arr[i])
    
    #now we got all the Kś and can thus update the solotion to the next tiem step
    #now update u_n
    for i in range(n_RK):
        u.x.array[:]+=dt*b_runge[i]*k_arr[i].x.array


    #TODO add protection against negative values
    #tvd_limiter(u)
    u.x.scatter_forward()

    #and also slope limiting here

    slope_limiter(u)

    t+=dt

#     if step%1000==0 and step>5000:
#         #make a plot to see whats going wrong
#         u_plot.interpolate(u)

#         vals = u_plot.x.array.reshape((-1,2))


#         A_vals=vals[:,0]
#         q_vals=vals[:,1]
#         plt.figure()
#         plt.title("A nt="+str(step))
#         plt.plot(A_vals)

#         plt.figure()
#         plt.title("Dischargent="+str(step))
#         plt.plot(q_vals)

#         plt.figure()
#         plt.title("Final velocities"+str(step))
#         plt.plot(q_vals/A_vals)
        

        
# plt.show()

# ============================================================
# FINAL PLOT
# ============================================================


u_plot.interpolate(u)

vals = u_plot.x.array.reshape((-1,2))


A_vals=vals[:,0]
q_vals=vals[:,1]



# plt.figure()
# plt.title("Initial Area")
# plt.plot(A_vals_init[::n_RK])

# plt.figure()
# plt.title("Initial Discharge")
# plt.plot(q_vals_init[::n_RK])

if artificial_viscosity:
    print("CURRENT SIMULATIONS DONE WITH ARTIFICIAL VISCOSITY")

plt.figure()
plt.title("Final Area")
plt.plot(A_vals)

plt.figure()
plt.title("Final Discharge")
plt.plot(q_vals)

plt.figure()
plt.title("Final velocities")
plt.plot(q_vals/A_vals)

print(len(u.x.array))
plt.show()