vertex.py

#!/usr/bin/env python3
 

 
from basf2 import register_module
from basf2 import B2WARNING
 
 
def fitVertex(
    list_name,
    conf_level,
    decay_string='',
    fitter='KFit',
    fit_type='vertex',
    constraint='',
    massConstraint=[],
    recoilMass=0,
    daughtersUpdate=False,
    smearing=0,
    path=None,
):
    """
    Perform the specified fit for each Particle in the given ParticleList.
 
    Info:
        Direct use of `fitVertex` is not recommended unless you know what you are doing.
        If you're unsure, you probably want to use `treeFit` or `kFit`.
 
    Parameters:
        list_name (str):        name of the input ParticleList
        conf_level (float):     minimum value of the confidence level to accept the fit.
            Setting this parameter to -1 selects all particle candidates.
            The value of 0 rejects the particle candidates with failed fit.
        decay_string (str):     select particles used for the vertex fit
        fitter (str):           Rave or KFit
        fit_type (str):         type of the kinematic fit (valid options are vertex/massvertex/mass/fourC/massfourC/recoilmass)
        constraint (str):       add additional constraint to the fit (valid options are empty string/ipprofile/iptube/mother)
        massConstraint (list(int) or list(str)): list of PDG ids or Names of the particles which are mass-constrained
        recoilMass (float):     invariant mass of recoil in GeV (valid only for recoilmass)
        daughtersUpdate (bool): make copy of the daughters and update them after the vertex fit
        smearing (float) :      IP tube width is smeared by this value (cm). meaningful only with 'KFit/vertex/iptube' option.
        path (basf2.Path):      modules are added to this path
    """
 
    warning = (
        "Direct use of fitVertex is not recommended unless you know what you are doing.\n"
        "Please use treeFit or kFit.\n"
        "See documentation at https://software.belle2.org"
    )
 
    B2WARNING(warning)
 
    _fitVertex(
        list_name, conf_level, decay_string, fitter, fit_type, constraint,
        massConstraint, recoilMass, daughtersUpdate, smearing, path)
 
 
def _fitVertex(
    list_name,
    conf_level,
    decay_string='',
    fitter='KFit',
    fit_type='vertex',
    constraint='',
    massConstraint=[],
    recoilMass=0,
    daughtersUpdate=False,
    smearing=0,
    path=None,
):
    """
    An internal function, performs the specified fit for each Particle in the given ParticleList.
 
    Warning:
        This is a private function, please use one of the aliases provided.
 
    Parameters:
        list_name (str):        name of the input ParticleList
        conf_level (float):     minimum value of the confidence level to accept the fit.
            Setting this parameter to -1 selects all particle candidates.
            The value of 0 rejects the particle candidates with failed fit.
        decay_string (str):     select particles used for the vertex fit
        fitter (str):           Rave or KFit
        fit_type (str):         type of the kinematic fit (valid options are vertex/massvertex/mass/fourC/massfourC/recoilmass)
        constraint (str):       add additional constraint to the fit (valid options are empty string/ipprofile/iptube/mother)
        massConstraint (list(int) or list(str)): list of PDG ids or Names of the particles which are mass-constrained
        recoilMass (float):     invariant mass of recoil in GeV (valid only for recoilmass)
        daughtersUpdate (bool): make copy of the daughters and update them after the vertex fit
        smearing (float) :      IP tube width is smeared by this value (cm). meaningful only with 'KFit/vertex/iptube' option.
        path (basf2.Path):      modules are added to this path
    """
 
    from pdg import from_names
 
    pvfit = register_module('ParticleVertexFitter')
    pvfit.set_name('ParticleVertexFitter_' + list_name)
    pvfit.param('listName', list_name)
    pvfit.param('confidenceLevel', conf_level)
    pvfit.param('vertexFitter', fitter)
    pvfit.param('fitType', fit_type)
    pvfit.param('withConstraint', constraint)
    pvfit.param('updateDaughters', daughtersUpdate)
    pvfit.param('decayString', decay_string)
    pvfit.param('recoilMass', recoilMass)
    pvfit.param('smearing', smearing)
    if massConstraint:
        pvfit.param('massConstraintList', from_names(massConstraint))
    path.add_module(pvfit)
 
 
