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

# In[ ]:


get_ipython().run_line_magic('matplotlib', 'inline')


# 
# Magnetostatic Fields
# =====================
# 
# An example of using PlasmaPy's `Magnetostatic` class in `physics` subpackage.
# 

# In[ ]:


import astropy.units as u
import matplotlib.pyplot as plt
import numpy as np

from plasmapy.formulary import magnetostatics
from plasmapy.plasma.sources import Plasma3D


# Some common magnetostatic fields can be generated and added to a plasma object.
# A dipole
# 
# 

# In[ ]:


dipole = magnetostatics.MagneticDipole(
    np.array([0, 0, 1]) * u.A * u.m * u.m, np.array([0, 0, 0]) * u.m
)
print(dipole)


# initialize a a plasma, where the magnetic field will be calculated on
# 
# 

# In[ ]:


plasma = Plasma3D(
    domain_x=np.linspace(-2, 2, 30) * u.m,
    domain_y=np.linspace(0, 0, 1) * u.m,
    domain_z=np.linspace(-2, 2, 20) * u.m,
)


# add the dipole field to it
# 
# 

# In[ ]:


plasma.add_magnetostatic(dipole)

X, Z = plasma.grid[0, :, 0, :], plasma.grid[2, :, 0, :]
U = plasma.magnetic_field[0, :, 0, :].value.T  # because grid uses 'ij' indexing
W = plasma.magnetic_field[2, :, 0, :].value.T  # because grid uses 'ij' indexing


# In[ ]:


plt.figure()
plt.axis("square")
plt.xlim(-2, 2)
plt.ylim(-2, 2)
plt.title(
    "Dipole field in x-z plane, generated by a dipole pointing in the z direction"
)
plt.streamplot(plasma.x.value, plasma.z.value, U, W)


# In[ ]:


cw = magnetostatics.CircularWire(
    np.array([0, 0, 1]), np.array([0, 0, 0]) * u.m, 1 * u.m, 1 * u.A
)
print(cw)


# initialize a a plasma, where the magnetic field will be calculated on
# 
# 

# In[ ]:


plasma = Plasma3D(
    domain_x=np.linspace(-2, 2, 30) * u.m,
    domain_y=np.linspace(0, 0, 1) * u.m,
    domain_z=np.linspace(-2, 2, 20) * u.m,
)


# add the circular coil field to it
# 
# 

# In[ ]:


plasma.add_magnetostatic(cw)

X, Z = plasma.grid[0, :, 0, :], plasma.grid[2, :, 0, :]
U = plasma.magnetic_field[0, :, 0, :].value.T  # because grid uses 'ij' indexing
W = plasma.magnetic_field[2, :, 0, :].value.T  # because grid uses 'ij' indexing

plt.figure()
plt.axis("square")
plt.xlim(-2, 2)
plt.ylim(-2, 2)
plt.title(
    "Circular coil field in x-z plane, generated by a circular coil in the x-y plane"
)
plt.streamplot(plasma.x.value, plasma.z.value, U, W)


# In[ ]:


gw_cw = cw.to_GeneralWire()

# the calculated magnetic field is close
print(gw_cw.magnetic_field([0, 0, 0]) - cw.magnetic_field([0, 0, 0]))


# In[ ]:


iw = magnetostatics.InfiniteStraightWire(
    np.array([0, 1, 0]), np.array([0, 0, 0]) * u.m, 1 * u.A
)
print(iw)


# initialize a a plasma, where the magnetic field will be calculated on
# 
# 

# In[ ]:


plasma = Plasma3D(
    domain_x=np.linspace(-2, 2, 30) * u.m,
    domain_y=np.linspace(0, 0, 1) * u.m,
    domain_z=np.linspace(-2, 2, 20) * u.m,
)

# add the infinite straight wire field to it
plasma.add_magnetostatic(iw)

X, Z = plasma.grid[0, :, 0, :], plasma.grid[2, :, 0, :]
U = plasma.magnetic_field[0, :, 0, :].value.T  # because grid uses 'ij' indexing
W = plasma.magnetic_field[2, :, 0, :].value.T  # because grid uses 'ij' indexing

plt.figure()
plt.title(
    "Dipole field in x-z plane, generated by a infinite straight wire "
    "pointing in the y direction"
)
plt.axis("square")
plt.xlim(-2, 2)
plt.ylim(-2, 2)
plt.streamplot(plasma.x.value, plasma.z.value, U, W)