siesta module

siesta

import subprocess
import os
import shutil
import glob

def list_variables():
  return [x.split() for x in  open("VARS").read().strip().split("\n")]

def new_input_file(parameters,var_names,file_name):
  if len(parameters) != len(var_names):
    print "ERROR: number of parameters not equal to number of input variables"
    quit()
  text=open("TEMPLATE").read()
  new_text=text[:]
  for i in range(len(var_names)):
    new_text=new_text.replace("$"+var_names[i],str(parameters[i]))
  fout=open(file_name,"w")
  fout.write(new_text)
  fout.close()

def run_siesta(input_file,output_file):
  fin=open(input_file,"r")
  siesta="/home/cathcart/code/trunk-367/Obj/siesta"
  run=subprocess.Popen([siesta],stdin=fin,stdout=subprocess.PIPE,stderr=subprocess.STDOUT,shell=True)
  output=run.communicate()[0]
  fin.close()
  fout=open(output_file,"w")
  fout.write(output)
  fout.close()
  return output

def parse(siesta_output):
  start=siesta_output.find("Total =")
  if start<0:
    print "ERROR: calculation failed. dumping last five lines"
    print "\n".join(siesta_output.split("\n")[-5::])
    return 9999
  end=siesta_output.find("\n",start)
  ans=siesta_output[start:end].split()[-1]
  return float(ans)

def siesta_function(name,parameters):
  var_file=list_variables()
  var_names=[x[0] for x in var_file]

  file_in=name+".fdf"
  file_out=name+".out"

  if os.path.exists(name+"/"+file_in):
    #input file already exists
    siesta_output_file=open(name+"/"+file_out,"r")
    siesta_output=siesta_output_file.read()
    siesta_output_file.close() 
  else:
    #set up folder
    os.mkdir(name)
    for pseudo in glob.glob('*.psf'):
      shutil.copy2(pseudo,name)#copy pseudopotentials across
    new_input_file(parameters,var_names,name+"/"+file_in)
    os.chdir(name)
    siesta_output=run_siesta(file_in,file_out)
    os.chdir("../")
  return parse(siesta_output)


if __name__=="__main__":
  var_file=list_variables()
  var_names=[x[0] for x in var_file]
  min_p_list=[float(x[1]) for x in var_file]
  max_p_list=[float(x[2]) for x in var_file]
  print siesta_function("test",min_p_list)

swarm

import math
import random

import siesta

class Vector:
  def __class__(self):
    return "Vector"
  def __init__(self, data):
    self.data = data
  def __repr__(self):
    return repr(self.data)  
  def __add__(self, other):
    data = []
    for j in range(len(self.data)):
      data.append(self.data[j] + other.data[j])
    return Vector(data)
  def __sub__(self,other):
    return self+(-1)*other
  def __len__(self):
    return len(self.data)
  def __rmul__(self,time):
    return self*time
  def __mul__(self,time):
    return Vector([time*self.data[i] for i in range(len(self.data))])
  def __getitem__(self,index):
    return self.data[index]
  def __setitem__(self,index,value):
    self.data[index]=value

class Particle:
  def __init__(self,data):
    self.position=Vector(data[0])
    self.velocity=Vector(data[1])
    self.local_min=Vector(data[0])
  def update_position(self):
    self.position=self.position+self.velocity
  def update_velocity(self,new_velocity):
    self.velocity=self.velocity+Vector(new_velocity)
  def update_local_min(self):
    if self.position < self.local_min:
      self.local_min=self.position
  def update(self,data):
    self.update_velocity(data)
    self.update_position()    
    self.update_local_min()
  def set_cost(self,value):
    self.cost=value

def distance(one,two):
  return math.sqrt(sum([(one[i]-two[i])**2 for i in range(len(one))]))

def grouping(vector_list,n):
  l=[]
  tmp_list=vector_list
  groups=[]
  t=len(vector_list)/n
  for one in vector_list[0::len(vector_list)/n]:
    sorted_list=sorted(tmp_list,key=lambda x: distance(one.position,x.position))
    groups.append(sorted_list[0:len(vector_list)/n])
    tmp_list=tmp_list[len(vector_list)/n::]
  return groups

def r():
  yield random.random()

def my_function(parameters):
  [x1,x2]=parameters
  return 100*(x2-x1**2)**2+(1-x1)**2

def run(N,iterations,grouping_N,min_p,max_p,function):
  #N=20#number of particles
  #grouping_N=6#number of groups to slip the swarm into for the local group mins
  #iterations=100#number of iterations to run
  #min_p=Vector([0,0])#min parameters
  #max_p=Vector([20,20])#max_parameters
  K_l=1.0#local vector constant divided by the particle mass
  K_g=1.0#global vector constant divided by the particle mass
  K_gv=1.0#group vector constant divided by the particle mass
  #[K_l,K_g,k_gv]=constants
  p=len(min_p)#number of parameters
  swarm=[]#list of swarm particles

  #initialise position and velocity for each particle
  for particle in range(N):
    position=[(min_p+(max_p-min_p)*next(r()))[i] for i in range(p)]
    velocity=[((max_p-min_p)*(2*next(r())-1))[i] for i in range(p)]
    swarm.append(Particle([position,velocity]))

  #distribute and calculate the value of the function at all of these points
  min_position=min_p

  for iteration in range(iterations):

#    print "%d %f" %(iteration,min_position[0])
    cost_results=map(function,[particle.position for particle in swarm])
    [swarm[i].set_cost(cost_results[i]) for i in range(len(swarm))]
#    for particle in swarm:
#      try:
#        particle.set_cost(function(particle.position))
#      except:
#        particle.set_cost(9999)

    #update all the local extrema and decide the new global extrema
    min_position=sorted(swarm,key=lambda x: x.cost)[0].position
    print "Iteration: %d, global min: %s"%(iteration,",".join([str(x) for x in min_position]))
    
    groups=grouping(swarm,grouping_N)
    for group in groups:
      group_position=sorted(group,key=lambda x: x.cost)[0].position
      for particle in group:

        global_vector=min_position-particle.position
        local_vector=particle.local_min-particle.position
        group_vector=group_position-particle.position

        [r_l,r_g,r_gv]=[next(r()) for i in range(3)]

        update_vector=(K_l*r_l*local_vector+K_g*r_g*global_vector+K_gv*r_gv*group_vector)
        particle.update(update_vector)

        #check min
        mi=particle.position-min_p
        for i in [y for y in mi if y <= 0]:
          x=mi.data.index(i)
          particle.position[x]=min_p[x]
          particle.velocity[x]=0

        #check max
        ma=particle.position-max_p
        for i in [y for y in ma if y >= 0]:
          x=ma.data.index(i)
          particle.position[x]=max_p[x]
          particle.velocity[x]=0

  p=sorted(swarm,key=lambda x: x.cost)[0].position
  return [math.sqrt(sum([i**2 for i in p])),p]

if __name__=="__main__":
  N=20
  iterations=10

  var_file=siesta.list_variables()
  min_p_list=[float(x[1]) for x in var_file]
  max_p_list=[float(x[2]) for x in var_file]
  min_p=Vector(min_p_list)
  max_p=Vector(max_p_list)

  ans=run(N,iterations,1,min_p,max_p,lambda x :siesta.siesta_function("_".join([str(y) for y in x]),x))[1]
  print ans
  out=open("final","w")
  out.write(str(ans))
  out.close()