test/tools/material.py

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