polarpy.polar2hdf5.polar_polarization_to_hdf5()

Complexity

Conditions

9

Size

Total Lines	98
Code Lines	44

Duplication

Lines	0
Ratio	0 %

Importance

Changes

0

import h5py
    def polar_polarization_to_hdf5(polarization_root_file, hdf5_out_file):
        """
        Converts the ROOT POLAR response into an HDF5 file so that users are not 
        dependent on ROOT.

        :param polarization_root_file: The ROOT file from which to build the response
        :param hdf5_out_file: The output HDF5 file name
        """

        # create a few lists so that we can hold the values

        energy = []
        degree = []
        angle = []

        energy_str = []
        degree_str = []
        angle_str = []

        # open the ROOT file

        with open_ROOT_file(polarization_root_file) as f:

            # This looks at all the info in the ROOT file
            # It is gross because ROOT is gross.

            tmp = [key.GetName() for key in f.GetListOfKeys()]
            tmp = filter(lambda x: 'sim' in x, tmp)
            for tmp2 in tmp:
                _, x, y, z = tmp2.split('_')

                energy.append(float(x))
                degree.append(float(y))
                angle.append(float(z))

                energy_str.append(x)
                degree_str.append(y)
                angle_str.append(z)

            # There are duplicates everywhere.
            # This makes sure we only grab what we need.

            energy = np.array(np.unique(energy))
            degree = np.array(np.unique(degree))
            angle = np.array(np.unique(angle))

            energy_str = np.array(np.unique(energy_str))
            degree_str = np.array(np.unique(degree_str))
            angle_str = np.array(np.unique(angle_str))

            # just to get the bins
            # must change this from ints later

            file_string = 'sim_%s_%s_%s' % (energy_str[1], degree_str[1], angle_str[1])

            bins, _, hist = th2_to_arrays(f.Get(file_string))

            out_matrix = np.zeros((len(energy), len(angle), len(degree), len(hist)))

            # Now we will build the HDF5 file. Much eaasier because the format is
            # beautiful.

            with h5py.File(hdf5_out_file, 'w', libver='latest') as database:

                for i, x in enumerate(energy_str):

                    for j, y in enumerate(angle_str):

                        for k, z in enumerate(degree_str):

                            file_string = 'sim_%s_%s_%s' % (x, z, y)

                            _, _, hist = th2_to_arrays(f.Get(file_string))

                            # Some beautiful matrix math

                            out_matrix[i, j, k, :] = hist

                # write to the matrix extension

                database.create_dataset('matrix', data=out_matrix, compression='lzf')

                if np.min(bins) < 0:
                    # we will try to automatically correct for the
                    # badly specified bins
                    bins = np.array(bins)

                    bins += -np.min(bins)

                    assert np.min(bins) >= 0, 'The scattering bins have egdes less than zero'
                    assert np.max(bins) <= 360, 'The scattering bins have egdes greater than 360'

                # Save all this out. We MUST write some docs describing the format at some point

                database.create_dataset('bins', data=bins, compression='lzf')
                database.create_dataset('pol_ang', data=angle, compression='lzf')
                database.create_dataset('pol_deg', data=degree, compression='lzf')
                database.create_dataset('energy', data=energy, compression='lzf')