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

# In[ ]:


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


# 
# 1D Maxwellian distribution function
# ===================================
# 
# We import the usual modules, and the hero of this notebook,
# the Maxwellian 1D distribution:
# 

# In[ ]:


import astropy.units as u
import matplotlib.pyplot as plt
import numpy as np
from astropy.constants import k_B, m_e


# Given we'll be plotting, import astropy's quantity support:
# 
# 

# In[ ]:


from astropy.visualization import quantity_support

from plasmapy.formulary import Maxwellian_1D

quantity_support()


# As a first example, let's get the probability density of
# finding an electron with a speed of 1 m/s if we have a
# plasma at a temperature of 30 000 K:
# 
# 

# In[ ]:


p_dens = Maxwellian_1D(
    v=1 * u.m / u.s, T=30000 * u.K, particle="e", v_drift=0 * u.m / u.s
)
print(p_dens)


# Note the units! Integrated over speed, this will give us a
# probability. Let's test that for a bunch of particles:
# 
# 

# In[ ]:


T = 3e4 * u.K
dv = 10 * u.m / u.s
v = np.arange(-5e6, 5e6, 10) * u.m / u.s


# Check that the integral over all speeds is 1
# (the particle has to be somewhere):
# 
# 

# In[ ]:


for particle in ["p", "e"]:
    pdf = Maxwellian_1D(v, T=T, particle=particle)
    integral = (pdf).sum() * dv
    print(f"Integral value for {particle}: {integral}")
    plt.plot(v, pdf, label=particle)
plt.legend()


# The standard deviation of this distribution should give us back the
# temperature:
# 
# 

# In[ ]:


std = np.sqrt((Maxwellian_1D(v, T=T, particle="e") * v**2 * dv).sum())
T_theo = (std**2 / k_B * m_e).to(u.K)

print("T from standard deviation:", T_theo)
print("Initial T:", T)