{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 5 - Vlasov-Ampère equations\n",
    "\n",
    "The equations we will solve are described in the model `VlasovAmpereOneSpecies`.\n",
    "To create the default parameter file from the console:\n",
    "\n",
    "```\n",
    "struphy params VlasovAmpereOneSpecies\n",
    "```\n",
    "\n",
    "Adapt the parameters and run the model with\n",
    "\n",
    "```\n",
    "python params_VlasovAmpereOneSpecies.py\n",
    "```\n",
    "\n",
    "In this notebook we shall re-create the launch file and perform some tests.\n",
    "\n",
    "## Weak Landau damping\n",
    "\n",
    "1. API imports:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from struphy import BaseUnits, EnvironmentOptions, Time\n",
    "from struphy import domains\n",
    "from struphy import grids\n",
    "from struphy import DerhamOptions\n",
    "from struphy import perturbations\n",
    "from struphy import maxwellians\n",
    "from struphy import BinningPlot, BoundaryParameters, LoadingParameters, WeightsParameters\n",
    "from struphy import Simulation\n",
    "\n",
    "from struphy.models import VlasovAmpereOneSpecies"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "2. Generic options:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# environment options\n",
    "env = EnvironmentOptions()\n",
    "\n",
    "# time stepping\n",
    "time_opts = Time(dt=0.05, Tend=0.5)  # , Tend = 3.5\n",
    "\n",
    "# geometry\n",
    "r1 = 12.56\n",
    "domain = domains.Cuboid(r1=r1)\n",
    "\n",
    "# fluid equilibrium (can be used as part of initial conditions)\n",
    "equil = None\n",
    "\n",
    "# grid\n",
    "grid = grids.TensorProductGrid(num_elements=(32, 1, 1))\n",
    "\n",
    "# derham options\n",
    "derham_opts = DerhamOptions()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "3. Model instance, info and species properties:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = VlasovAmpereOneSpecies(alpha=1.0, epsilon=1.0)\n",
    "model.info()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "4. Instance of simulation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sim = Simulation(model,\n",
    "                 env=env,\n",
    "                 time_opts=time_opts,\n",
    "                 domain=domain,\n",
    "                 equil=equil,\n",
    "                 grid=grid,\n",
    "                 derham_opts=derham_opts,)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "4. Kinetic species parameters:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loading_params = LoadingParameters(ppc=10000)\n",
    "weights_params = WeightsParameters(control_variate=True)\n",
    "boundary_params = BoundaryParameters()\n",
    "model.kinetic_ions.set_markers(\n",
    "    loading_params=loading_params, weights_params=weights_params, boundary_params=boundary_params\n",
    ")\n",
    "model.kinetic_ions.set_sorting_boxes()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# particle binning\n",
    "binplot_1 = BinningPlot(slice=\"e1\", n_bins=128, ranges=(0.0, 1.0))\n",
    "binplot_2 = BinningPlot(slice=\"v1\", n_bins=128, ranges=(-5.0, 5.0))\n",
    "binplot_3 = BinningPlot(slice=\"e1_v1\", n_bins=(128, 128), ranges=((0.0, 1.0), (-5.0, 5.0)))\n",
    "\n",
    "binning_plots = (binplot_1, binplot_2, binplot_3)\n",
    "\n",
    "model.kinetic_ions.set_save_data(binning_plots=binning_plots)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "5. Propagator options:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.propagators.push_eta.options = model.propagators.push_eta.Options()\n",
    "model.propagators.coupling_va.options = model.propagators.coupling_va.Options()\n",
    "model.initial_poisson.options = model.initial_poisson.Options(stab_eps=1e-12)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "6. Initial conditions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "background = maxwellians.Maxwellian3D(n=(1.0, None))\n",
    "model.kinetic_ions.var.add_background(background)\n",
    "\n",
    "perturbation = perturbations.ModesCos(ls=[1], amps=[0.001])\n",
    "init = maxwellians.Maxwellian3D(n=(1.0, perturbation))\n",
    "model.kinetic_ions.var.add_initial_condition(init)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us run the simulation. However, depending on the configuration, running in a notebook might be very slow. In order to get fast execution, run from the console. First, create the default parameter file and rename it \n",
    "\n",
    "```\n",
    "struphy params VlasovAmpereOneSpecies\n",
    "mv params_VlasovAmpereOneSpecies.py landau.py\n",
    "```\n",
    "\n",
    "Adapt it with the parameters from this notebook. Start the run with\n",
    "\n",
    "```\n",
    "python landau.py\n",
    "```\n",
    "\n",
    "or \n",
    "\n",
    "```\n",
    "mpirun -n 2 python landau.py\n",
    "```\n",
    "\n",
    "for a run on two threads. Line profiling can be enabled with \n",
    "\n",
    "```\n",
    "LINE_PROFILE=1 mpirun -n 2 landau.py\n",
    "```\n",
    "\n",
    "Let us look at the slower run in the notebook:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sim.run(verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sim.pproc(celldivide=8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sim.load_plotting_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot in v1\n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "v1_bins = sim.f.kinetic_ions.v1_density.grid_v1\n",
    "f_v1_init = sim.f.kinetic_ions.v1_density.f_binned[0]\n",
    "\n",
    "plt.plot(v1_bins, f_v1_init)\n",
    "plt.xlabel(\"vx\")\n",
    "plt.title(\"Initial Maxwellian\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot in e1\n",
    "\n",
    "e1_bins = sim.f.kinetic_ions.e1_density.grid_e1\n",
    "df_e1_init = sim.f.kinetic_ions.e1_density.delta_f_binned[0]\n",
    "\n",
    "plt.plot(e1_bins, df_e1_init)\n",
    "plt.xlabel(\"$\\eta_1$\")\n",
    "plt.title(\"Initial spatial perturbation\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot in e1-v1\n",
    "\n",
    "e1_bins = sim.f.kinetic_ions.e1_v1_density.grid_e1\n",
    "v1_bins = sim.f.kinetic_ions.e1_v1_density.grid_v1\n",
    "f_init = sim.f.kinetic_ions.e1_v1_density.f_binned[0]\n",
    "df_init = sim.f.kinetic_ions.e1_v1_density.delta_f_binned[0]\n",
    "f_end = sim.f.kinetic_ions.e1_v1_density.f_binned[-1]\n",
    "df_end = sim.f.kinetic_ions.e1_v1_density.delta_f_binned[-1]\n",
    "\n",
    "plt.figure(figsize=(14, 10))\n",
    "\n",
    "plt.subplot(2, 2, 1)\n",
    "plt.pcolor(e1_bins, v1_bins, f_init.T)\n",
    "plt.xlabel(\"$\\eta_1$\")\n",
    "plt.ylabel(\"$v_x$\")\n",
    "plt.title(\"Initial Maxwellian\")\n",
    "plt.colorbar()\n",
    "\n",
    "plt.subplot(2, 2, 2)\n",
    "plt.pcolor(e1_bins, v1_bins, df_init.T)\n",
    "plt.xlabel(\"$\\eta_1$\")\n",
    "plt.ylabel(\"$v_x$\")\n",
    "plt.title(\"Initial perturbation\")\n",
    "plt.colorbar()\n",
    "\n",
    "plt.subplot(2, 2, 3)\n",
    "plt.pcolor(e1_bins, v1_bins, f_end.T)\n",
    "plt.xlabel(\"$\\eta_1$\")\n",
    "plt.ylabel(\"$v_x$\")\n",
    "plt.title(\"Final Maxwellian\")\n",
    "plt.colorbar()\n",
    "\n",
    "plt.subplot(2, 2, 4)\n",
    "plt.pcolor(e1_bins, v1_bins, df_end.T)\n",
    "plt.xlabel(\"$\\eta_1$\")\n",
    "plt.ylabel(\"$v_x$\")\n",
    "plt.title(\"Final perturbation\")\n",
    "plt.colorbar()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# electric field\n",
    "e1, e2, e3 = sim.grids_log\n",
    "e_field = sim.spline_values.em_fields.e_field_log\n",
    "print(e_field)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "e_vals = e_field.data[0][0]\n",
    "plt.plot(e1, e_vals[:, 0, 0], label=\"E\")\n",
    "plt.xlabel(\"$\\eta_1$\")\n",
    "plt.title(\"Initial electric field\")\n",
    "plt.legend()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "env (3.12.3)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}