test for making movies with Firedrake

firedrake-movie-test.py

import argparse
import firedrake
from firedrake import (
    max_value, sqrt, Constant, inner, as_vector, grad, dx, ds, dS
)
import numpy as np
from numpy import pi as π

parser = argparse.ArgumentParser()
parser.add_argument('--method', choices=['good', 'bad'])
args = parser.parse_args()

nx, ny = 32, 32
mesh = firedrake.UnitSquareMesh(nx, ny, diagonal='crossed')
x = firedrake.SpatialCoordinate(mesh)
y = Constant((0.5, 0.5))
r = x - y
u = as_vector((-r[1], r[0]))

x_c = as_vector((5 / 8, 5 / 8))
R_c = Constant(1 / 8)

x_b = as_vector((3 / 8, 3 / 8))
R_b = Constant(1 / 8)

q_expr = (
    max_value(0, 1 - sqrt(inner(x - x_c, x - x_c) / R_c**2)) +
    max_value(0, 1 - inner(x - x_b, x - x_b) / R_b**2)
)

Q1 = firedrake.FunctionSpace(mesh, family='DG', degree=1)
q0 = firedrake.project(q_expr, Q1)

q, φ = firedrake.TrialFunction(Q1), firedrake.TestFunction(Q1)
m = q * φ * dx

q = q0.copy(deepcopy=True)
cell_flux = -inner(grad(φ), q * u) * dx

n = firedrake.FacetNormal(mesh)
u_n = max_value(inner(u, n), 0)
f = q * u_n
facet_flux = (f('+') - f('-')) * (φ('+') - φ('-')) * dS

outflux = q * max_value(0, inner(u, n)) * φ * ds

dq_dt = -(cell_flux + facet_flux + outflux)

δq = firedrake.Function(Q1)
q1 = firedrake.Function(Q1)
q2 = firedrake.Function(Q1)

F2 = firedrake.replace(dq_dt, {q: q1})
F3 = firedrake.replace(dq_dt, {q: q2})

final_time = 2 * π
min_diameter = 1 / nx
max_speed = np.sqrt(2)
cfl_timestep = min_diameter / max_speed / 3
num_steps = 4 * int(final_time / cfl_timestep)
dt = Constant(final_time / num_steps)

problems = [
    firedrake.LinearVariationalProblem(m, dt * dq_dt, δq),
    firedrake.LinearVariationalProblem(m, dt * F2, δq),
    firedrake.LinearVariationalProblem(m, dt * F3, δq)
]

parameters = {
    'solver_parameters': {
        'ksp_type': 'preonly',
        'pc_type': 'bjacobi',
        'sub_pc_type': 'ilu'
    }
}
solvers = [
    firedrake.LinearVariationalSolver(problem, **parameters)
    for problem in problems
]

qs = [q.copy(deepcopy=True)]
output_freq = 20

print('Starting simulation.')

limiter = firedrake.VertexBasedLimiter(Q1)
for step in range(num_steps):
    solvers[0].solve()
    q1.assign(q + δq)

    solvers[1].solve()
    q2.assign(3 * q / 4 + (q1 + δq) / 4)

    solvers[2].solve()
    q.assign(q / 3 + 2 * (q2 + δq) / 3)
    limiter.apply(q)

    if (step + 1) % output_freq == 0:
        qs.append(q.copy(deepcopy=True))

print('Done simulation.')

print('Starting animation.')
import time
import matplotlib.pyplot as plt
from firedrake.plot import FunctionPlotter

start = time.time()

num_sample_points = 16
fn_plotter = FunctionPlotter(mesh, num_sample_points)

fig, axes = plt.subplots()
axes.set_aspect('equal')
axes.get_xaxis().set_visible(False)
axes.get_yaxis().set_visible(False)
kwargs = {'num_sample_points': num_sample_points, 'vmin': 0.0, 'vmax': 1.0}
colors = firedrake.tripcolor(q, axes=axes, **kwargs)

from matplotlib.animation import FuncAnimation

if args.method == 'good':
    def animate(q):
        values = fn_plotter(q)
        colors.set_array(values)
else:
    def animate(q):
        firedrake.tripcolor(q, axes=axes, **kwargs)

interval = 1e3 * output_freq * float(dt)
animation = FuncAnimation(fig, animate, frames=qs, interval=interval)
animation.save('firedrake-movie.mpg')

end = time.time()
print(f'Animation done; elapsed time: {end - start:5.1f}s')