Ejemplo n.º 1
0
def addSeeding(
    s: acts.examples.Sequencer,
    trackingGeometry: acts.TrackingGeometry,
    field: acts.MagneticFieldProvider,
    geoSelectionConfigFile: Optional[Union[Path, str]] = None,
    seedingAlgorithm: SeedingAlgorithm = SeedingAlgorithm.Default,
    truthSeedRanges: TruthSeedRanges = TruthSeedRanges(),
    particleSmearingSigmas: ParticleSmearingSigmas = ParticleSmearingSigmas(),
    initialVarInflation: Optional[list] = None,
    seedfinderConfigArg: SeedfinderConfigArg = SeedfinderConfigArg(),
    trackParamsEstimationConfig:
    TrackParamsEstimationConfig = TrackParamsEstimationConfig(),
    inputParticles: str = "particles_initial",
    outputDirRoot: Optional[Union[Path, str]] = None,
    logLevel: Optional[acts.logging.Level] = None,
    rnd: Optional[acts.examples.RandomNumbers] = None,
) -> acts.examples.Sequencer:
    """This function steers the seeding

    Parameters
    ----------
    s: Sequencer
        the sequencer module to which we add the Seeding steps (returned from addSeeding)
    trackingGeometry : tracking geometry
    field : magnetic field
    geoSelectionConfigFile : Path|str, path, None
        Json file for space point geometry selection. Not required for SeedingAlgorithm.TruthSmeared.
    seedingAlgorithm : SeedingAlgorithm, Default
        seeding algorithm to use: one of Default (no truth information used), TruthSmeared, TruthEstimated
    truthSeedRanges : TruthSeedRanges(rho, z, phi, eta, absEta, pt, nHits)
        TruthSeedSelector configuration. Each range is specified as a tuple of (min,max).
        Defaults of no cuts specified in Examples/Algorithms/TruthTracking/ActsExamples/TruthTracking/TruthSeedSelector.hpp
    particleSmearingSigmas : ParticleSmearingSigmas(d0, d0PtA, d0PtB, z0, z0PtA, z0PtB, t0, phi, theta, pRel)
        ParticleSmearing configuration.
        Defaults specified in Examples/Algorithms/TruthTracking/ActsExamples/TruthTracking/ParticleSmearing.hpp
    initialVarInflation : list
        List of 6 scale factors to inflate the initial covariance matrix
        Defaults (all 1) specified in Examples/Algorithms/TruthTracking/ActsExamples/TruthTracking/ParticleSmearing.hpp
    seedfinderConfigArg : SeedfinderConfigArg(maxSeedsPerSpM, cotThetaMax, sigmaScattering, radLengthPerSeed, minPt, bFieldInZ, impactMax, deltaR, collisionRegion, r, z, beamPos)
        SeedfinderConfig settings. deltaR, collisionRegion, r, z are ranges specified as a tuple of (min,max). beamPos is specified as (x,y).
        Defaults specified in Core/include/Acts/Seeding/SeedfinderConfig.hpp
    trackParamsEstimationConfig : TrackParamsEstimationConfig(deltaR)
        TrackParamsEstimationAlgorithm configuration. Currently only deltaR=(min,max) range specified here.
        Defaults specified in Examples/Algorithms/TrackFinding/include/ActsExamples/TrackFinding/TrackParamsEstimationAlgorithm.hpp
    inputParticles : str, "particles_initial"
        input particles name in the WhiteBoard
    outputDirRoot : Path|str, path, None
        the output folder for the Root output, None triggers no output
    logLevel : acts.logging.Level, None
        logging level to override setting given in `s`
    rnd : RandomNumbers, None
        random number generator. Only used by SeedingAlgorithm.TruthSmeared.
    """
    def customLogLevel(custom: acts.logging.Level = acts.logging.INFO):
        """override logging level"""
        if logLevel is None:
            return s.config.logLevel
        return acts.logging.Level(max(custom.value, logLevel.value))

    if int(customLogLevel()) <= int(acts.logging.DEBUG):
        acts.examples.dump_args_calls(locals())
    logger = acts.logging.getLogger("addSeeding")

    if truthSeedRanges is not None:
        selAlg = acts.examples.TruthSeedSelector(
            **acts.examples.defaultKWArgs(
                ptMin=truthSeedRanges.pt[0],
                ptMax=truthSeedRanges.pt[1],
                etaMin=truthSeedRanges.eta[0],
                etaMax=truthSeedRanges.eta[1],
                nHitsMin=truthSeedRanges.nHits[0],
                nHitsMax=truthSeedRanges.nHits[1],
                rhoMin=truthSeedRanges.rho[0],
                rhoMax=truthSeedRanges.rho[1],
                zMin=truthSeedRanges.z[0],
                zMax=truthSeedRanges.z[1],
                phiMin=truthSeedRanges.phi[0],
                phiMax=truthSeedRanges.phi[1],
                absEtaMin=truthSeedRanges.absEta[0],
                absEtaMax=truthSeedRanges.absEta[1],
            ),
            level=customLogLevel(),
            inputParticles=inputParticles,
            inputMeasurementParticlesMap="measurement_particles_map",
            outputParticles="truth_seeds_selected",
        )
        s.addAlgorithm(selAlg)
        selectedParticles = selAlg.config.outputParticles
    else:
        selectedParticles = inputParticles

    # Create starting parameters from either particle smearing or combined seed
    # finding and track parameters estimation
    if seedingAlgorithm == SeedingAlgorithm.TruthSmeared:
        rnd = rnd or acts.examples.RandomNumbers(seed=42)
        logger.info("Using smeared truth particles for seeding")
        # Run particle smearing
        ptclSmear = acts.examples.ParticleSmearing(
            level=customLogLevel(),
            inputParticles=selectedParticles,
            outputTrackParameters="estimatedparameters",
            randomNumbers=rnd,
            # gaussian sigmas to smear particle parameters
            **acts.examples.defaultKWArgs(
                sigmaD0=particleSmearingSigmas.d0,
                sigmaD0PtA=particleSmearingSigmas.d0PtA,
                sigmaD0PtB=particleSmearingSigmas.d0PtB,
                sigmaZ0=particleSmearingSigmas.z0,
                sigmaZ0PtA=particleSmearingSigmas.z0PtA,
                sigmaZ0PtB=particleSmearingSigmas.z0PtB,
                sigmaT0=particleSmearingSigmas.t0,
                sigmaPhi=particleSmearingSigmas.phi,
                sigmaTheta=particleSmearingSigmas.theta,
                sigmaPRel=particleSmearingSigmas.pRel,
                initialVarInflation=initialVarInflation,
            ),
        )
        s.addAlgorithm(ptclSmear)
    else:

        spAlg = acts.examples.SpacePointMaker(
            level=customLogLevel(),
            inputSourceLinks="sourcelinks",
            inputMeasurements="measurements",
            outputSpacePoints="spacepoints",
            trackingGeometry=trackingGeometry,
            geometrySelection=acts.examples.readJsonGeometryList(
                str(geoSelectionConfigFile)),
        )
        s.addAlgorithm(spAlg)

        # Run either: truth track finding or seeding
        if seedingAlgorithm == SeedingAlgorithm.TruthEstimated:
            logger.info(
                "Using truth track finding from space points for seeding")
            # Use truth tracking
            truthTrackFinder = acts.examples.TruthTrackFinder(
                level=customLogLevel(),
                inputParticles=selectedParticles,
                inputMeasurementParticlesMap=selAlg.config.
                inputMeasurementParticlesMap,
                outputProtoTracks="prototracks",
            )
            s.addAlgorithm(truthTrackFinder)
            inputProtoTracks = truthTrackFinder.config.outputProtoTracks
            inputSeeds = ""
        elif seedingAlgorithm == SeedingAlgorithm.Default:
            logger.info("Using default seeding")
            # Use seeding
            seedFinderConfig = acts.SeedfinderConfig(
                **acts.examples.defaultKWArgs(
                    rMin=seedfinderConfigArg.r[0],
                    rMax=seedfinderConfigArg.r[1],
                    deltaRMin=seedfinderConfigArg.deltaR[0],
                    deltaRMax=seedfinderConfigArg.deltaR[1],
                    deltaRMinTopSP=seedfinderConfigArg.deltaR[0],
                    deltaRMinBottomSP=seedfinderConfigArg.deltaR[0],
                    deltaRMaxTopSP=seedfinderConfigArg.deltaR[1],
                    deltaRMaxBottomSP=seedfinderConfigArg.deltaR[1],
                    collisionRegionMin=seedfinderConfigArg.collisionRegion[0],
                    collisionRegionMax=seedfinderConfigArg.collisionRegion[1],
                    zMin=seedfinderConfigArg.z[0],
                    zMax=seedfinderConfigArg.z[1],
                    maxSeedsPerSpM=seedfinderConfigArg.maxSeedsPerSpM,
                    cotThetaMax=seedfinderConfigArg.cotThetaMax,
                    sigmaScattering=seedfinderConfigArg.sigmaScattering,
                    radLengthPerSeed=seedfinderConfigArg.radLengthPerSeed,
                    minPt=seedfinderConfigArg.minPt,
                    bFieldInZ=seedfinderConfigArg.bFieldInZ,
                    impactMax=seedfinderConfigArg.impactMax,
                    beamPos=(
                        None if seedfinderConfigArg.beamPos is None or all(
                            [x is None for x in seedfinderConfigArg.beamPos])
                        else acts.Vector2(
                            seedfinderConfigArg.beamPos[0] or 0.0,
                            seedfinderConfigArg.beamPos[1] or 0.0,
                        )),
                ), )

            seedFilterConfig = acts.SeedFilterConfig(
                maxSeedsPerSpM=seedFinderConfig.maxSeedsPerSpM,
                deltaRMin=seedFinderConfig.deltaRMin,
            )

            gridConfig = acts.SpacePointGridConfig(
                bFieldInZ=seedFinderConfig.bFieldInZ,
                minPt=seedFinderConfig.minPt,
                rMax=seedFinderConfig.rMax,
                zMax=seedFinderConfig.zMax,
                zMin=seedFinderConfig.zMin,
                deltaRMax=seedFinderConfig.deltaRMax,
                cotThetaMax=seedFinderConfig.cotThetaMax,
            )

            seedingAlg = acts.examples.SeedingAlgorithm(
                level=customLogLevel(acts.logging.VERBOSE),
                inputSpacePoints=[spAlg.config.outputSpacePoints],
                outputSeeds="seeds",
                outputProtoTracks="prototracks",
                gridConfig=gridConfig,
                seedFilterConfig=seedFilterConfig,
                seedFinderConfig=seedFinderConfig,
            )
            s.addAlgorithm(seedingAlg)
            inputProtoTracks = seedingAlg.config.outputProtoTracks
            inputSeeds = seedingAlg.config.outputSeeds
        elif seedingAlgorithm == SeedingAlgorithm.Orthogonal:
            logger.info("Using orthogonal seeding")
            # Use seeding
            seedFinderConfig = acts.SeedFinderOrthogonalConfig(
                **acts.examples.defaultKWArgs(
                    rMin=seedfinderConfigArg.r[0],
                    rMax=seedfinderConfigArg.r[1],
                    deltaRMin=seedfinderConfigArg.deltaR[0],
                    deltaRMax=seedfinderConfigArg.deltaR[1],
                    collisionRegionMin=seedfinderConfigArg.collisionRegion[0],
                    collisionRegionMax=seedfinderConfigArg.collisionRegion[1],
                    zMin=seedfinderConfigArg.z[0],
                    zMax=seedfinderConfigArg.z[1],
                    maxSeedsPerSpM=seedfinderConfigArg.maxSeedsPerSpM,
                    cotThetaMax=seedfinderConfigArg.cotThetaMax,
                    sigmaScattering=seedfinderConfigArg.sigmaScattering,
                    radLengthPerSeed=seedfinderConfigArg.radLengthPerSeed,
                    minPt=seedfinderConfigArg.minPt,
                    bFieldInZ=seedfinderConfigArg.bFieldInZ,
                    impactMax=seedfinderConfigArg.impactMax,
                    beamPos=(
                        None if seedfinderConfigArg.beamPos is None or all(
                            [x is None for x in seedfinderConfigArg.beamPos])
                        else acts.Vector2(
                            seedfinderConfigArg.beamPos[0] or 0.0,
                            seedfinderConfigArg.beamPos[1] or 0.0,
                        )),
                ), )

            seedFilterConfig = acts.SeedFilterConfig(
                maxSeedsPerSpM=seedFinderConfig.maxSeedsPerSpM,
                deltaRMin=seedFinderConfig.deltaRMin,
            )

            seedingAlg = acts.examples.SeedingOrthogonalAlgorithm(
                level=customLogLevel(acts.logging.VERBOSE),
                inputSpacePoints=[spAlg.config.outputSpacePoints],
                outputSeeds="seeds",
                outputProtoTracks="prototracks",
                seedFilterConfig=seedFilterConfig,
                seedFinderConfig=seedFinderConfig,
            )
            s.addAlgorithm(seedingAlg)
            inputProtoTracks = seedingAlg.config.outputProtoTracks
            inputSeeds = seedingAlg.config.outputSeeds
        else:
            logger.fatal("unknown seedingAlgorithm %s", seedingAlgorithm)

        parEstimateAlg = acts.examples.TrackParamsEstimationAlgorithm(
            level=customLogLevel(acts.logging.VERBOSE),
            inputSeeds=inputSeeds,
            inputProtoTracks=inputProtoTracks,
            inputSpacePoints=[spAlg.config.outputSpacePoints],
            inputSourceLinks=spAlg.config.inputSourceLinks,
            outputTrackParameters="estimatedparameters",
            outputProtoTracks="prototracks_estimated",
            trackingGeometry=trackingGeometry,
            magneticField=field,
            **acts.examples.defaultKWArgs(
                initialVarInflation=initialVarInflation,
                deltaRMin=trackParamsEstimationConfig.deltaR[0],
                deltaRMax=trackParamsEstimationConfig.deltaR[1],
            ),
        )
        s.addAlgorithm(parEstimateAlg)

        if outputDirRoot is not None:
            outputDirRoot = Path(outputDirRoot)
            if not outputDirRoot.exists():
                outputDirRoot.mkdir()
            s.addWriter(
                acts.examples.TrackFinderPerformanceWriter(
                    level=customLogLevel(),
                    inputProtoTracks=inputProtoTracks,
                    inputParticles=
                    selectedParticles,  # the original selected particles after digitization
                    inputMeasurementParticlesMap=selAlg.config.
                    inputMeasurementParticlesMap,
                    filePath=str(outputDirRoot /
                                 "performance_seeding_trees.root"),
                ))

            s.addWriter(
                acts.examples.SeedingPerformanceWriter(
                    level=customLogLevel(acts.logging.DEBUG),
                    inputProtoTracks=inputProtoTracks,
                    inputParticles=selectedParticles,
                    inputMeasurementParticlesMap=selAlg.config.
                    inputMeasurementParticlesMap,
                    filePath=str(outputDirRoot /
                                 "performance_seeding_hists.root"),
                ))

            s.addWriter(
                acts.examples.RootTrackParameterWriter(
                    level=customLogLevel(acts.logging.VERBOSE),
                    inputTrackParameters=parEstimateAlg.config.
                    outputTrackParameters,
                    inputProtoTracks=parEstimateAlg.config.outputProtoTracks,
                    inputParticles=inputParticles,
                    inputSimHits="simhits",
                    inputMeasurementParticlesMap=selAlg.config.
                    inputMeasurementParticlesMap,
                    inputMeasurementSimHitsMap="measurement_simhits_map",
                    filePath=str(outputDirRoot / "estimatedparams.root"),
                    treeName="estimatedparams",
                ))

    return s
