# Licensed under a 3-clause BSD style license - see LICENSE.rst
"""Utility functions to convert ROOT data to numpy / FITS data.
"""
from __future__ import absolute_import, division, print_function, unicode_literals
import warnings
from collections import OrderedDict
import numpy as np
from astropy.io import fits
from astropy.table import Table
__all__ = ['hist1d_to_table',
'graph1d_to_table',
'TH2_to_FITS_header',
'TH2_to_FITS_data',
'TH2_to_FITS',
]
__doctest_skip__ = ['TH2_to_FITS']
[docs]def hist1d_to_table(hist):
"""Convert 1D ROOT histogram into astropy table.
Parameters
----------
hist : ROOT.TH1
ROOT histogram
Returns
-------
table : `~astropy.table.Table`
Astropy table
"""
bins = range(1, hist.GetNbinsX() + 1)
data = OrderedDict()
names = [('x', 'GetBinCenter'),
('x_bin_lo', 'GetBinLowEdge'),
('x_bin_width', 'GetBinWidth'),
('y', 'GetBinContent'),
('y_err', 'GetBinError'),
('y_err_lo', 'GetBinErrorLow'),
('y_err_hi', 'GetBinErrorUp'),
]
for column, method in names:
try:
getter = getattr(hist, method)
data[column] = [getter(i) for i in bins]
# Note: `GetBinErrorLow` is not available in old ROOT versions!?
except AttributeError:
pass
table = Table(data)
table['x_bin_hi'] = table['x_bin_lo'] + table['x_bin_width']
return table
[docs]def graph1d_to_table(graph):
"""Convert ROOT TGraph to an astropy Table.
Parameters
----------
graph : ROOT.TGraph
ROOT graph
Returns
-------
table : `~astropy.table.Table`
Astropy table
"""
bins = range(0, graph.GetN())
data = OrderedDict()
names = [('x', 'GetX'),
('x_err', 'GetEX'),
('x_err_lo', 'GetEXlow'),
('x_err_hi', 'GetEXhigh'),
('y', 'GetY'),
('y_err', 'GetEY'),
('y_err_lo', 'GetEYlow'),
('y_err_hi', 'GetEYhigh'),
]
for column, method in names:
try:
buffer_ = getattr(graph, method)()
data[column] = [buffer_[i] for i in bins]
except IndexError:
pass
table = Table(data)
return table
[docs]def TH2_to_FITS_data(hist, flipx=True):
"""Convert TH2 bin values into a numpy array.
Parameters
----------
hist : ROOT.TH2
ROOT histogram
Returns
-------
data : array
Histogram data
"""
# Note: Numpy array index order is (y, x), whereas ROOT TH2 has (x, y)
nx, ny = hist.GetNbinsX(), hist.GetNbinsY()
# TODO: This doesn't work properly:
# dtype = type(hist.GetBinContent(0))
dtype = 'float32'
array = np.empty((ny, nx), dtype=dtype)
for ix in range(nx):
for iy in range(ny):
array[iy, ix] = hist.GetBinContent(ix, iy)
if flipx:
array = array[:, ::-1]
return array
[docs]def TH2_to_FITS(hist, flipx=True):
"""Convert ROOT 2D histogram to FITS format.
Parameters
----------
hist : ROOT.TH2
2-dim ROOT histogram
Returns
-------
hdu : `~astropy.io.fits.ImageHDU`
Histogram in FITS format.
Examples
--------
>>> import ROOT
>>> from gammapy.utils.root import TH2_to_FITS
>>> root_hist = ROOT.TH2F()
>>> fits_hdu = TH2_to_FITS(root_hist)
>>> fits_hdu.writetofits('my_image.fits')
"""
header = TH2_to_FITS_header(hist, flipx)
if header['CDELT1'] > 0:
warnings.warn('CDELT1 > 0 might not be handled properly.'
'A TH2 representing an astro image should have '
'a reversed x-axis, i.e. xlow > xhi')
data = TH2_to_FITS_data(hist, flipx)
hdu = fits.ImageHDU(data=data, header=header)
return hdu
def tree_to_table(tree, tree_name):
"""Convert a ROOT TTree to an astropy Table.
Parameters
----------
tree : ROOT.TTree
ROOT TTree
Returns
-------
table : `~astropy.table.Table`
ROOT tree data as an astropy table.
"""
from rootpy import asrootpy
from rootpy.io import open
from rootpy.root2array import tree_to_recarray
tree = asrootpy(tree)
array = tree.to_array()
file = open(infile)
tree_name = 'TableAllMu_WithoutNan' # 'ParTree_Postselect'
tree = file.get(tree_name, ignore_unsupported=True)
array = tree_to_recarray(tree)
table = recarray_to_table(array)
# Remove empty columns
empty_columns = []
for col in table.columns:
#if table['col'].min() == table['col'].max()
if (table['col'] == 0).all():
empty_columns.append(col)
print(empty_columns)
table.remove_columns(empty_columns)
for name in array.dtype.names:
# FITS can't save these types.
data = array[name]
if data.dtype == np.dtype('uint64'):
data = data.copy().astype('int64')
table.add_column(name, data)
return table