def kFit(list_name,
         conf_level,
         fit_type='vertex',
         constraint='',
         daughtersUpdate=False,
         decay_string='',
         massConstraint=[],
         recoilMass=0,
         smearing=0,
         path=None):
    """
    Perform KFit for each Particle in the given ParticleList.
 
    Parameters:
        list_name (str):        name of the input ParticleList
        conf_level (float):     minimum value of the confidence level to accept the fit
            Setting this parameter to -1 selects all particle candidates.
            The value of 0 rejects particle candidates with a failed fit.
        fit_type (str):         type of the kinematic fit. Valid options are
 
          * ``mass`` for a mass-constrained fit
          * ``vertex`` for a vertex fit
          * ``massvertex`` for a vertex fit with a mass constraint on the mother particle
          * ``fourC`` for a vertex fit in which the mother particle's four-momentum is constrained to the beam four-momentum
          * ``massfourC`` for a vertex fit with a 4-momentum constraint and mass constraints on the specified daughter particles
          * ``recoilmass`` for kinematic fit in which the mass of the mother particle's recoil four-momentum with respect
            to the beam four-momentum is constrained
 
        constraint (str):       add an additional constraint to the fit (valid options are ipprofile or iptube)
        massConstraint (list(int) or list(str)): list of PDG ids or Names of the particles which are mass-constrained
        recoilMass (float):     invariant mass of recoil in GeV (valid only for recoilmass)
        daughtersUpdate (bool): make copy of the daughters and update them after the KFit
        decay_string (str):     select particles used for the KFit
        smearing (float) :      IP tube width is smeared by this value (cm). meaningful only with 'iptube' constraint.
        path (basf2.Path):      modules are added to this path
    """
 
    _fitVertex(
        list_name, conf_level, decay_string, 'KFit', fit_type, constraint,
        massConstraint, recoilMass, daughtersUpdate, smearing, path)
 
 
def raveFit(
    list_name,
    conf_level,
    fit_type='vertex',
    decay_string='',
    constraint='',
    daughtersUpdate=False,
    path=None,
    silence_warning=False,
):
    """
    Performs a vertex fit using `RAVE <https://github.com/rave-package>`_
    for each Particle in the given ParticleList.
 
    Warning:
        `RAVE <https://github.com/rave-package>`_ is deprecated since it is not maintained.
        Whilst we will not remove RAVE, it is not recommended for analysis use, other than benchmarking or legacy studies.
        Instead, we recommend :doc:`TreeFitter` (`vertex.treeFit`) or `vertex.kFit`.
 
    Parameters:
        list_name (str):    name of the input ParticleList
        conf_level (float): minimum value of the confidence level to accept the fit.
            Setting this parameter to -1 selects all particle candidates.
            The value of 0 rejects the particle candidates with failed fit.
        fit_type (str): type of the RAVE vertex fit.Valid options are
 
          * ``mass`` for a mass-constrained fit (the 7x7 error matrix of the mother particle has to be defined)
          * ``vertex`` for a vertex fit without any mass constraint
          * ``massvertex`` for a mass-constrained vertex fit
 
        decay_string (str): select particles used for the vertex fit
        constraint (str):   add additional constraint to the fit
            (valid options are ipprofile or iptube).
        daughtersUpdate (bool): make copy of the daughters and update them after the Rave vertex fit
        path (basf2.Path):  modules are added to this path
        silence_warning (bool): silence the warning advertising TreeFitter use
    """
 
    # verbose deprecation message
    message_a = (
        "RAVE is deprecated since it is not maintained.\n"
        "Whilst we will not remove RAVE, it is not recommended for analysis use, other than benchmarking or legacy studies.\n"
        "Instead, we recommend TreeFitter (treeFit) or KFit.\n"
        "Try: \n  treeFit(\'" + list_name + "\'," + str(conf_level) +
        ", updateAllDaughters=False, path=mypath)\n"
    )
    message_b = "To silence this warning, add silence_warning=True when you call this function."
 
    # if the user wants a constraint, they should check the doc, or send a ticket if it's not implemented
    if constraint == '':
        message_if = ""
    else:
        message_if = (
            "Please consult the documentation at  https://software.belle2.org \n"
            "(search for TreeFitter) for special constraints.\n"
        )
 
    if not silence_warning:
        B2WARNING(message_a + message_if + message_b)
 
    _fitVertex(list_name, conf_level, decay_string, 'Rave', fit_type, constraint, None, 0, daughtersUpdate, 0, path)
 
 
