Skip to content

Instantly share code, notes, and snippets.

@finsberg
Last active April 6, 2022 12:10
Show Gist options
  • Select an option

  • Save finsberg/f9dcaa8965de675219b9e5472bd4964c to your computer and use it in GitHub Desktop.

Select an option

Save finsberg/f9dcaa8965de675219b9e5472bd4964c to your computer and use it in GitHub Desktop.
Example of solving and ODE and a PDE independently at the same time using FEniCS
import matplotlib.pyplot as plt
import dolfin
from pulse import Constant
def boundary(x):
return x[0] < dolfin.DOLFIN_EPS or x[0] > 1.0 - dolfin.DOLFIN_EPS
def rhs(x, y):
"""Right hand side of a simple oscillator"""
return dolfin.as_vector([-y, x])
def poisson(u, v):
f = dolfin.Expression(
"10*exp(-(pow(x[0] - 0.5, 2) + pow(x[1] - 0.5, 2)) / 0.02)", degree=2
)
a = dolfin.inner(dolfin.grad(u), dolfin.grad(v)) * dolfin.dx
L = f * v * dolfin.dx
return a - L
mesh = dolfin.UnitSquareMesh(24, 24)
# Make a mixed function space with first subspace being the PDE
# and the next subspaces being the states in the ODE
P_pde = dolfin.FiniteElement("Lagrange", mesh.ufl_cell(), 1)
P_ode = dolfin.VectorElement("Lagrange", mesh.ufl_cell(), 1, dim=2)
V = dolfin.FunctionSpace(mesh, dolfin.MixedElement([P_pde, P_ode]))
# Define boundary condition for the PDE
bc = (dolfin.DirichletBC(V.sub(0), dolfin.Constant(0.0), boundary),)
state = dolfin.Function(V)
state_test = dolfin.TestFunction(V)
u, s = dolfin.split(state)
v, w = dolfin.split(state_test)
s1, s2 = dolfin.split(s)
w1, w2 = dolfin.split(w)
# Keep track of the ODE solution in the previous iteration
V_s1 = V.sub(1).sub(0).collapse()
V_s2 = V.sub(1).sub(1).collapse()
s1_ = dolfin.Function(V_s1)
s2_ = dolfin.Function(V_s2)
s1_assigner = dolfin.FunctionAssigner(V_s1, V.sub(1).sub(0))
s2_assigner = dolfin.FunctionAssigner(V_s2, V.sub(1).sub(1))
# Set initial conditions
s1_.assign(dolfin.Constant(1.0))
s2_.assign(dolfin.Constant(0.0))
# Residual for the PDE
R_pde = poisson(u, v)
#
dt = 0.01
T = 7.0
Dt = Constant(dt)
# Get the right hand side of the ODE
F = rhs(s1, s2)
# The derivative of the states (forward difference)
s1_t = (s1 - s1_) / Dt
s2_t = (s2 - s2_) / Dt
# Residual of the ODE
R_ode = (s1_t - F[0]) * w1 * dolfin.dx + (s2_t - F[1]) * w2 * dolfin.dx
# Total residual
R = R_pde + R_ode
# Keep track of the ode solution (just pick the first dof)
xs = [s1_.vector().get_local()[0]]
ys = [s2_.vector().get_local()[0]]
t = 0.0
ts = [t]
while t < T:
t += float(dt)
ts.append(t)
dolfin.solve(R == 0, state, bc)
s1_assigner.assign(s1_, state.sub(1).sub(0))
s2_assigner.assign(s2_, state.sub(1).sub(1))
xs.append(s1_.vector().get_local()[0])
ys.append(s2_.vector().get_local()[0])
u, s = state.split()
dolfin.File("u.pvd") << u
fig, ax = plt.subplots()
ax.plot(ts, xs, label="x")
ax.plot(ts, ys, label="y")
ax.legend()
fig.savefig("oscilator_ode.png")
plt.show()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment