DevelopmentPage/BeyondProjects: helasamp_lib.py

File helasamp_lib.py, 51.4 KB (added by Olivier Mattelaer, 14 years ago)
Line 
1################################################################################
2#
3# Copyright (c) 2010 The MadGraph Development team and Contributors
4#
5# This file is a part of the MadGraph 5 project, an application which
6# automatically generates Feynman diagrams and matrix elements for arbitrary
7# high-energy processes in the Standard Model and beyond.
8#
9# It is subject to the MadGraph license which should accompany this
10# distribution.
11#
12# For more information, please visit: http://madgraph.phys.ucl.ac.be
13#
14################################################################################
15## Diagram of Class
16##
17## Variable <--- ScalarVariable
18## |
19## +- LorentzObject
20##
21##
22## list <--- AddVariable
23## |
24## +- MultVariable <--- MultLorentz
25##
26## list <--- LorentzObjectRepresentation <-- ConstantObject
27##
28################################################################################
29"""This module develop basic structure for the symbolic treatment of the Helas
30Amplitude. This also authorize to write the Helas Routine in whatever languages.
31
32The Symbolic treatment are based on a series of (scalar) variable of the form
33a * x + b. 'a' and 'b' are (complex) number and x is an abstract variable.
34Those variable can be multiply/add to each other and will then create a list of
35those variable in 'MultVariable','AddVariable'.
36
37As We will work on two different level, we have two concrete representation of
38Variable:
39 - 'ScalarVariable' in order to represent a number (mass, component of
40 impulsion,...
41 - 'LorentzObject' in order to store HighLevel object with lorentz/spin
42 indices (as Gamma, Impulsion,...)
43
44In order to pass from HighLevel Object (deriving from LorentzObject) to concrete
45definition based on ScalarVariable, we define a new class
46'LorentzObjectRepresentation'
47"""
48import copy
49from numbers import Number
50#===============================================================================
51# FracVariable
52#===============================================================================
53class FracVariable(object):
54 """ A Fraction of Variable
55 """
56
57 def __init__(self, numerator, denominator):
58 self.numerator = numerator
59 self.denominator = denominator
60 if isinstance(self.numerator, Number):
61 self.tag = self.denominator.tag[:]
62 else:
63 self.tag = self.numerator.tag + self.denominator.tag
64
65 def simplify(self):
66 """apply rule of simplification"""
67 if not isinstance(self.numerator, Number):
68 self.numerator = self.numerator.simplify()
69 self.denominator = self.denominator.simplify()
70
71 if isinstance(self.denominator, ConstantObject):
72 self.numerator /= self.denominator.constant_term
73 return self.numerator
74 else:
75 return self
76
77 def factorize(self):
78 """made the factorization"""
79 self.numerator = self.numerator.factorize()
80 self.denominator = self.denominator.factorize()
81
82
83 def expand(self):
84 """Expand the content information"""
85 if isinstance(self.numerator, Number):
86 return FracVariable(self.numerator, self.denominator.expand())
87 return FracVariable(self.numerator.expand(), self.denominator.expand())
88
89 def __eq__(self, obj):
90 """Define the Equality of two Object"""
91
92 return (self.numerator == obj.numerator) and \
93 (self.denominator == obj.denominator)
94
95 def __mul__(self, obj):
96 """multiply by an object"""
97 # Multiply the denominator by one to ensure that the new object didn't
98 #have a pointer through the old one
99 if isinstance(obj, FracVariable):
100 return FracVariable(self.numerator * obj.numerator, \
101 self.denominator * obj.denominator)
102 return FracVariable(self.numerator * obj, self.denominator * 1)
103
104 __rmul__ = __mul__
105
106 def __div__(self, obj):
107 """Deal with division"""
108 # Multiply the denominator by one to ensure that the new object didn't
109 #have a pointer through the old one
110 return FracVariable(self.numerator * 1, self.denominator * obj)
111
112
113
114 __truediv__ = __div__
115
116 def __add__(self,obj):
117 """We don't need to deal with addition-substraction on this type of Object"""
118 if isinstance(obj, Number):
119 if obj:
120 self.numerator += obj * self.denominator
121 return self
122 assert(isinstance(obj, FracVariable))
123
124 if self.denominator != obj.denominator:
125 return NotImplemented('The Denominator should be the Same')
126 else:
127 new = FracVariable(self.numerator + obj.numerator, self.denominator)
128 return new
129
130 def __str__(self):
131 if isinstance(self.numerator, LorentzObjectRepresentation):
132 text = 'number of lorentz index :'+ str(len(self.numerator.lorentz_ind))+ '\n'
133 text += 'number of spin index :'+ str(len(self.numerator.spin_ind))+ '\n'
134 text += 'other info '+ str(self.numerator.tag) + '\n'
135 for ind in self.numerator.listindices():
136 ind=tuple(ind)
137 text += str(ind)+' --> '
138 if self.numerator.get_rep(ind) == 0:
139 text+= '0\n'
140 else:
141 text+='[ '+str(self.numerator.get_rep(ind)) + ' ] / [ '+ \
142 str(self.denominator.get_rep([0]))+ ']\n'
143 return text
144 else:
145 return '%s / %s' % (self.numerator, self.denominator)
146
147#===============================================================================
148# AddVariable
149#===============================================================================
150class AddVariable(list):
151 """ A list of Variable with addition as operator between themselves.
152 """
153 def __init__(self, old_data=[], prefactor=1, constant=0):
154 """ initialization of the object with default value """
155
156 self.prefactor = prefactor
157 self.constant_term = constant
158 self.tag = []
159 list.__init__(self, old_data)
160
161 def simplify(self):
162 """ apply rule of simplification """
163
164 new = self[0].simplify() + self.constant_term
165 #simplify each term of the sum
166 for i,term in enumerate(self[1:]):
167 new += term.simplify()
168 new *= self.prefactor
169
170 if isinstance(new, AddVariable):
171 return new.full_simplify()
172 else:
173 return new.simplify()
174
175 def full_simplify(self):
176
177 #look for identical term
178 i=-1
179 while i<len(self):
180 i+=1
181 for j in range(len(self)-1,i,-1):
182 if self[i] == self[j]:
183 self[i].prefactor += self[j].prefactor
184 del self[j]
185
186 #supress term with zero prefactor
187 for i in range(len(self)-1,-1,-1):
188 if self[i].prefactor == 0:
189 del self[i]
190
191 #avoid AddVariable if only one term
192 if len(self) == 1:
193 new = self[0] +self.constant_term
194 return new
195 elif len(self) == 0:
196 return ConstantObject(self.constant_term)
197
198 return self
199
200 def expand(self):
201 """Expand the content information"""
202
203 try:
204 new = self[0]
205 except:
206 new = 0
207 else:
208 new = new.expand()
209
210 for item in self[1:]:
211 new += item.expand()
212 if self.constant_term:
213 new += self.constant_term
214 return new
215
216 def __mul__(self,obj):
217
218 if isinstance(obj, Number):
219 new = self.__class__([], self.prefactor, self.constant_term)
220 new.constant_term *= obj
221 for term in self:
222 new += obj * term
223 return new
224
225 elif isinstance(obj, (Variable,MultVariable)):
226 return obj * self
227
228 elif isinstance(obj, AddVariable):
229 new = AddVariable()
230 for term in self:
231 new += term * obj
232 if self.constant_term:
233 new += self.constant_term * obj
234 return new
235 else:
236 return obj * self
237
238 def __add__(self, obj):
239 """Define all the different addition."""
240
241 if isinstance(obj, Number):
242 new = AddVariable(self, self.prefactor, self.constant_term)
243 new.constant_term += obj
244 return new
245 elif isinstance(obj, Variable):
246 new = AddVariable(self, self.prefactor, self.constant_term)
247 obj = copy.copy(obj)
248 new.append(obj)
249 return new
250 elif isinstance(obj, MultVariable):
251 new = AddVariable(self, self.prefactor, self.constant_term)
252 obj = MultVariable(obj, obj.prefactor, obj.constant_term)
253 new.append(obj)
254 return new
255 elif isinstance(obj, AddVariable):
256 new = AddVariable(list.__add__(self,obj), 1, self.constant_term + \
257 obj.constant_term)
258 return new
259 else:
260 return NotImplemented
261
262 def __div__(self, obj):
263 """ Implement division"""
264
265 if isinstance(obj, Number):
266 factor = 1 / obj
267 return factor * self
268 else:
269 return FracVariable(1 * self, 1 * obj)
270
271 def __rdiv__(self, obj):
272 """Deal division in a inverse way"""
273
274 new = AddVariable(self, self.prefactor, self.constant_term)
275 if isinstance(obj, Number):
276 return FracVariable( obj, new)
277 else:
278 return NotImplemented
279
280 def __sub__(self, obj):
281 return self + (-1)* obj
282
283 __radd__ = __add__
284 __iadd__ = __add__
285 __rmul__ = __mul__
286 __truediv__ = __div__
287 __rtruediv__ = __rdiv__
288
289 def __rsub__(self, obj):
290 return (-1) * self + obj
291
292 def append(self,obj):
293 """ add a newVariable in the Multiplication list """
294 if obj.prefactor:
295 self.constant_term += obj.constant_term
296 obj.constant_term = 0
297 list.append(self,obj)
298
299 def __eq__(self, obj):
300 """Define The Equality"""
301
302 if self.__class__ != obj.__class__:
303 return False
304
305 for term in self:
306 self_prefactor = [term2.prefactor for term2 in self if term2==term]
307 obj_prefactor = [term2.prefactor for term2 in obj if term2==term]
308 if len(self_prefactor) != len(obj_prefactor):
309 return False
310 self_prefactor.sort()
311 obj_prefactor.sort()
312 if self_prefactor != obj_prefactor:
313 return False
314
315 #Pass all the test
316 return True
317
318 def __str__(self):
319 text=''
320 if self.constant_term:
321 text+='['
322 if self.prefactor!=1:
323 text += str(self.prefactor) + ' * '
324 text+='( '
325 text+=' + '.join([str(item) for item in self])
326 text+=' )'
327 if self.constant_term!=0:
328 text += ' + '+str(self.constant_term)+' ]'
329 return text
330
331 def factorize(self):
332 """ try to foctorize as much as possible the expression """
333
334 count = {}
335 max, maxvar = 0, None
336 for term in self:
337 if isinstance(term, MultVariable):
338 for var in term:
339 count[var.variable] = count.setdefault(var.variable, 0) + 1
340 if count[var.variable] > max:
341 max, maxvar = max +1, var
342 else:
343 count[term.variable] = count.setdefault(term.variable, 0) + 1
344 if count[term.variable] > max:
345 max, maxvar = max +1, term
346
347 if max <= 1:
348 return self
349 else:
350 newadd = AddVariable()
351 #mult = MultVariable([maxvar,newadd])
352 constant = AddVariable([], 1, self.constant_term)
353 for term in self:
354 if maxvar == term:
355 term.power -= 1
356 newadd.append(term.simplify())
357 elif isinstance(term, MultVariable):
358 if maxvar in term:
359 pos = term.index(maxvar)
360 if term[pos].power == 1:
361 del term[pos]
362 newadd.append(term.simplify())
363 else:
364 term[pos].power -= 1
365 newadd.append(term.simplify())
366 else:
367 constant.append(term)
368 else:
369 constant.append(term)
370
371 newadd = newadd.factorize()
372 if len(constant):
373 constant = constant.factorize()
374 out = AddVariable([MultVariable([maxvar, newadd]),constant])
375 else:
376 if isinstance(newadd, MultVariable):
377 newadd.append(maxvar)
378 out = newadd
379 else:
380 out = MultVariable([maxvar,newadd])
381 return out
382
383
384#===============================================================================
385# MultVariable
386#===============================================================================
387class MultVariable(list):
388 """ A list of Variable with multiplication as operator between themselves.
389 """
390
391 add_class = AddVariable
392
393 def __init__(self,old_data=[], prefactor=1, constant=0):
394 """ initialization of the object with default value """
395
396 self.prefactor = prefactor
397 self.constant_term = constant
398 self.tag = []
399 list.__init__(self, old_data)
400
401 def simplify(self):
402 #Ask to Simplify each term
403 for i,fact in enumerate(self):
404 self[i] = fact.simplify()
405
406
407 #Check that we have at least two different factor:
408 if len(self) == 1:
409 return self.prefactor * self[0]
410 return self
411
412 def factorize(self):
413 """Try to factorize this (nothing to do)"""
414 return self
415
416 def expand(self):
417 """ expand each part of the product and combine them """
418
419 out = self[0].expand()
420 for fact in self[1:]:
421 out *= fact.expand()
422 out *= self.prefactor
423 return out
424
425 #Defining rule of Multiplication, Addition, Substraction
426 def __mul__(self, obj):
427
428 if isinstance(obj, Number):
429 return self.__class__(self, self.prefactor * obj, \
430 self.constant_term * obj)
431
432 elif isinstance(obj, Variable):
433 return obj *self
434
435 elif isinstance(obj, MultVariable):
436 if self.constant_term == 0 and obj.constant_term ==0:
437 new = self.__class__(self,self.prefactor)
438 for fact in obj:
439 new.append(fact)
440 new.prefactor *= obj.prefactor
441 return new
442 else:
443 new = self.add_class()
444 new += self.constant_term * obj
445 new += obj.constant_term * self
446 new.constant_term /=2
447 new2 = self.__class__(self, self.prefactor)
448 for fact in obj:
449 new2.append(fact)
450 new2.prefactor *= obj.prefactor
451 new += new2
452 return new
453
454
455 elif isinstance(obj, AddVariable):
456 if self.constant_term:
457 self2 = self.__class__(self,self.prefactor)
458 new = self.add_class(obj, obj.prefactor)
459 for i,data in enumerate(new):
460 new[i] *= self2
461 if obj.constant_term:
462 new += obj.constant_term * self2
463 new += self.constant_term * obj
464 else:
465 new = self.add_class(obj, obj.prefactor)
466 for i,data in enumerate(new):
467 new[i] *= self
468 if obj.constant_term:
469 new += obj.constant_term * self
470 return new
471 else:
472 return NotImplemented
473
474
475 def __add__(self,obj):
476
477 if isinstance(obj, Number):
478 new= self.__class__(self, self.prefactor, self.constant_term)
479 new.constant_term += obj
480 return new
481
482 elif isinstance(obj, MultVariable):
483 new = self.add_class()
484 new.append(self.__class__(self, self.prefactor, self.constant_term))
485 new.append(self.__class__(obj, obj.prefactor, obj.constant_term))
486 return new
487 else:
488 #call the implementation of addition implemented in obj
489 return obj + self
490
491 def __sub__(self, obj):
492 return self + (-1) * obj
493
494 def __rsub__(self, obj):
495 return (-1) * self + obj
496
497 def __div__(self, obj):
498 """ define the division """
499 if isinstance(obj, Number):
500 factor = 1 / obj
501 return factor * self
502 else:
503 return FracVariable(1 * self, 1 * obj)
504
505 def __rdiv__(self, obj):
506 """Deal division in a inverse way"""
507
508 new = self.__class__(self, self.prefactor, self.constant_term)
509 if isinstance(obj, Number):
510 return FracVariable( obj, new)
511 else:
512 return NotImplemented
513
514 __radd__ = __add__
515 __iadd__ = __add__
516 __rmul__ = __mul__
517 __truediv__ = __div__
518 __rtruediv__ = __rdiv__
519
520
521
522 def append(self,obj):
523 """ add a newVariable in the Multiplication list """
524
525 self.prefactor *= obj.prefactor
526 if obj in self:
527 index =self.index(obj)
528 self[index].power += 1
529 else:
530 obj.prefactor = 1
531 obj.constant_term = 0
532 list.append(self,obj)
533
534 def __eq__(self, obj):
535 """Define When two MultVariable are identical"""
536
537 if self.__class__ != obj.__class__ or len(self) != len(obj):
538 return False
539 else:
540 out1 = [var.variable for var in self]
541 out1.sort()
542 out2 = [var.variable for var in obj]
543 out2.sort()
544 if out1 != out2:
545 return False
546 #for var in self:
547 # if self.count(var) != obj.count(var):
548 # return False
549
550 return True
551
552 def __str__(self):
553 """ String representation """
554
555 text=''
556 if self.prefactor!=1:
557 text += str(self.prefactor) + ' * '
558 text+='( '
559 text+=' * '.join([str(item) for item in self])
560 text+=' )'
561 return text
562 __rep__ = __str__
563
564#===============================================================================
565# Variable
566#===============================================================================
567class Variable(object):
568 """This is the standard object for all the variable linked to expression.
569 the variable 'X' is associte to a prefactor and a constant term. Such that
570 this object can be reprensentet mathematicaly by a * X + b
571 """
572
573 mult_class = MultVariable
574 add_class = AddVariable
575
576 class VariableError(Exception):
577 """class for error in Variable object"""
578 pass
579
580 def __init__(self, prefactor=1, constant=0, variable='x'):
581 """ [prefactor] * Variable + [constant]"""
582
583 self.prefactor = prefactor
584 self.constant_term = constant
585 self.variable = variable
586 self.power = 1
587
588 def simplify(self):
589 """Define How to simplify this object."""
590 return self
591
592 def expand(self):
593 """Return a more basic representation of this variable."""
594 return self
595
596 def factorize(self):
597 """ try to factorize this"""
598 return self
599
600 # Define basic operation (multiplication, addition, soustraction)
601 def __mul__(self, obj):
602 """ How to multiply object together
603 product of Variable -> MultVariable
604 """
605
606 if isinstance(obj, Number):
607 if not obj:
608 return 0
609 out = copy.copy(self)
610 out.prefactor *= obj
611 out.constant_term *= obj
612 return out
613
614 elif isinstance(obj, Variable):
615 if self.constant_term == 0 and obj.constant_term ==0:
616 new = self.mult_class()
617 new.append(copy.copy(self))
618 new.append(copy.copy(obj))
619 return new
620 else:
621 new = self.add_class()
622 new += obj.constant_term * self
623 new += self.constant_term * obj
624 prod_part = self.mult_class()
625 prod_part.append(copy.copy(self))
626 prod_part.append(copy.copy(obj))
627 new += prod_part
628 new.constant_term /= 2
629
630 return new
631
632 elif isinstance(obj, MultVariable):
633 if self.constant_term !=0:
634 new = self.add_class()
635 #store prefactor-constant since mult-add will change those
636 prefactor_self, prefactor_obj = self.prefactor, obj.prefactor
637 constant_self, constant_obj = self.constant_term, obj.constant_term
638 #repass to default
639 self.prefactor, obj.prefactor = 1, 1
640 self.constant_term, obj.constant_term = 0, 0
641 new.append(prefactor_self * self * prefactor_obj * obj)
642 new.append(prefactor_self * self * constant_obj)
643 new.append(constant_self * prefactor_obj * obj)
644 new += constant_obj * constant_self
645 return new
646 else:
647 out = self.mult_class()
648 for factor in obj:
649 out.append(copy.copy(factor))
650 out.append(copy.copy(self))
651 out.prefactor *= obj.prefactor
652 return out
653
654 elif isinstance(obj, AddVariable):
655 new = AddVariable()
656
657 for term in obj:
658 new += self * term
659 new += obj.constant_term *self
660 return new
661 else:
662 return NotImplemented
663
664
665 def __add__(self, obj):
666 """ How to make an addition
667 Addition of Variable -> AddVariable
668 """
669
670 if isinstance(obj, Number):
671 out = copy.copy(self)
672 out.constant_term += obj
673 return out
674 #return self.__class__(*self.data(), prefactor=prefactor,
675 # constant=constant)
676
677 elif isinstance(obj, (Variable)):
678 new = self.add_class()
679 new.append(copy.copy(self))
680 new.append(copy.copy(obj))
681 return new
682
683 elif isinstance(obj, MultVariable):
684 new = self.add_class()
685 new.append(copy.copy(self))
686 new.append(self.mult_class(obj, obj.prefactor, obj.constant_term))
687 return new
688
689 elif isinstance(obj, AddVariable):
690 new = self.add_class(obj, obj.prefactor, obj.constant_term)
691 new.append(self)
692 return new
693 else:
694 return NotImplemented
695
696 def __sub__(self, obj):
697 return self + -1 * obj
698
699 def __rsub__(self, obj):
700 return (-1) * self + obj
701
702 def __div__(self, obj):
703 """ define the division """
704
705 if isinstance(obj, Number):
706 factor = 1 / obj
707 return factor * self
708 else:
709 return FracVariable(1 * self, 1 * obj)
710
711 def __rdiv__(self, obj):
712 """Deal division in a inverse way"""
713
714 new = copy.copy(self)
715 if isinstance(obj, Number):
716 return FracVariable( obj, new)
717 else:
718 return NotImplemented
719
720 __radd__ = __add__
721 __iadd__ = __add__
722 __rmul__ = __mul__
723 __truediv__ = __div__
724 __rtruediv__ = __rdiv__
725
726 def __eq__(self,obj):
727 """ identical if the variable is the same """
728 if isinstance(obj, Variable):
729 return self.variable == obj.variable
730 else:
731 return False
732
733 def __str__(self):
734 text = ''
735 if self.constant_term:
736 text +='['
737 if self.prefactor != 1:
738 text += str(self.prefactor)+' * '
739 text +=str(self.variable)
740 if self.power !=1:
741 text += '**%s' % self.power
742 if self.constant_term != 0:
743 text+= ' + '+str(self.constant_term)+']'
744 return text
745
746 __repr__ = __str__
747
748#===============================================================================
749# ScalarVariable
750#===============================================================================
751class ScalarVariable(Variable):
752 """ A concrete symbolic scalar variable
753 """
754
755 def __init__(self, type, indices='', prefactor=1, constant=0):
756 """ initialization of the object with default value """
757
758 self.type = type[0]
759 self.indices = indices
760 Variable.__init__(self, prefactor, constant, type)
761
762# def __str__(self):
763# text = ''
764# if self.prefactor != 1:
765# text += str(self.prefactor)+' * '
766# text += self.variable
767# if self.constant_term != 0:
768# text+= ' + '+str(self.constant_term)
769# return text
770
771#===============================================================================
772# AddLorentz
773#===============================================================================
774# Not needed at present stage
775
776
777#===============================================================================
778# MultLorentz
779#===============================================================================
780class MultLorentz(MultVariable):
781 """Specific class for LorentzObject Multiplication"""
782
783 add_class = AddVariable
784
785 def find_lorentzcontraction(self):
786 """ """
787 out={}
788 for i, fact in enumerate(self):
789 for j in range(len(fact.lorentz_ind)):
790 for k, fact2 in enumerate(self):
791 if i == k:
792 continue
793 try:
794 l = fact2.lorentz_ind.index(fact.lorentz_ind[j])
795 except:
796 pass
797 else:
798 out[(i,j)] = (k,l)
799 return out
800
801 def find_spincontraction(self):
802 """ """
803 out={}
804 for i, fact in enumerate(self):
805 for j in range(len(fact.spin_ind)):
806 for k, fact2 in enumerate(self):
807 if i == k:
808 continue
809 try:
810 l = fact2.spin_ind.index(fact.spin_ind[j])
811 except:
812 pass
813 else:
814 out[(i,j)] = (k,l)
815 return out
816
817 def check_equivalence(self, obj, i, j, lorentzcontractself, lorentzcontractobj, \
818 spincontractself, spincontractobj, map):
819 """check if i and j are compatible up to sum up indices"""
820
821 # Fail if not the same class
822 if self[i].__class__ != obj[j].__class__:
823 return False
824
825 # Check if an assignement already exist for any of the factor consider
826 if map.has_key(i):
827 if map[i] != j:
828 # The self factor does't point on the obj factor under focus
829 return False
830 elif j in map.values():
831 # the obj factor is already mapped by another variable
832 return False
833
834 # Check if the tag information is identical
835 if self[i].tag != obj[j].tag:
836 return False
837
838 # Check all lorentz indices
839 for k in range(len(self[i].lorentz_ind)):
840 # Check if the same indices
841 if self[i].lorentz_ind[k] == obj[j].lorentz_ind[k]:
842 continue
843 # Check if those indices are contracted
844 if (i,k) not in lorentzcontractself or \
845 (j,k) not in lorentzcontractobj:
846 return False
847
848 #return object-indices of contraction
849 i2, k2 = lorentzcontractself[(i,k)]
850 j2, l2 = lorentzcontractobj[(j,k)]
851
852 # Check that both contract at same position
853 if k2 != l2:
854 return False
855
856 # Check that both object are the same type
857 if self[i2].__class__ != obj[j2].__class__:
858 return False
859
860 # Check if one of the term is already map
861 if map.has_key(i2) and map[i2] != j2:
862 if map[i2] != j2:
863 return False
864 else:
865 # if not mapped, map it
866 map[i2] = j2
867
868
869 # Do the same but for spin indices
870 for k in range(len(self[i].spin_ind)):
871 # Check if the same indices
872 if self[i].spin_ind[k] == obj[j].spin_ind[k]:
873 continue
874 #if not check if this indices is sum
875 if (i,k) not in spincontractself or \
876 (j,k) not in spincontractobj:
877
878 return False
879
880 #return correspondance
881 i2, k2 = spincontractself[(i,k)]
882 j2, l2 = spincontractobj[(j,k)]
883
884 # Check that both contract at same position
885 if k2 != l2:
886 return False
887
888 # Check that both object are the same type
889 if self[i2].__class__ != obj[j2].__class__:
890 return False
891
892 # Check if one of the term is already map
893 if map.has_key(i2) and map[i2] != j2:
894 if map[i2] != j2:
895 return False
896 else:
897 map[i2] = j2
898
899 # If pass all the test then this is a possibility
900 return True
901
902 def find_equivalence(self, obj, pos, lorentzcontractself, lorentzcontractobj, \
903 spincontractself, spincontractobj, map):
904
905 possibility = []
906
907
908 for pos2 in range(len(obj)):
909 init_map = dict(map)
910 if self.check_equivalence(obj, pos, pos2, lorentzcontractself, \
911 lorentzcontractobj, spincontractself, \
912 spincontractobj, init_map):
913 init_map[pos] = pos2
914 possibility.append(init_map)
915
916 if pos+1 ==len(self) and possibility:
917 return True
918 for map in possibility:
919 if self.find_equivalence(obj, pos+1, lorentzcontractself, \
920 lorentzcontractobj, spincontractself, \
921 spincontractobj, map):
922 return True
923 return False
924
925 def append(self,obj):
926 """ add a newVariable in the Multiplication list """
927
928 self.prefactor *= obj.prefactor
929 obj.prefactor = 1
930 obj.constant_term = 0
931 list.append(self,obj)
932
933
934 def __eq__(self, obj):
935
936 # Check Standard Equality
937 if MultVariable.__eq__(self, obj):
938 return True
939
940 if len(self) != len(obj):
941 return False
942
943 # Check that the number and the type of contraction are identical both
944 #for spin and lorentz indices
945 spin_contract_self = self.find_spincontraction()
946 spin_contract_obj = obj.find_spincontraction()
947
948 #check basic consitency
949 if len(spin_contract_self) != len(spin_contract_obj):
950 return False
951
952 lorentz_contract_self = self.find_lorentzcontraction()
953 lorentz_contract_obj = obj.find_lorentzcontraction()
954 #check basic consitency
955 if len(lorentz_contract_self) != len(lorentz_contract_obj):
956 return False
957
958 # Try to achieve to have a mapping between the two representations
959 mapping = {}
960 a= self.find_equivalence(obj, 0, lorentz_contract_self, \
961 lorentz_contract_obj, spin_contract_self, \
962 spin_contract_obj, mapping)
963 return a
964
965#===============================================================================
966# LorentzObject
967#===============================================================================
968class LorentzObject(Variable):
969 """ A symbolic Object for All Helas object. All Helas Object Should
970 derivated from this class"""
971
972 mult_class = MultLorentz
973 add_class = AddVariable
974
975 def __init__(self,lorentz_indices, spin_indices, other_indices, prefactor=1,
976 constant=0, variable=''):
977 """ initialization of the object with default value """
978
979 self.lorentz_ind = lorentz_indices
980 self.spin_ind = spin_indices
981 self.tag = other_indices
982 if not variable:
983 variable = self.__class__.__name__
984 if hasattr(self, 'particle'):
985 variable+=str(self.particle)
986 if lorentz_indices:
987 variable+='^{%s}' % ','.join([str(ind) for ind in lorentz_indices])
988 if spin_indices:
989 variable+='_{%s}' % ','.join([str(ind) for ind in spin_indices])
990
991
992 Variable.__init__(self, prefactor, constant, variable)
993
994 def expand(self):
995 """Expand the content information into LorentzObjectRepresentation."""
996 try:
997 return self.representation
998 except:
999 self.create_representation()
1000 return self.prefactor * self.representation
1001
1002 def create_representation(self):
1003 raise self.VariableError("This Object doesn't have define representation")
1004
1005 def __eq__(self, obj):
1006 """redifine equality"""
1007
1008 return (self.__class__ == obj.__class__ and \
1009 self.lorentz_ind == obj.lorentz_ind and \
1010 self.spin_ind == obj.spin_ind and \
1011 self.tag == obj.tag and
1012 self.variable == obj.variable)
1013
1014
1015#===============================================================================
1016# IndicesIterator
1017#===============================================================================
1018class IndicesIterator:
1019 """Class needed for the iterator"""
1020
1021 def __init__(self,len):
1022 self.len = len
1023 self.max_value = 3
1024 if len:
1025 self.data = [-1] + [0] * (len - 1)
1026 else:
1027 self.data = 0
1028 self.next = self.nextscalar
1029
1030 def __iter__(self):
1031 return self
1032
1033 def next(self):
1034 for i in range(self.len):
1035 if self.data[i] < self.max_value:
1036 self.data[i] += 1
1037 return self.data
1038 else:
1039 self.data[i] = 0
1040 raise StopIteration
1041
1042 def nextscalar(self):
1043 if self.data:
1044 raise StopIteration
1045 else:
1046 self.data += 1
1047 return [0]
1048
1049
1050
1051
1052#===============================================================================
1053# LorentzObjectRepresentation
1054#===============================================================================
1055class LorentzObjectRepresentation(list):
1056 """A concrete representation of the LorentzObject."""
1057
1058 class LorentzObjectRepresentationError(Exception):
1059 """Specify error for LorentzObjectRepresentation"""
1060
1061 def __init__(self, representation, lorentz_indices, spin_indices, other_indices):
1062 """ initialize """
1063 self.lorentz_ind = lorentz_indices
1064 self.spin_ind = spin_indices
1065 self.tag = other_indices
1066 if self.lorentz_ind or self.spin_ind:
1067 list.__init__(self, representation)
1068 else:
1069 list.append(self,representation)
1070
1071
1072 def __add__(self, obj, fact=1):
1073 assert(isinstance(obj, LorentzObjectRepresentation))
1074
1075 if self.lorentz_ind != obj.lorentz_ind or self.spin_ind != obj.spin_ind:
1076 switch_order=[]
1077 for value in self.lorentz_ind:
1078 try:
1079 index = obj.lorentz_ind.index(value)
1080 except:
1081 raise self.LorentzObjectRepresentationError("Invalid" + \
1082 "addition. Object doen't have the same indices")
1083 else:
1084 switch_order.append(index)
1085 for value in self.spin_ind:
1086 try:
1087 index = obj.spin_ind.index(value)
1088 except:
1089 raise self.LorentzObjectRepresentationError("Invalid" + \
1090 "addition. Object doen't have the same indices")
1091 else:
1092 switch_order.append(len(self.lorentz_ind) + index)
1093
1094 switch = lambda ind : [ind[switch_order[i]] for i in range(len(ind))]
1095 else:
1096 switch = lambda ind : ind
1097
1098 new = LorentzObjectRepresentation([], self.lorentz_ind, self.spin_ind , \
1099 self.tag + obj.tag)
1100
1101 for ind in self.listindices():
1102 #permute index for the second object
1103 new.set_rep(ind, self.get_rep(ind) + fact * obj.get_rep(switch(ind)))
1104
1105
1106 return new.simplify()
1107 __iadd__ = __add__
1108
1109 def __sub__(self,obj):
1110 return self.__add__(obj,fact=-1)
1111
1112 def __rsub__(self, obj):
1113 return obj.__add__(self, fact=-1)
1114
1115 def __isub__(self, obj):
1116 return self.__add__(obj,fact=-1)
1117
1118
1119 def __mul__(self,obj):
1120 """New multiplication performing directly the einstein/spin sommation.
1121 """
1122 if isinstance(obj, (Number,Variable)):
1123 out= LorentzObjectRepresentation([], self.lorentz_ind, self.spin_ind,
1124 self.tag)
1125 for ind in out.listindices():
1126 out.set_rep(ind, obj * self.get_rep(ind))
1127 return out
1128 elif isinstance(obj, FracVariable):
1129 out = self * obj.numerator
1130 out /= obj.denominator
1131 return out
1132
1133
1134 assert(obj.__class__ == LorentzObjectRepresentation)
1135
1136 l_ind, sum_l_ind = self.compare_indices(self.lorentz_ind, \
1137 obj.lorentz_ind)
1138 s_ind, sum_s_ind = self.compare_indices(self.spin_ind, \
1139 obj.spin_ind)
1140
1141 if not( sum_l_ind or sum_s_ind):
1142 return self.tensor_product(obj)
1143
1144 new_object = LorentzObjectRepresentation([], l_ind, s_ind, \
1145 self.tag + obj.tag)
1146
1147
1148 for indices in new_object.listindices():
1149 dict_l_ind = self.pass_ind_in_dict(indices[:len(l_ind)], l_ind)
1150 dict_s_ind = self.pass_ind_in_dict(indices[len(l_ind):], s_ind)
1151 new_object.set_rep(indices, \
1152 self.contraction(obj,sum_l_ind, sum_s_ind,\
1153 dict_l_ind, dict_s_ind))
1154 return new_object
1155
1156 @staticmethod
1157 def pass_ind_in_dict(indices,key):
1158
1159 if not key:
1160 return {}
1161 out={}
1162 for i,ind in enumerate(indices):
1163 out[key[i]]=ind
1164 return out
1165
1166 @staticmethod
1167 def compare_indices(list1, list2):
1168 """return the each of the list."""
1169
1170 are_unique = []
1171 are_sum = []
1172 for indice in list1:
1173 if indice in list2:
1174 are_sum.append(indice)
1175 else:
1176 are_unique.append(indice)
1177 for indice in list2:
1178 if indice not in are_sum:
1179 are_unique.append(indice)
1180
1181 return are_unique, are_sum
1182
1183 def contraction(self, obj, l_sum, s_sum, l_dict, s_dict):
1184
1185 out = 0
1186 for l_value in IndicesIterator(len(l_sum)):
1187 l_dict.update( self.pass_ind_in_dict(l_value, l_sum))
1188 for s_value in IndicesIterator(len(s_sum)):
1189 s_dict_final = s_dict.copy()
1190 s_dict_final.update( self.pass_ind_in_dict(s_value, s_sum))
1191
1192 self_ind = self.combine_indices(l_dict, s_dict_final)
1193 obj_ind = obj.combine_indices(l_dict, s_dict_final)
1194
1195 factor = self.get_rep(self_ind)
1196 factor *= obj.get_rep(obj_ind)
1197 if factor:
1198 factor *= (-1)**(len(l_value)-l_value.count(0))
1199 out += factor
1200 return out
1201
1202 def combine_indices(self, l_dict, s_dict):
1203
1204 out=[]
1205 for value in self.lorentz_ind:
1206 out.append(l_dict[value])
1207 for value in self.spin_ind:
1208 out.append(s_dict[value])
1209 return out
1210
1211# # def insert_in_indices(self, indices, position, value):
1212# ## assert(len(position) == len(value))
1213# out=[]
1214# old_value = 0
1215# for i, value in enumerate(position):
1216# out = indices[old_value: value]
1217# out.append(value[i])
1218# old_value = value
1219# out += indices[old_value:]
1220# return out
1221
1222 def einstein_sum(self, obj, indices, indices2, ind1, ind2):
1223
1224 prov_obj=LorentzObjectRepresentation([], range(len(ind1)), [], [])
1225
1226 out = 0
1227 for value in prov_obj.listindices():
1228 self_ind = self.insert_in_indices(indices, ind1, value)
1229 obj_ind = obj.insert_in_indices(indices2, ind2, value)
1230 out += (-1)**(len(value)-value.count(0)) * self.get_rep(self_ind) * \
1231 obj.get_rep(obj_ind)
1232
1233 return out
1234
1235
1236
1237
1238 def tensor_product(self, obj):
1239 assert(isinstance(obj, LorentzObjectRepresentation))
1240
1241 new_object = LorentzObjectRepresentation([], \
1242 self.lorentz_ind + obj.lorentz_ind, \
1243 self.spin_ind+obj.spin_ind, \
1244 self.tag+obj.tag)
1245
1246 lor1 = len(self.lorentz_ind)
1247 lor2 = len(obj.lorentz_ind)
1248 spin1 = len(self.spin_ind)
1249 spin2 = len(obj.spin_ind)
1250
1251 selfind = lambda indices: indices[:lor1] + \
1252 indices[lor1+lor2: lor1 + lor2 + spin1]
1253 objind = lambda indices: indices[lor1: lor1 + lor2] +\
1254 indices[lor1 + lor2 + spin1:]
1255 if lor1 == 0 and spin1 == 0:
1256 selfind = lambda indices: [0]
1257 elif lor2 ==0 and spin2 ==0:
1258 objind = lambda indices: [0]
1259
1260 for indices in new_object.listindices():
1261 new_object.set_rep(indices, self.get_rep(selfind(indices))*
1262 obj.get_rep(objind(indices)))
1263 return new_object
1264
1265
1266 def __div__(self, obj):
1267 """ define division
1268 Only division by scalar!!!"""
1269 out= LorentzObjectRepresentation([], self.lorentz_ind, self.spin_ind,
1270 self.tag)
1271 for ind in out.listindices():
1272 out.set_rep(ind, self.get_rep(ind) / obj)
1273 return out
1274
1275 __rmul__ = __mul__
1276 __imul__ = __mul__
1277 __truediv__ = __div__
1278 __rtruediv__ = __div__
1279 __rdiv__ = __div__
1280
1281 def factorize(self):
1282 """Try to factorize each component"""
1283
1284 for ind in self.listindices():
1285 self.set_rep(ind, self.get_rep(ind).factorize())
1286 return self
1287
1288 def simplify(self):
1289 """Check if we can simplify the object (check for non treated Sum)"""
1290
1291 #look for Einstein Sum
1292 for i,ind in enumerate(self.lorentz_ind):
1293 try:
1294 j= i + 1 + self.lorentz_ind[i+1:].index(ind)
1295 except ValueError:
1296 pass
1297 else:
1298 obj = self.auto_contraction(i, j)
1299 obj = obj.simplify()
1300 return obj
1301
1302 # Look for spin contraction
1303 for i,ind in enumerate(self.spin_ind):
1304 try:
1305 j = i + 1 + self.spin_ind[i+1:].index(ind)
1306 except ValueError:
1307 pass
1308 else:
1309 obj = self.spin_sum(i, j)
1310 obj = obj.simplify()
1311 return obj
1312
1313 #Look for internal simplification
1314 for ind in self.listindices():
1315 term = self.get_rep(ind)
1316 try:
1317 term = term.simplify()
1318 except:
1319 continue
1320 else:
1321 self.set_rep(ind, term)
1322
1323
1324 #no additional simplification
1325 return self
1326
1327
1328 def listindices(self):
1329 """Return an iterator in order to be able to loop easily on all the
1330 indices of the object.
1331 """
1332
1333 return IndicesIterator(len(self.lorentz_ind)+len(self.spin_ind))
1334
1335
1336 def get_rep(self,indices):
1337 """return the value/Variable associate to the indices"""
1338
1339 data=self
1340 for value in indices:
1341 data = data[value]
1342 return data
1343
1344 def set_rep(self,indices,value):
1345 """assign 'value' at the indices position"""
1346
1347 data = self
1348 for ind in indices[:-1]:
1349 try:
1350 data = data[ind]
1351 except IndexError:
1352 for i in range(len(data),ind + 1):
1353 data.append([])
1354 data = data[ind]
1355
1356 #now we can assign the value
1357 try:
1358 data[indices[-1]] = value
1359 except IndexError:
1360 for i in range(len(data),indices[-1]+1):
1361 data.append(0)
1362 data[indices[-1]] = value
1363
1364
1365
1366
1367 def auto_contraction(self, ind1, ind2):
1368 """ perform the contraction of ind1 and ind2 """
1369 new_lorentz = list(self.lorentz_ind)
1370 if ind1 > ind2:
1371 ind1, ind2 = ind2, ind1
1372 del new_lorentz[ind2]
1373 del new_lorentz[ind1]
1374
1375 new_object = LorentzObjectRepresentation([], new_lorentz, \
1376 self.spin_ind, self.tag)
1377
1378 if new_object.lorentz_ind or self.spin_ind:
1379 old_ind = lambda indices, i: tuple(indices[:ind1]) + tuple([i]) +\
1380 tuple(indices[ind1:ind2-1]) + tuple([i]) + \
1381 tuple(indices[ind2-1:])
1382 else:
1383 old_ind = lambda indices, i: (i, i)
1384
1385 for ind in new_object.listindices():
1386 self.get_rep(old_ind(ind, 2)),self.get_rep(old_ind(ind, 3))
1387 new_object.set_rep(ind,
1388 self.get_rep(old_ind(ind, 0)) \
1389 - self.get_rep(old_ind(ind, 1)) \
1390 - self.get_rep(old_ind(ind, 2)) \
1391 - self.get_rep(old_ind(ind, 3)))
1392
1393
1394 return new_object
1395
1396 def spin_sum(self, ind1, ind2):
1397 """Perform spin contraction """
1398
1399 #create the spin list for the object
1400 new_spin = list(self.spin_ind)
1401 if ind1 > ind2:
1402 ind1, ind2 = ind2, ind1
1403 del new_spin[ind2]
1404 del new_spin[ind1]
1405
1406 # Define a New Empty Object
1407 new_object = LorentzObjectRepresentation([],self.lorentz_ind, \
1408 new_spin, self.tag)
1409
1410 #flip to indices to the position in the full list of indices
1411 ind1 , ind2 = len(self.lorentz_ind) + ind1, len(self.lorentz_ind) + ind2
1412
1413 if self.lorentz_ind or new_spin:
1414 old_ind = lambda indices, i: tuple(indices[:ind1]) + tuple([i]) +\
1415 tuple(indices[ind1:ind2-1]) + tuple([i]) + \
1416 tuple(indices[ind2-1:])
1417 else:
1418 old_ind = lambda indices, i: (i, i)
1419
1420
1421 for ind in new_object.listindices():
1422 # Run on all the indices of the new Object and I will define the value
1423 #accordingly
1424
1425 new_object.set_rep(ind,
1426 self.get_rep(old_ind(ind, 0)) \
1427 + self.get_rep(old_ind(ind, 1)) \
1428 + self.get_rep(old_ind(ind, 2)) \
1429 + self.get_rep(old_ind(ind, 3)))
1430
1431 return new_object
1432
1433 def create_output(self, language, output_path):
1434 """create the Helas routine linked to this object"""
1435
1436 try:
1437 Writerclass = eval('HelasWriterFor'+language)
1438 except NameError:
1439 raise Exception('No class define for defining Helas Routine in %s' \
1440 % language)
1441
1442 writing_class = Writerclass(self, output_path)
1443 writing_class.write()
1444
1445 def __eq__(self, obj):
1446 """Check that two representation are identical"""
1447
1448 if self.__class__ != obj.__class__:
1449 return False
1450 if self.lorentz_ind != obj.lorentz_ind:
1451 return False
1452 if self.spin_ind != obj.spin_ind:
1453 return False
1454
1455 for ind in self.listindices():
1456 self_comp = self.get_rep(ind)
1457 try:
1458 obj_comp = obj.get_rep(ind)
1459 except:
1460 return False
1461
1462 if self_comp != obj_comp:
1463 return False
1464 if not isinstance(self_comp, Number):
1465 if self_comp.prefactor != obj_comp.prefactor:
1466 return False
1467 if self_comp.constant_term != obj_comp.constant_term:
1468 return False
1469
1470 #Pass all the test
1471 return True
1472
1473 def __str__(self):
1474 """ string representation """
1475 text = 'number of lorentz index :'+ str(len(self.lorentz_ind))+ '\n'
1476 text += 'number of spin index :'+ str(len(self.spin_ind))+ '\n'
1477 text += 'other info '+ str(self.tag) + '\n'
1478 for ind in self.listindices():
1479 ind=tuple(ind)
1480 text += str(ind)+' --> '
1481 text+=str(self.get_rep(ind))+'\n'
1482 return text
1483
1484#===============================================================================
1485# ConstantObject
1486#===============================================================================
1487class ConstantObject(LorentzObjectRepresentation):
1488
1489 def __init__(self, value):
1490 pass
1491 self.lorentz_ind = []
1492 self.spin_ind = []
1493 self.tag = []
1494 self.constant_term = value
1495 self.prefactor = 1
1496
1497 def __add__(self, obj):
1498 """Zero is neutral"""
1499 return self.constant_term + obj
1500 __radd__ = __add__
1501
1502 def __mul__(self, obj):
1503 """Zero is absorbant"""
1504 return self.constant_term * obj
1505
1506 def get_rep(self,ind):
1507 """return representation"""
1508 return self.constant_term
1509
1510 def __eq__(self,obj):
1511 if isinstance(obj, ConstantObject):
1512 return True
1513 else:
1514 return self.constant_term == obj
1515
1516 def expand(self):
1517 return self
1518
1519 def __str__(self):
1520 return '0'