Ejemplo n.º 1
0
def classify_etc(loadfile, savefile, v):
    """
    1. Moments algorithm
    2. HybridTrack algorithm
    3. Classify
    """

    progressflag = not file_vars()[0]

    pn = 'pix10_5noise15'
    HTname = 'python HT v1.52'
    MTname = 'moments v1.0'

    vprint(v, 1, 'Starting {} at {} with {}% mem usage'.format(
        loadfile, time.ctime(), psutil.virtual_memory().percent))

    pyobj_to_h5 = {}

    try:
        with h5py.File(loadfile, 'a', driver='core') as f, h5py.File(
                savefile, 'a', driver='core') as h5save:
            #
            n = 0
            if progressflag:
                pbar = progressbar.ProgressBar(
                    widgets=[progressbar.Percentage(), ' ',
                             progressbar.Bar(), ' ',
                             progressbar.ETA()], maxval=len(f))
                pbar.start()

            keylist = f.keys()
            keylist.sort()
            for ind in keylist:
                vprint(v, 3, 'Beginning track {} in {}'.format(ind, loadfile))
                if int(ind) % 50 == 0:
                    vprint(v, 2,
                           'Beginning track {} in {}'.format(ind, loadfile))

                # evt is for the G4Track object
                evt = f[ind]
                # shouldn't get a KeyError because we are looping over ind

                # trk is for the Track object
                try:
                    trk = evt[pn]
                except KeyError:
                    vprint(v, 1,
                           '**Missing key {} in {}{}, skipping'.format(
                               pn, loadfile, evt.name))
                    continue

                trkpath = ind + '/' + pn
                errorcode = 0
                n += 1

                if n > 50:
                    # testing
                    # vprint(v, 1, 'Finished 50 files, exiting')
                    # break
                    pass

                # load G4Track
                vprint(v, 3, 'Loading G4Track {} in {}'.format(ind, loadfile))
                try:
                    this_g4track = G4Track.from_dth5(evt)
                except trackio.InterfaceError:
                    vprint(v, 2, 'InterfaceError in G4Track at {}, {}'.format(
                        loadfile, ind))
                    continue

                # load Track
                vprint(v, 3, 'Loading Track {} in {}'.format(ind, loadfile))
                try:
                    this_track = Track.from_dth5(trk, g4track=this_g4track)
                except trackio.InterfaceError:
                    vprint(v, 2, 'InterfaceError in Track at {}, {}'.format(
                        loadfile, ind))
                    continue

                # check for error codes - Track.from_dth5 may return an int
                if isinstance(this_track, int):
                    this_save = h5save.create_group(trkpath)
                    this_save.attrs.create(
                        'errorcode', this_track, shape=np.shape(this_track))
                    continue

                # run moments algorithm
                vprint(v, 3, 'Running moments on track {} in {}'.format(
                    ind, loadfile))
                try:
                    mom = tm.MomentsReconstruction(this_track.image)
                    mom.reconstruct()
                except ht.NoEndsFound:
                    mom.alpha = None
                    errorcode = 4
                    this_save = h5save.create_group('mom_' + trkpath)
                    this_save.attrs.create(
                        'errorcode', errorcode, shape=np.shape(errorcode))
                else:
                    # write into savefile
                    trackio.write_object_to_hdf5(
                        mom, h5save, 'mom_' + trkpath, pyobj_to_h5=pyobj_to_h5)
                # write into track object
                if mom.alpha:
                    this_track.add_algorithm(
                        MTname,
                        alpha_deg=mom.alpha * 180 / np.pi,
                        beta_deg=np.nan, info=None)
                else:
                    this_track.add_algorithm(
                        MTname,
                        alpha_deg=np.nan,
                        beta_deg=np.nan, info=None)

                # run HT algorithm (v1.52)
                vprint(v, 3, 'Running HT on track {} in {}'.format(
                    ind, loadfile))
                try:
                    _, HTinfo = ht.reconstruct(this_track)
                except (ht.InfiniteLoop, ht.NoEndsFound):
                    # write empty into track object
                    this_track.add_algorithm(
                        HTname,
                        alpha_deg=np.nan,
                        beta_deg=np.nan,
                        info=None)
                else:
                    # trim memory usage
                    if hasattr(HTinfo, 'ridge'):
                        if HTinfo.ridge:
                            for ridgept in HTinfo.ridge:
                                ridgept.cuts = None
                                ridgept.best_cut = None
                    # write into track object
                    this_track.add_algorithm(
                        HTname,
                        alpha_deg=HTinfo.alpha_deg,
                        beta_deg=HTinfo.beta_deg,
                        info=HTinfo)

                # write py HDF5 format into savefile (including alg outputs)
                try:
                    trackio.write_object_to_hdf5(
                        this_track, h5save, trkpath,
                        pyobj_to_h5=pyobj_to_h5)
                except trackio.InterfaceError:
                    vprint(v, 1, 'InterfaceError writing to file at ' +
                           '{}, {}'.format(savefile, ind))
                    continue

                # run classifier
                vprint(v, 3, 'Running MC classifier on track {} in {}'.format(
                    ind, loadfile))
                classifier = cl.Classifier(this_track.g4track)
                try:
                    classifier.mc_classify()
                except cl.TrackTooShortError:
                    classifier.error = 'TrackTooShortError'
                else:
                    vprint(v, 3,
                           'Running ends classifier on track {} in {}'.format(
                               ind, loadfile))
                    try:
                        classifier.end_classify(this_track, mom=mom)
                    except tp.G4TrackTooBigError:
                        classifier.error = 'G4TrackTooBigError'
                    except cl.NoEnds:
                        classifier.error = 'NoEnds'
                # write into savefile
                vprint(v, 3, 'Writing classifier into {} for track {}'.format(
                    savefile, ind))
                trackio.write_object_to_hdf5(
                    classifier, h5save, 'cl_' + trkpath,
                    pyobj_to_h5=pyobj_to_h5)

                if progressflag:
                    pbar.update(n)
            if progressflag:
                pbar.finish()
        # f gets closed here
    except NotImplementedError:
        pass