#!/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>
# 
# # IDMaxwellNRPy: An Einstein Toolkit Initial Data Thorn for Maxwell's equations
# 
# ## Author: Zach Etienne
# ### Formatting improvements courtesy Brandon Clark
# 
# ### NRPy+ Source Code for this module: [Maxwell/MaxwellCartesian_ID.py](../edit/Maxwell/MaxwellCartesian_ID.py) [\[tutorial\]](Tutorial-MaxwellCartesian.ipynb) Constructs these Maxwell initial data as SymPy expressions.
# 
# ## Introduction:
# In this part of the tutorial, we will construct an Einstein Toolkit (ETK) thorn (module) that will set up *initial data* for Maxwell's equations. In the [Tutorial-MaxwellCartesian](Tutorial-MaxwellCartesian.ipynb), tutorial notebook we used NRPy+ to contruct the SymPy expressions for plane-wave initial data. 
# 
# We will construct this thorn in two steps.
# 
# 1. Call on NRPy+ to convert the SymPy expressions for the initial data into one C-code kernel.
# 2. Write the C code and linkages to the Einstein Toolkit infrastructure (i.e., the .ccl files) to complete this Einstein Toolkit module.

# <a id='toc'></a>
# 
# # Table of Contents
# $$\label{toc}$$
# 
# This notebook is organized as follows
# 
# 1. [Step 1](#initializenrpy): Call on NRPy+ to convert the SymPy expression for the Maxwell's equations initial data into a C-code kernel
# 1. [Step 2](#einstein): Interfacing with the Einstein Toolkit
#     1. [Step 2.a](#einstein_c): Constructing the Einstein Toolkit C-code calling functions that include the C code kernels
#     1. [Step 2.b](#einstein_ccl): CCL files - Define how this module interacts and interfaces with the larger Einstein Toolkit infrastructure
#     1. [Step 2.c](#einstein_list): Add the C code to the Einstein Toolkit compilation list
# 1. [Step 3](#latex_pdf_output): Output this notebook to $\LaTeX$-formatted PDF
# 

# <a id='initializenrpy'></a>
# 
# # Step 1: Call on NRPy+ to convert the SymPy expression for the Maxwell's equations initial data into a C-code kernel \[Back to [top](#toc)\]
# $$\label{initializenrpy}$$
# 
# After importing the core modules, since we are writing an ETK thorn, we'll need to set `"grid::GridFuncMemAccess"` to `"ETK"`. SymPy expressions for plane wave initial data are written inside [Maxwell/MaxwellCartesian_ID.py](../edit/Maxwell/MaxwellCartesian_ID.py), and we simply import them for use here.

# In[1]:


# Step 1: Call on NRPy+ to convert the SymPy expression
#         for the Maxwell's equations initial data into a C-code kernel

# Step 1a: Import needed NRPy+ core modules:
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 loop as lp                # NRPy+: Generate C code loops
import sympy as sp               # SymPy: The Python computer algebra package upon which NRPy+ depends
from outputC import lhrh         # NRPy+: Core C code output module
import finite_difference as fin  # NRPy+: Finite difference C code generation module

# Step 1b: This is an Einstein Toolkit (ETK) thorn. Here we
#          tell NRPy+ that gridfunction memory access will
#          therefore be in the "ETK" style.
par.set_parval_from_str("grid::GridFuncMemAccess","ETK")
#Set the spatial dimension parameter to 3.
par.set_parval_from_str("grid::DIM", 3)
DIM = par.parval_from_str("grid::DIM")

import Maxwell.MaxwellCartesian_Evol as mwrhs # This declare the parameter System_to_use needed for MaxwellCartesian_ID
# Step 1c: Call the MaxwellCartesian_ID() function from within the
#          Maxwell/MaxwellCartesian_ID.py.py module.
import Maxwell.MaxwellCartesian_ID as mwid
# Step 3: Set up the plane wave initial data. This sets uu_ID and vv_ID.
mwid.MaxwellCartesian_ID()

# Step 4: Register gridfunctions so they can be written to by NRPy.
# System I:
AID = ixp.register_gridfunctions_for_single_rank1("EVOL","AID")
EID = ixp.register_gridfunctions_for_single_rank1("EVOL","EID")
psiI = gri.register_gridfunctions("EVOL","psiI")

# Step 5: FIXME: ADD DESCRIPTION
for i in range(DIM):
    AID[i] = mwid.AidD[i]
    EID[i] = mwid.EidD[i]
psiI = mwid.psi_ID

# Step 4: Register gridfunctions so they can be written to by NRPy.
# System II:
AIID = ixp.register_gridfunctions_for_single_rank1("EVOL","AIID")
EIID = ixp.register_gridfunctions_for_single_rank1("EVOL","EIID")
psiII = gri.register_gridfunctions("EVOL","psiII")
Gamma = gri.register_gridfunctions("EVOL","Gamma")

