#!/usr/bin/env python
# coding: utf-8

# <script async src="https://www.googletagmanager.com/gtag/js?id=UA-59152712-8"></script>
# <script>
#   window.dataLayer = window.dataLayer || [];
#   function gtag(){dataLayer.push(arguments);}
#   gtag('js', new Date());
# 
#   gtag('config', 'UA-59152712-8');
# </script>
# 
# # Generating C code for the right-hand sides of Maxwell's equations, in ***curvilinear*** coordinates, using a reference metric formalism
# 
# ## Author: Ian Ruchlin
# ### Formatting improvements courtesy Brandon Clark
# 
# ## This notebook outlines the generation of C code for Maxwell's equations in curvilinear coordinates, using reference metric formalism. It explores two systems: the first with fixed constraint-violating numerical error and the second employing an auxiliary variable to enhance stability by allowing constraint violations to propagate away.
# 
# ### The following formulations of Maxwell's equations, called System I and System II, are described in [Illustrating Stability Properties of Numerical Relativity in Electrodynamics](https://arxiv.org/abs/gr-qc/0201051) by Knapp et al.
# 
# **Notebook Status:** <font color='red'><b> In progress </b></font>
# 
# **Validation Notes:** This module has not yet undergone validation testing. Do ***not*** use it until after appropriate validation testing has been performed.
# 
# ## Introduction:
# [Maxwell's equations](https://en.wikipedia.org/wiki/Maxwell%27s_equations) are subject to the Gauss' law constraint
# $$\mathcal{C} \equiv \hat{D}_{i} E^{i} - 4 \pi \rho = 0 \; ,$$
# where $E^{i}$ is the electric vector field, $\hat{D}_{i}$ is the [covariant derivative](https://en.wikipedia.org/wiki/Covariant_derivative) associated with the reference metric $\hat{\gamma}_{i j}$ (which is taken to represent flat space), and $\rho$ is the electric charge density. We use $\mathcal{C}$ as a measure of numerical error. Maxwell's equations are also required to satisfy $\hat{D}_{i} B^{i} = 0$, where $B^{i}$ is the magnetic vector field. The magnetic constraint implies that the magnetic field can be expressed as
# $$B_{i} = \epsilon_{i j k} \hat{D}^{j} A^{k} \; ,$$
# where $\epsilon_{i j k}$ is the totally antisymmetric [Levi-Civita tensor](https://en.wikipedia.org/wiki/Levi-Civita_symbol) and $A^{i}$ is the vector potential field. Together with the scalar potential $\psi$, the electric field can be expressed in terms of the potential fields as
# $$E_{i} = -\hat{D}_{i} \psi - \partial_{t} A_{i} \; .$$
# For now, we work in vacuum, where the electric charge density and the electric current density vector both vanish ($\rho = 0$ and $j_{i} = 0$).
# 
# In addition to the Gauss constraints, the electric and magnetic fields obey two independent [electromagnetic invariants](https://en.wikipedia.org/wiki/Classification_of_electromagnetic_fields#Invariants)
# \begin{align}
# \mathcal{P} &\equiv B_{i} B^{i} - E_{i} E^{i} \; , \\
# \mathcal{Q} &\equiv E_{i} B^{i} \; .
# \end{align}
# In vacuum, these satisfy $\mathcal{P} = \mathcal{Q} = 0$.

# <a id='toc'></a>
# 
# # Table of Contents:
# $$\label{toc}$$
# 
# This notebook is organized as follows
# 
# 1. [Step 1](#sys1): System I
# 1. [Step 2](#sys2): System II
# 1. [Step 3](#latex_pdf_output): Output this notebook to $\LaTeX$-formatted PDF file

# <a id='sys1'></a>
# 
# # Step 1: System I \[Back to [top](#toc)\]
# $$\label{sys1}$$
# 
# In terms of the above definitions, the evolution Maxwell's equations take the form
# \begin{align}
# \partial_{t} A_{i} &= -E_{i} - \hat{D}_{i} \psi \; , \\
# \partial_{t} E_{i} &= -\hat{D}_{j} \hat{D}^{j} A_{i} + \hat{D}_{i} \hat{D}_{j} A^{j}\; , \\
# \partial_{t} \psi &= -\hat{D}_{i} A^{i} \; .
# \end{align}
# Note that this coupled system contains mixed second derivatives in the second term on the right hand side of the $E^{i}$ evolution equation. We will revisit this fact when building System II.
# 
# It can be shown that the Gauss constraint satisfies the evolution equation
# $$\partial_{t} \mathcal{C} = 0 \; .$$
# This implies that any constraint violating numerical error remains fixed in place during the evolution. This becomes problematic when the violations grow large and spoil the physics of the simulation.

# In[1]:


import NRPy_param_funcs as par   # NRPy+: parameter interface
import indexedexp as ixp         # NRPy+: Symbolic indexed expression (e.g., tensors, vectors, etc.) support
import grid as gri               # NRPy+: Functions having to do with numerical grids
import finite_difference as fin  # NRPy+: Finite difference C code generation module
import reference_metric as rfm   # NRPy+: Reference metric support
from outputC import lhrh         # NRPy+: Core C code output module

