Height data as input into DOLFINx

I’m currently working on a DG model to solve the shallow water equations, in both the dynamical wave form and the kinematic wave approximation. In the kinematic wave approximation the equations are simplified to:

\frac{\partial h}{\partial t}+ \nabla \cdot (h \mathbf{u}) = r^*
\mathbf{u}=\frac{\mathbf{S}_0 h^{2/3}}{n_{m}\sqrt{|\mathbf{S}_0|}}

with r^* the net influx from rain, infiltration and snow melt etc., n_m the manning friction constant and \mathbf{S}_0 the surface gradient.

I would like to have an array with elevation data (and the same for all the other data) as input, and then solve the equations for this. Currently I do this:

def s0sqrts0(x):
    #THIS SHOULD RETURN S0/ABS(SQRTS0)
    return ufl.as_vector([np.sqrt(0.005)*(0.5*x[0]+1.),-np.sqrt(0.005)])


def Fc(U):
    #in the kinematic wave approximation Q=sqrtabs(so)H^5/3/n
    return s0sqrts0(x)*U**(5/3)/n_manning


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

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

where I have hard coded \frac{\mathbf{S}_0 }{\sqrt{|\mathbf{S}_0|}}. But I would prefer this to be some input array instead, with heights etc. How can this be effectively done in DOLFINx? Should I create an other function on which I interpolate the height data and have this as input into the function? Or are there more effective ways to do this?

Should I create an other function on which I interpolate the height data and have this as input into the function?

This is likely the most effective way. You can control the numerical discretsiation of your height data by choosing an appropriate function space. A simple first step is choosing DG0, which will give you piecewise constant values of height for each cell in your mesh.