{
 "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",
    "function varnames(::typeof(cons2cons), ::NonconservativeLinearAdvectionEquation)\n",
    "    return (\"scalar\", \"advection_velocity\")\n",
    "end\n",
    "\n",
    "default_analysis_integrals(::NonconservativeLinearAdvectionEquation) = ()\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,\n",
    "                             ::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",
    "# 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 OrdinaryDiffEqTsit5\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",
    "    return 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",
    "            ode_default_options()..., callback = callbacks)\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",
    "            ode_default_options()..., callback = callbacks);"
   ],
   "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": [
    "For non-trivial boundary conditions involving non-conservative terms,\n",
    "please refer to the section on [Other available example elixirs with non-trivial BC](https://trixi-framework.github.io/TrixiDocumentation/stable/tutorials/non_periodic_boundaries/#Other-available-example-elixirs-with-non-trivial-BC)."
   ],
   "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",
    "function varnames(::typeof(cons2cons), ::NonconservativeLinearAdvectionEquation)\n",
    "    return (\"scalar\", \"advection_velocity\")\n",
    "end\n",
    "\n",
    "default_analysis_integrals(::NonconservativeLinearAdvectionEquation) = ()\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,\n",
    "                             ::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",
    "# 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",
    "# Create a simulation setup\n",
    "import .NonconservativeLinearAdvection\n",
    "using Trixi\n",
    "using OrdinaryDiffEqTsit5\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,\n",
    "                                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",
    "    return 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",
    "            ode_default_options()...);\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\", \"OrdinaryDiffEqTsit5\", \"Plots\"],\n",
    "           mode = PKGMODE_MANIFEST)"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Additional modifications"
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "When one carries auxiliary variable(s) in the solution vector, e.g., for non-constant\n",
    "coefficient advection problems some routines may require modification to avoid adding\n",
    "dissipation to the variable coefficient quantity `a` that is carried as an auxiliary variable in\n",
    "the solution vector. In particular, a specialized `DissipationLocalLaxFriedrichs` term\n",
    "used together with the numerical surface flux `flux_lax_friedrichs` prevents \"smearing\"\n",
    "the variable coefficient `a` artificially."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "# Specialized dissipation term for the Lax-Friedrichs surface flux\n",
    "@inline function (dissipation::DissipationLocalLaxFriedrichs)(u_ll, u_rr,\n",
    "                                                              orientation::Integer,\n",
    "                                                              equation::NonconservativeLinearAdvectionEquation)\n",
    "    λ = dissipation.max_abs_speed(u_ll, u_rr, orientation, equation)\n",
    "\n",
    "    diss = -0.5 * λ * (u_rr - u_ll)\n",
    "    # do not add dissipation to the variable coefficient a used as last entry of u\n",
    "    return SVector(diss[1], zero(u_ll))\n",
    "end"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "Another modification is necessary if one wishes to use the stage limiter `PositivityPreservingLimiterZhangShu`\n",
    "during the time integration. This limiter takes in a `variable` (or set of variables) to limit and ensure positivity.\n",
    "However, these variables are used to compute the limiter quantities that are then applied to every\n",
    "variable in the solution vector `u`. To avoid artificially limiting (and in turn changing) the variable coefficient\n",
    "quantity that should remain unchanged, a specialized implementation of the `limiter_zhang_shu!` function is required.\n",
    "For the example equation given in this tutorial, this new function for the limiting would take the form"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "# Specialized positivity limiter that avoids modification of the auxiliary variable `a`\n",
    "function Trixi.limiter_zhang_shu!(u, threshold, variable, mesh,\n",
    "                                  equations::NonconservativeLinearAdvectionEquation,\n",
    "                                  dg, cache)\n",
    "    weights = dg.basis\n",
    "\n",
    "    for element in eachelement(dg, cache)\n",
    "        # determine minimum value\n",
    "        value_min = typemax(eltype(u))\n",
    "        for i in eachnode(dg)\n",
    "            u_node = get_node_vars(u, equations, dg, i, element)\n",
    "            value_min = min(value_min, variable(u_node, equations))\n",
    "        end\n",
    "\n",
    "        # detect if limiting is necessary\n",
    "        value_min < threshold || continue\n",
    "\n",
    "        # compute mean value\n",
    "        u_mean = zero(get_node_vars(u, equations, dg, 1, element))\n",
    "        for i in eachnode(dg)\n",
    "            u_node = get_node_vars(u, equations, dg, i, element)\n",
    "            u_mean += u_node * weights[i]\n",
    "        end\n",
    "        # note that the reference element is [-1,1]^ndims(dg), thus the weights sum to 2\n",
    "        u_mean = u_mean / 2^ndims(mesh)\n",
    "\n",
    "        # Compute the value directly with the mean values, as we assume that\n",
    "        # Jensen's inequality holds.\n",
    "        value_mean = variable(u_mean, equations)\n",
    "        theta = (value_mean - threshold) / (value_mean - value_min)\n",
    "        for i in eachnode(dg)\n",
    "            u_node = get_node_vars(u, equations, dg, i, element)\n",
    "\n",
    "            _, a_node = u_node\n",
    "            scalar_mean, _ = u_mean\n",
    "\n",
    "            # mean values of variable coefficient not used as it must not be overwritten\n",
    "            u_mean = SVector(scalar_mean, a_node)\n",
    "\n",
    "            set_node_vars!(u, theta * u_node + (1 - theta) * u_mean,\n",
    "                           equations, dg, i, element)\n",
    "        end\n",
    "    end\n",
    "\n",
    "    return nothing\n",
    "end"
   ],
   "metadata": {},
   "execution_count": null
  }
 ],
 "nbformat_minor": 3,
 "metadata": {
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.10.10"
  },
  "kernelspec": {
   "name": "julia-1.10",
   "display_name": "Julia 1.10.10",
   "language": "julia"
  }
 },
 "nbformat": 4
}