#!/usr/bin/env python3
##########################################################################
# basf2 (Belle II Analysis Software Framework) #
# Author: The Belle II Collaboration #
# #
# See git log for contributors and copyright holders. #
# This file is licensed under LGPL-3.0, see LICENSE.md. #
##########################################################################
"""
D* Veto reconstruction for ModeSelector.
This module reconstructs D* candidates from D mesons (daughters of B candidates)
combined with soft pions or pi0s from the Rest of Event. The reconstructed
D* mass difference (deltaMassDiff) and vertex fit chi2 probability are stored
as ExtraInfo on the B candidates.
This helps identify B -> D* X decays that were reconstructed as B -> D X
by the FEI, which the ModeSelector uses as input features (feature blocks 5-8).
"""
import math
import basf2 as b2
import modularAnalysis as ma
from ROOT import Belle2
from variables import variables as vm
from vertex import kFit, treeFit
class _SetDstarVetoDefaults(b2.Module):
"""Set missing D* veto ExtraInfo keys to NaN on all B candidates.
This ensures every B candidate has Dstp_deltaMassDiff and Dst0_deltaMassDiff
set (even when the corresponding D* type is not its first daughter and no veto
candidate is reconstructed).
"""
def __init__(self, particle_lists, keys):
"""Set up with lists of particle list names and ExtraInfo key names."""
super().__init__()
#: Particle list names to iterate over
self._particle_lists = particle_lists
#: ExtraInfo key names to set to NaN if missing
self._keys = keys
def event(self):
"""Set missing ExtraInfo keys to NaN for every candidate in each list."""
for list_name in self._particle_lists:
plist = Belle2.PyStoreObj(list_name)
if not plist.isValid():
continue
for i in range(plist.obj().getListSize()):
p = plist.obj().getParticle(i)
for key in self._keys:
if not p.hasExtraInfo(key):
p.addExtraInfo(key, math.nan)
def add_dstar_veto_aliases():
"""
Add variable aliases needed for D* veto reconstruction.
"""
# D* - D mass difference using invariant mass
vm.addAlias('trueM', 'M - dM')
vm.addAlias('massDiffInvM', 'formula(InvM - daughter(0, M))')
vm.addAlias('trueMassDiff', 'trueM - daughter(0,trueM)')
vm.addAlias('deltaMassDiffInvM', 'massDiffInvM - trueMassDiff')
vm.addAlias('deltaMassDiff', 'massDifference(0) - trueMassDiff')
vm.addAlias('dmID', 'extraInfo(decayModeID)')
[docs]
def addDstarVeto(
particleLists,
path: b2.Path = None,
deltaMassDiffCut: tuple = (-0.02, 0.02),
dMassCut: tuple = (-0.03, 0.03),
writeExtraInfo: bool = True,
skipTreeFit: bool = True
):
"""
Add D* veto reconstruction to the path for B meson particle lists.
This function reconstructs D* candidates by combining D mesons (first daughter
of B candidates) with soft pions or pi0s from the Rest of Event, to identify
cases where the FEI reconstructed B -> D X but the true decay was B -> D* X.
For B candidates with D0 as first daughter:
- D*+ -> D0 pi+ (from ROE)
- D*0 -> D0 pi0 (from ROE)
For B candidates with D+ as first daughter:
- D*+ -> D+ pi0 (from ROE)
The veto builds its own ROE on private particles that wrap the B candidates, so it
neither uses nor creates an ROE related to the B candidates. An ROE built by the user
on the input lists, before or after this function, is independent of the veto.
Parameters:
particleLists (str or list): Name(s) of B meson particle list(s)
(e.g., 'B+:feiHadronic' or ['B+:feiHadronic', 'B0:feiHadronic'])
path (basf2.Path): The basf2 path to add modules to.
deltaMassDiffCut (tuple): Cut on deltaMassDiff (D* mass diff - true mass diff) in GeV.
dMassCut (tuple): Cut on D and D* mass deviation (dM) in GeV.
writeExtraInfo (bool): Whether to write ExtraInfo to particles.
skipTreeFit (bool): If True, skip the vertex TreeFit (significant speedup). The
deltaMassDiff will use InvM-based computation instead of fit-based, and chiProb
will not be available (stored as NaN). Candidates are ranked by
abs(deltaMassDiffInvM) instead of chiProb. Default: True
The following ExtraInfo fields are added to B candidates:
For D0 daughter (D*+ and D*0 veto):
- Dstp_deltaMassDiff: Delta mass difference for D*+ -> D0 pi+
- Dstp_chiProb: Vertex fit chi2 probability for D*+ (NaN if skipTreeFit)
- Dst0_deltaMassDiff: Delta mass difference for D*0 -> D0 pi0
- Dst0_chiProb: Vertex fit chi2 probability for D*0 (NaN if skipTreeFit)
For D+ daughter (D*+ veto only):
- Dstp_deltaMassDiff: Delta mass difference for D*+ -> D+ pi0
- Dstp_chiProb: Vertex fit chi2 probability for D*+ (NaN if skipTreeFit)
"""
if path is None:
b2.B2FATAL("Path is required for addDstarVeto")
if isinstance(particleLists, str):
particleLists = [particleLists]
add_dstar_veto_aliases()
# ExtraInfo variable mappings
extra_info_dstp = {
'daughter(0,chiProb)': 'Dstp_chiProb',
'daughter(0,deltaMassDiff)': 'Dstp_deltaMassDiff',
}
extra_info_dst0 = {
'daughter(0,chiProb)': 'Dst0_chiProb',
'daughter(0,deltaMassDiff)': 'Dst0_deltaMassDiff',
}
# For veto reconstruction, use InvM-based deltaMassDiff
extra_info_veto_dstp = {
'daughter(0,deltaMassDiffInvM)': 'Dstp_deltaMassDiff',
}
extra_info_veto_dst0 = {
'daughter(0,deltaMassDiffInvM)': 'Dst0_deltaMassDiff',
}
if not skipTreeFit:
extra_info_veto_dstp['daughter(0,chiProb)'] = 'Dstp_chiProb'
extra_info_veto_dst0['daughter(0,chiProb)'] = 'Dst0_chiProb'
# Create pi+ list for D*+ -> D0 pi+ reconstruction
from_ip = "[[dr < 2] and [abs(dz) < 4]]"
p_cut = " and [p > 0.05] and [useCMSFrame(p) < 0.5]"
ma.fillParticleList('pi+:dstVeto', from_ip + p_cut, path=path)
# Create pi0 list for D* veto reconstruction. The lists are private to the veto,
# so they cannot collide with a list of the same name created elsewhere, e.g. one
# without the photon MVA, whose suppression cuts would then silently reject all pi0s.
# The selection reproduces the 50% efficiency pi0 selection optimised in May 2020
# (photon and pi0 cuts, mass-constrained fit) that the models were trained with.
photon_cuts = '[clusterNHits > 1.5] and thetaInCDCAcceptance and ' \
'[[clusterReg == 1 and E > 0.025] or [clusterReg == 2 and E > 0.025] or [clusterReg == 3 and E > 0.040]]'
ma.fillParticleList('gamma:dstVeto', photon_cuts, path=path)
# Photon MVA weights the training was done with (MC16rd)
ma.getBeamBackgroundProbability('gamma:dstVeto', weight='MC16rd', path=path)
ma.getFakePhotonProbability('gamma:dstVeto', weight='MC16rd', path=path)
ma.reconstructDecay('pi0:dstVetoAll -> gamma:dstVeto gamma:dstVeto', '0.105 < InvM < 0.150', dmID=1, path=path)
kFit('pi0:dstVetoAll', 0.0, 'mass', path=path)
# Apply additional pi0 cuts (matching training preprocessing)
pi0Cuts = '[useCMSFrame(p) < 0.5]'
pi0Cuts += ' and [daughter(0,beamBackgroundSuppression) > 0.5] and [daughter(0,fakePhotonSuppression) > 0.1]'
pi0Cuts += ' and [daughter(1,beamBackgroundSuppression) > 0.5] and [daughter(1,fakePhotonSuppression) > 0.1]'
ma.cutAndCopyList('pi0:dstVeto', 'pi0:dstVetoAll', pi0Cuts, path=path)
# --- Process each particle list ---
for particleList in particleLists:
particle_type = particleList.split(':')[0] # e.g., 'B+' or 'B0'
list_label = particleList.split(':')[1] if ':' in particleList else ''
# --- Process B candidates with D*+ or D*0 as first daughter ---
# These already have the correct mass difference, just store it
dstp_daughter_list = f'{particle_type}:dstVeto_Dstp_{list_label}'
dst0_daughter_list = f'{particle_type}:dstVeto_Dst0_{list_label}'
ma.cutAndCopyList(dstp_daughter_list, particleList,
'[abs(daughter(0,PDG)) == 413]', path=path)
ma.cutAndCopyList(dst0_daughter_list, particleList,
'[abs(daughter(0,PDG)) == 423]', path=path)
if writeExtraInfo:
ma.variablesToExtraInfo(dstp_daughter_list, extra_info_dstp, option=0, path=path)
ma.variablesToExtraInfo(dst0_daughter_list, extra_info_dst0, option=0, path=path)
# --- Process B candidates with D0 or D+ as first daughter (veto) ---
for d_pdg, d_str in [(421, 'D0'), (411, 'Dp')]:
channel_name = f'{particle_type}:dstVeto_{d_str}_{list_label}'
d_particle = 'D0' if d_pdg == 421 else 'D+'
ma.cutAndCopyList(channel_name, particleList,
f'abs(daughter(0,PDG)) == {d_pdg}', path=path)
# Wrap each B candidate in a new single-daughter particle and build the ROE for
# the wrappers. A particle can only have one ROE, so building it on the B
# candidates would make the veto reuse an ROE built by the user on the same
# candidates (with possibly different input lists), and an ROE built here would
# in turn be reused by the user. The wrappers are private to the veto, and their
# ROE contains the same tracks and clusters as an ROE of the B candidate.
wrapper_list = f'Xsd:dstVeto_{d_str}_{"Bp" if particle_type == "B+" else "B0"}_{list_label}'
ma.reconstructDecay(f'{wrapper_list} -> {channel_name}', '', allowChargeViolation=True, path=path)
path.modules()[-1].set_log_level(b2.LogLevel.ERROR)
ma.buildRestOfEvent(wrapper_list, path=path)
# Create ROE path
roe_path = b2.Path()
dead_end_path = b2.Path()
ma.signalSideParticleFilter(wrapper_list, '', roe_path, dead_end_path)
# Get particles from ROE or direct B daughters
roe_condition = f'[isInRestOfEvent == 1] or [isDescendantOfList({channel_name},1) == 1]'
ma.cutAndCopyList('pi0:dstVetoROE', 'pi0:dstVeto', roe_condition, path=roe_path)
if d_str == 'D0':
# D0 can form D*+ (with pi+) or D*0 (with pi0)
dst_daughters = ['pi+', 'pi0']
ma.cutAndCopyList('pi+:dstVetoROE', 'pi+:dstVeto', roe_condition, path=roe_path)
else:
# D+ can only form D*+ (with pi0)
dst_daughters = ['pi0']
# Fill signal side D
ma.fillSignalSideParticleList(f'{d_particle}:dstVetoSig', f'Xsd -> [{particle_type} -> ^{d_particle}]',
path=roe_path)
dstp_lists = []
dst0_lists = []
for i, dst_daughter in enumerate(dst_daughters):
if d_str == 'Dp' or (d_str == 'D0' and dst_daughter == 'pi+'):
dst_list = f'D*+:dstVeto_{i}'
else:
dst_list = f'D*0:dstVeto_{i}'
ma.reconstructDecay(f'{dst_list} -> {d_particle}:dstVetoSig {dst_daughter}:dstVetoROE',
'', dmID=i, path=roe_path)
# Mass window cuts
cut_str = f'[{deltaMassDiffCut[0]} < deltaMassDiffInvM < {deltaMassDiffCut[1]}]'
cut_str += f' and [{dMassCut[0]} < dM < {dMassCut[1]}]'
cut_str += f' and [{dMassCut[0]} < daughter(0,dM) < {dMassCut[1]}]'
ma.applyCuts(dst_list, cut_str, path=roe_path)
# Rank by best candidate
if dst_daughter == 'pi0':
ma.rankByHighest(dst_list, 'daughter(1,chiProb)', 1, path=roe_path)
elif dst_daughter == 'pi+':
if skipTreeFit:
ma.rankByLowest(dst_list, 'abs(deltaMassDiffInvM)', 1, path=roe_path)
# else: will rank by vertex fit quality after fit
if not skipTreeFit:
# Vertex fit with mass constraints
treeFit(
list_name=dst_list,
conf_level=0,
ipConstraint=False,
updateAllDaughters=False,
massConstraint=["D*+", "D*0", "D+", "D0", "K_S0", "pi0"],
path=roe_path,
)
if dst_daughter == 'pi+':
ma.rankByHighest(dst_list, 'chiProb', 1, path=roe_path)
if 'D*+:dstVeto' in dst_list:
dstp_lists.append(dst_list)
else:
dst0_lists.append(dst_list)
# Merge D* lists
if dstp_lists:
ma.copyLists('D*+:dstVeto', dstp_lists, writeOut=False, path=roe_path)
ma.applyCuts('D*+:dstVeto', 'useCMSFrame(p) < 3', path=roe_path)
ma.rankByLowest('D*+:dstVeto', 'dmID', 1, path=roe_path)
if dst0_lists:
ma.copyLists('D*0:dstVeto', dst0_lists, writeOut=False, path=roe_path)
ma.applyCuts('D*0:dstVeto', 'useCMSFrame(p) < 3', path=roe_path)
ma.rankByLowest('D*0:dstVeto', 'dmID', 1, path=roe_path)
# Create dummy particles to transfer ExtraInfo back to signal side
if dstp_lists:
ma.reconstructDecay('Xsd:dstVetoDstp -> D*+:dstVeto', '', allowChargeViolation=True, path=roe_path)
roe_path.modules()[-1].set_log_level(b2.LogLevel.ERROR)
if writeExtraInfo:
ma.variableToSignalSideExtraInfo('Xsd:dstVetoDstp', extra_info_veto_dstp, path=roe_path)
if dst0_lists:
ma.reconstructDecay('Xsd:dstVetoDst0 -> D*0:dstVeto', '', allowChargeViolation=True, path=roe_path)
roe_path.modules()[-1].set_log_level(b2.LogLevel.ERROR)
if writeExtraInfo:
ma.variableToSignalSideExtraInfo('Xsd:dstVetoDst0', extra_info_veto_dst0, path=roe_path)
# Execute ROE path
path.for_each('RestOfEvent', 'RestOfEvents', roe_path)
# The ROE loop writes the results to the wrappers; copy them to the B candidates
if writeExtraInfo:
veto_keys = list(extra_info_veto_dstp.values())
if d_str == 'D0':
veto_keys += list(extra_info_veto_dst0.values())
ma.variablesToDaughterExtraInfo(
wrapper_list, f'Xsd -> ^{particle_type}',
{f'extraInfo({key})': key for key in veto_keys}, path=path)
if writeExtraInfo:
# After all real values are set, fill any remaining missing deltaMassDiff
# keys with NaN. This ensures every B candidate has these keys so that
# ModeSelectorModule's hasExtraInfo check does not FATAL. NaN is converted
# to None in ModeSelectorModule and stored as 0 in the sparse feature matrix,
# matching the behaviour of a genuinely absent veto candidate.
path.add_module(_SetDstarVetoDefaults(
[particleList],
['Dstp_deltaMassDiff', 'Dst0_deltaMassDiff'],
))
b2.B2INFO(f"DstarVeto: Added D* veto reconstruction for {particleLists}")
