<reponame>icecube/simweights
from glob import glob

import pylab as plt
from icecube import dataio, simclasses

import simweights

filedir = "/data/sim/scratch/IceSim/dynstack/21682/0000000-0000999/"
filelist = sorted(glob(filedir + "Level2_IC86.2016_corsika.021682.00000?.i3.zst"))

# SimWeights iterates over the info table row-wise so to mimic it we need a list
I3PrimaryInjectorInfo = []

# However, the weight object is iterated column-wise so we need to create a dictionary of lists
I3CorsikaWeight = {"energy": [], "type": [], "zenith": [], "weight": []}

# loop over all the files we want to read
for filename in filelist:

    # open the i3 files with the dataio interface
    infile = dataio.I3File(filename)
    print("Reading " + filename)

    # loop over the frames in the file
    while infile.more():

        # get the frame
        frame = infile.pop_frame()

        # if this is an S-Frame
        if frame.Stop == frame.Simulation:
            # get the info from the frame
            info = frame["I3PrimaryInjectorInfo"]

            # convert the info object into a dictionary which represents a row
            inforow = {x: getattr(info, x) for x in dir(info) if x[0] != "_"}

            # add the row to to the list
            I3PrimaryInjectorInfo.append(inforow)

        # if this is a physics event in the right sub-event stream
        elif frame.Stop == frame.Physics and frame["I3EventHeader"].sub_event_stream == "InIceSplit":

            # get the weighting object
            w = frame["I3CorsikaWeight"]

            # for each of the columns we need get it from the frame object
            # and put it in the correct column
            I3CorsikaWeight["energy"].append(w.primary.energy)
            I3CorsikaWeight["type"].append(w.primary.type)
            I3CorsikaWeight["zenith"].append(w.primary.dir.zenith)
            I3CorsikaWeight["weight"].append(w.weight)

# make a dictionary object to mimic the file structure of a pandas file
fileobj = dict(I3PrimaryInjectorInfo=I3PrimaryInjectorInfo, I3CorsikaWeight=I3CorsikaWeight)

# create the weighter object
weighter = simweights.CorsikaWeighter(fileobj)

# create an object to represent our cosmic-ray primary flux model
flux = simweights.GaisserH4a()

# get the weights by passing the flux to the weighter
weights = weights = weighter.get_weights(flux)

# print some info about the weighting object
print(weighter.tostring(flux))

# create equal spaced bins in log space
bins = plt.geomspace(3e4, 1e6, 50)

# get energy of the primary cosmic-ray from `PolyplopiaPrimary`
primary_energy = weighter.get_weight_column("energy")

# histogram the primary energy with the weights
plt.hist(primary_energy, weights=weights, bins=bins)

# make the plot look good
plt.loglog()
plt.xlabel("Primary Energy [GeV]")
plt.ylabel("Event Rate [Hz]")
plt.xlim(bins[0], bins[-1])
plt.ylim(0.1, 10)
plt.show()

from glob import glob

import pylab as plt
from icecube import dataio, simclasses

import simweights

filedir = "/data/sim/scratch/IceSim/dynstack/21682/0000000-0000999/"
filelist = sorted(glob(filedir + "Level2_IC86.2016_corsika.021682.00000?.i3.zst"))

# SimWeights iterates over the info table row-wise so to mimic it we need a list
I3PrimaryInjectorInfo = []

# However, the weight object is iterated column-wise so we need to create a dictionary of lists
I3CorsikaWeight = {"energy": [], "type": [], "zenith": [], "weight": []}

# loop over all the files we want to read
for filename in filelist:

    # open the i3 files with the dataio interface
    infile = dataio.I3File(filename)
    print("Reading " + filename)

    # loop over the frames in the file
    while infile.more():

        # get the frame
        frame = infile.pop_frame()

        # if this is an S-Frame
        if frame.Stop == frame.Simulation:
            # get the info from the frame
            info = frame["I3PrimaryInjectorInfo"]

            # convert the info object into a dictionary which represents a row
            inforow = {x: getattr(info, x) for x in dir(info) if x[0] != "_"}

            # add the row to to the list
            I3PrimaryInjectorInfo.append(inforow)

        # if this is a physics event in the right sub-event stream
        elif frame.Stop == frame.Physics and frame["I3EventHeader"].sub_event_stream == "InIceSplit":

            # get the weighting object
            w = frame["I3CorsikaWeight"]

            # for each of the columns we need get it from the frame object
            # and put it in the correct column
            I3CorsikaWeight["energy"].append(w.primary.energy)
            I3CorsikaWeight["type"].append(w.primary.type)
            I3CorsikaWeight["zenith"].append(w.primary.dir.zenith)
            I3CorsikaWeight["weight"].append(w.weight)

# make a dictionary object to mimic the file structure of a pandas file
fileobj = dict(I3PrimaryInjectorInfo=I3PrimaryInjectorInfo, I3CorsikaWeight=I3CorsikaWeight)

# create the weighter object
weighter = simweights.CorsikaWeighter(fileobj)

# create an object to represent our cosmic-ray primary flux model
flux = simweights.GaisserH4a()

# get the weights by passing the flux to the weighter
weights = weights = weighter.get_weights(flux)

# print some info about the weighting object
print(weighter.tostring(flux))

# create equal spaced bins in log space
bins = plt.geomspace(3e4, 1e6, 50)

# get energy of the primary cosmic-ray from `PolyplopiaPrimary`
primary_energy = weighter.get_weight_column("energy")

# histogram the primary energy with the weights
plt.hist(primary_energy, weights=weights, bins=bins)

# make the plot look good
plt.loglog()
plt.xlabel("Primary Energy [GeV]")
plt.ylabel("Event Rate [Hz]")
plt.xlim(bins[0], bins[-1])
plt.ylim(0.1, 10)
plt.show()
# SimWeights iterates over the info table row-wise so to mimic it we need a list # However, the weight object is iterated column-wise so we need to create a dictionary of lists # loop over all the files we want to read # open the i3 files with the dataio interface # loop over the frames in the file # get the frame # if this is an S-Frame # get the info from the frame # convert the info object into a dictionary which represents a row # add the row to to the list # if this is a physics event in the right sub-event stream # get the weighting object # for each of the columns we need get it from the frame object # and put it in the correct column # make a dictionary object to mimic the file structure of a pandas file # create the weighter object # create an object to represent our cosmic-ray primary flux model # get the weights by passing the flux to the weighter # print some info about the weighting object # create equal spaced bins in log space # get energy of the primary cosmic-ray from `PolyplopiaPrimary` # histogram the primary energy with the weights # make the plot look good