Ejemplo n.º 2
0
def runITkSeeding(field, csvInputDir, outputDir, s=None):

    # Read input space points from input csv files
    evReader = CsvSpacePointReader(
        level=acts.logging.INFO,
        inputStem="spacepoints",
        inputCollection="pixel",
        inputDir=csvInputDir,
        outputSpacePoints="PixelSpacePoints",
    )

    gridConfig = acts.SpacePointGridConfig(
        bFieldInZ=2 * u.T,
        minPt=900 * u.MeV,
        rMax=320 * u.mm,  # pixel: 320 mm, strip: 1000 mm
        zMax=3000 * u.mm,
        zMin=-3000 * u.mm,
        deltaRMax=280 * u.mm,  # pixel: 280 mm, strip: 600 mm
        cotThetaMax=27.2899,  # pixel: 27.2899 (4 eta)
        impactMax=2 * u.mm,  # pixel: 2 mm, strip: 20 mm
        zBinEdges=[
            -3000.0,
            -2500.0,
            -1400.0,
            -925.0,
            -450.0,
            -250.0,
            250.0,
            450.0,
            925.0,
            1400.0,
            2500.0,
            3000.0,
        ],  # zBinEdges enables non-equidistant binning in z, in case the binning is not defined the edges are evaluated automatically using equidistant binning
        phiBinDeflectionCoverage=3,
    )

    seedFinderConfig = acts.SeedfinderConfig(
        rMax=gridConfig.rMax,
        deltaRMin=20 * u.mm,
        deltaRMax=gridConfig.deltaRMax,
        deltaRMinTopSP=6 * u.mm,  # pixel: 6 mm, strip: 20 mm
        deltaRMinBottomSP=6 * u.mm,  # pixel: 6 mm, strip: 20 mm
        deltaRMaxTopSP=280 * u.mm,  # pixel: 280 mm, strip: 3000 mm
        deltaRMaxBottomSP=120 * u.mm,  # pixel: 120 mm, strip: 3000 mm
        collisionRegionMin=-200 * u.mm,
        collisionRegionMax=200 * u.mm,
        zMin=gridConfig.zMin,
        zMax=gridConfig.zMax,
        maxSeedsPerSpM=4,
        cotThetaMax=gridConfig.cotThetaMax,
        sigmaScattering=2,
        radLengthPerSeed=0.1,
        minPt=gridConfig.minPt,
        bFieldInZ=gridConfig.bFieldInZ,
        beamPos=acts.Vector2(0 * u.mm, 0 * u.mm),
        impactMax=gridConfig.impactMax,
        maxPtScattering=float("inf") * u.GeV,
        deltaZMax=900 * u.mm,
        interactionPointCut=True,
        zBinEdges=gridConfig.zBinEdges,
        skipPreviousTopSP=True,
        rRangeMiddleSP=[
            [40.0, 90.0],
            [40.0, 200.0],
            [46.0, 200.0],
            [46.0, 200.0],
            [46.0, 250.0],
            [46.0, 250.0],
            [46.0, 250.0],
            [46.0, 200.0],
            [46.0, 200.0],
            [40.0, 200.0],
            [40.0, 90.0],
        ],  # if useVariableMiddleSPRange is set to false, the vector rRangeMiddleSP can be used to define a fixed r range for each z bin: {{rMin, rMax}, ...}. If useVariableMiddleSPRange is set to false and the vector is empty, the cuts won't be applied
        useVariableMiddleSPRange=True,  # if useVariableMiddleSPRange is true, the values in rRangeMiddleSP will be calculated based on r values of the SPs and deltaRMiddleSPRange
        deltaRMiddleSPRange=10,
        seedConfirmation=True,
        centralSeedConfirmationRange=acts.SeedConfirmationRange(
            zMinSeedConf=250 * u.mm,
            zMaxSeedConf=250 * u.mm,
            rMaxSeedConf=140 * u.mm,
            nTopForLargeR=1,
            nTopForSmallR=2,
        ),  # contains parameters for seed confirmation
        forwardSeedConfirmationRange=acts.SeedConfirmationRange(
            zMinSeedConf=3000 * u.mm,
            zMaxSeedConf=-3000 * u.mm,
            rMaxSeedConf=140 * u.mm,
            nTopForLargeR=1,
            nTopForSmallR=2,
        ),
    )

    seedFilterConfig = acts.SeedFilterConfig(
        maxSeedsPerSpM=seedFinderConfig.maxSeedsPerSpM,
        deltaRMin=seedFinderConfig.deltaRMin,
        impactWeightFactor=100,
        compatSeedWeight=100,
        compatSeedLimit=3,
        curvatureSortingInFilter=True,
    )

    seedingAlg = acts.examples.SeedingAlgorithm(
        level=acts.logging.VERBOSE,
        inputSpacePoints=[evReader.config.outputSpacePoints],
        outputSeeds="PixelSeeds",
        outputProtoTracks="prototracks",
        allowSeparateRMax=False,
        gridConfig=gridConfig,
        seedFinderConfig=seedFinderConfig,
        seedFilterConfig=seedFilterConfig,
        zBinNeighborsTop=[
            [0, 0],
            [-1, 0],
            [-1, 0],
            [-1, 0],
            [-1, 0],
            [-1, 1],
            [0, 1],
            [0, 1],
            [0, 1],
            [0, 1],
            [0, 0],
        ],  # allows to specify the number of neighbors desired for each bin, [-1,1] means one neighbor on the left and one on the right, if the vector is empty the algorithm returns the 8 surrounding bins
        zBinNeighborsBottom=[
            [0, 1],
            [0, 1],
            [0, 1],
            [0, 1],
            [0, 1],
            [0, 0],
            [-1, 0],
            [-1, 0],
            [-1, 0],
            [-1, 0],
            [-1, 0],
        ],
        numPhiNeighbors=1,
    )

    s = s or acts.examples.Sequencer(
        events=1, numThreads=-1, logLevel=acts.logging.INFO
    )

    s.addReader(evReader)
    s.addAlgorithm(seedingAlg)

    return s
