Analyze fen list with Stockfish

analyze.py

#!/usr/bin/python
import time 
import sh
import sys
from multiprocessing import Pool

def get_fens():
  with open('pgn_fen') as f:
    next = False
    for line in f.read().split('\n'):
      if line.startswith('FEN:'):
        next = True
        cur = line.replace('FEN:', '')
      elif next:
        if line.startswith(' }'):
          next = False
          print cur
        else:
          if len(cur) > 0:
            cur += ' '
          cur += line

def run_sf(args):
  fen, depth = args

  state = {'started': False, 'scores': [], 'done': False}
  def process_output(line, stdin, process):
    if not state['started']:
      stdin.put('position fen %s\ngo depth %d\n' % (fen, depth))
      state['started'] = True
    if 'score cp' in line and 'upperbound' not in line and 'lowerbound' not in line:
      state['scores'].append(int(line.split('score cp')[1].split()[0]))
      if ' depth %d' % (depth) in line:
        state['done'] = True
        process.kill()

  sf = sh.Command('./stockfish')(_out=process_output)
  iter = 0
  while iter < 30 and not state['done']:
    time.sleep(0.1)
    iter += 1
  sf.kill()

  return (fen, state['scores'])

def chunks(l, n):
  return [l[i:i+n] for i in range(0, len(l), n)]

def multi_probcut(fens):
  searches = []
  for fen in fens[10:]:
    searches.append((fen, 15))

  p = Pool(3)
  for chunk in chunks(searches, 12):
    results = p.map(run_sf, chunk)
    for result in results:
      print result[0], result[1] 

def fit(all, low_depth, high_depth):
  # Generate the difference
  x = []
  y = []
  for sample in all:
    x.append(sample[low_depth])
    y.append(sample[high_depth])

  import numpy
  x = numpy.array(x)
  y = numpy.array(y)
  A = numpy.vstack([x, numpy.ones(len(x))]).T
  m, c = numpy.linalg.lstsq(A, y)[0]

  """
  import matplotlib.pyplot as plt
  plt.plot(x, y, 'o', markersize=1)
  plt.plot(x, m*x + c, 'r')
  plt.show()
  """

  computed = []
  for idx in range(len(x)):
    computed.append(x[idx] * m + c - y[idx])
    #print '%5d %5d %5d %5d' % (x[idx], y[idx], int(x[idx] * m + c), int(computed[idx]))
  print '%2d %2d - %5.1f %5.1f %5.1f' % (low_depth+1, high_depth+1, m, c, numpy.std(computed))

def prediction():
  with open('output') as f:
    lines = f.read().split('\n')

  all = []
  for line in lines:
    try:
      scores = line.split('[')[1][:-1].split(',')
      scores = [int(s.strip()) for s in scores]
      if len(scores) == 15:
        all.append(scores)
    except:
      pass

  """
  for low_depth in range(14):
    for high_depth in range(low_depth+1, 15):
      fit(all, low_depth, high_depth)
  """
  for high_depth in range(4, 15):
    fit(all, high_depth - 4, high_depth)

def main():
  prediction()
  return

  with open('fens') as f:
    fens = f.read().split('\n')

  #run_sf(('rnbqkbnr/pppppppp/8/8/4P3/8/PPPP1PPP/RNBQKBNR b KQkq e3 0 1', 15))
  multi_probcut(fens)
  return

  evals = []
  evals8 = []

  state = {'next': True, 'fen_idx': 0}

  def process_output(line, stdin, process):
    if line.startswith('info depth 8'):
      eval = float(line.split()[7]) / 100.0
      if fens[state['fen_idx']].split()[1] == 'b':
        eval = -eval
      evals8.append(eval)
      state['next'] = True
      state['fen_idx'] += 1

    if line.startswith('Total'):
      evals.append(float(line.split()[2]))

    if state['next']:
      stdin.put('position fen %s\n' % (fens[state['fen_idx']]))
      stdin.put('eval\n')
      stdin.put('go depth 8\n')
      state['next'] = False
      if state['fen_idx'] >= 150: #len(fens):
        process.kill()
        return 

  sf = sh.Command('./stockfish')(_out=process_output)
  sf.wait()

  for idx in range(len(evals)):
    print evals[idx], evals8[idx]

main()