def treeFit(
    list_name,
    conf_level=0.001,
    massConstraint=[],
    ipConstraint=False,
    updateAllDaughters=False,
    massConstraintDecayString='',
    massConstraintMassValues=[],
    customOriginConstraint=False,
    customOriginVertex=[0.001, 0, 0.0116],
    customOriginCovariance=[0.0048, 0, 0, 0, 0.003567, 0, 0, 0, 0.0400],
    originDimension=3,
    treatAsInvisible='',
    ignoreFromVertexFit='',
    path=None,
):
    """
    Perform a :doc:`TreeFitter` fit for each Particle in the given ParticleList.
 
    :Example:
        An example of usage for the decay chain :math:`B^0\\to\\pi^+\\pi^-\\pi^0`  is the following:
 
    ::
 
      reconstructDecay('pi0:A -> gamma:pi0 gamma:pi0', '0.130 < InvM < 0.14', path=mypath)
      reconstructDecay('B0:treefit -> pi+:my pi-:my pi0:A ', '', path=mypath)
      treeFit('B0:treefit', ipConstraint=True, path=mypath)
 
 
    .. note::
        The Particle object loaded from the KLMCluster does not have proper covariance matrix yet (by release-06 at least).
        The performance of TreeFit with these particles is not guaranteed.
        Alternatively, one can perform the TreeFit ignoring these particles from the vertex fit with
        the option ``ignoreFromVertexFit``.
 
 
    Parameters:
        list_name (str):     name of the input ParticleList
        conf_level (float):  minimum value of the confidence level to accept the fit.
            Setting this parameter to -1 selects all particle candidates.
            The value of 0 rejects the particle candidates with failed fit.
        massConstraint (list(int, str)): list of PDG ids (positive and negative values are allowed) or names of the
            (anti-)particles which are mass-constrained
        ipConstraint (bool): constrain head production vertex to IP (x-y-z) constraint
        updateAllDaughters (bool): if true the entire tree will be updated with the fitted values
            for momenta and vertex position. Otherwise only the momenta of the head of the tree will be updated,
            however for all daughters we also update the vertex position with the fit results as this would
            otherwise be set to {0, 0, 0} contact us if this causes any hardship/confusion.
        massConstraintDecayString (str): Decay string to select which particles' mass should be constrained
        massConstraintMassValues (list(float, int, str)): list of invariant masses to be used for the mass constraints
            of the particles selected via the decay string massConstraintDecayString. A floating point value is taken directly
            while an integer or string is interpreted as defining the PDG id or name of a particle
            and the corresponding PDG value is taken
        customOriginConstraint (bool): use a custom origin vertex as the production vertex of your particle.
            This is useful when fitting D*/D without wanting to fit a B but constraining the process to be B-decay-like.
            (think of semileptonic modes and stuff with a neutrino in the B decay).
        customOriginVertex (list(float)): 3d vector of the vertex coordinates you want to use as custom origin.
            Default numbers are taken for B-mesons
        customOriginCovariance (list(float)): 3x3 covariance matrix for the custom vertex (type: vector).
            Default numbers extracted from generator distribution width of B-mesons.
        originDimension (int): If the origin or IP constraint (``customOriginVertex`` or ``ipConstraint``) are used,
            this specifies the dimension of the constraint (3D or 2D).
        treatAsInvisible (str): Decay string to select one particle that will be treated as invisible in the fit.
        ignoreFromVertexFit (str): Decay string to select particles that will be ignored to determine the vertex position.
        path (basf2.Path): modules are added to this path
    """
 
    import pdg
 
    treeFitter = register_module("TreeFitter")
    treeFitter.set_name('TreeFitter_' + list_name)
    if massConstraint:
        treeFitter.param('massConstraintList', pdg.from_names(massConstraint))
    treeFitter.param('particleList', list_name)
    treeFitter.param('confidenceLevel', conf_level)
    treeFitter.param('ipConstraint', ipConstraint)
    treeFitter.param('updateAllDaughters', updateAllDaughters)
    treeFitter.param('customOriginConstraint', customOriginConstraint)
    treeFitter.param('customOriginVertex', customOriginVertex)
    treeFitter.param('customOriginCovariance', customOriginCovariance)
    treeFitter.param('originDimension', originDimension)
    treeFitter.param('treatAsInvisible', treatAsInvisible)
    treeFitter.param('ignoreFromVertexFit', ignoreFromVertexFit)
    treeFitter.param('massConstraintDecayString', massConstraintDecayString)
    if massConstraintMassValues:
        treeFitter.param('massConstraintMassValues', [
                pdg.get(item).Mass() if isinstance(item, (int, str)) else item for item in massConstraintMassValues])
    path.add_module(treeFitter)
 
 
