397a398,528
> 
> ################################
> ## insert polarization filter ##
> ################################
> # 0. make declarations/definitions 
> # 1. find ID of tree containing a vertex with a "bad" leg
> # 2. initiate loop over loop_diagrams
> # 3. check whether any part of loop_diagram is connected to 'good/bad' tree
> 
> ##################################################################
> # PID reference:
> # 22=a
> # 23=z
> # 24=w+
> # 25=h
> # 230 = Z0
> # 231 = ZT
> # 232 = ZA
> # 240 = W0
> # 241 = WT
> # 242 = WA 
> 
>         # 0. make declarations/definitions 
>         doFilterDebug=False
>         isTreeAlreadyVetoed=False
>         gotPol1=False
>         gotPol2=False
>         doTagSnglPol=True
>         valid_diag = False # assume bad diagram
> 
>         # (anti)tag list
>         redList=[22,23,24,-24]
>         antiTagList=[232,230] # removes ZX-lep-lep vertex
>         pol1TagList=[231] # requires ZX-lep-lep vertex
>         pol2TagList=[231] # requires ZX-lep-lep vertex
> 
>         snglTreeList=[] # loops with one tree
>         pol1TreeList=[] # loops with many trees
>         pol2TreeList=[] # loops with many trees
>         vetoTreeList=[] # loops with `bad` trees
> 
>         
> 
>         # 1. find ID of tree containing a vertex with a "bad" leg
>         #   Note: a "tree in structs" loop runs over all tree subgraphs
>         #   -- a "tree" here is any tree-level subgraph attached to a loop
>         #   -- structure contains all "tree" diagrams (non-loops)
>         #   -- a tree contains one or more vertices
>         #   -- a vertex contains one or more "legs"
>         #   -- a leg has a PID
>         #
>         #   case 1. single tree with two ZX
>         #   case 2. two trees, each with one ZX
> 
> 
>         for tree in structs:
>             isTreeAlreadyVetoed=False
> 
>             # does the tree have only one vertex? (box/pentagon/etc diagram!)
>             #   -- check if vertex has a ``bad'' leg (veto)
>             #   -- check if vertex has a ``good'' leg (tag)
>             if len(tree.get('vertices')) == 1:
>                 vertex = tree.get('vertices')[0]
>                 for leg in vertex.get('legs'): 
>                     if (leg.get('id') in redList) or \
>                        (leg.get('id') in antiTagList):
>                         # automatically reject trees with red flags
>                         vetoTreeList.append(tree.get('id'))
>                         # no need to explore vertex any further
>                         break # out of loop
> 
>                     elif (leg.get('id') in pol1TagList):
>                         # tag the simple tree as ``good''
>                         pol1TreeList.append(tree.get('id'))
> 
>                     elif (leg.get('id') in pol2TagList):
>                         # tag the simple tree as ``good''
>                         pol2TreeList.append(tree.get('id'))
> 
> 
> 
>             # does the tree have multiple verticies? (triangle diagram!)
>             #   -- check if vertex has a ``bad'' leg (veto)
>             #   -- check if vertex has a ``good'' leg (tag)
>             else:
>                 # reset conditions for new tree subgraph
>                 gotPol1=False
>                 gotPol2=False
> 
>                 for vertex in tree.get('vertices'):
>                     if(isTreeAlreadyVetoed): break # out of loop
> 
>                     for leg in vertex.get('legs'):
>                         if (leg.get('id') in redList) or \
>                            (leg.get('id') in antiTagList):
>                             # automatically reject trees with red flags
>                             vetoTreeList.append(tree.get('id'))
>                             # no need to explore vertex any further
>                             # (note: really, one should skip to next tree)
>                             isTreeAlreadyVetoed=True
>                             break # out of loop
> 
>                         elif leg.get('id') in pol1TagList:
>                             # check if vertex contains "good" legs anywhere
>                             gotPol1=True
>                             # for case of pol1=pol2
>                             if doFilterDebug or doTagSnglPol:
>                                 gotPol2=True
>                         
>                         elif leg.get('id') in pol2TagList:
>                             # check if vertex contains "good" legs anywhere
>                             gotPol2=True
>                             # for case of pol1=pol2
>                             if doFilterDebug or doTagSnglPol:
>                                 gotPol1=True 
> 
> 
>                 # do we want to tag this tree as special?
>                 # note: we want to ignore/keep benign trees, 
>                 #       e.g., external gluon legs
>                 if(gotPol1 and gotPol2 and not isTreeAlreadyVetoed):
>                     snglTreeList.append(tree.get('id'))
>                 
>         
>         if doFilterDebug:
>             print(f"{snglTreeList=}")
>             print(f"{pol1TreeList=}")
>             print(f"{pol2TreeList=}")
>             print(f"{vetoTreeList=}")
>             
>         # 2. initiate loop over loop_diagrams
402c533,535
<             valid_diag = True
---
> 
>             # assume diagram will be rejected
>             valid_diag = False
412a546,626
>             
>             # get list of PIDs attached to loop
>             connectedPIDList = diag.get_pdgs_attached_to_loop(structs)
>             if doFilterDebug:
>                 print('new diagram')
> 
>             # case 1: box diagram
>             if(len(connectedPIDList)==4):
>                 # reset for new diagram
>                 valid_diag = False
>                 gotPol1=False
>                 gotPol2=False
>                 isTreeAlreadyVetoed=False
> 
>                 if doFilterDebug:
>                     #valid_diag = True
>                     print('box diagram')
>                     treeMapList = diag.get('canonical_tag')
>                     print(f"{len(treeMapList)=},\t{treeMapList=}")
> 
>                 for boxLegID in connectedPIDList:
>                     if (boxLegID in redList) or\
>                        (boxLegID in antiTagList):
>                         isTreeAlreadyVetoed = True
>                         valid_diag = False
>                         break # out of loop
> 
>                     # not elif in order to recycle filter for case of pol1=pol2
>                     if boxLegID in pol1TagList:
>                         gotPol1=True
> 
>                     # not elif in order to recycle filter for case of pol1=pol2
>                     if boxLegID in pol2TagList:
>                         gotPol2=True
>                 
> 
>                 if(gotPol1 and gotPol2 and not isTreeAlreadyVetoed):
>                     # must pass all three conditions
>                     valid_diag = True
>                 else:
>                     valid_diag = False
> 
>             # case 2: triangle diagram
>             elif(len(connectedPIDList)==3):
>                 # 'canonical_tag' contains the list of trees attached to a loop
>                 # this includes: initial-state trees and final-state trees.
>                 # the following loop will: 
>                 #   -- go through each treeMapping in treeMapList = diag.get('canonical_tag')
>                 #   -- the [0] element of a treeMap is PID of loop object (+/- = CCW/CW loop) 
>                 #   -- the [1] element of a treeMap is the treeIDList
>                 #   -- the [0] element of treeMapList *is* the treeID
> 
>                 # reset for new diagram
>                 valid_diag = False
>                 treeMapList = diag.get('canonical_tag')
>                 if doFilterDebug:
>                     #valid_diag = True
>                     print('triangle diagram')
>                     print(f"{len(treeMapList)=},\t{treeMapList=}")
>                     
>                 
>                 # 3. check whether loop_diagram is connected to "bad" tree
>                 #   -- 'canonical_tag' links "loop" subgraph and "tree" subgraph
>                 for treeMapping in treeMapList:
>                     #treeMapList = treeMapping[1]
>                     canonTreeID = treeMapping[1][0]
> 
>                     if canonTreeID in vetoTreeList:
>                         # is the tree ``bad''?
>                         valid_diag = False
>                         break # out of loop
> 
>                     elif canonTreeID in snglTreeList:
>                         # is the tree ``good''?
>                         valid_diag = True
> 
>             else:
>                 print("number of connected legs = %i" % len(connectedPIDList))
>                 print("number of connected legs is not 3 or 4; exiting")
>                 return
> 
477a692,693
>                 if doFilterDebug:
>                     print('keeping diagram')
479a696,697
>                 if doFilterDebug:
>                     print('rejecting diagram')
