tinypede.py

Go to the documentation of this file.

#!/usr/bin/env python3
 
 
 
import numpy as np
from scipy.stats import chi2
import array
import math
import time
 
 
 
class MilleRecord(object):
 
  
  def __init__(self, doublePrec=False):
    
    self.__doublePrecision = doublePrec
    
    self.__position = 1 
    
    self.__numData = 0  
    
    self.__recLen = 0 
    
    self.__iMeas = 0    
    
    self.__iErr = 0     
    
    self.__inder = array.array('i') 
    
    self.__glder = array.array('d' if doublePrec else 'f') 
     
  
  def addData(self, dataList):
    if (self.__numData == 0):  # first word is error counter
      self.__inder.append(0)
      self.__glder.append(0.) 
    self.__numData += 1     
 
    aMeas, aPrec, indLocal, derLocal, labGlobal, derGlobal = dataList
    self.__inder.append(0)
    self.__glder.append(aMeas)
    self.__inder.fromlist(indLocal)
    self.__glder.fromlist(derLocal)    
    self.__inder.append(0)
    self.__glder.append(1.0 / math.sqrt(aPrec))  # convert to error
    self.__inder.fromlist(labGlobal)
    self.__glder.fromlist(derGlobal)
 
  
  def getData(self):
    aMeas = self.__glder[self.__iMeas]
    indLocal = []
    derLocal = []
    for i in range(self.__iMeas + 1, self.__iErr):
      indLocal.append(self.__inder[i])
      derLocal.append(self.__glder[i])
    aPrec = 1.0 / self.__glder[self.__iErr] ** 2  # convert to precision 
    indGlobal = []
    derGlobal = []
    for i in range(self.__iErr + 1, self.__position):
      indGlobal.append(self.__inder[i])
      derGlobal.append(self.__glder[i])    
    return aMeas, aPrec, indLocal, derLocal, indGlobal, derGlobal   
 
  
  def printRecord(self):
    print(" MilleRecord, len: ", len(self.__inder))
    print(self.__inder)
    print(self.__glder)
 
  
  def writeRecord(self, aFile):
    header = array.array('i')  # header with number of words
    header.append(-len(self.__inder) * 2 if self.__doublePrecision else len(self.__inder) * 2)
    header.tofile(aFile)
    self.__glder.tofile(aFile)
    self.__inder.tofile(aFile)
 
  
  def readRecord(self, aFile, aType):
    if aType == 'F':
      # FOrtran record info
      lenf = array.array('i')
      lenf.fromfile(aFile, 1)
    header = array.array('i')  # header with number of words
    header.fromfile(aFile, 1)
    self.__recLen = abs(header[0] // 2)
    if header[0] < 0:
      self.__glder = array.array('d')
    else: 
      self.__glder = array.array('f')
    self.__glder.fromfile(aFile, self.__recLen)
    self.__inder.fromfile(aFile, self.__recLen)
    if aType == 'F':
      # FOrtran record info
      lenf = array.array('i')
      lenf.fromfile(aFile, 1)
    
  
  def moreData(self):
    if (self.__position < self.__recLen):
      while (self.__position < self.__recLen and self.__inder[self.__position] != 0):
        self.__position += 1
      self.__iMeas = self.__position 
      self.__position += 1
      while (self.__position < self.__recLen and self.__inder[self.__position] != 0):
        self.__position += 1
      self.__iErr = self.__position 
      self.__position += 1
      # special data?
      if (self.__iMeas + 1 == self.__iErr and self.__glder[self.__iErr] < 0):
        self.__position += int(-self.__glder[self.__iErr])
      else: 
        while (self.__position < self.__recLen and self.__inder[self.__position] != 0):
          self.__position += 1
      self.__numData += 1  
      return True
    else:
      return False
 
  
  def specialDataTag(self):
    aTag = -1
    if (self.__iMeas + 1 == self.__iErr and self.__glder[self.__iErr] < 0):
      aTag = int(-self.__glder[self.__iErr] * 10. + 0.5) % 10
    return aTag
 
 
 
class Pede(object):
 
  
  def __init__(self):
    
    self.__binaryFiles = []
    
    self.__numPar = 0
    
    self.__parCounters = {}
    
    self.__parIndices = {}
    
    self.__matrix = None
    
    self.__vector = None
    
    self.__sumChi2 = 0.
    
    self.__sumNdf = 0
    
    self.__numRec = 0
    
    self.__numRejects = 0
    
    self.__numBad = 0
    
    self.__numCons = 0
    
    self.__listCons = []
    
    #
    print()
    print(" TinyPede - a simple PEDE implementation in python")
    print()
    
  
  def addBinaryFile(self, binaryFileName, maxRec=1000000, fileType='C'):
    self.__binaryFiles.append((binaryFileName, maxRec, fileType))
    
  
  def defineParameters(self, entriesCut=25, fixedPar=[]):
    for fileName, maxRecord, fileType in self.__binaryFiles:
      # open binary file
      if fileType == 'F':
        print(" Reading Fortran binary file ", fileName)
      else: 
        print(" Reading binary file ", fileName)
      binaryFile = open(fileName, "rb")
      nRec = 0
      try:
    
        while nRec < maxRecord:
          nRec += 1
          # read record
          rec = MilleRecord()
          rec.readRecord(binaryFile, fileType)
          nData = 0
          while (rec.moreData()):
            aTag = rec.specialDataTag()
            if (aTag >= 0):
              continue
            # get data
            nData += 1
            indGlobal = rec.getData()[4]
            for label in indGlobal:
              if label not in self.__parCounters:
                self.__parCounters[label] = 0
              self.__parCounters[label] += 1        
  
      except EOFError:
        pass
      
      print("   records found ", nRec)
      binaryFile.close()
    # fix parameters
    for label in fixedPar:
      if label in self.__parCounters:
        self.__parCounters[label] *= -1
    # active parameters
    for l in sorted(self.__parCounters):
      if self.__parCounters[l] >= entriesCut:
        self.__parIndices[l] = self.__numPar
        self.__numPar += 1
 
    print(" Found in binary files")     
    print("   total number of parameters ", len(self.__parCounters))
    print("   minimum number of entries  ", entriesCut)
    print("   number of free parameters  ", self.__numPar)
    print()
 
  
  def addConstraints(self, textFile):
    if self.__numPar == 0:
      print(" no variable parameters defined ")
      return
    # read text file
    print(" Reading constraint file ", textFile)
    nline = 0
    nCon = 0
    pairs = []  # (index, coefficient)
    tf = open(textFile)
    for line in iter(tf):
      nline += 1
      fields = line.split()
      if len(fields) < 2:
        continue
      if fields[0] == 'Constraint':
        # next constraint
        nCon += 1
        # previous constraint had variable parameters (pairs)?
        if len(pairs) > 0:
          self.__numCons += 1
          self.__listCons.append(pairs)
          pairs = []
        continue
      if fields[0] == '*':
        continue
      # label and coefficient
      label = int(fields[0])
      value = float(fields[1])
      if label in self.__parIndices:
        # variable parameter (label)
        pairs.append((self.__parIndices[label], value))
        
    tf.close()
    # last constraint
    if len(pairs) > 0:
      self.__numCons += 1
      self.__listCons.append(pairs)
      
    print("   number of all constraints  ", nCon)
    print("   number of accepted constr. ", self.__numCons)
    #print self.__listCons      
       
  
  def construct(self, chi2Factor=50.):
    # create global matrix and vector 
    # size: number of parameters + number of Lagrange multipliers for constraints
    self.__vector = np.zeros(self.__numPar + self.__numCons)
    self.__matrix = np.zeros((self.__numPar + self.__numCons, self.__numPar + self.__numCons))
    print()
    print(" Constructing linear equations system")
    print("   local chi2 scaling factor  ", chi2Factor)
    for fileName, maxRecord, fileType in self.__binaryFiles:
      # open binary file
      binaryFile = open(fileName, "rb")
      nRec = 0
      try:
        while nRec < maxRecord:
          # read record
          rec = MilleRecord()
          rec.readRecord(binaryFile, fileType)
          nRec += 1
          dataBlocks = [] 
          # number of local parameters
          numLoc = 0
          # number of measurements
          numMeas = 0
          
          while (rec.moreData()):
            aTag = rec.specialDataTag()
            if (aTag >= 0):
              continue
            # get data
            dataBlocks.append(rec.getData())
            numLoc = max(numLoc, dataBlocks[-1][2][-1])
            numMeas += 1
            
          #print(" numLoc ", nRec, numLoc, numMeas)
          vecMeas = np.zeros(numMeas)
          vecMeasW = np.zeros(numMeas)
          locDer = np.zeros((numMeas, numLoc))
          locDerW = np.zeros((numMeas, numLoc))
          gloDer = np.zeros((numMeas, self.__numPar))
          gloDerW = np.zeros((numMeas, self.__numPar))
          # loop over data blocks    
          for iData, data in enumerate(dataBlocks):
            #aRes = data.getResidual()[0]
            aMeas, aPrec, indLocal, derLocal, indGlobal, derGlobal = data
            # update local-local matrix
            for i, der in zip(indLocal, derLocal):
              locDer[iData, i - 1] = der
              locDerW[iData, i - 1] = der * aPrec
              vecMeas[iData] = aMeas       
              vecMeasW[iData] = aMeas * aPrec       
            # update local-global matrix
            for l, der in zip(indGlobal, derGlobal):
              if l not in self.__parIndices:
                continue
              i = self.__parIndices[l]
              gloDer[iData, i] = der
              gloDerW[iData, i] = der * aPrec
  
          # local fit
          matLocLoc = np.dot(locDerW.T, locDer)
          vecLoc = np.dot(locDerW.T, vecMeas)
          try: 
            locCov = np.linalg.inv(matLocLoc)         
          except np.linalg.LinAlgError as e:
            print("   local fit failed:", e)
            self.__numBad += 1
            continue          
          locSol = np.dot(locCov, vecLoc)
          # chi2, ndf
          locNdf = numMeas - numLoc
          locChi2 = np.dot(vecMeas - np.dot(locDer, locSol), vecMeasW - np.dot(locDerW, locSol))
          prob = 1. - chi2.cdf(locChi2 / chi2Factor, locNdf)
          # check rank
          nRank = np.linalg.matrix_rank(matLocLoc)
          if nRank < numLoc:
            self.__numBad += 1
            print(" ??? bad local fit ", nRec, locNdf, locChi2, numLoc, nRank)
            for i in range(numLoc):
              gcor2 = 1.0 - 1.0 / (matLocLoc[i, i] * locCov[i, i])  # global correlation (squared)
              print("    locPar ", i + 1, matLocLoc[i, i], locCov[i, i], gcor2)
            continue  
          # reject (at 3 sigma)
          if prob < 0.0027:
            self.__numRejects += 1
            continue
          # accepted
          self.__sumNdf += locNdf
          self.__sumChi2 += locChi2
          # update global matrix, vector (parameter part)
          self.__matrix[:self.__numPar,:self.__numPar] += np.dot(gloDerW.T, gloDer)
          self.__vector[:self.__numPar] += np.dot(gloDerW.T, vecMeas)
          # account for correlations (via local parameters)
          matGloLoc = np.dot(gloDerW.T, locDer)
          self.__matrix[:self.__numPar,:self.__numPar] -= np.dot(matGloLoc, np.dot(locCov, matGloLoc.T))
          self.__vector[:self.__numPar] -= np.dot(matGloLoc, locSol)
  
      except EOFError:
        pass
      
      binaryFile.close()
      self.__numRec += nRec 
    
    # add Lagrange multipliers
    indexCons = self.__numPar
    for pairs in self.__listCons:
      for idx, val in pairs:
        self.__matrix[indexCons, idx] = val
        self.__matrix[idx, indexCons] = val
      indexCons += 1  
    
    # record statistics
    print("   total number of records    ", self.__numRec)
    print("   number of rejected records ", self.__numRejects)
    print("   number of bad records      ", self.__numBad)
    # matrix information
    matSize = self.__numPar + self.__numCons
    nonZero = np.count_nonzero(self.__matrix)
    print("   size of global matrix      ", matSize)
    print("   fraction of elements <> 0. ", int(100.*nonZero / matSize ** 2), '%')
    
  
  def dump(self):
    print(" Global matrix, non-zero elements")
    for i in range(self.__numPar + self.__numCons):
      for j in range(i, self.__numPar + self.__numCons):
        if self.__matrix[i, j] != 0.:
          print("  ", i, j, self.__matrix[i, j]) 
        
  
  def solve(self):
    print()
    print(" Solving linear equations system")
    # solve equation system
    try: 
      aCov = np.linalg.inv(self.__matrix)      
    except np.linalg.LinAlgError as e:
      print("   matrix inversion failed:", e)
      print("   >>> improve input data <<<")
      return      
    aSol = np.dot(aCov, self.__vector)
    print("   sum(ndf)                 ", self.__sumNdf - self.__numPar)
    print("   initial sum(chi2)        ", self.__sumChi2)
    deltaChi2 = np.dot(self.__vector, aSol)
    print("   chi2 reduction by fit    ", deltaChi2)
    print("   final sum(chi2)          ", self.__sumChi2 - deltaChi2)
    print("   final sum(chi2)/sum(ndf) ", (self.__sumChi2 - deltaChi2) / (self.__sumNdf - self.__numPar))
    print()
    print(" Solution")
    print("   parameter label, #entries, correction, error")
    for l in sorted(self.__parCounters):
      if l in self.__parIndices:
        i = self.__parIndices[l]
        print("  ", l, self.__parCounters[l], aSol[i], math.sqrt(aCov[i, i]) if aCov[i, i] > 0. else -math.sqrt(-aCov[i, i]))
      else:
        print("  ", l, self.__parCounters[l], " fixed")
 
          
if __name__ == '__main__':
  print(time.asctime())
  #
  # Some steering parameters
  # 
  # maximum number of records to read (for each file)
  maxRec = 10000
  # minimum number of entries (appearances in binary files) for variable global parameters
  entries = 25
 
  #
  # get pede instance
  #
  p = Pede()
  
  #
  # MP2 test case (from pede -t)
  #
  # add (Fortran) binary file(s)
  p.addBinaryFile("mp2tst.bin", maxRec, 'F')
  # define parameters
  p.defineParameters(entries, fixedPar=[])
  # add constraints
  p.addConstraints("mp2con.txt")
  # construct linear equation system
  p.construct()
  #p.dump()
  # solve linear equation system
  p.solve()
    
  print()
  print(time.asctime())