Ejemplo n.º 3
0
def runSeeding(trackingGeometry, field, outputDir, s=None):

    srcdir = Path(__file__).resolve().parent.parent.parent.parent

    csv_dir = os.path.join(outputDir, "csv")
    if not os.path.exists(csv_dir):
        os.mkdir(csv_dir)

    # Input
    rnd = acts.examples.RandomNumbers(seed=42)
    evGen = acts.examples.EventGenerator(
        level=acts.logging.INFO,
        generators=[
            acts.examples.EventGenerator.Generator(
                multiplicity=acts.examples.FixedMultiplicityGenerator(n=2),
                vertex=acts.examples.GaussianVertexGenerator(
                    stddev=acts.Vector4(0, 0, 0, 0),
                    mean=acts.Vector4(0, 0, 0, 0)),
                particles=acts.examples.ParametricParticleGenerator(
                    p=(1 * u.GeV, 10 * u.GeV),
                    eta=(-2, 2),
                    phi=(0, 360 * u.degree),
                    randomizeCharge=True,
                    numParticles=4,
                ),
            )
        ],
        outputParticles="particles_input",
        randomNumbers=rnd,
    )

    # Read input from input collection (e.g. Pythia8 output)

    # evGen = acts.examples.RootParticleReader(
    #     level=acts.logging.INFO,
    #     particleCollection="particles_input",
    #     inputDir="output",
    #     inputFile="pythia8_particles.root",
    # )

    # Simulation
    simAlg = acts.examples.FatrasSimulation(
        level=acts.logging.INFO,
        inputParticles=evGen.config.outputParticles,
        # inputParticles=evGen.config.particleCollection,
        outputParticlesInitial="particles_initial",
        outputParticlesFinal="particles_final",
        outputSimHits="simhits",
        randomNumbers=rnd,
        trackingGeometry=trackingGeometry,
        magneticField=field,
        generateHitsOnSensitive=True,
    )

    # Digitization
    digiCfg = acts.examples.DigitizationConfig(
        acts.examples.readDigiConfigFromJson(
            str(srcdir /
                "Examples/Algorithms/Digitization/share/default-smearing-config-generic.json"
                )),
        trackingGeometry=trackingGeometry,
        randomNumbers=rnd,
        inputSimHits=simAlg.config.outputSimHits,
    )
    digiAlg = acts.examples.DigitizationAlgorithm(digiCfg, acts.logging.INFO)

    selAlg = acts.examples.TruthSeedSelector(
        level=acts.logging.INFO,
        ptMin=1 * u.GeV,
        eta=(-2.5, 2.5),
        nHitsMin=9,
        inputParticles=simAlg.config.outputParticlesFinal,
        inputMeasurementParticlesMap=digiCfg.outputMeasurementParticlesMap,
        outputParticles="particles_selected",
    )

    inputParticles = selAlg.config.outputParticles

    spAlg = acts.examples.SpacePointMaker(
        level=acts.logging.INFO,
        inputSourceLinks=digiCfg.outputSourceLinks,
        inputMeasurements=digiCfg.outputMeasurements,
        outputSpacePoints="spacepoints",
        trackingGeometry=trackingGeometry,
        geometrySelection=acts.examples.readJsonGeometryList(
            str(srcdir /
                "Examples/Algorithms/TrackFinding/share/geoSelection-genericDetector.json"
                )),
    )

    gridConfig = acts.SpacePointGridConfig(
        bFieldInZ=1.99724 * u.T,
        minPt=500 * u.MeV,
        rMax=200 * u.mm,
        zMax=2000 * u.mm,
        zMin=-2000 * u.mm,
        deltaRMax=60 * u.mm,
        cotThetaMax=7.40627,  # 2.7 eta
    )

    seedFilterConfig = acts.SeedFilterConfig(maxSeedsPerSpM=1,
                                             deltaRMin=1 * u.mm)

    seedFinderConfig = acts.SeedfinderConfig(
        rMax=gridConfig.rMax,
        deltaRMin=seedFilterConfig.deltaRMin,
        deltaRMax=gridConfig.deltaRMax,
        collisionRegionMin=-250 * u.mm,
        collisionRegionMax=250 * u.mm,
        zMin=gridConfig.zMin,
        zMax=gridConfig.zMax,
        maxSeedsPerSpM=seedFilterConfig.maxSeedsPerSpM,
        cotThetaMax=gridConfig.cotThetaMax,
        sigmaScattering=50,
        radLengthPerSeed=0.1,
        minPt=gridConfig.minPt,
        bFieldInZ=gridConfig.bFieldInZ,
        beamPos=acts.Vector2(0 * u.mm, 0 * u.mm),
        impactMax=3 * u.mm,
    )

    seedingAlg = acts.examples.SeedingAlgorithm(
        level=acts.logging.VERBOSE,
        inputSpacePoints=[spAlg.config.outputSpacePoints],
        outputSeeds="seeds",
        outputProtoTracks="prototracks",
        gridConfig=gridConfig,
        seedFilterConfig=seedFilterConfig,
        seedFinderConfig=seedFinderConfig,
    )

    parEstimateAlg = acts.examples.TrackParamsEstimationAlgorithm(
        level=acts.logging.VERBOSE,
        inputProtoTracks=seedingAlg.config.outputProtoTracks,
        inputSpacePoints=[spAlg.config.outputSpacePoints],
        inputSourceLinks=digiCfg.outputSourceLinks,
        outputTrackParameters="estimatedparameters",
        outputProtoTracks="prototracks_estimated",
        trackingGeometry=trackingGeometry,
        magneticField=field,
    )

    s = s or acts.examples.Sequencer(
        events=100, numThreads=-1, logLevel=acts.logging.INFO)

    s.addReader(evGen)
    s.addAlgorithm(simAlg)
    s.addAlgorithm(digiAlg)
    s.addAlgorithm(selAlg)
    s.addAlgorithm(spAlg)
    s.addAlgorithm(seedingAlg)
    s.addAlgorithm(parEstimateAlg)

    s.addWriter(
        acts.examples.RootParticleWriter(
            level=acts.logging.INFO,
            inputParticles=evGen.config.outputParticles,
            filePath=outputDir + "/evgen_particles.root",
        ))
    s.addWriter(
        acts.examples.CsvParticleWriter(
            level=acts.logging.INFO,
            inputParticles=evGen.config.outputParticles,
            outputStem="evgen_particles",
            outputDir=outputDir + "/csv",
        ))

    s.addWriter(
        acts.examples.RootParticleWriter(
            level=acts.logging.INFO,
            inputParticles=simAlg.config.outputParticlesFinal,
            filePath=outputDir + "/fatras_particles_final.root",
        ))
    s.addWriter(
        acts.examples.CsvParticleWriter(
            level=acts.logging.INFO,
            inputParticles=simAlg.config.outputParticlesFinal,
            outputStem="fatras_particles_final",
            outputDir=outputDir + "/csv",
        ))

    s.addWriter(
        acts.examples.RootParticleWriter(
            level=acts.logging.INFO,
            inputParticles=simAlg.config.outputParticlesInitial,
            filePath=outputDir + "/fatras_particles_initial.root",
        ))
    s.addWriter(
        acts.examples.CsvParticleWriter(
            level=acts.logging.INFO,
            inputParticles=simAlg.config.outputParticlesInitial,
            outputStem="fatras_particles_initial",
            outputDir=outputDir + "/csv",
        ))

    s.addWriter(
        acts.examples.TrackFinderPerformanceWriter(
            level=acts.logging.INFO,
            inputProtoTracks=seedingAlg.config.outputProtoTracks,
            inputParticles=inputParticles,
            inputMeasurementParticlesMap=digiCfg.outputMeasurementParticlesMap,
            filePath=outputDir + "/performance_seeding_trees.root",
        ))

    s.addWriter(
        acts.examples.SeedingPerformanceWriter(
            level=acts.logging.DEBUG,
            inputProtoTracks=seedingAlg.config.outputProtoTracks,
            inputParticles=inputParticles,
            inputMeasurementParticlesMap=digiCfg.outputMeasurementParticlesMap,
            filePath=outputDir + "/performance_seeding_hists.root",
        ))

    s.addWriter(
        acts.examples.RootTrackParameterWriter(
            level=acts.logging.VERBOSE,
            inputTrackParameters=parEstimateAlg.config.outputTrackParameters,
            inputProtoTracks=parEstimateAlg.config.outputProtoTracks,
            inputParticles=simAlg.config.outputParticlesFinal,
            inputSimHits=simAlg.config.outputSimHits,
            inputMeasurementParticlesMap=digiCfg.outputMeasurementParticlesMap,
            inputMeasurementSimHitsMap=digiCfg.outputMeasurementSimHitsMap,
            filePath=outputDir + "/estimatedparams.root",
            treeName="estimatedparams",
        ))

    return s