par.set_parval_from_str("reference_metric::CoordSystem", "Spherical")
par.set_parval_from_str("grid::DIM", 3)

rfm.reference_metric()

# The name of this module ("maxwell") is given by __name__:
thismodule = __name__

# Step 0: Read the spatial dimension parameter as DIM.
DIM = par.parval_from_str("grid::DIM")

# Step 1: Set the finite differencing order to 4.
par.set_parval_from_str("finite_difference::FD_CENTDERIVS_ORDER", 4)

# Step 2: Register gridfunctions that are needed as input.
psi = gri.register_gridfunctions("EVOL", ["psi"])

# Step 3a: Declare the rank-1 indexed expressions E_{i}, A_{i},
#          and \partial_{i} \psi. Derivative variables like these
#          must have an underscore in them, so the finite
#          difference module can parse the variable name properly.
ED = ixp.register_gridfunctions_for_single_rank1("EVOL", "ED")
AD = ixp.register_gridfunctions_for_single_rank1("EVOL", "AD")
psi_dD = ixp.declarerank1("psi_dD")

# Step 3b: Declare the rank-2 indexed expression \partial_{j} A_{i},
#          which is not symmetric in its indices.
#          Derivative variables like these must have an underscore
#          in them, so the finite difference module can parse the
#          variable name properly.
AD_dD = ixp.declarerank2("AD_dD", "nosym")

# Step 3c: Declare the rank-3 indexed expression \partial_{jk} A_{i},
#          which is symmetric in the two {jk} indices.
AD_dDD = ixp.declarerank3("AD_dDD", "sym12")

# Step 4: Calculate first and second covariant derivatives, and the
#         necessary contractions.
# First covariant derivative
# D_{j} A_{i} = A_{i,j} - \Gamma^{k}_{ij} A_{k}
AD_dHatD = ixp.zerorank2()
for i in range(DIM):
    for j in range(DIM):
        AD_dHatD[i][j] = AD_dD[i][j]
        for k in range(DIM):
            AD_dHatD[i][j] -= rfm.GammahatUDD[k][i][j] * AD[k]

# Second covariant derivative
# D_{k} D_{j} A_{i} = \partial_{k} D_{j} A_{i} - \Gamma^{l}_{jk} D_{l} A_{i}
#                    - \Gamma^{l}_{ik} D_{j} A_{l}
#                   =  A_{i,jk}
#                    - \Gamma^{l}_{ij,k} A_{l}
#                    - \Gamma^{l}_{ij}   A_{l,k}
#                    - \Gamma^{l}_{jk}   A_{i;\hat{l}}
#                    - \Gamma^{l}_{ik}   A_{l;\hat{j}}
AD_dHatDD = ixp.zerorank3()
for i in range(DIM):
    for j in range(DIM):
        for k in range(DIM):
            AD_dHatDD[i][j][k] = AD_dDD[i][j][k]
            for l in range(DIM):
                AD_dHatDD[i][j][k] += - rfm.GammahatUDDdD[l][i][j][k]  * AD[l]  \
                                      - rfm.GammahatUDD[l][i][j] *    AD_dD[l][k] \
                                      - rfm.GammahatUDD[l][j][k] * AD_dHatD[i][l] \
                                      - rfm.GammahatUDD[l][i][k] * AD_dHatD[l][j]

# Covariant divergence
# D_{i} A^{i} = ghat^{ij} D_{j} A_{i}
DivA = 0
# Gradient of covariant divergence
# DivA_dD_{i} = ghat^{jk} A_{k;\hat{j}\hat{i}}
DivA_dD = ixp.zerorank1()
# Covariant Laplacian
# LapAD_{i} = ghat^{jk} A_{i;\hat{j}\hat{k}}
LapAD = ixp.zerorank1()
for i in range(DIM):
    for j in range(DIM):
        DivA += rfm.ghatUU[i][j] * AD_dHatD[i][j]
        for k in range(DIM):
            DivA_dD[i] += rfm.ghatUU[j][k] * AD_dHatDD[k][j][i]
            LapAD[i]   += rfm.ghatUU[j][k] * AD_dHatDD[i][j][k]

# Step 5: Define right-hand sides for the evolution.
AD_rhs = ixp.zerorank1()
ED_rhs = ixp.zerorank1()
for i in range(DIM):
    AD_rhs[i] = -ED[i] - psi_dD[i]
    ED_rhs[i] = -LapAD[i] + DivA_dD[i]
psi_rhs = -DivA

# Step 6: Generate C code for System I Maxwell's evolution equations,
#         print output to the screen (standard out, or stdout).
lhrh_list = []
for i in range(DIM):
    lhrh_list.append(lhrh(lhs=gri.gfaccess("rhs_gfs", "AD" + str(i)), rhs=AD_rhs[i]))
    lhrh_list.append(lhrh(lhs=gri.gfaccess("rhs_gfs", "ED" + str(i)), rhs=ED_rhs[i]))
