Ejemplo n.º 1
0
def main(cfg, run_number, scratch):
    with open(cfg, 'r') as stream:
        cfg = yaml.load(stream)
    cfg['run_number'] = run_number
    cfg['run_folder'] = get_run_folder(run_number)

    tray = I3Tray()

    random_service, random_service_prop, _ = create_random_services(
        dataset_number=cfg['dataset_number'],
        run_number=cfg['run_number'],
        seed=cfg['seed'])

    tray.context['I3RandomService'] = random_service

    model = MuonGun.load_model(cfg['muongun_model'])
    model.flux.min_multiplicity = cfg['muongun_min_multiplicity']
    model.flux.max_multiplicity = cfg['muongun_max_multiplicity']
    spectrum = MuonGun.OffsetPowerLaw(  cfg['gamma'],
                                        cfg['e_min']*icetray.I3Units.TeV,
                                        cfg['e_min']*icetray.I3Units.TeV,
                                        cfg['e_max']*icetray.I3Units.TeV)
    surface = MuonGun.Cylinder(1600, 800,
                            dataclasses.I3Position(31.25, 19.64, 0))

    if cfg['muongun_generator'] == 'energy':
        scale = MuonGun.BasicSurfaceScalingFunction()
        scale.SetSideScaling(4., 17266, 3.41, 1.74)
        scale.SetCapScaling(4., 23710, 3.40, 1.88)
        generator = MuonGun.EnergyDependentSurfaceInjector(surface,
                                                            model.flux,
                                                            spectrum,
                                                            model.radius,
                                                            scale)
    elif cfg['muongun_generator'] == 'static':
        generator = MuonGun.StaticSurfaceInjector(surface,
                                                    model.flux,
                                                    spectrum,
                                                    model.radius)
    elif cfg['muongun_generator'] =='floodlight':
        generator = MuonGun.Floodlight(surface = surface,
                                       energyGenerator=spectrum,
                                       cosMin=cfg['muongun_floodlight_min_cos'],
                                       cosMax=cfg['muongun_floodlight_max_cos'],
                                       )
    else:
        err_msg = 'MuonGun generator {} is not known.'
        err_msg += " Must be 'energy','static' or 'floodlight"
        raise ValueError(err_msg.format(cfg['muongun_generator']))

    tray.Add(MuonGun.segments.GenerateBundles, 'MuonGenerator',
                            Generator=generator,
                            NEvents=cfg['n_events_per_run'],
                            GCDFile=cfg['gcd'])

    tray.Add("Rename", keys=["I3MCTree", "I3MCTree_preMuonProp"])

    tray.AddSegment(
        segments.PropagateMuons,
        "PropagateMuons",
        RandomService=random_service_prop)
    if scratch:
        outfile = cfg['scratchfile_pattern'].format(**cfg)
    else:
        outfile = cfg['outfile_pattern'].format(**cfg)
    outfile = outfile.replace(' ', '0')
    outfile = outfile.replace('.bz2', '')
    tray.AddModule("I3Writer", "writer",
                   Filename=outfile,
                   Streams=[icetray.I3Frame.DAQ,
                            icetray.I3Frame.Physics,
                            icetray.I3Frame.Stream('S'),
                            icetray.I3Frame.Stream('M')])

    tray.AddModule("TrashCan", "the can")
    tray.Execute()
    tray.Finish()
