First attempt to parse the cheetah CXIdb files from LCLS

parse_cxi.py

import tables
import numpy as np


class CheetahCXI(object):
    """
    A parser for the CXIdb files generated by cheetah at LCLS.
    """
    
    def __init__(self, filename):
        
        self._fhandle = tables.File(filename)
        self._ds1_data = self._fhandle.root.entry_1.instrument_1.detector_1.data
        
        return


    def __del__(self):
        self._fhandle.close()
        return
    
        
    @property
    def num_shots(self):
        return self._ds1_data.shape[0]
    
        
    def energy(self, mean=True):
        """
        Returns the energy in eV.
        
        Parameters
        ----------
        mean : bool
            Return the mean. If `False`, returns the measured energy for each
            shot.
        """
        n = self._get_nodes('photon_energy_eV', strict=True)[0]
        if mean:
            return n.read().mean()
        else:
            return n.read()
    
        
    def shot(self, shot_number, detector='ds1'):
        """
        Read the intensity data for a single shot.
        
        Parameters
        ----------
        shot_number : int
            Which shot to read.
            
        detector : str
            Which detector to read the data for.
            
        Returns
        -------
        shot_data : np.ndarray
            An array of intensity data. The first dimension indexes shots, the
            next two are a single image in cheetah's CSPAD format.
        """
        if detector == 'ds1':
            return self._ds1_data.read(shot_number)
        else:
            raise NotImplementedError('only ds1 support in right now')
    
        
    def shot_range(self, start, stop, detector='ds1'):
        """
        Read a range of shots.
        
        Parameters
        ----------
        start : int
            Shot to start at (zero indexed).
            
        stop : int
            Where to stop.
            
        detector : str
            Which detector to read the data for.
            
        Returns
        -------
        shot_data : np.ndarray
            An array of intensity data. The first dimension indexes shots, the
            next two are a single image in cheetah's CSPAD format.
        """
        if detector == 'ds1':
            return self._ds1_data.read(start, stop)
        else:
            raise NotImplementedError('only ds1 support in right now')
    
        
    def shot_iterator(self, detector='ds1'):
        """
        Return an iterable that enables looping over shot data on disk. This
        object lets python access shots one at a time without loading all of
        them into memory
        
        Parameters
        ----------
        detector : str
            Which detector to read the data for.
            
        Returns
        -------
        shotiter : tables.earray.EArray
            An iterable element array .
        """
        
        if detector == 'ds1':
            return self._ds1_data.iterrows()
        else:
            raise NotImplementedError('only ds1 support in right now')
    
            
    def _get_groups(self, name, root='/'):
        """
        Locates groups in the HDF5 file structure, beneath `root`, with name
        matching `name`.

        Returns
        -------
        groups : list
            A list of pytables group objects
        """
        groups = []
        for g in self._fhandle.walkGroups(root):                
            gname = g.__str__().split('/')[-1].split()[0]
            if gname.find(name) == 0:
                groups.append(g)
        return groups
    

    def _get_nodes(self, name, root='/', strict=False):
        """
        Locates nodes in the HDF5 file structure, beneath `root`, with name
        matching `name`.

        Returns
        -------
        nodes : list
            A list of pytables nodes objects
        """
        nodes = []
        for n in self._fhandle.walkNodes(root):                
            nname = n.__str__().split('/')[-1].split()[0]
            if not isinstance(n, tables.link.SoftLink):
                if strict:
                    if nname == name:
                        nodes.append(n)
                else:
                    if nname.find(name) == 0:
                        nodes.append(n)
        return nodes
    
    
if __name__ == '__main__':
    
    t = CheetahCXI('test.cxi')
    
    print t.energy()
    print t.num_shots
    print t.shot(0)
    print t.shot_range(0, 2)
    
    for s in t.shot_iterator():
        print s
        break
   
    del t