nueana/selection.py at master · lynnt20/nueana · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
import numpy as np

def InSpill(df,spill_start=0.2, spill_end=2.2):
    return (df.slc.barycenterFM.flashTime > spill_start) & (df.slc.barycenterFM.flashTime < spill_end)

def InScore(df,score_cut=0.02):
    return (df.slc.barycenterFM.score > score_cut)

def InRealisticFV(df):
    """
    Filter events based on realistic fiducial volume criteria.

    This function applies spatial cuts to determine if events fall within the realistic
    fiducial volume of the detector based on z and y coordinates.

    Parameters
    ----------
    df : pandas.DataFrame
        DataFrame containing event data with 'z' and 'y' coordinate columns.

    Returns
    -------
    pandas.Series or numpy.ndarray
        Boolean mask where True indicates the event is within the realistic fiducial volume.
        The condition is satisfied when either:
        - z >= 250 cm AND y < 100 cm, OR
        - z < 250 cm

    Notes
    -----
    The realistic fiducial volume is defined by the logical OR of two conditions:
    1. Events with z-coordinate >= 250 cm must have y-coordinate < 100 cm
    2. All events with z-coordinate < 250 cm are included regardless of y-coordinate
    """

    return (((df.z >= 250) & (df.y < 100)) | (df.z < 250))

def select(indf,
           spring=False,
           realisticFV=True,
           spill_start=0.2,
           spill_end=2.2,
           score_cut=0.02,
           min_shower_energy=0.5,
           max_track_length=200,
           max_conversion_gap=2,
           min_dedx=1,
           max_dedx=2.5,
           min_opening_angle=0.03,
           max_opening_angle=0.2):
    """
    Apply selection cuts to neutrino interaction data.

    Parameters
    ----------
    indf : pandas.DataFrame
        Input DataFrame with neutrino interaction data
    realisticFV : bool, optional
        Whether to apply realistic active volume cut (default: True)
    spill_start : float, optional
        Minimum flash time for beam spill (default: 0.2)
    spill_end : float, optional
        Maximum flash time for beam spill (default: 2.2)
    score_cut : float, optional
        Minimum flash matching score (default: 0.02)
    min_shower_energy : float, optional
        Minimum primary shower energy in GeV (default: 0.5)
    max_track_length : float, optional
        Maximum track length in cm for muon rejection (default: 200)
    max_conversion_gap : float, optional
        Maximum conversion gap (default: 2)
    min_dedx : float, optional
        Minimum dE/dx on best plane (default: 1)
    max_dedx : float, optional
        Maximum dE/dx on best plane (default: 2.5)
    min_opening_angle : float, optional
        Minimum shower opening angle (default: 0.03)
    max_opening_angle : float, optional
        Maximum shower opening angle (default: 0.2)

    Returns
    -------
    dict
        Dictionary of DataFrames after each selection cut
    """
    df_dict = {}
    df = indf.copy()

    # ** these cuts done already in makedf
    # * require nuscore > 0.5
    # * require not clear cosmic
    # * require reco vertex in AV
    # * require that there is a primary shower (at least one pfp w/ trackScore < 0.5)
    if realisticFV:
        df = df[InRealisticFV(df.slc.vertex)]
    df_dict['preselection'] = df

    # * require that the matched (many-to-many) is inside the beam spill
    df = df[InSpill(df, spill_start, spill_end) & InScore(df, score_cut)]
    df_dict['flash matching'] = df

    # * require that primary shower > min_shower_energy
    if spring:
        df[('primshw','shw','reco_energy','','','')] = df.primshw.shw.maxplane_energy * 1.25
        df = df[df.primshw.shw.maxplane_energy > min_shower_energy]
    else:
        df[('primshw','shw','reco_energy','','','')] = df.primshw.shw.bestplane_energy * 1.25
        df = df[df.primshw.shw.bestplane_energy > min_shower_energy]
    df_dict['shower energy'] = df

    # * require track length < max_track_length cm
    df = df[np.isnan(df.primtrk.trk.len) | (df.primtrk.trk.len < max_track_length)]
    # df = df.drop('primtrk',axis=1,level=0)
    df_dict['muon rejection'] = df

    df = df[df.primshw.shw.conversion_gap < max_conversion_gap]
    df_dict['conversion gap'] = df

    df = df[(df.primshw.shw.bestplane_dEdx > min_dedx) & (df.primshw.shw.bestplane_dEdx < max_dedx)]
    df_dict['dEdx'] = df

    df = df[(df.primshw.shw.open_angle > min_opening_angle) &
            (df.primshw.shw.open_angle < max_opening_angle)]
    df_dict['opening angle'] = df

    return df_dict