def main(cfg, run_number, scratch):
    with open(cfg, 'r') as stream:
        if int(yaml.__version__[0]) < 5:
            # backwards compatibility for yaml versions before version 5
            cfg = yaml.load(stream)
        else:
            cfg = yaml.full_load(stream)
    cfg['run_number'] = run_number
    cfg['run_folder'] = get_run_folder(run_number)

    tray = I3Tray()

    random_services, _ = create_random_services(
        dataset_number=cfg['dataset_number'],
        run_number=cfg['run_number'],
        seed=cfg['seed'],
        n_services=2)

    random_service, random_service_prop = random_services

    tray.context['I3RandomService'] = random_service

    model = MuonGun.load_model(cfg['muongun_model'])
    model.flux.min_multiplicity = cfg['muongun_min_multiplicity']
    model.flux.max_multiplicity = cfg['muongun_max_multiplicity']
    spectrum = MuonGun.OffsetPowerLaw(cfg['gamma'],
                                      cfg['e_min'] * icetray.I3Units.GeV,
                                      cfg['e_min'] * icetray.I3Units.GeV,
                                      cfg['e_max'] * icetray.I3Units.GeV)
    surface = MuonGun.Cylinder(1600, 800,
                               dataclasses.I3Position(31.25, 19.64, 0))

    if cfg['muongun_generator'] == 'energy':
        scale = MuonGun.BasicSurfaceScalingFunction()
        scale.SetSideScaling(4., 17266, 3.41, 1.74)
        scale.SetCapScaling(4., 23710, 3.40, 1.88)
        generator = MuonGun.EnergyDependentSurfaceInjector(
            surface, model.flux, spectrum, model.radius, scale)
    elif cfg['muongun_generator'] == 'static':
        generator = MuonGun.StaticSurfaceInjector(surface, model.flux,
                                                  spectrum, model.radius)
    elif cfg['muongun_generator'] == 'floodlight':
        generator = MuonGun.Floodlight(
            surface=surface,
            energyGenerator=spectrum,
            cosMin=cfg['muongun_floodlight_min_cos'],
            cosMax=cfg['muongun_floodlight_max_cos'],
        )
    else:
        err_msg = 'MuonGun generator {} is not known.'
        err_msg += " Must be 'energy','static' or 'floodlight"
        raise ValueError(err_msg.format(cfg['muongun_generator']))

    tray.Add(MuonGun.segments.GenerateBundles,
             'MuonGenerator',
             Generator=generator,
             NEvents=cfg['n_events_per_run'],
             GCDFile=cfg['gcd'])

    tray.Add("Rename", keys=["I3MCTree", "I3MCTree_preMuonProp"])

    tray.AddSegment(segments.PropagateMuons,
                    "PropagateMuons",
                    RandomService=random_service_prop,
                    **cfg['muon_propagation_config'])
    if scratch:
        outfile = cfg['scratchfile_pattern'].format(**cfg)
    else:
        outfile = cfg['outfile_pattern'].format(**cfg)
    outfile = outfile.replace(' ', '0')
    if cfg['distance_splits'] is not None:
        click.echo('SplittingDistance: {}'.format(cfg['distance_splits']))
        distance_splits = np.atleast_1d(cfg['distance_splits'])
        dom_limits = np.atleast_1d(cfg['threshold_doms'])
        if len(dom_limits) == 1:
            dom_limits = np.ones_like(distance_splits) * cfg['threshold_doms']
        oversize_factors = np.atleast_1d(cfg['oversize_factors'])
        order = np.argsort(distance_splits)

        distance_splits = distance_splits[order]
        dom_limits = dom_limits[order]
        oversize_factors = oversize_factors[order]

        stream_objects = generate_stream_object(distance_splits, dom_limits,
                                                oversize_factors)
        tray.AddModule(OversizeSplitterNSplits,
                       "OversizeSplitterNSplits",
                       thresholds=distance_splits,
                       thresholds_doms=dom_limits,
                       oversize_factors=oversize_factors)
        for stream_i in stream_objects:
            outfile_i = stream_i.transform_filepath(outfile)
            tray.AddModule("I3Writer",
                           "writer_{}".format(stream_i.stream_name),
                           Filename=outfile_i,
                           Streams=[
                               icetray.I3Frame.DAQ, icetray.I3Frame.Physics,
                               icetray.I3Frame.Stream('S'),
                               icetray.I3Frame.Stream('M')
                           ],
                           If=stream_i)
            click.echo('Output ({}): {}'.format(stream_i.stream_name,
                                                outfile_i))
    else:
        click.echo('Output: {}'.format(outfile))
        tray.AddModule("I3Writer",
                       "writer",
                       Filename=outfile,
                       Streams=[
                           icetray.I3Frame.DAQ, icetray.I3Frame.Physics,
                           icetray.I3Frame.Stream('S'),
                           icetray.I3Frame.Stream('M')
                       ])
    click.echo('Scratch: {}'.format(scratch))
    tray.AddModule("TrashCan", "the can")
    tray.Execute()
    tray.Finish()