# Step 5: FIXME: ADD DESCRIPTION
for i in range(DIM):
    AIID[i] = mwid.AidD[i]
    EIID[i] = mwid.EidD[i]
psiII = mwid.psi_ID
Gamma = mwid.Gamma_ID

# Step 6: Create the C code output kernel.
Maxwell_ID_to_print = [\
                        lhrh(lhs=gri.gfaccess("out_gfs","AID0"),rhs=AID[0]),\
                        lhrh(lhs=gri.gfaccess("out_gfs","AID1"),rhs=AID[1]),\
                        lhrh(lhs=gri.gfaccess("out_gfs","AID2"),rhs=AID[2]),\
                        lhrh(lhs=gri.gfaccess("out_gfs","EID0"),rhs=EID[0]),\
                        lhrh(lhs=gri.gfaccess("out_gfs","EID1"),rhs=EID[1]),\
                        lhrh(lhs=gri.gfaccess("out_gfs","EID2"),rhs=EID[2]),\
                        lhrh(lhs=gri.gfaccess("out_gfs","psiI"),rhs=psiI),\
                        lhrh(lhs=gri.gfaccess("out_gfs","AIID0"),rhs=AIID[0]),\
                        lhrh(lhs=gri.gfaccess("out_gfs","AIID1"),rhs=AIID[1]),\
                        lhrh(lhs=gri.gfaccess("out_gfs","AIID2"),rhs=AIID[2]),\
                        lhrh(lhs=gri.gfaccess("out_gfs","EIID0"),rhs=EIID[0]),\
                        lhrh(lhs=gri.gfaccess("out_gfs","EIID1"),rhs=EIID[1]),\
                        lhrh(lhs=gri.gfaccess("out_gfs","EIID2"),rhs=EIID[2]),\
                        lhrh(lhs=gri.gfaccess("out_gfs","psiII"),rhs=psiII),\
                        lhrh(lhs=gri.gfaccess("out_gfs","Gamma"),rhs=Gamma),]
Maxwell_ID_CcodeKernel = fin.FD_outputC("returnstring",Maxwell_ID_to_print)

Maxwell_ID_looped = lp.loop(["i2","i1","i0"],["0","0","0"],["cctk_lsh[2]","cctk_lsh[1]","cctk_lsh[0]"],\
                            ["1","1","1"],["#pragma omp parallel for","",""],"",\
                            Maxwell_ID_CcodeKernel.replace("time","cctk_time"))

# Step 7: Create directories for the thorn if they don't exist.
get_ipython().system("mkdir MaxwellID     2>/dev/null # 2>/dev/null: Don't throw an error if the directory already exists.")
get_ipython().system("mkdir MaxwellID/src 2>/dev/null # 2>/dev/null: Don't throw an error if the directory already exists.")

# Step 8: Write the C code kernel to file.
with open("MaxwellID/src/Maxwell_ID.h", "w") as file:
    file.write(str(Maxwell_ID_looped))


# <a id='einstein'></a>
# 
# # Step 2: Interfacing with the Einstein Toolkit \[Back to [top](#toc)\]
# $$\label{einstein}$$
# 

# <a id='einstein_c'></a>
# 
# ## Step 2.a: Constructing the Einstein Toolkit C-code calling functions that include the C code kernels \[Back to [top](#toc)\]
# $$\label{einstein_c}$$
# 
# We will write another C file with the functions we need here.

# In[2]:


get_ipython().run_cell_magic('writefile', 'MaxwellID/src/InitialData.c', '#include <math.h>\n#include <stdio.h>\n\n#include "cctk.h"\n#include "cctk_Parameters.h"\n#include "cctk_Arguments.h"\n\nvoid Maxwell_InitialData(CCTK_ARGUMENTS)\n{\n  DECLARE_CCTK_ARGUMENTS\n  DECLARE_CCTK_PARAMETERS\n\n  const CCTK_REAL *xGF = x;\n  const CCTK_REAL *yGF = y;\n  const CCTK_REAL *zGF = z;\n  const CCTK_REAL *gammaDD00GF = gxx;\n  const CCTK_REAL *gammaDD01GF = gxy;\n  const CCTK_REAL *gammaDD02GF = gxz;\n  const CCTK_REAL *gammaDD11GF = gyy;\n  const CCTK_REAL *gammaDD12GF = gyz;\n  const CCTK_REAL *gammaDD22GF = gzz;\n#include "Maxwell_ID.h"\n}\n')


