{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# 12: Adding a non-conservative equation"
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "**Note:** To improve responsiveness via caching, the notebooks are updated only once a week. They are only\n",
    "available for the latest stable release of Trixi.jl at the time of caching."
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "If you want to use Trixi.jl for your own research, you might be interested in\n",
    "a new physics model that is not present in Trixi.jl. In this tutorial,\n",
    "we will implement the nonconservative linear advection equation in a periodic domain\n",
    "$$\n",
    "\\left\\{\n",
    "\\begin{aligned}&\\partial_t u(t,x) + a(x) \\partial_x u(t,x) = 0 \\\\\n",
    "&u(0,x)=\\sin(x) \\\\\n",
    "&u(t,-\\pi)=u(t,\\pi)\n",
    "\\end{aligned}\n",
    "\\right.\n",
    "$$\n",
    "where $a(x) = 2 + \\cos(x)$. The analytic solution is\n",
    "$$\n",
    "u(t,x)=-\\sin \\left(2 \\tan ^{-1}\\left(\\sqrt{3} \\tan \\left(\\frac{\\sqrt{3} t}{2}-\\tan ^{-1}\\left(\\frac{1}{\\sqrt{3}}\\tan \\left(\\frac{x}{2}\\right)\\right)\\right)\\right)\\right)\n",
    "$$\n",
    "In Trixi.jl, such a mathematical model\n",
    "is encoded as a subtype of `Trixi.AbstractEquations`."
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Basic setup"
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "Since there is no native support for variable coefficients, we need to transform the PDE to the following system:\n",
    "$$\n",
    "\\left\\{\n",
    "\\begin{aligned}&\\partial_t \\begin{pmatrix}u(t,x)\\\\a(t,x) \\end{pmatrix} +\\begin{pmatrix} a(t,x) \\partial_x u(t,x) \\\\ 0 \\end{pmatrix} = 0 \\\\\n",
    "&u(0,x)=\\sin(x) \\\\\n",
    "&a(0,x)=2+\\cos(x) \\\\\n",
    "&u(t,-\\pi)=u(t,\\pi)\n",
    "\\end{aligned}\n",
    "\\right.\n",
    "$$"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "# Define new physics\n",
    "using Trixi\n",
    "using Trixi: AbstractEquations, get_node_vars\n",
    "import Trixi: varnames, default_analysis_integrals, flux, max_abs_speed_naive,\n",
    "              have_nonconservative_terms\n",
    "\n",
    "# Since there is no native support for variable coefficients, we use two\n",
    "# variables: one for the basic unknown `u` and another one for the coefficient `a`\n",
    "struct NonconservativeLinearAdvectionEquation <: AbstractEquations{1 #= spatial dimension =#,\n",
    "                                                                   2 #= two variables (u,a) =#}\n",
    "end\n",
    "\n",
    "varnames(::typeof(cons2cons), ::NonconservativeLinearAdvectionEquation) = (\"scalar\", \"advection_velocity\")\n",
    "\n",
    "default_analysis_integrals(::NonconservativeLinearAdvectionEquation) = ()\n",
    "\n",
    "\n",
    "# The conservative part of the flux is zero\n",
    "flux(u, orientation, equation::NonconservativeLinearAdvectionEquation) = zero(u)\n",
    "\n",
    "# Calculate maximum wave speed for local Lax-Friedrichs-type dissipation\n",
    "function max_abs_speed_naive(u_ll, u_rr, orientation::Integer, ::NonconservativeLinearAdvectionEquation)\n",
    "    _, advection_velocity_ll = u_ll\n",
    "    _, advection_velocity_rr = u_rr\n",
    "\n",
    "    return max(abs(advection_velocity_ll), abs(advection_velocity_rr))\n",
    "end\n",
    "\n",
    "\n",
    "# We use nonconservative terms\n",
    "have_nonconservative_terms(::NonconservativeLinearAdvectionEquation) = Trixi.True()\n",
    "\n",
    "# This \"nonconservative numerical flux\" implements the nonconservative terms.\n",
    "# In general, nonconservative terms can be written in the form\n",
    "#   g(u) ∂ₓ h(u)\n",
    "# Thus, a discrete difference approximation of this nonconservative term needs\n",
    "# - `u mine`:  the value of `u` at the current position (for g(u))\n",
    "# - `u_other`: the values of `u` in a neighborhood of the current position (for ∂ₓ h(u))\n",
    "function flux_nonconservative(u_mine, u_other, orientation,\n",
    "                              equations::NonconservativeLinearAdvectionEquation)\n",
    "    _, advection_velocity = u_mine\n",
    "    scalar, _             = u_other\n",
    "\n",
    "    return SVector(advection_velocity * scalar, zero(scalar))\n",
    "end"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "The implementation of nonconservative terms uses a single \"nonconservative flux\"\n",
    "function `flux_nonconservative`. It will basically be applied in a loop of the\n",
    "form\n",
    "```julia\n",
    "du_m(D, u) = sum(D[m, l] * flux_nonconservative(u[m], u[l], 1, equations)) # orientation 1: x\n",
    "```\n",
    "where `D` is the derivative matrix and `u` contains the nodal solution values."
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "Now, we can run a simple simulation using a DGSEM discretization."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "# Create a simulation setup\n",
    "using Trixi\n",
    "using OrdinaryDiffEq\n",
    "\n",
    "equation = NonconservativeLinearAdvectionEquation()\n",
    "\n",
    "# You can derive the exact solution for this setup using the method of\n",
    "# characteristics\n",
    "function initial_condition_sine(x, t, equation::NonconservativeLinearAdvectionEquation)\n",
    "    x0 = -2 * atan(sqrt(3) * tan(sqrt(3) / 2 * t - atan(tan(x[1] / 2) / sqrt(3))))\n",
    "    scalar = sin(x0)\n",
    "    advection_velocity = 2 + cos(x[1])\n",
    "    SVector(scalar, advection_velocity)\n",
    "end\n",
    "\n",
    "# Create a uniform mesh in 1D in the interval [-π, π] with periodic boundaries\n",
    "mesh = TreeMesh(-Float64(π), Float64(π), # min/max coordinates\n",
    "                initial_refinement_level=4, n_cells_max=10^4)\n",
    "\n",
    "# Create a DGSEM solver with polynomials of degree `polydeg`\n",
    "# Remember to pass a tuple of the form `(conservative_flux, nonconservative_flux)`\n",
    "# as `surface_flux` and `volume_flux` when working with nonconservative terms\n",
    "volume_flux  = (flux_central, flux_nonconservative)\n",
    "surface_flux = (flux_lax_friedrichs, flux_nonconservative)\n",
    "solver = DGSEM(polydeg=3, surface_flux=surface_flux,\n",
    "               volume_integral=VolumeIntegralFluxDifferencing(volume_flux))\n",
    "\n",
    "# Setup the spatial semidiscretization containing all ingredients\n",
    "semi = SemidiscretizationHyperbolic(mesh, equation, initial_condition_sine, solver)\n",
    "\n",
    "# Create an ODE problem with given time span\n",
    "tspan = (0.0, 1.0)\n",
    "ode = semidiscretize(semi, tspan)\n",
    "\n",
    "# Set up some standard callbacks summarizing the simulation setup and computing\n",
    "# errors of the numerical solution\n",
    "summary_callback = SummaryCallback()\n",
    "analysis_callback = AnalysisCallback(semi, interval=50)\n",
    "callbacks = CallbackSet(summary_callback, analysis_callback)\n",
    "\n",
    "# OrdinaryDiffEq's `solve` method evolves the solution in time and executes\n",
    "# the passed callbacks\n",
    "sol = solve(ode, Tsit5(), abstol=1.0e-6, reltol=1.0e-6,\n",
    "            save_everystep=false, callback=callbacks)\n",
    "\n",
    "# Print the timer summary\n",
    "summary_callback()\n",
    "\n",
    "# Plot the numerical solution at the final time\n",
    "using Plots: plot\n",
    "plot(sol)"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "You see a plot of the final solution."
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "We can check whether everything fits together by refining the grid and comparing\n",
    "the numerical errors. First, we look at the error using the grid resolution\n",
    "above."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "error_1 = analysis_callback(sol).l2 |> first"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "Next, we increase the grid resolution by one refinement level and run the\n",
    "simulation again."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "mesh = TreeMesh(-Float64(π), Float64(π), # min/max coordinates\n",
    "                initial_refinement_level=5, n_cells_max=10^4)\n",
    "\n",
    "semi = SemidiscretizationHyperbolic(mesh, equation, initial_condition_sine, solver)\n",
    "\n",
    "tspan = (0.0, 1.0)\n",
    "ode = semidiscretize(semi, tspan);\n",
    "\n",
    "summary_callback = SummaryCallback()\n",
    "analysis_callback = AnalysisCallback(semi, interval=50)\n",
    "callbacks = CallbackSet(summary_callback, analysis_callback);\n",
    "\n",
    "sol = solve(ode, Tsit5(), abstol=1.0e-6, reltol=1.0e-6,\n",
    "            save_everystep=false, callback=callbacks);\n",
    "summary_callback()"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "error_2 = analysis_callback(sol).l2 |> first"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "error_1 / error_2"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "As expected, the new error is roughly reduced by a factor of 16, corresponding\n",
    "to an experimental order of convergence of 4 (for polynomials of degree 3)."
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Summary of the code"
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "Here is the complete code that we used (without the callbacks since these\n",
    "create a lot of unnecessary output in the doctests of this tutorial).\n",
    "In addition, we create the `struct` inside the new module `NonconservativeLinearAdvection`.\n",
    "That ensures that we can re-create `struct`s defined therein without having to\n",
    "restart Julia."
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "Define new physics"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "module NonconservativeLinearAdvection\n",
    "\n",
    "using Trixi\n",
    "using Trixi: AbstractEquations, get_node_vars\n",
    "import Trixi: varnames, default_analysis_integrals, flux, max_abs_speed_naive,\n",
    "              have_nonconservative_terms\n",
    "\n",
    "# Since there is not yet native support for variable coefficients, we use two\n",
    "# variables: one for the basic unknown `u` and another one for the coefficient `a`\n",
    "struct NonconservativeLinearAdvectionEquation <: AbstractEquations{1 #= spatial dimension =#,\n",
    "                                                                   2 #= two variables (u,a) =#}\n",
    "end\n",
    "\n",
    "varnames(::typeof(cons2cons), ::NonconservativeLinearAdvectionEquation) = (\"scalar\", \"advection_velocity\")\n",
    "\n",
    "default_analysis_integrals(::NonconservativeLinearAdvectionEquation) = ()\n",
    "\n",
    "\n",
    "# The conservative part of the flux is zero\n",
    "flux(u, orientation, equation::NonconservativeLinearAdvectionEquation) = zero(u)\n",
    "\n",
    "# Calculate maximum wave speed for local Lax-Friedrichs-type dissipation\n",
    "function max_abs_speed_naive(u_ll, u_rr, orientation::Integer, ::NonconservativeLinearAdvectionEquation)\n",
    "    _, advection_velocity_ll = u_ll\n",
    "    _, advection_velocity_rr = u_rr\n",
    "\n",
    "    return max(abs(advection_velocity_ll), abs(advection_velocity_rr))\n",
    "end\n",
    "\n",
    "\n",
    "# We use nonconservative terms\n",
    "have_nonconservative_terms(::NonconservativeLinearAdvectionEquation) = Trixi.True()\n",
    "\n",
    "# This \"nonconservative numerical flux\" implements the nonconservative terms.\n",
    "# In general, nonconservative terms can be written in the form\n",
    "#   g(u) ∂ₓ h(u)\n",
    "# Thus, a discrete difference approximation of this nonconservative term needs\n",
    "# - `u mine`:  the value of `u` at the current position (for g(u))\n",
    "# - `u_other`: the values of `u` in a neighborhood of the current position (for ∂ₓ h(u))\n",
    "function flux_nonconservative(u_mine, u_other, orientation,\n",
    "                              equations::NonconservativeLinearAdvectionEquation)\n",
    "    _, advection_velocity = u_mine\n",
    "    scalar, _            = u_other\n",
    "\n",
    "    return SVector(advection_velocity * scalar, zero(scalar))\n",
    "end\n",
    "\n",
    "end # module\n",
    "\n",
    "\n",
    "\n",
    "# Create a simulation setup\n",
    "import .NonconservativeLinearAdvection\n",
    "using Trixi\n",
    "using OrdinaryDiffEq\n",
    "\n",
    "equation = NonconservativeLinearAdvection.NonconservativeLinearAdvectionEquation()\n",
    "\n",
    "# You can derive the exact solution for this setup using the method of\n",
    "# characteristics\n",
    "function initial_condition_sine(x, t, equation::NonconservativeLinearAdvection.NonconservativeLinearAdvectionEquation)\n",
    "    x0 = -2 * atan(sqrt(3) * tan(sqrt(3) / 2 * t - atan(tan(x[1] / 2) / sqrt(3))))\n",
    "    scalar = sin(x0)\n",
    "    advection_velocity = 2 + cos(x[1])\n",
    "    SVector(scalar, advection_velocity)\n",
    "end\n",
    "\n",
    "# Create a uniform mesh in 1D in the interval [-π, π] with periodic boundaries\n",
    "mesh = TreeMesh(-Float64(π), Float64(π), # min/max coordinates\n",
    "                initial_refinement_level=4, n_cells_max=10^4)\n",
    "\n",
    "# Create a DGSEM solver with polynomials of degree `polydeg`\n",
    "# Remember to pass a tuple of the form `(conservative_flux, nonconservative_flux)`\n",
    "# as `surface_flux` and `volume_flux` when working with nonconservative terms\n",
    "volume_flux  = (flux_central, NonconservativeLinearAdvection.flux_nonconservative)\n",
    "surface_flux = (flux_lax_friedrichs, NonconservativeLinearAdvection.flux_nonconservative)\n",
    "solver = DGSEM(polydeg=3, surface_flux=surface_flux,\n",
    "               volume_integral=VolumeIntegralFluxDifferencing(volume_flux))\n",
    "\n",
    "# Setup the spatial semidiscretization containing all ingredients\n",
    "semi = SemidiscretizationHyperbolic(mesh, equation, initial_condition_sine, solver)\n",
    "\n",
    "# Create an ODE problem with given time span\n",
    "tspan = (0.0, 1.0)\n",
    "ode = semidiscretize(semi, tspan);\n",
    "\n",
    "# Set up some standard callbacks summarizing the simulation setup and computing\n",
    "# errors of the numerical solution\n",
    "summary_callback = SummaryCallback()\n",
    "analysis_callback = AnalysisCallback(semi, interval=50)\n",
    "callbacks = CallbackSet(summary_callback, analysis_callback);\n",
    "\n",
    "# OrdinaryDiffEq's `solve` method evolves the solution in time and executes\n",
    "# the passed callbacks\n",
    "sol = solve(ode, Tsit5(), abstol=1.0e-6, reltol=1.0e-6,\n",
    "            save_everystep=false);\n",
    "\n",
    "# Plot the numerical solution at the final time\n",
    "using Plots: plot\n",
    "plot(sol);"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Package versions"
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "These results were obtained using the following versions."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "using InteractiveUtils\n",
    "versioninfo()\n",
    "\n",
    "using Pkg\n",
    "Pkg.status([\"Trixi\", \"OrdinaryDiffEq\", \"Plots\"],\n",
    "           mode=PKGMODE_MANIFEST)"
   ],
   "metadata": {},
   "execution_count": null
  }
 ],
 "nbformat_minor": 3,
 "metadata": {
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.10.3"
  },
  "kernelspec": {
   "name": "julia-1.10",
   "display_name": "Julia 1.10.3",
   "language": "julia"
  }
 },
 "nbformat": 4
}