{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# 20: Custom semidiscretizations"
   ],
   "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": [
    "As described in the overview section,\n",
    "semidiscretizations are high-level descriptions of spatial discretizations\n",
    "in Trixi.jl. Trixi.jl's main focus is on hyperbolic conservation\n",
    "laws represented in a `SemidiscretizationHyperbolic`.\n",
    "Hyperbolic-parabolic problems based on the advection-diffusion equation or\n",
    "the compressible Navier-Stokes equations can be represented in a\n",
    "`SemidiscretizationHyperbolicParabolic`. This is described in the\n",
    "basic tutorial on parabolic terms and its extension to\n",
    "custom parabolic terms.\n",
    "In this tutorial, we will describe how these semidiscretizations work and how\n",
    "they can be used to create custom semidiscretizations involving also other tasks."
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Overview of the right-hand side evaluation"
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "The semidiscretizations provided by Trixi.jl are set up to create `ODEProblem`s from the\n",
    "[SciML ecosystem for ordinary differential equations](https://diffeq.sciml.ai/latest/).\n",
    "In particular, a spatial semidiscretization can be wrapped in an ODE problem\n",
    "using `semidiscretize`, which returns an `ODEProblem`. This `ODEProblem`\n",
    "bundles an initial condition, a right-hand side (RHS) function, the time span,\n",
    "and possible parameters. The `ODEProblem`s created by Trixi.jl use the semidiscretization\n",
    "passed to `semidiscretize` as a parameter.\n",
    "For a `SemidiscretizationHyperbolic`, the `ODEProblem` wraps\n",
    "`Trixi.rhs!` as ODE RHS.\n",
    "For a `SemidiscretizationHyperbolicParabolic`,  Trixi.jl\n",
    "uses a `SplitODEProblem` combining `Trixi.rhs_parabolic!` for the\n",
    "(potentially) stiff part and `Trixi.rhs!` for the other part."
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Standard Trixi.jl setup"
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "In this tutorial, we will consider the linear advection equation\n",
    "with source term\n",
    "$$\n",
    "\\partial_t u(t,x) + \\partial_x u(t,x) = -\\exp(-t) \\sin\\bigl(\\pi (x - t) \\bigr)\n",
    "$$\n",
    "with periodic boundary conditions in the domain `[-1, 1]` as a\n",
    "model problem.\n",
    "The initial condition is\n",
    "$$\n",
    "u(0,x) = \\sin(\\pi x).\n",
    "$$\n",
    "The source term results in some damping and the analytical solution\n",
    "$$\n",
    "u(t,x) = \\exp(-t) \\sin\\bigl(\\pi (x - t) \\bigr).\n",
    "$$\n",
    "First, we discretize this equation using the standard functionality\n",
    "of Trixi.jl."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "using Trixi, OrdinaryDiffEq, Plots"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "The linear scalar advection equation is already implemented in\n",
    "Trixi.jl as `LinearScalarAdvectionEquation1D`. We construct\n",
    "it with an advection velocity `1.0`."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "equations = LinearScalarAdvectionEquation1D(1.0)"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "Next, we use a standard `DGSEM` solver."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "solver = DGSEM(polydeg = 3)"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "We create a simple `TreeMesh` in 1D."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "coordinates_min = (-1.0,)\n",
    "coordinates_max = (+1.0,)\n",
    "mesh = TreeMesh(coordinates_min, coordinates_max;\n",
    "                initial_refinement_level = 4,\n",
    "                n_cells_max = 10^4,\n",
    "                periodicity = true)"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "We wrap everything in in a semidiscretization and pass the source\n",
    "terms as a standard Julia function. Please note that Trixi.jl uses\n",
    "`SVector`s from\n",
    "[StaticArrays.jl](https://github.com/JuliaArrays/StaticArrays.jl)\n",
    "to store the conserved variables `u`. Thus, the return value of the\n",
    "source terms must be wrapped in an `SVector` - even if we consider\n",
    "just a scalar problem."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "function initial_condition(x, t, equations)\n",
    "    return SVector(exp(-t) * sinpi(x[1] - t))\n",
    "end\n",
    "\n",
    "function source_terms_standard(u, x, t, equations)\n",
    "    return -initial_condition(x, t, equations)\n",
    "end\n",
    "\n",
    "semi = SemidiscretizationHyperbolic(mesh, equations, initial_condition,\n",
    "                                    solver;\n",
    "                                    source_terms = source_terms_standard)"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "Now, we can create the `ODEProblem`, solve the resulting ODE\n",
    "using a time integration method from\n",
    "[OrdinaryDiffEq.jl](https://github.com/SciML/OrdinaryDiffEq.jl),\n",
    "and visualize the numerical solution at the final time using\n",
    "[Plots.jl](https://github.com/JuliaPlots/Plots.jl)."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "tspan = (0.0, 3.0)\n",
    "ode = semidiscretize(semi, tspan)\n",
    "\n",
    "sol = solve(ode, RDPK3SpFSAL49(); ode_default_options()...)\n",
    "\n",
    "plot(sol; label = \"numerical sol.\", legend = :topright)"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "We can also plot the analytical solution for comparison.\n",
    "Since Trixi.jl uses `SVector`s for the variables, we take their `first`\n",
    "(and only) component to get the scalar value for manual plotting."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "let\n",
    "   x = range(-1.0, 1.0; length = 200)\n",
    "   plot!(x, first.(initial_condition.(x, sol.t[end], equations)),\n",
    "         label = \"analytical sol.\", linestyle = :dash, legend = :topright)\n",
    "end"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "We can also add the initial condition to the plot."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "plot!(sol.u[1], semi, label = \"u0\", linestyle = :dot, legend = :topleft)"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "You can of course also use some\n",
    "[callbacks](https://trixi-framework.github.io/Trixi.jl/stable/callbacks/)\n",
    "provided by Trixi.jl as usual."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "summary_callback = SummaryCallback()\n",
    "analysis_interval = 100\n",
    "analysis_callback = AnalysisCallback(semi; interval = analysis_interval)\n",
    "alive_callback = AliveCallback(; analysis_interval)\n",
    "callbacks = CallbackSet(summary_callback,\n",
    "                        analysis_callback,\n",
    "                        alive_callback)\n",
    "\n",
    "sol = solve(ode, RDPK3SpFSAL49();\n",
    "            ode_default_options()..., callback = callbacks)\n",
    "summary_callback()"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Using a custom ODE right-hand side function"
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "Next, we will solve the same problem but use our own ODE RHS function.\n",
    "To demonstrate this, we will artificially create a global variable\n",
    "containing the current time of the simulation."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "const GLOBAL_TIME = Ref(0.0)\n",
    "\n",
    "function source_terms_custom(u, x, t, equations)\n",
    "    t = GLOBAL_TIME[]\n",
    "    return -initial_condition(x, t, equations)\n",
    "end"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "Next, we create our own RHS function to update the global time of\n",
    "the simulation before calling the RHS function from Trixi.jl."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "function rhs_source_custom!(du_ode, u_ode, semi, t)\n",
    "    GLOBAL_TIME[] = t\n",
    "    Trixi.rhs!(du_ode, u_ode, semi, t)\n",
    "end"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "Next, we create an `ODEProblem` manually copying over the data from\n",
    "the one we got from `semidiscretize` earlier."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "ode_source_custom = ODEProblem(rhs_source_custom!,\n",
    "                               ode.u0,\n",
    "                               ode.tspan,\n",
    "                               ode.p #= semi =#)\n",
    "sol_source_custom = solve(ode_source_custom, RDPK3SpFSAL49();\n",
    "                          ode_default_options()...)\n",
    "\n",
    "plot(sol_source_custom; label = \"numerical sol.\")\n",
    "let\n",
    "    x = range(-1.0, 1.0; length = 200)\n",
    "    plot!(x, first.(initial_condition.(x, sol_source_custom.t[end], equations)),\n",
    "          label = \"analytical sol.\", linestyle = :dash, legend = :topleft)\n",
    "end\n",
    "plot!(sol_source_custom.u[1], semi, label = \"u0\", linestyle = :dot, legend = :topleft)"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "This also works with callbacks as usual."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "summary_callback = SummaryCallback()\n",
    "analysis_interval = 100\n",
    "analysis_callback = AnalysisCallback(semi; interval = analysis_interval)\n",
    "alive_callback = AliveCallback(; analysis_interval)\n",
    "callbacks = CallbackSet(summary_callback,\n",
    "                        analysis_callback,\n",
    "                        alive_callback)\n",
    "\n",
    "sol = solve(ode_source_custom, RDPK3SpFSAL49();\n",
    "            ode_default_options()..., callback = callbacks)\n",
    "summary_callback()"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Setting up a custom semidiscretization"
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "Using a global constant is of course not really nice from a software\n",
    "engineering point of view. Thus, it can often be useful to collect\n",
    "additional data in the parameters of the `ODEProblem`. Thus, it is\n",
    "time to create our own semidiscretization. Here, we create a small\n",
    "wrapper of a standard semidiscretization of Trixi.jl and the current\n",
    "global time of the simulation."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "struct CustomSemidiscretization{Semi, T} <: Trixi.AbstractSemidiscretization\n",
    "    semi::Semi\n",
    "    t::T\n",
    "end\n",
    "\n",
    "semi_custom = CustomSemidiscretization(semi, Ref(0.0))"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "To get pretty printing in the REPL, you can consider specializing\n",
    "\n",
    "- `Base.show(io::IO, parameters::CustomSemidiscretization)`\n",
    "- `Base.show(io::IO, ::MIME\"text/plain\", parameters::CustomSemidiscretization)`\n",
    "\n",
    "for your custom semidiscretiation."
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "Next, we create our own source terms that use the global time stored\n",
    "in the custom semidiscretiation."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "source_terms_custom_semi = let semi_custom = semi_custom\n",
    "    function source_terms_custom_semi(u, x, t, equations)\n",
    "        t = semi_custom.t[]\n",
    "        return -initial_condition(x, t, equations)\n",
    "    end\n",
    "end"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "We also create a custom ODE RHS to update the current global time\n",
    "stored in the custom semidiscretization. We unpack the standard\n",
    "semidiscretization created by Trixi.jl and pass it to `Trixi.rhs!`."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "function rhs_semi_custom!(du_ode, u_ode, semi_custom, t)\n",
    "    semi_custom.t[] = t\n",
    "    Trixi.rhs!(du_ode, u_ode, semi_custom.semi, t)\n",
    "end"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "Finally, we set up an `ODEProblem` and solve it numerically."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "ode_semi_custom = ODEProblem(rhs_semi_custom!,\n",
    "                             ode.u0,\n",
    "                             ode.tspan,\n",
    "                             semi_custom)\n",
    "sol_semi_custom = solve(ode_semi_custom, RDPK3SpFSAL49();\n",
    "                        ode_default_options()...)"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "If we want to make use of additional functionality provided by\n",
    "Trixi.jl, e.g., for plotting, we need to implement a few additional\n",
    "specializations. In this case, we forward everything to the standard\n",
    "semidiscretization provided by Trixi.jl wrapped in our custom\n",
    "semidiscretization."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "Base.ndims(semi::CustomSemidiscretization) = ndims(semi.semi)\n",
    "function Trixi.mesh_equations_solver_cache(semi::CustomSemidiscretization)\n",
    "    Trixi.mesh_equations_solver_cache(semi.semi)\n",
    "end"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "Now, we can plot the numerical solution as usual."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "plot(sol_semi_custom; label = \"numerical sol.\")\n",
    "let\n",
    "    x = range(-1.0, 1.0; length = 200)\n",
    "    plot!(x, first.(initial_condition.(x, sol_semi_custom.t[end], equations)),\n",
    "          label = \"analytical sol.\", linestyle = :dash, legend = :topleft)\n",
    "end\n",
    "plot!(sol_semi_custom.u[1], semi, label = \"u0\", linestyle = :dot, legend = :topleft)"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "This also works with many callbacks as usual. However, the\n",
    "`AnalysisCallback` requires some special handling since it\n",
    "makes use of a performance counter contained in the standard\n",
    "semidiscretizations of Trixi.jl to report some\n",
    "performance metrics.\n",
    "Here, we forward all accesses to the performance counter to the\n",
    "wrapped semidiscretization."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "function Base.getproperty(semi::CustomSemidiscretization, s::Symbol)\n",
    "    if s === :performance_counter\n",
    "        wrapped_semi = getfield(semi, :semi)\n",
    "        wrapped_semi.performance_counter\n",
    "    else\n",
    "        getfield(semi, s)\n",
    "    end\n",
    "end"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "Moreover, the `AnalysisCallback` also performs some error\n",
    "calculations. We also need to forward them to the wrapped\n",
    "semidiscretization."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "function Trixi.calc_error_norms(func, u, t, analyzer,\n",
    "                                semi::CustomSemidiscretization,\n",
    "                                cache_analysis)\n",
    "    Trixi.calc_error_norms(func, u, t, analyzer,\n",
    "                           semi.semi,\n",
    "                           cache_analysis)\n",
    "end"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "Now, we can work with the callbacks used before as usual."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "summary_callback = SummaryCallback()\n",
    "analysis_interval = 100\n",
    "analysis_callback = AnalysisCallback(semi_custom;\n",
    "                                     interval = analysis_interval)\n",
    "alive_callback = AliveCallback(; analysis_interval)\n",
    "callbacks = CallbackSet(summary_callback,\n",
    "                        analysis_callback,\n",
    "                        alive_callback)\n",
    "\n",
    "sol = solve(ode_semi_custom, RDPK3SpFSAL49();\n",
    "            ode_default_options()..., callback = callbacks)\n",
    "summary_callback()"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "For even more advanced usage of custom semidiscretizations, you\n",
    "may look at the source code of the ones contained in Trixi.jl, e.g.,\n",
    "- `SemidiscretizationHyperbolicParabolic`\n",
    "- `SemidiscretizationEulerGravity`\n",
    "- `SemidiscretizationEulerAcoustics`\n",
    "- `SemidiscretizationCoupled`"
   ],
   "metadata": {}
  },
  {
   "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
}