Ejemplo n.º 4
0
def runCKFTracks(
    trackingGeometry,
    decorators,
    geometrySelection: Path,
    digiConfigFile: Path,
    field,
    outputDir: Path,
    truthSmearedSeeded=False,
    truthEstimatedSeeded=False,
    outputCsv=True,
    inputParticlePath: Optional[Path] = None,
    s=None,
):
    s = s or Sequencer(events=100, numThreads=-1)

    logger = acts.logging.getLogger("CKFExample")

    for d in decorators:
        s.addContextDecorator(d)

    rnd = acts.examples.RandomNumbers(seed=42)

    if inputParticlePath is None:
        logger.info("Generating particles using particle gun")

        evGen = acts.examples.EventGenerator(
            level=acts.logging.INFO,
            generators=[
                acts.examples.EventGenerator.Generator(
                    multiplicity=acts.examples.FixedMultiplicityGenerator(n=2),
                    vertex=acts.examples.GaussianVertexGenerator(
                        stddev=acts.Vector4(0, 0, 0, 0),
                        mean=acts.Vector4(0, 0, 0, 0)),
                    particles=acts.examples.ParametricParticleGenerator(
                        p=(1 * u.GeV, 10 * u.GeV),
                        eta=(-2, 2),
                        phi=(0, 360 * u.degree),
                        randomizeCharge=True,
                        numParticles=4,
                    ),
                )
            ],
            outputParticles="particles_input",
            randomNumbers=rnd,
        )
        s.addReader(evGen)
        inputParticles = evGen.config.outputParticles
    else:
        logger.info("Reading particles from %s", inputParticlePath.resolve())
        assert inputParticlePath.exists()
        inputParticles = "particles_read"
        s.addReader(
            RootParticleReader(
                level=acts.logging.INFO,
                filePath=str(inputParticlePath.resolve()),
                particleCollection=inputParticles,
                orderedEvents=False,
            ))

    # Selector
    selector = acts.examples.ParticleSelector(
        level=acts.logging.INFO,
        inputParticles=inputParticles,
        outputParticles="particles_selected",
    )
    s.addAlgorithm(selector)

    # Simulation
    simAlg = acts.examples.FatrasSimulation(
        level=acts.logging.INFO,
        inputParticles=selector.config.outputParticles,
        outputParticlesInitial="particles_initial",
        outputParticlesFinal="particles_final",
        outputSimHits="simhits",
        randomNumbers=rnd,
        trackingGeometry=trackingGeometry,
        magneticField=field,
        generateHitsOnSensitive=True,
    )
    s.addAlgorithm(simAlg)

    # Run the sim hits smearing
    digiCfg = acts.examples.DigitizationConfig(
        acts.examples.readDigiConfigFromJson(str(digiConfigFile)),
        trackingGeometry=trackingGeometry,
        randomNumbers=rnd,
        inputSimHits=simAlg.config.outputSimHits,
    )
    digiAlg = acts.examples.DigitizationAlgorithm(digiCfg, acts.logging.INFO)
    s.addAlgorithm(digiAlg)

    # Run the particle selection
    # The pre-selection will select truth particles satisfying provided criteria
    # from all particles read in by particle reader for further processing. It
    # has no impact on the truth hits themselves
    selAlg = acts.examples.TruthSeedSelector(
        level=acts.logging.INFO,
        ptMin=500 * u.MeV,
        nHitsMin=9,
        inputParticles=simAlg.config.outputParticlesInitial,
        inputMeasurementParticlesMap=digiCfg.outputMeasurementParticlesMap,
        outputParticles="particles_seed_selected",
    )
    s.addAlgorithm(selAlg)

    inputParticles = selAlg.config.outputParticles

    # Create starting parameters from either particle smearing or combined seed
    # finding and track parameters estimation
    if truthSmearedSeeded:
        logger.info("Using smeared truth particles for seeding")
        # Run particle smearing
        ptclSmear = acts.examples.ParticleSmearing(
            level=acts.logging.INFO,
            inputParticles=inputParticles,
            outputTrackParameters="smearedparameters",
            randomNumbers=rnd,
            # gaussian sigmas to smear particle parameters
            sigmaD0=20 * u.um,
            sigmaD0PtA=30 * u.um,
            sigmaD0PtB=0.3 / 1 * u.GeV,
            sigmaZ0=20 * u.um,
            sigmaZ0PtA=30 * u.um,
            sigmaZ0PtB=0.3 / 1 * u.GeV,
            sigmaPhi=1 * u.degree,
            sigmaTheta=1 * u.degree,
            sigmaPRel=0.01,
            sigmaT0=1 * u.ns,
            initialVarInflation=[1, 1, 1, 1, 1, 1],
        )
        outputTrackParameters = ptclSmear.config.outputTrackParameters
        s.addAlgorithm(ptclSmear)
    else:
        # Create space points
        spAlg = acts.examples.SpacePointMaker(
            level=acts.logging.INFO,
            inputSourceLinks=digiAlg.config.outputSourceLinks,
            inputMeasurements=digiAlg.config.outputMeasurements,
            outputSpacePoints="spacepoints",
            trackingGeometry=trackingGeometry,
            geometrySelection=acts.examples.readJsonGeometryList(
                str(geometrySelection)),
        )
        s.addAlgorithm(spAlg)

        # Run either: truth track finding or seeding
        if truthEstimatedSeeded:
            logger.info(
                "Using truth track finding from space points for seeding")
            # Use truth tracking
            truthTrackFinder = acts.examples.TruthTrackFinder(
                level=acts.logging.INFO,
                inputParticles=inputParticles,
                inputMeasurementParticlesMap=digiAlg.config.
                outputMeasurementParticlesMap,
                outputProtoTracks="prototracks",
            )
            s.addAlgorithm(truthTrackFinder)
            inputProtoTracks = truthTrackFinder.config.outputProtoTracks
            inputSeeds = ""
        else:
            logger.info("Using seeding")
            # Use seeding
            gridConfig = acts.SpacePointGridConfig(
                bFieldInZ=1.99724 * u.T,
                minPt=500 * u.MeV,
                rMax=200 * u.mm,
                zMax=2000 * u.mm,
                zMin=-2000 * u.mm,
                deltaRMax=60 * u.mm,
                cotThetaMax=7.40627,  # 2.7 eta
            )

            seedFilterConfig = acts.SeedFilterConfig(maxSeedsPerSpM=1,
                                                     deltaRMin=1 * u.mm)

            seedFinderConfig = acts.SeedfinderConfig(
                rMax=gridConfig.rMax,
                deltaRMin=seedFilterConfig.deltaRMin,
                deltaRMax=gridConfig.deltaRMax,
                collisionRegionMin=-250 * u.mm,
                collisionRegionMax=250 * u.mm,
                zMin=gridConfig.zMin,
                zMax=gridConfig.zMax,
                maxSeedsPerSpM=seedFilterConfig.maxSeedsPerSpM,
                cotThetaMax=gridConfig.cotThetaMax,
                sigmaScattering=50,
                radLengthPerSeed=0.1,
                minPt=gridConfig.minPt,
                bFieldInZ=gridConfig.bFieldInZ,
                beamPos=acts.Vector2(0 * u.mm, 0 * u.mm),
                impactMax=3 * u.mm,
            )
            seeding = acts.examples.SeedingAlgorithm(
                level=acts.logging.INFO,
                inputSpacePoints=[spAlg.config.outputSpacePoints],
                outputSeeds="seeds",
                outputProtoTracks="prototracks",
                gridConfig=gridConfig,
                seedFilterConfig=seedFilterConfig,
                seedFinderConfig=seedFinderConfig,
            )
            s.addAlgorithm(seeding)
            inputProtoTracks = seeding.config.outputProtoTracks
            inputSeeds = seeding.config.outputSeeds

        # Write truth track finding / seeding performance
        trackFinderPerformanceWriter = acts.examples.TrackFinderPerformanceWriter(
            level=acts.logging.INFO,
            inputProtoTracks=inputProtoTracks,
            inputParticles=
            inputParticles,  # the original selected particles after digitization
            inputMeasurementParticlesMap=digiAlg.config.
            outputMeasurementParticlesMap,
            filePath=str(outputDir / "performance_seeding_trees.root"),
        )
        s.addWriter(trackFinderPerformanceWriter)

        # Estimate track parameters from seeds
        paramEstimation = acts.examples.TrackParamsEstimationAlgorithm(
            level=acts.logging.INFO,
            inputSeeds=inputSeeds,
            inputProtoTracks=inputProtoTracks,
            inputSpacePoints=[spAlg.config.outputSpacePoints],
            inputSourceLinks=digiCfg.outputSourceLinks,
            outputTrackParameters="estimatedparameters",
            outputProtoTracks="prototracks_estimated",
            trackingGeometry=trackingGeometry,
            magneticField=field,
            bFieldMin=0.1 * u.T,
            deltaRMax=100.0 * u.mm,
            deltaRMin=10.0 * u.mm,
            sigmaLoc0=25.0 * u.um,
            sigmaLoc1=100.0 * u.um,
            sigmaPhi=0.02 * u.degree,
            sigmaTheta=0.02 * u.degree,
            sigmaQOverP=0.1 / 1.0 * u.GeV,
            sigmaT0=1400.0 * u.s,
            initialVarInflation=[1, 1, 1, 1, 1, 1],
        )
        s.addAlgorithm(paramEstimation)
        outputTrackParameters = paramEstimation.config.outputTrackParameters

    # Setup the track finding algorithm with CKF
    # It takes all the source links created from truth hit smearing, seeds from
    # truth particle smearing and source link selection config
    trackFinder = acts.examples.TrackFindingAlgorithm(
        level=acts.logging.INFO,
        measurementSelectorCfg=acts.MeasurementSelector.Config([
            (acts.GeometryIdentifier(), ([], [15.0], [10]))
        ]),
        inputMeasurements=digiAlg.config.outputMeasurements,
        inputSourceLinks=digiAlg.config.outputSourceLinks,
        inputInitialTrackParameters=outputTrackParameters,
        outputTrajectories="trajectories",
        findTracks=acts.examples.TrackFindingAlgorithm.makeTrackFinderFunction(
            trackingGeometry, field),
    )
    s.addAlgorithm(trackFinder)

    # write track states from CKF
    trackStatesWriter = acts.examples.RootTrajectoryStatesWriter(
        level=acts.logging.INFO,
        inputTrajectories=trackFinder.config.outputTrajectories,
        # @note The full particles collection is used here to avoid lots of warnings
        # since the unselected CKF track might have a majority particle not in the
        # filtered particle collection. This could be avoided when a seperate track
        # selection algorithm is used.
        inputParticles=selector.config.outputParticles,
        inputSimHits=simAlg.config.outputSimHits,
        inputMeasurementParticlesMap=digiAlg.config.
        outputMeasurementParticlesMap,
        inputMeasurementSimHitsMap=digiAlg.config.outputMeasurementSimHitsMap,
        filePath=str(outputDir / "trackstates_ckf.root"),
        treeName="trackstates",
    )
    s.addWriter(trackStatesWriter)

    # write track summary from CKF
    trackSummaryWriter = acts.examples.RootTrajectorySummaryWriter(
        level=acts.logging.INFO,
        inputTrajectories=trackFinder.config.outputTrajectories,
        # @note The full particles collection is used here to avoid lots of warnings
        # since the unselected CKF track might have a majority particle not in the
        # filtered particle collection. This could be avoided when a seperate track
        # selection algorithm is used.
        inputParticles=selector.config.outputParticles,
        inputMeasurementParticlesMap=digiAlg.config.
        outputMeasurementParticlesMap,
        filePath=str(outputDir / "tracksummary_ckf.root"),
        treeName="tracksummary",
    )
    s.addWriter(trackSummaryWriter)

    # Write CKF performance data
    ckfPerfWriter = acts.examples.CKFPerformanceWriter(
        level=acts.logging.INFO,
        inputParticles=inputParticles,
        inputTrajectories=trackFinder.config.outputTrajectories,
        inputMeasurementParticlesMap=digiAlg.config.
        outputMeasurementParticlesMap,
        # The bottom seed could be the first, second or third hits on the truth track
        nMeasurementsMin=selAlg.config.nHitsMin - 3,
        ptMin=0.4 * u.GeV,
        filePath=str(outputDir / "performance_ckf.root"),
    )
    s.addWriter(ckfPerfWriter)

    if outputCsv:
        csv_dir = outputDir / "csv"
        csv_dir.mkdir(parents=True, exist_ok=True)
        logger.info("Writing CSV files")
        csvMTJWriter = acts.examples.CsvMultiTrajectoryWriter(
            level=acts.logging.INFO,
            inputTrajectories=trackFinder.config.outputTrajectories,
            inputMeasurementParticlesMap=digiAlg.config.
            outputMeasurementParticlesMap,
            outputDir=str(csv_dir),
        )
        s.addWriter(csvMTJWriter)

    return s