# -*- coding: utf-8 -*-
# get the modules needed for mathematical functions and plotting
#
#import math
from numpy import *
#%matplotlib inline
import matplotlib.pyplot as plt
#
# enter any constants needed
#
m_e = 511 # keV
print("mass of electron ",m_e," keV.")
#
# Compton scattering, energy of outgoing photon as function of angle
#
def comptonE(energy_in,theta):
    energy_out = energy_in/(1+energy_in*(1-math.cos(theta))/m_e)
    return energy_out
#
# plot energy over some angular range
#
nIter = 100
theta_bot = 0
theta_top = 3.1415926
theta_step = (theta_top-theta_bot)/nIter
energy_in = 0.4*m_e
theta = zeros((nIter+1))
photonE = zeros((nIter+1))
#
# do the calculation
#
for i in range(0,nIter+1):
    theta[i] = i*theta_step
    photonE[i] = comptonE(energy_in,theta[i])
#    print("Scattering angle ",theta[i]," photon energy ",photonE[i])
#
# fill a plot
#
plt.plot(theta,photonE)
plt.title('Energy of Compton scattered photon');
plt.xlabel('Angle (radians)')
plt.ylabel('Energy (keV)')