# <a id='einstein_ccl'></a>
# 
# ## Step 2.b: CCL files - Define how this module interacts and interfaces with the larger Einstein Toolkit infrastructure \[Back to [top](#toc)\]
# $$\label{einstein_ccl}$$
# 
# Writing a module ("thorn") within the Einstein Toolkit requires that three "ccl" files be constructed, all in the root directory of the thorn:
# 
# 1. `interface.ccl`: defines the gridfunction groups needed, and provides keywords denoting what this thorn provides and what it should inherit from other thorns. Specifically, this file governs the interaction between this thorn and others; more information can be found in the [official Einstein Toolkit documentation](http://einsteintoolkit.org/usersguide/UsersGuidech12.html#x17-178000D2.2). 
# With "implements", we give our thorn its unique name. By "inheriting" other thorns, we tell the Toolkit that we will rely on variables that exist and are declared "public" within those functions.

# In[3]:


get_ipython().run_cell_magic('writefile', 'MaxwellID/interface.ccl', 'implements: MaxwellID\ninherits: admbase MaxwellEvol grid\n')


# 2. `param.ccl`: specifies free parameters within the thorn, enabling them to be set at runtime. It is required to provide allowed ranges and default values for each parameter. More information on this file's syntax can be found in the [official Einstein Toolkit documentation](http://einsteintoolkit.org/usersguide/UsersGuidech12.html#x17-183000D2.3).

# In[4]:


get_ipython().run_cell_magic('writefile', 'MaxwellID/param.ccl', 'shares: grid\n\nUSES KEYWORD type\n\nrestricted:\nCCTK_KEYWORD initial_data "Type of initial data"\n{\n  "toroid"      :: "Toroidal Dipole field"\n} "toroid"\n\nrestricted:\nCCTK_REAL amp "The amplitude of the initial wave packet"\n{\n 0.0:* :: "Should be positive"\n} 1.0\n\nrestricted:\nCCTK_REAL lam "A parameter describing the size of the wave package"\n{\n 0.0:* :: "Should be positive"\n} 1.0\n')


# 3. `schedule.ccl`: allocates storage for gridfunctions, defines how the thorn's functions should be scheduled in a broader simulation, and specifies the regions of memory written to or read from gridfunctions. `schedule.ccl`'s official documentation may be found [here](http://einsteintoolkit.org/usersguide/UsersGuidech12.html#x17-186000D2.4). 
# 
# We specify here the standardized ETK "scheduling bins" in which we want each of our thorn's functions to run.

# In[5]:


get_ipython().run_cell_magic('writefile', 'MaxwellID/schedule.ccl', '\nschedule Maxwell_InitialData at CCTK_INITIAL as Maxwell_InitialData\n{\n  STORAGE:       MaxwellEvol::system_I[3]\n  STORAGE:       MaxwellEvol::system_II[3]\n  LANG:          C\n  READS: admbase::gxx(Everywhere)\n  READS: admbase::gxy(Everywhere)\n  READS: admbase::gxz(Everywhere)\n  READS: admbase::gyy(Everywhere)\n  READS: admbase::gyz(Everywhere)\n  READS: admbase::gzz(Everywhere)\n  READS: grid::x(Everywhere)\n  READS: grid::y(Everywhere)\n  READS: grid::y(Everywhere)\n  WRITES: AI0GF(Everywhere)\n  WRITES: AI1GF(Everywhere)\n  WRITES: AI2GF(Everywhere)\n  WRITES: EI0GF(Everywhere)\n  WRITES: EI1GF(Everywhere)\n  WRITES: EI2GF(Everywhere)\n  WRITES: psiIGF(Everywhere)\n  WRITES: AII0GF(Everywhere)\n  WRITES: AII1GF(Everywhere)\n  WRITES: AII2GF(Everywhere)\n  WRITES: EII0GF(Everywhere)\n  WRITES: EII1GF(Everywhere)\n  WRITES: EII2GF(Everywhere)\n  WRITES: psiIIGF(Everywhere)\n  WRITES: GammaGF(Everywhere)\n} "Initial data for Maxwell\'s equations"\n')


# <a id='einstein_list'></a>
# 
# ## Step 2.c: Add the C code to the Einstein Toolkit compilation list \[Back to [top](#toc)\]
# $$\label{einstein_list}$$
# 
# We will also need `make.code.defn`, which indicates the list of files that need to be compiled. This thorn only has the one C file to compile.

# In[6]:


get_ipython().run_cell_magic('writefile', 'MaxwellID/src/make.code.defn', 'SRCS = InitialData.c\n')


# <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-ETK_thorn-IDMaxwellNRPy.pdf](Tutorial-ETK_thorn-IDMaxwellNRPy.pdf) (Note that clicking on this link may not work; you may need to open the PDF file through another means.)

# In[7]:


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