Ejemplo n.º 3
0
def main(cfg, run_number, scratch):
    with open(cfg, 'r') as stream:
        cfg = yaml.load(stream, Loader=yaml.Loader)
    cfg['run_number'] = run_number
    cfg['run_folder'] = get_run_folder(run_number)

    if scratch:
        outfile = cfg['scratchfile_pattern'].format(**cfg)
    else:
        outfile = cfg['outfile_pattern'].format(**cfg)
    outfile = outfile.replace(' ', '0')

    if cfg['distance_splits'] is not None:
        click.echo('SplittingDistances: {}'.format(cfg['distance_splits']))
        click.echo('Oversizefactors: {}'.format(cfg['oversize_factors']))
    click.echo('NEvents: {}'.format(cfg['n_events_per_run']))
    click.echo('EMin: {}'.format(cfg['e_min']))
    click.echo('EMax: {}'.format(cfg['e_max']))
    click.echo('EBreak: {}'.format(cfg['muongun_e_break']))
    click.echo('Gamma: {}'.format(cfg['gamma']))
    click.echo('ZenithMin: {}'.format(cfg['zenith_min']))
    click.echo('ZenithMax: {}'.format(cfg['zenith_max']))

    # create convex hull
    if 'use_convex_hull' in cfg and cfg['use_convex_hull']:

        # hardcode icecube corner points
        # ToDo: read from geometry file
        points = [
            [-570.90002441, -125.13999939, 501],  # string 31
            [-256.14001465, -521.08001709, 501],  # string 1
            [361., -422.82998657, 501],  # string 6
            [576.36999512, 170.91999817, 501],  # string 50
            [338.44000244, 463.72000122, 501],  # string 74
            [101.04000092, 412.79000854, 501],  # string 72
            [22.11000061, 509.5, 501],  # string 78
            [-347.88000488, 451.51998901, 501],  # string 75
            [-570.90002441, -125.13999939, -502],  # string 31
            [-256.14001465, -521.08001709, -502],  # string 1
            [361., -422.82998657, -502],  # string 6
            [576.36999512, 170.91999817, -502],  # string 50
            [338.44000244, 463.72000122, -502],  # string 74
            [101.04000092, 412.79000854, -502],  # string 72
            [22.11000061, 509.5, -502],  # string 78
            [-347.88000488, 451.51998901, -502],  # string 75
        ]
        convex_hull = ConvexHull(points)
    else:
        convex_hull = None

    if 'extend_past_hull' not in cfg:
        cfg['extend_past_hull'] = 0.0

    random_services, _ = create_random_services(
        dataset_number=cfg['dataset_number'],
        run_number=cfg['run_number'],
        seed=cfg['seed'],
        n_services=2)

    random_service, random_service_prop = random_services

    # create muon
    muon = create_muon(
        azimuth_range=[cfg['azimuth_min'], cfg['azimuth_max']],
        zenith_range=[cfg['zenith_min'], cfg['zenith_max']],
        energy_range=[cfg['e_min'], cfg['e_max']],
        anchor_time_range=cfg['anchor_time_range'],
        anchor_x_range=cfg['anchor_x_range'],
        anchor_y_range=cfg['anchor_y_range'],
        anchor_z_range=cfg['anchor_z_range'],
        length_to_go_back=cfg['length_to_go_back'],
        convex_hull=convex_hull,
        extend_past_hull=cfg['extend_past_hull'],
        random_service=random_services[0],
    )

    tray = I3Tray()

    tray.context['I3RandomService'] = random_service

    tray.AddModule("I3InfiniteSource",
                   "TheSource",
                   Prefix=cfg['gcd'],
                   Stream=icetray.I3Frame.DAQ)

    if cfg['MuonGenerator'] == 'MuonGunSinglemuons':
        tray.AddSegment(segments.GenerateSingleMuons,
                        "GenerateCosmicRayMuons",
                        NumEvents=cfg['n_events_per_run'],
                        FromEnergy=cfg['e_min'] * icetray.I3Units.GeV,
                        ToEnergy=cfg['e_max'] * icetray.I3Units.GeV,
                        BreakEnergy=cfg['muongun_e_break'] *
                        icetray.I3Units.GeV,
                        GammaIndex=cfg['gamma'],
                        ZenithRange=[
                            cfg['zenith_min'] * icetray.I3Units.deg,
                            cfg['zenith_max'] * icetray.I3Units.deg
                        ])

    elif cfg['MuonGenerator'] == 'MuonGunGeneral':
        model = MuonGun.load_model(cfg['muongun_model'])
        model.flux.min_multiplicity = cfg['muongun_min_multiplicity']
        model.flux.max_multiplicity = cfg['muongun_max_multiplicity']
        spectrum = MuonGun.OffsetPowerLaw(cfg['gamma'],
                                          cfg['e_min'] * icetray.I3Units.GeV,
                                          cfg['e_min'] * icetray.I3Units.GeV,
                                          cfg['e_max'] * icetray.I3Units.GeV)
        surface = MuonGun.Cylinder(1600, 800,
                                   dataclasses.I3Position(31.25, 19.64, 0))

        if cfg['muongun_generator'] == 'energy':
            scale = MuonGun.BasicSurfaceScalingFunction()
            scale.SetSideScaling(4., 17266, 3.41, 1.74)
            scale.SetCapScaling(4., 23710, 3.40, 1.88)
            generator = MuonGun.EnergyDependentSurfaceInjector(
                surface, model.flux, spectrum, model.radius, scale)
        elif cfg['muongun_generator'] == 'static':
            generator = MuonGun.StaticSurfaceInjector(surface, model.flux,
                                                      spectrum, model.radius)
        elif cfg['muongun_generator'] == 'floodlight':
            generator = MuonGun.Floodlight(
                surface=surface,
                energyGenerator=spectrum,
                cosMin=cfg['muongun_floodlight_min_cos'],
                cosMax=cfg['muongun_floodlight_max_cos'],
            )
        else:
            err_msg = 'MuonGun generator {} is not known.'
            err_msg += " Must be 'energy','static' or 'floodlight"
            raise ValueError(err_msg.format(cfg['muongun_generator']))

        tray.Add(MuonGun.segments.GenerateBundles,
                 'MuonGenerator',
                 Generator=generator,
                 NEvents=cfg['n_events_per_run'],
                 GCDFile=cfg['gcd'])

        tray.Add("Rename", keys=["I3MCTree", "I3MCTree_preMuonProp"])

    elif cfg['MuonGenerator'] == 'MuonResimulation':
        tray.AddModule(ParticleMultiplier,
                       'make_particles',
                       num_events=cfg['n_events_per_run'],
                       primary=muon)
    else:
        err_msg = 'MuonGenerator {} is not known.'
        err_msg += " Must be 'MuonGunSinglemuons','MuonGunGeneral' or 'MuonResimulation"
        raise ValueError(err_msg.format(cfg['MuonGenerator']))

    # --------------------------------------
    # Propagate Muons
    # --------------------------------------
    tray.AddSegment(segments.PropagateMuons,
                    "PropagateMuons",
                    RandomService=random_service_prop,
                    **cfg['muon_propagation_config'])

    # --------------------------------------
    # Distance Splits
    # --------------------------------------
    if cfg['distance_splits'] is not None:
        click.echo('SplittingDistance: {}'.format(cfg['distance_splits']))
        distance_splits = np.atleast_1d(cfg['distance_splits'])
        dom_limits = np.atleast_1d(cfg['threshold_doms'])
        if len(dom_limits) == 1:
            dom_limits = np.ones_like(distance_splits) * cfg['threshold_doms']
        oversize_factors = np.atleast_1d(cfg['oversize_factors'])
        order = np.argsort(distance_splits)

        distance_splits = distance_splits[order]
        dom_limits = dom_limits[order]
        oversize_factors = oversize_factors[order]

        stream_objects = generate_stream_object(distance_splits, dom_limits,
                                                oversize_factors)
        tray.AddModule(OversizeSplitterNSplits,
                       "OversizeSplitterNSplits",
                       thresholds=distance_splits,
                       thresholds_doms=dom_limits,
                       oversize_factors=oversize_factors)
        for stream_i in stream_objects:
            outfile_i = stream_i.transform_filepath(outfile)
            tray.AddModule("I3Writer",
                           "writer_{}".format(stream_i.stream_name),
                           Filename=outfile_i,
                           Streams=[
                               icetray.I3Frame.DAQ, icetray.I3Frame.Physics,
                               icetray.I3Frame.Stream('S'),
                               icetray.I3Frame.Stream('M')
                           ],
                           If=stream_i)
            click.echo('Output ({}): {}'.format(stream_i.stream_name,
                                                outfile_i))
    else:
        click.echo('Output: {}'.format(outfile))
        tray.AddModule("I3Writer",
                       "writer",
                       Filename=outfile,
                       Streams=[
                           icetray.I3Frame.DAQ, icetray.I3Frame.Physics,
                           icetray.I3Frame.Stream('S'),
                           icetray.I3Frame.Stream('M')
                       ])

    click.echo('Scratch: {}'.format(scratch))

    tray.Execute()

    del tray