# Example: Linbasex¶

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

import numpy as np
import abel
import os
import bz2

import matplotlib.pylab as plt

# This example demonstrates linbasex inverse Abel transform
# of a velocity-map image of photoelectrons from O2- photodetachment at 454 nm.
# Measured at  The Australian National University
# J. Chem. Phys. 133, 174311 (2010) DOI: 10.1063/1.3493349

# Load image as a numpy array - numpy handles .gz, .bz2
imagefile = bz2.BZ2File('data/O2-ANU1024.txt.bz2')

IM = IM[::2,::2]
# the [::2, ::2] reduces the image size x1/2, decreasing processing memory load

# Image center should be mid-pixel and the image square,
# center=convolution takes care of this

un = [0, 2]  # spherical harmonic orders
proj_angles = np.arange(0, 2*np.pi, np.pi/20) # projection angles
# adjust these parameter to 'improve' the look
smoothing = 0.9  # smoothing Gaussian 1/e width
threshold = 0.01 # exclude small amplitude Newton spheres
# no need to change these
clip = 0

# linbasex inverse Abel transform
LIM = abel.Transform(IM, method="linbasex", origin="convolution",
center_options=dict(square=True),
transform_options=dict(basis_dir=None, return_Beta=True,
legendre_orders=un,
proj_angles=proj_angles,
smoothing=smoothing,
threshold=threshold))

# angular, and radial integration - direct from linbasex transform
# as class attributes
speed  = LIM.Beta[0]
anisotropy = LIM.Beta[1]

# normalize to max intensity peak i.e. max peak height = 1
speed /= speed[200:].max()  # exclude transform noise near centerline of image

# plots of the analysis
fig = plt.figure(figsize=(11, 5))
ax1 = plt.subplot2grid((1, 2), (0, 0))
ax2 = plt.subplot2grid((1, 2), (0, 1))

# join 1/2 raw data : 1/2 inversion image
inv_IM = LIM.transform
cols = inv_IM.shape[1]
c2 = cols//2
vmax = IM[:, :c2-100].max()
inv_IM *= vmax/inv_IM[:, c2+100:].max()
JIM = np.concatenate((IM[:, :c2], inv_IM[:, c2:]), axis=1)

# raw data
im1 = ax1.imshow(JIM, origin='upper', aspect='auto', vmin=0, vmax=vmax)
ax1.set_xlabel('column (pixels)')
ax1.set_ylabel('row (pixels)')
ax1.set_title('VMI, inverse Abel: {:d}x{:d}'.format(*inv_IM.shape),
fontsize='small')

# Plot the 1D speed distribution and anisotropy parameter ("looks" better
# if multiplied by the intensity)
ax2.plot(radial, speed*anisotropy, label=r'anisotropy $\times$ speed')
row, cols = IM.shape
ax2.axis(xmin=100*cols/1024, xmax=500*cols/1024, ymin=-1.5, ymax=1.8)
ax2.set_title("speed, anisotropy parameter", fontsize='small')
ax2.set_ylabel('intensity')
r'linbasex inverse Abel transform of O$_{2}{}^{-}$ electron velocity-map image',