test/tools/material.py

0001 import sys
0002 import re
0003 import bisect
0004 import copy
0005 import math
0006 from python25 import *
0007
0008 class ElementLevel:
0009   name = ""
0010   copy = 0
0011
0012   def parse(self, name):
0013     pattern  = re.compile(r'([A-Za-z0-9]+)\[([0-9]+)\]')
0014     match = pattern.match(name)
0015     self.name = match.group(1)
0016     self.copy = int(match.group(2))
0017
0018   def full_name(self):
0019     if self.copy == 0:
0020       return self.name
0021     else:
0022       return "%s[%d]" % (self.name, self.copy)
0023
0024   # self matches pattern iff:
0025   #  - they have the same name
0026   #  - they have the same index, or pattern has index 0
0027   def match(self, pattern):
0028     return (self.name == pattern.name) and ((pattern.copy == 0) or (self.copy == pattern.copy))
0029
0030   def __init__(self, name, copy = None):
0031     if copy is None:
0032       self.parse(name)
0033     else:
0034       self.name = name
0035       self.copy = copy
0036
0037   def __str__(self):
0038     return self.full_name()
0039
0040
0041 # base class for Element and ElementFilter
0042 # defines __hash__ for speed, but it's NOT generically immutable
0043 class ElementBase:
0044   name = []
0045
0046   def parse(self, name):
0047     self.name = [ElementLevel(item) for item in name.lstrip('/').split('/')]
0048
0049   def full_name(self):
0050     return "//" + "/".join([item.full_name() for item in self.name])
0051
0052   def levels(self):
0053     return len(self.name)
0054
0055   def __init__(self, other = None):
0056     if other is None:
0057       self.name = []
0058     elif isinstance(other, ElementBase):
0059       self.name = other.name
0060     elif isinstance(other, str):
0061       self.parse(other)
0062     else:
0063       raise TypeError("Cannot initialize an ElementBase from type %s" % type(other))
0064
0065   def __str__(self):
0066     return self.full_name()
0067
0068   def __eq__(self, other):
0069     return self.name == other.name
0070
0071   def __ne__(self, other):
0072     return self.name != other.name
0073
0074   def __hash__(self):
0075     return hash(self.full_name())
0076
0077
0078 class ElementFilter(ElementBase):
0079
0080   def match(self, element):
0081     if self.levels() != element.levels():
0082       return False
0083     return all( [element.name[i].match( self.name[i] ) for i in range(self.levels())] )
0084
0085
0086 class Element(ElementBase):
0087   directions = {
0088     'none'      : 0,
0089     'r'         : 1,
0090     'z'         : 2,
0091     'abs(z)'    : 3,
0092     '|z|'       : 3,
0093     'eta'       : 4,
0094     'abs(eta)'  : 5,
0095     '|eta|'     : 5,
0096     'phi'       : 6
0097   }
0098   dir_labels = {
0099     'none'      : 'None',
0100     'r'         : 'R',
0101     'z'         : 'Z',
0102     'abs(z)'    : 'Z',
0103     '|z|'       : 'Z',
0104     'eta'       : 'Eta',
0105     'abs(eta)'  : 'Eta',
0106     '|eta|'     : 'Eta',
0107     'phi'       : 'Phi'
0108   }
0109   position = (0, 0, 0, 0, 0, 0, 0)
0110
0111   def match(self, pattern):
0112     if self.levels() != pattern.levels():
0113       return False
0114     return all( [self.name[i].match( pattern.name[i] ) for i in range(self.levels())] )
0115
0116   def __init__(self, position, other):
0117     ElementBase.__init__(self, other)
0118     self.position = position
0119
0120
0121 def parse(source):
0122   pattern = re.compile(r'(.*) *\(([0-9-.]+) *, *([0-9-.]+) *, *([0-9-.]+)\)')
0123   elements = []
0124   for line in source:
0125     match = pattern.match(line)
0126     if not match:
0127       print('Warning: the following line does not match the parsing rules:')
0128       print(line)
0129       print()
0130       continue
0131     r   = float(match.group(2))
0132     z   = float(match.group(3))
0133     phi = float(match.group(4))
0134     eta = -math.log(r / (z + math.sqrt(r*r + z*z)))
0135     position = (0, r, z, abs(z), eta, abs(eta), phi)
0136     name = match.group(1)
0137     elements.append(Element(position, name))
0138   return elements
0139
0140
0141 # collapse the elements into a set of filters
0142 def collapse(elements):
0143   size = elements[0].levels()
0144   names = [{} for i in range(size)]
0145   filters = [ ElementFilter(element) for element in elements ]
0146   for i in range(size):
0147     for filter in filters:
0148       name = filter.name[i].name
0149       copy = filter.name[i].copy
0150       if name not in names[i]:
0151         # new name
0152         names[i][name] = copy
0153       elif names[i][name] and names[i][name] != copy:
0154         # name with different copy number
0155         names[i][name] = 0
0156     for filter in filters:
0157       name = filter.name[i].name
0158       copy = filter.name[i].copy
0159       if names[i][name] == copy:
0160         filter.name[i].copy = 0
0161   return filters
0162
0163
0164 # elements is as returned by parse()
0165 # cuts must be sorted in ascending order
0166 def split_along(direction, elements, cuts):
0167   filters = collapse(elements)
0168   groups  = [[] for i in range(len(cuts)+1)]
0169   for element in elements:
0170     i = bisect.bisect(cuts, element.position[direction])
0171     matching_filters = [ filter for filter in filters if filter.match(element) ]
0172     if len(matching_filters) == 0:
0173       print("Error: no matches for element %s" % element.full_name())
0174     elif len(matching_filters) > 1:
0175       print("Error: too many matches for element %s" % element.full_name())
0176     else:
0177       groups[i].append( matching_filters[0] )
0178   return groups
0179
0180 # return True if and only if all elements match one and only one filter group
0181 def check_groups(elements, groups):
0182   for element in elements:
0183     matching_groups = set()
0184     for (index, group) in enumerate(groups):
0185       for filter in group:
0186         if element.match(filter):
0187           matching_groups.add(index)
0188           break
0189     if len(matching_groups) == 0:
0190       # filters have lost validity
0191       return False
0192     elif len(matching_groups) > 1:
0193       # filters have lost discriminating power
0194       return False
0195   # all elements match one and only one filter group
0196   return True
0197
0198
0199 def remove_copy_number(elements, old_groups):
0200   levels = elements[0].levels()
0201   for level in range(levels):
0202     # if this level has no copy number, skip it
0203     if not any([filter.name[level].copy for group in old_groups for filter in group]):
0204       continue
0205     # try to remove the copy number for this level
0206     new_groups = [[] for group in old_groups]
0207     for (new, old) in zip(new_groups, old_groups):
0208       cache = []
0209       for filter in old:
0210         new_filter = copy.deepcopy(filter)
0211         new_filter.name[level].copy = 0
0212         new_hash = hash(new_filter.full_name())
0213         if new_hash not in cache:
0214           new.append(new_filter)
0215           cache.append(new_hash)
0216
0217     # check if the filter are still working (each element must belong to one and only one group)
0218     if check_groups(elements, new_groups):
0219       # the copy number can be safely removed
0220       old_groups = new_groups
0221
0222   return old_groups
0223