def TagV(
    list_name,
    MCassociation='',
    confidenceLevel=0.,
    trackFindingType="standard_PXD",
    constraintType="tube",
    askMCInfo=False,
    reqPXDHits=0,
    maskName='all',
    fitAlgorithm='KFit',
    kFitReqReducedChi2=5.0,
    useTruthInFit=False,
    useRollBack=False,
    path=None,
):
    """
    For each Particle in the given Breco ParticleList:
    perform the fit of tag side using the track list from the RestOfEvent dataobject
    save the MC Btag in case of signal MC
 
    Parameters:
 
        list_name (str): name of the input Breco ParticleList
        MCassociation (str): use standard MC association or the internal one
        confidenceLevel (float): minimum value of the ConfidenceLevel to accept the fit. 0 selects CL > 0
        constraintType (str): choose the constraint used in the fit. Can be set to
 
          * noConstraint;
          * IP: **default**, tag B constrained to be on the IP;
          * tube: tube along the tag B line of flight, only for fully reconstructed signal B;
          * boost: long tube along the boost direction;
          * (breco): deprecated, but similar to tube;
 
        trackFindingType (str): choose how to look for tag tracks. Can be set to
 
          * standard: all tracks except from Kshorts;
          * standard_PXD: **default**, same as above but consider only tracks with at least 1 PXD hit.
              If the fit fails, attempt again with the standard option;
 
        fitAlgorithm (str):     Fitter used for the tag vertex fit: Rave or KFit (default)
        kFitReqReducedChi2 (float): The required chi2/ndf to accept the kFit result, if it is higher, iteration procedure is applied
        askMCInfo (bool): True when requesting MC Information from the tracks performing the vertex fit
        reqPXDHits (int): minimum N PXD hits for a track (default is 0)
        maskName (str): get particles from a specified ROE mask
        useTruthInFit (bool): True when the tag vertex fit is performed with the true momentum and
            position of the tracks (default is false). The variable :b2:var:`TagVFitTruthStatus` is set to 1
            if the truth-matching succeeds and 2 otherwise.
        useRollBack (bool): True when the tag vertex fit is performed with position of tracks rolled back to
            position of the mother B (default is false). The variable :b2:var:`TagVRollBackStatus` is set to 1
            if the truth-matching succeeds and 2 otherwise.
        path (basf2.Path): modules are added to this path
 
    Warning:
        Note that the useFitAlgorithm (str) parameter is deprecated and replaced by constraintType (str)
        and trackFindingType (str)
 
    Warning:
        The trackFindingType ``singleTrack`` and ``singleTrack_PXD`` are broken and **cannot** be used any more.
    """
 
    tvfit = register_module('TagVertex')
    tvfit.set_name('TagVertex_' + list_name)
    tvfit.param('listName', list_name)
    tvfit.param('maskName', maskName)
    tvfit.param('confidenceLevel', confidenceLevel)
    tvfit.param('MCAssociation', MCassociation)
    tvfit.param('trackFindingType', trackFindingType)
    tvfit.param('constraintType', constraintType)
    tvfit.param('askMCInformation', askMCInfo)
    tvfit.param('reqPXDHits', reqPXDHits)
    tvfit.param('fitAlgorithm', fitAlgorithm)
    tvfit.param('kFitReqReducedChi2', kFitReqReducedChi2)
    tvfit.param('useTruthInFit', useTruthInFit)
    tvfit.param('useRollBack', useRollBack)
    path.add_module(tvfit)
 
 
def fitPseudo(
    list_name,
    path,
):
    """
    Add a pseudo \"vertex fit\" which adds a covariance matrix from the combination of the four-vectors of the daughters.
    This is similar to BaBar's "Add4" function.
    It is commonly used for :math:`\\pi^0\\to\\gamma\\gamma` reconstruction where a vertex fit is not possible.
 
    Here is the basic usage:
 
    ::
 
        from modularAnalysis import fitPseudo
        from stdPi0s import stdPi0s
        stdPi0s("eff40_May2020", path=mypath)
        fitPseudo("pi0:eff40_May2020", path=mypath)
 
    Parameters:
        list_name (str): the name of the list to add the covariance matrix to
        path (basf2.Path): modules are added to this path
    """
    pseudofit = register_module('PseudoVertexFitter')
    pseudofit.set_name('PseudoVertexFitter_' + list_name)
    pseudofit.param('listName', list_name)
    path.add_module(pseudofit)
 
 
if __name__ == '__main__':
    from basf2.utils import pretty_print_module
    pretty_print_module(__name__, "vertex")