Test direct methods with step function.py

# -*- coding: utf-8 -*-
from __future__ import division
from __future__ import print_function

import numpy as np
import abel
import matplotlib.pyplot as plt
import scipy

def direct1d(n=51):
    x = np.linspace(0, 101, n)
    
    origS = np.zeros(n)

    r0, r1 = 50, 85
    step = np.logical_and(x>r0, x<r1)
    origS[step] = 1

    projS = np.zeros(n)
    projS[step] = np.sqrt(r1**2-x[step]**2)
    projS[x<r0] = np.sqrt(r1**2-x[x<r0]**2) - np.sqrt(r0**2-x[x<r0]**2)
    
    
    
    C = abel.direct.direct_transform(origS, dr=x[1]-x[0],
                                            direction='forward', backend='C', correction=True)
    py = abel.direct.direct_transform(origS, dr=x[1]-x[0],
                                            direction='forward', backend='Python', correction=True)
                                            
    py_simps = abel.direct.direct_transform(origS, dr=x[1]-x[0],
                                            direction='forward', backend='Python', correction=True, int_func=scipy.integrate.simps)
    
    py_trapz = abel.direct.direct_transform(origS, dr=x[1]-x[0],
                                            direction='forward', backend='Python', correction=True, int_func=scipy.integrate.trapz)

    ratioS = projS.max()/py_trapz.max()

    plt.title("direct forward Abel transform - 1d Gaussian and 1d step")

    plt.plot(projS,    'b', label="Step projection")
    plt.plot(C*ratioS, 'y', label="direct_C", marker='o', ms=3, mec='none')
    plt.plot(py*ratioS,'g', label="direct_Py", marker='o', ms=3, mec='none')
    # plt.plot(py_simps*ratioS,'m', label="direct_Py w/simps", marker='o', ms=3, mec='none')
    plt.plot(py_trapz*ratioS,'r', label="direct_Py w/trapz", marker='o', ms=3, mec='none')
    
    
    
    plt.legend(loc=0, labelspacing=0.1, frameon=False)
    plt.savefig("direct-1d-gauss-step.png", dpi=200)
    plt.show()

if __name__ == "__main__":
    direct1d()