lhrh_list.append(lhrh(lhs=gri.gfaccess("rhs_gfs", "psi"), rhs=psi_rhs))

fin.FD_outputC("stdout", lhrh_list)


# <a id='sys2'></a>
# 
# # Step 2: System II \[Back to [top](#toc)\]
# $$\label{sys2}$$
# 
# Define the auxiliary variable
# $$\Gamma \equiv \hat{D}_{i} A^{i} \; .$$
# Substituting this into Maxwell's equations yields the system
# \begin{align}
# \partial_{t} A_{i} &= -E_{i} - \hat{D}_{i} \psi \; , \\
# \partial_{t} E_{i} &= -\hat{D}_{j} \hat{D}^{j} A_{i} + \hat{D}_{i} \Gamma \; , \\
# \partial_{t} \psi &= -\Gamma \; , \\
# \partial_{t} \Gamma &= -\hat{D}_{i} \hat{D}^{i} \psi \; .
# \end{align}
# 
# 
# 
# It can be shown that the Gauss constraint now satisfies the wave equation
# $$\partial_{t}^{2} \mathcal{C} = \hat{D}_{i} \hat{D}^{i} \mathcal{C} \; .$$
# Thus, any constraint violation introduced by numerical error propagates away at the speed of light. This property increases the stability of the simulation, compared to System I above. A similar trick is used in the [BSSN formulation](Tutorial-BSSNCurvilinear.ipynb) of Einstein's equations.

# In[2]:


# We inherit here all of the definitions from System I, above

# Step 7a: Register the scalar auxiliary variable \Gamma
Gamma = gri.register_gridfunctions("EVOL", ["Gamma"])

# Step 7b: Declare the ordinary gradient \partial_{i} \Gamma
Gamma_dD = ixp.declarerank1("Gamma_dD")

# Step 8a: Construct the second covariant derivative of the scalar \psi
# \psi_{;\hat{i}\hat{j}} = \psi_{,i;\hat{j}}
#                        = \psi_{,ij} - \Gamma^{k}_{ij} \psi_{,k}
psi_dDD = ixp.declarerank2("psi_dDD", "sym01")
psi_dHatDD = ixp.zerorank2()
for i in range(DIM):
    for j in range(DIM):
        psi_dHatDD[i][j] = psi_dDD[i][j]
        for k in range(DIM):
            psi_dHatDD[i][j] += - rfm.GammahatUDD[k][i][j] * psi_dD[k]

# Step 8b: Construct the covariant Laplacian of \psi
# Lappsi = ghat^{ij} D_{j} D_{i} \psi
Lappsi = 0
for i in range(DIM):
    for j in range(DIM):
        Lappsi += rfm.ghatUU[i][j] * psi_dHatDD[i][j]

# Step 9: Define right-hand sides for the evolution.
AD_rhs = ixp.zerorank1()
ED_rhs = ixp.zerorank1()
for i in range(DIM):
    AD_rhs[i] = -ED[i] - psi_dD[i]
    ED_rhs[i] = -LapAD[i] + Gamma_dD[i]
psi_rhs = -Gamma
Gamma_rhs = -Lappsi

# Step 10: Generate C code for System II Maxwell's evolution equations,
#          print output to the screen (standard out, or stdout).
lhrh_list = []
for i in range(DIM):
    lhrh_list.append(lhrh(lhs=gri.gfaccess("rhs_gfs", "AD" + str(i)), rhs=AD_rhs[i]))
    lhrh_list.append(lhrh(lhs=gri.gfaccess("rhs_gfs", "ED" + str(i)), rhs=ED_rhs[i]))
lhrh_list.append(lhrh(lhs=gri.gfaccess("rhs_gfs", "psi"), rhs=psi_rhs))
lhrh_list.append(lhrh(lhs=gri.gfaccess("rhs_gfs", "Gamma"), rhs=Gamma_rhs))

fin.FD_outputC("stdout", lhrh_list)


# <a id='latex_pdf_output'></a>
# 
# # Step 3: Output this notebook to $\LaTeX$-formatted PDF file \[Back to [top](#toc)\]
# $$\label{latex_pdf_output}$$
# 
# The following code cell converts this Jupyter notebook into a proper, clickable $\LaTeX$-formatted PDF file. After the cell is successfully run, the generated PDF may be found in the root NRPy+ tutorial directory, with filename
# [Tutorial-MaxwellCurvilinear.pdf](Tutorial-MaxwellCurvilinear.pdf) (Note that clicking on this link may not work; you may need to open the PDF file through another means.)

# In[3]:


import cmdline_helper as cmd    # NRPy+: Multi-platform Python command-line interface
cmd.output_Jupyter_notebook_to_LaTeXed_PDF("Tutorial-MaxwellCurvilinear")