diff --git a/demos/test_demos_run.py b/demos/test_demos_run.py new file mode 100644 index 0000000000..c055389218 --- /dev/null +++ b/demos/test_demos_run.py @@ -0,0 +1,143 @@ +import glob +import importlib +import os +import subprocess +import sys +from collections import namedtuple +from pathlib import Path +from os.path import abspath, basename, dirname, join, splitext + +import pyadjoint +import pytest + + +Demo = namedtuple("Demo", ["loc", "requirements"]) + + +CWD = abspath(dirname(__file__)) +DEMO_DIR = join(CWD, "..", "..", "..", "demos") + +SERIAL_DEMOS = [ + Demo(("benney_luke", "benney_luke"), ["vtk"]), + Demo(("boussinesq", "boussinesq"), []), + Demo(("burgers", "burgers"), ["vtk"]), + Demo(("camassa-holm", "camassaholm"), ["vtk"]), + Demo(("DG_advection", "DG_advection"), ["matplotlib"]), + Demo(("eigenvalues_QG_basinmodes", "qgbasinmodes"), ["matplotlib", "slepc", "vtk"]), + Demo(("extruded_continuity", "extruded_continuity"), []), + Demo(("helmholtz", "helmholtz"), ["vtk"]), + Demo(("higher_order_mass_lumping", "higher_order_mass_lumping"), ["vtk"]), + Demo(("immersed_fem", "immersed_fem"), []), + Demo(("linear_fluid_structure_interaction", "linear_fluid_structure_interaction"), ["vtk"]), + Demo(("linear-wave-equation", "linear_wave_equation"), ["vtk"]), + Demo(("ma-demo", "ma-demo"), ["vtk"]), + Demo(("matrix_free", "navier_stokes"), ["mumps", "vtk"]), + Demo(("matrix_free", "poisson"), []), + Demo(("matrix_free", "rayleigh-benard"), ["hypre", "mumps", "vtk"]), + Demo(("matrix_free", "stokes"), ["hypre", "mumps", "vtk"]), + Demo(("multigrid", "geometric_multigrid"), ["vtk"]), + Demo(("netgen", "netgen_mesh"), ["mumps", "netgen", "slepc", "vtk"]), + Demo(("nonlinear_QG_winddrivengyre", "qg_winddrivengyre"), ["vtk"]), + Demo(("parallel-printing", "parprint"), []), + Demo(("poisson", "poisson_mixed"), ["vtk"]), + Demo(("quasigeostrophy_1layer", "qg_1layer_wave"), ["hypre", "vtk"]), + Demo(("saddle_point_pc", "saddle_point_systems"), ["hypre", "mumps"]), + Demo(("vlasov_poisson_1d", "vp1d"), []) +] +PARALLEL_DEMOS = [ + Demo(("full_waveform_inversion", "full_waveform_inversion"), ["adjoint"]), +] + + +@pytest.fixture +def env(): + env = os.environ.copy() + env["MPLBACKEND"] = "pdf" + return env + + +def test_no_missing_demos(): + all_demo_locs = { + demo.loc + for demos in [SERIAL_DEMOS, PARALLEL_DEMOS] + for demo in demos + } + for rst_file in glob.glob(f"{DEMO_DIR}/*/*.py.rst"): + rst_path = Path(rst_file) + demo_dir = rst_path.parent.name + demo_name, _, _ = rst_path.name.split(".") + demo_loc = (demo_dir, demo_name) + assert demo_loc in all_demo_locs + all_demo_locs.remove(demo_loc) + assert not all_demo_locs, "Unrecognised demos listed" + + +def _maybe_skip_demo(demo, skip_dependency): + skip_dep, dependency_skip_markers_and_reasons = skip_dependency + # Add pytest skips for missing imports or packages + + for dep, _, reason in dependency_skip_markers_and_reasons: + if dep in demo.requirements and skip_dep(dep): + pytest.skip(reason) + + +def _prepare_demo(demo, monkeypatch, tmpdir): + # Change to the temporary directory (monkeypatch ensures that this + # is undone when the fixture usage disappears) + monkeypatch.chdir(tmpdir) + + demo_dir, demo_name = demo.loc + rst_file = f"{DEMO_DIR}/{demo_dir}/{demo_name}.py.rst" + + # Check if we need to generate any meshes + geos = glob.glob("%s/*.geo" % dirname(rst_file)) + for geo in geos: + name = "%s.msh" % splitext(basename(geo))[0] + if os.path.exists(name): + # No need to generate if it's already there + continue + try: + subprocess.check_call(["gmsh", geo, "-format", "msh2", "-3", "-o", str(tmpdir.join(name))]) + except (subprocess.CalledProcessError, OSError): + # Skip if unable to make mesh + pytest.skip("Unable to generate mesh file, skipping test") + + # Get the name of the python file that pylit will make + name = splitext(basename(rst_file))[0] + py_file = str(tmpdir.join(name)) + # Convert rst demo to runnable python file + subprocess.check_call(["pylit", rst_file, py_file]) + return Path(py_file) + + +def _exec_file(py_file): + # To execute a file we import it. We therefore need to modify sys.path so the + # tempdir can be found. + sys.path.insert(0, str(py_file.parent)) + importlib.import_module(py_file.with_suffix("").name) + sys.path.pop(0) # cleanup + + +@pytest.mark.skipcomplex +@pytest.mark.parametrize("demo", SERIAL_DEMOS, ids=["/".join(d.loc) for d in SERIAL_DEMOS]) +def test_serial_demo(demo, env, monkeypatch, tmpdir, skip_dependency): + _maybe_skip_demo(demo, skip_dependency) + py_file = _prepare_demo(demo, monkeypatch, tmpdir) + _exec_file(py_file) + + if "adjoint" in demo.requirements: + pyadjoint.pause_annotation() + pyadjoint.get_working_tape().clear_tape() + + +@pytest.mark.parallel(2) +@pytest.mark.skipcomplex +@pytest.mark.parametrize("demo", PARALLEL_DEMOS, ids=["/".join(d.loc) for d in PARALLEL_DEMOS]) +def test_parallel_demo(demo, env, monkeypatch, tmpdir, skip_dependency): + _maybe_skip_demo(demo, skip_dependency) + py_file = _prepare_demo(demo, monkeypatch, tmpdir) + _exec_file(py_file) + + if "adjoint" in demo.requirements: + pyadjoint.pause_annotation() + pyadjoint.get_working_tape().clear_tape() diff --git a/demos/vlasov_poisson_1d/vlasov_0s_HR.png b/demos/vlasov_poisson_1d/vlasov_0s_HR.png new file mode 100644 index 0000000000..0e0e3176fe Binary files /dev/null and b/demos/vlasov_poisson_1d/vlasov_0s_HR.png differ diff --git a/demos/vlasov_poisson_1d/vlasov_0s_LR.png b/demos/vlasov_poisson_1d/vlasov_0s_LR.png new file mode 100644 index 0000000000..e28376b5d4 Binary files /dev/null and b/demos/vlasov_poisson_1d/vlasov_0s_LR.png differ diff --git a/demos/vlasov_poisson_1d/vlasov_15_HR.png b/demos/vlasov_poisson_1d/vlasov_15_HR.png new file mode 100644 index 0000000000..d5de92a08f Binary files /dev/null and b/demos/vlasov_poisson_1d/vlasov_15_HR.png differ diff --git a/demos/vlasov_poisson_1d/vlasov_15_LR.png b/demos/vlasov_poisson_1d/vlasov_15_LR.png new file mode 100644 index 0000000000..7e8b3faea5 Binary files /dev/null and b/demos/vlasov_poisson_1d/vlasov_15_LR.png differ diff --git a/demos/vlasov_poisson_1d/vp1d.py.rst b/demos/vlasov_poisson_1d/vp1d.py.rst new file mode 100644 index 0000000000..d21a7e0d4d --- /dev/null +++ b/demos/vlasov_poisson_1d/vp1d.py.rst @@ -0,0 +1,390 @@ +1D Vlasov-Poisson Equation +=========================== + +This tutorial was contributed by `Colin Cotter +`__ and Werner Bauer. + +A plasma is a continuum of moving particles with nonunique velocity +at each point in space. In :math:`d` dimensions, the plasma is +described by a density :math:`f(x,v,t)` where :math:`x\in \mathbb{R}^d` +are the physical coordinates and :math:`v \in \mathbb{R}^d` are velocity +coordinates. Hence, in :math:`d` dimensions, a :math:`2d` +dimensional mesh is required. To deal with this curse of +dimensionality, particle-in-cell methods are usually used. However, +in 1 dimension, it is tractable to simulate the plasma on a 2 +dimensional mesh. + +The Vlasov equation models the (collisionless) conservation of plasma +particles, according to + +.. math:: + f_t + \nabla_{\vec{x}} \cdot (\vec{v}f) + \nabla_{\vec{v}} \cdot (\vec{a}f) = 0, + +where + +.. math:: + \nabla_{\vec{x}} = (\partial_{x_1},\ldots, \partial_{x_d}), \quad + \nabla_{\vec{v}} = (\partial_{v_1},\ldots, \partial_{v_d}). + +To close the system, we need a formula for the acceleration :math:`\vec{a}`. +In the (single species) Vlasov-Poisson model, the acceleration is +determined by the electrostatic force, + +.. math:: + \vec{a} = -\frac{1}{m}\nabla\phi, + +where :math:`m` +is the mass per plasma particle, and :math:`\phi` is the electrostatic +potential determined by the Poisson equation, + +.. math:: + -\nabla^2\phi = q_0\int_{\mathbb{R}^d} f(\vec{x},\vec{v},t)\,\mathrm{d} v, + +where :math:`q_0` is the electric charge per plasma particle. + +In this demo we specialise to :math:`d=1`, and the equations become + +.. math:: + f_t + (fv)_x + (-f\phi_x/m)_v = 0, \quad + -\phi_{xx} = q_0\int f(x,v,t)\,\mathrm{d} v, + +with coordinates :math:`(x,v)\in \mathbb{R}^2`. From now on we will +relabel these coordinates :math:`(x,v)\mapsto (x_1,x_2)`, obtaining +the equivalent form, + +.. math:: + f_t + \nabla\cdot(\vec{u}f) = 0, \quad \vec{u} = (v,-\phi_x/m), \quad + -\phi_{x_1x_1} = q_0\int f(x_1,x_2,t)\,\mathrm{d} x_2, + +where :math:`\nabla=(\partial_{x_1},\partial_{x_2})`. From now on we will +choose units such that :math:`q_0,m` are absorbed into the definition of +:math:`f`. + +To proceed, we need to develop variational formulations of these +equations. + +For the density we will use a discontinuous Galerkin formulation, +and the continuity equation becomes + +.. math:: + + \int_\Omega \! q \frac{\partial f}{\partial t} \, \mathrm{d} x + &= \int_\Omega \! f \nabla \cdot (q \vec{u}) \, \mathrm{d} x\\ + &\quad- \int_{\Gamma_\mathrm{int}} \! \widetilde{f}(q_+ \vec{u} \cdot \vec{n}_+ + + q_- \vec{u} \cdot \vec{n}_-) \, \mathrm{d} S\\ + &\quad- \int_{\Gamma_I} q f_\mathrm{in} \vec{u} \cdot + \vec{n} \, \mathrm{d} s\\ + &\quad- \int_{\Gamma_O} q f \vec{u} \cdot + \vec{n} \, \mathrm{d} s + \qquad \forall q \in V, + +where :math:`\Omega` is the computational domain in :math:`(x,v)` +space, :math:`V` is the discontinuous finite element space, +:math:`\Gamma_\mathrm{int}` is the set of interior cell edges, +:math:`\Gamma_I` is the inlet part of +exterior boundary where :math:`\vec{u}\cdot\vec{n}<0`, +:math:`\Gamma_O` is the outlet part of +exterior boundary where :math:`\vec{u}\cdot\vec{n}>0`, :math:`n` is +the normal to each edge, :math:`\tilde{f}` is the upwind value of +:math:`f`, and :math:`f_{\mathrm{in}}` is the inflow boundary value +for :math:`f`. See the Discontinuous Galerkin advection +:doc:`demo` for more details. The unapproximated +problem should have :math:`-\infty < x_2 < \infty`, i.e. unbounded velocities, but we approximate +the problem by solving in the domain :math:`\Omega=I_1\times [-H/2, H/2]`, +where :math:`I` is some chosen interval in the spatial dimension. + +For the Poisson equation, we will use a regular Galerkin formulation. +The difficulty in the formulation is the integral over :math:`x_2`. We +deal with this by considering a space :math:`\bar{W}` which is restricted +to functions that are constant in the vertical. Multiplying by a +test function :math:`\psi\in \bar{W}` and integrating by parts gives + +.. math:: + + \int \psi_{x_1}\phi_{x_1}\, \mathrm{d} x_1 + = \int \int f(x_1,x_2,t) \psi\, \mathrm{d} x_1\,\mathrm{d} x_2, \quad + \forall \psi \in \bar{W}. + +Since the left hand side integrand is independent of :math:`v=x_2`, we +can integrate over :math:`x_2` and divide by :math:`H`, to obtain + +.. math:: + + \int_\Omega \psi_{x_1}\phi_{x_1}/H\, \mathrm{d} x + = \int f \psi\, \mathrm{d} x, \quad + \forall \psi \in \bar{W}, + +which is now in a form which we can implement easily in Firedrake. One +final issue is that this problem only has a solution up to an additive +constant, so we further restrict :math:`\phi \in \mathring{\bar{W}}`, +where + +.. math:: + \mathring{\bar{W}} = \{ w\in \bar{W}: \bar{w}=0\}, + +where here the bar indicates a spatial average, + +.. math:: + + \bar{w} = \frac{\int_{\Omega} w\, \mathrm{d} x}{\int_{\Omega} 1 \mathrm{d} x}. + +Then we seek the solution of + +.. math:: + + \int_\Omega \psi_{x_1}\phi_{x_1}\,\mathrm{d} x + = \int H(f-\bar{f}) \psi\, \mathrm{d} x, \quad + \forall \psi \in \mathring{\bar{W}}. + +To discretise in time, we will use an SSPRK3 time discretisation, similar to the DG advection :doc:`demo`. At +each Runge-Kutta stage, we must solve for the electrostatic potential, +and then use it to compute :math:`\vec{u}`, in order to compute +:math:`\partial f/\partial t`. + +As usual, to implement this problem, we start by importing the +Firedrake namespace. :: + + from firedrake import * + +We build the mesh by constructing a 1D mesh, which will be extruded in +the vertical. Here we will use periodic boundary conditions in the +:math:`x_1` direction, :: + + ncells = 50 + L = 8*pi + base_mesh = PeriodicIntervalMesh(ncells, L) + +The mesh is then extruded upwards in the "velocity" direction. :: + + H = 10.0 + nlayers = 50 + mesh = ExtrudedMesh(base_mesh, layers=nlayers, layer_height=H/nlayers) + +We want to have :math:`v=0` in the middle of the domain, so that we +can have negative and positive velocities. This requires to edit the +coordinate field. :: + + x, v = SpatialCoordinate(mesh) + mesh.coordinates.interpolate(as_vector([x, v-H/2])) + +Now we build a discontinuous finite element space for the density, :: + + V = FunctionSpace(mesh, 'DQ', 1) + +and a continuous finite element space for the electostatic potential. +The space is continuous in the horizontal and constant in the vertical, +specified through the ``vfamily``. :: + + Wbar = FunctionSpace(mesh, 'CG', 1, vfamily='R', vdegree=0) + +We create a :class:`~.Function` to store the solution at the current +time, and then set its initial condition, + +.. math:: + + f(x,v,0) = \frac{1}{\sqrt{2\pi}}v^2\exp(-v^2/2)(1+ A\cos(kx)), + \quad A=0.05, \quad k=0.5. + +:: + + fn = Function(V, name="density") + A = Constant(0.05) + k = Constant(0.5) + fn.interpolate(v**2*exp(-v**2/2)*(1 + A*cos(k*x))/(2*pi)**0.5) + +We will need the (conserved) average :math:`\bar{f}` for the Poisson +equation. :: + + One = Function(V).assign(1.0) + fbar = assemble(fn*dx)/assemble(One*dx) + +We create a :class:`~.Function` to store the electrostatic potential. :: + + phi = Function(Wbar, name="potential") + +The next task is to create the solver for the electrostatic potential, which +will be called every timestep. + +We create a :class:`~.Function` to store the intermediate densities at each +Runge-Kutta stage. The right hand side of the Poisson equation will be +evaluated using this :class:`~.Function` to obtain the potential at each +stage. Defining this beforehand will enable us to reuse the solver. :: + + fstar = Function(V) + +Now we express the Poisson equation in UFL. :: + + psi = TestFunction(Wbar) + dphi = TrialFunction(Wbar) + phi_eqn = dphi.dx(0)*psi.dx(0)*dx - H*(fstar-fbar)*psi*dx + +To deal mathematically with the null space of the potential, we expressed the +problem in :math:`\mathring{\bar{W}}`. Programmatically we will express the +problem in :math:`\bar{W}` and deal with the null space by defining a basis +for the space of globally constant functions, which we will later pass to PETSc +so the solver can remove it from the solution. :: + + nullspace = VectorSpaceBasis(constant=True, comm=COMM_WORLD) + +However, the null space also means that the assembled matrix of the +Poisson problem will be singular, which will prevent us from using a +direct solver. To deal with this, we will precondition the Poisson problem +with a version shifted by :math:`\int_{\Omega}\phi\psi\mathrf{d}x`. The +shifted problem is well-posed on :math:`\bar{W}`, so the assembled matrix +will be non-singular and so solvable with direct methods. :: + + shift_eqn = dphi.dx(0)*psi.dx(0)*dx + dphi*psi*dx + +We use these to define a :class:`~.LinearVariationalProblem`. :: + + phi_problem = LinearVariationalProblem(lhs(phi_eqn), rhs(phi_eqn), + phi, aP=shift_eqn) + +Now we build the :class:`~.LinearVariationalSolver`. The problem +is preconditioned by the shifted operator which is solved using a direct +solver, and we pass the nullspace of globally constant functions to +the solver. :: + + params = { + 'ksp_type': 'gmres', + 'pc_type': 'lu', + 'ksp_rtol': 1.0e-8, + } + phi_solver = LinearVariationalSolver(phi_problem, + nullspace=nullspace, + solver_parameters=params) + +Now we move onto the solver to compute :math:`\partial f/\partial t`. We +define a symbolic :math:`\Delta t` which we will update later. :: + + dtc = Constant(0) + +At each stage, the solver will take in the intermediate solution ``fstar`` and +return the stage increment :math:`\Delta t\partial f/\partial t` in ``df_out``. :: + + df_out = Function(V) + +Now we express the equation in UFL, starting with the left hand side +bilinear form :: + + q = TestFunction(V) + u = as_vector([v, -phi.dx(0)]) + n = FacetNormal(mesh) + un = 0.5*(dot(u, n) + abs(dot(u, n))) + df = TrialFunction(V) + df_a = q*df*dx + +The problem is defined on an extruded mesh, so the interior facets are +separated into horizontal and vertical ones. :: + + dS = dS_h + dS_v + +Now we build the right hand side linear form. A conditional operator +is used to deal with the inflow and outflow parts of the exterior +boundary. Due to the periodic boundary conditions in :math:`x_1`, the only exterior boundary is at the top and bottom of the domain, with measure `ds_tb`. :: + + df_L = dtc*(div(u*q)*fstar*dx + - (q('+') - q('-'))*(un('+')*fstar('+') - un('-')*fstar('-'))*dS + - conditional(dot(u, n) > 0, q*dot(u, n)*fstar, 0.)*ds_tb + ) + +We then use this to build a solver. :: + + df_problem = LinearVariationalProblem(df_a, df_L, df_out) + df_solver = LinearVariationalSolver(df_problem) + +We are getting close to the time loop. We set up some timestepping +parameters. :: + + T = 50.0 # maximum timestep + t = 0. # model time + ndump = 100 # frequency of file dumps + dumpn = 0 # dump counter + nsteps = 5000 + dt = T/nsteps + dtc.assign(dt) + +We set up some :class:`~.Function`\s to store Runge-Kutta stage variables. :: + + f1 = Function(V) + f2 = Function(V) + +We set up a ``VTKFile`` object to write output every ``ndump`` +timesteps. :: + + outfile = VTKFile("vlasov.pvd") + +We want to output the initial condition, so need to solve for the electrostatic +potential that corresponds to the initial density. :: + + fstar.assign(fn) + phi_solver.solve() + outfile.write(fn, phi) + phi.assign(.0) + +Now we start the timeloop using a lovely progress bar. Note that +we have 5000 timesteps so this may take a few minutes to run:: + + for step in ProgressBar("Timestep").iter(range(nsteps)): + +Each Runge-Kutta stage involves solving for :math:`\phi` before solving +for :math:`\partial f/\partial t`. Here is the first stage. :: + + # + fstar.assign(fn) + phi_solver.solve() + df_solver.solve() + f1.assign(fn + df_out) + +The second stage. :: + + # + fstar.assign(f1) + phi_solver.solve() + df_solver.solve() + f2.assign(3*fn/4 + (f1 + df_out)/4) + +The third stage. :: + + # + fstar.assign(f2) + phi_solver.solve() + df_solver.solve() + fn.assign(fn/3 + 2*(f2 + df_out)/3) + t += dt + +Finally we output to the VTK file if it is time to do that. :: + + # + dumpn += 1 + if dumpn % ndump == 0: + dumpn = 0 + outfile.write(fn, phi) + +Images of the solution at shown below. + +.. figure:: vlasov_0s_LR.png + :align: center + + Solution at :math:`t = 0.` + +.. figure:: vlasov_15_LR.png + :align: center + + Solution at :math:`t = 15.` + +We also present solutions at double the resolution, by doubling the number +of horizontal cells and the number of layers, halving the timestep (by doubling the number of steps), and doubling ``nsteps``. + +.. figure:: vlasov_0s_HR.png + :align: center + + Solution at :math:`t = 0.` + +.. figure:: vlasov_15_HR.png + :align: center + + Solution at :math:`t = 15.` + + +A Python script version of this demo can be found :demo:`here `. diff --git a/docs/source/advanced_tut.rst b/docs/source/advanced_tut.rst index 6794207abc..f2b4dfd44f 100644 --- a/docs/source/advanced_tut.rst +++ b/docs/source/advanced_tut.rst @@ -24,4 +24,5 @@ element systems. Rayleigh-Benard convection. Netgen support. Full-waveform inversion: spatial and wave sources parallelism. + 1D Vlasov-Poisson equation using vertical independent function spaces. Steady Boussinesq problem with integral constraints. diff --git a/tests/firedrake/demos/test_demos_run.py b/tests/firedrake/demos/test_demos_run.py index cc0938272b..3bd846f121 100644 --- a/tests/firedrake/demos/test_demos_run.py +++ b/tests/firedrake/demos/test_demos_run.py @@ -42,6 +42,7 @@ Demo(("poisson", "poisson_mixed"), ["vtk"]), Demo(("quasigeostrophy_1layer", "qg_1layer_wave"), ["hypre", "vtk"]), Demo(("saddle_point_pc", "saddle_point_systems"), ["hypre", "mumps"]), + Demo(('vlasov_poisson_1d', 'vp1d'), []), ] PARALLEL_DEMOS = [ Demo(("full_waveform_inversion", "full_waveform_inversion"), ["adjoint"]),