ornlenv.py

import sys
import numpy as np
import scipy.optimize as optimize
from bumps.names import *
import bumps.fitters as fitters
import gym.spaces as spaces
import gym
import random 

from stable_baselines3 import PPO, A2C
from stable_baselines3.common.env_util import make_vec_env
print("orderenv")
class OrderMethods():

  def bsol(temp,p):
    Tn,Jt,Nf,Bk=p
    t=4.0*(Jt/(Jt+1.0))*Tn/temp
    if (Tn<=0) or (Jt<=0) or temp>=Tn:
      xout=0.0
    else:
      xout=optimize.brentq(OrderMethods.bfun,0.0,t,args=(temp,p),xtol=1e-6)

    return xout


  def bfun(x,T,p):
    Tn,Jt,Nf,Bk=p
    if x==0.0:
      B=-1.0 # so that it wont find solution at zero
    else:
      B=(x-3*OrderMethods.brill(Jt,x)*(Jt/(Jt+1))*(Tn/T))
    return B

  def brill(j,x):
    temp=(2*j+1.0)/2/j
    if x==0:
      Br=0.0
    else:
      Br=temp/np.tanh(temp*x)-1.0/np.tanh(x/2/j)/2/j
    return Br

  def Intensity(T,p):
    Tn,Jt,Nf,Bk=p
    br=OrderMethods.brill(Jt,OrderMethods.bsol(T,p))
    bout=Bk+Nf*br**2
    return bout

  def orderparameter(T, Tn, Jt, Nf, Bk):
    p=[Tn, Jt, Nf, Bk]
    I=[]
    for t in T:
      I.append(OrderMethods.Intensity(t,p))
    return np.array(I)

  def fit(model):
    # nllfs = []
    # zin = []
    # for tns in np.arange(50,200, 5):
      # #print("zs", zs)
      # model.Tn.value = tns
      # model.update()
      # schi=model.nllf()
      # nllfs.append(schi)
      # zin.append(tns)
    # plt.scatter(zin, nllfs)
    # plt.xlabel("tn in")
    # plt.ylabel("chi") 
    # plt.show()
    # plt.close()

    # model.Tn.value = 200
    model.update()

    problem = FitProblem(model)
    result = fitters.fit(problem, method='dream', name='order', store='/work/kmm11/orderout/dreamOut')
    for p, v in zip(problem._parameters, result.dx):
      p.dx = v
    return result.x, result.dx, problem.chisq(), problem._parameters

class OrnlEnv(gym.Env):

  def __init__(self):
    self.reward_scale = 200
    self.episodeNum = 0
    self.steps = 0
    self.T = np.array([])
    #self.Bk = .3
    self.error = []

    self.fixedTn = 128.817
    self.fixedNf = 148.449
    self.fixedJt = .403
    self.fixedBk = .1935
    self.print = False

    # self.observation_space = spaces.Box(low = np.array([3]), high = np.array([340]))
    # self.action_space = spaces.Box(low = np.array([3]), high = np.array([340]))
    self.action_space = spaces.Discrete(10)
    self.action_options = [1, 5, 10, 15, 20, 25, 30, 35, 40, 50]
    self.curTemp = 3
    self.x = 150 #startTn value

    #logging arrays and vars
    self.rewards = [] #interepisodic
    self.chisqds = []
    self.convergsTn = []
    self.convergsNf = []
    self.convergsJt = []
    self.convergsBk = []
    self.transTemps = []
    self.acts = []
    self.Jts = []
    self.Nfs = []
    self.Bks = []
    self.totReward = 0 
    self.info = {}   
  

  def step(self, action): 
    if self.print : print("stepping: ", self.episodeNum)

    self.curTemp += self.action_options[action]
    #self.curTemp = action
    self.curTemp = self.round_to(self.curTemp, 0.5)

    self.steps += 1
    reward = -self.reward_scale
    self.T = np.append(self.T, self.curTemp)
    if self.print : print("Temperature: ", self.T)
    
    if self.steps > 4:
      self.I = self.getData(self.T)
      #self.I = OrderMethods.getData(self.T, self.Bk)
      #print("Intensity: ", self.I)
      self.error= np.sqrt(self.I)
      M = Curve(OrderMethods.orderparameter, self.T, self.I, self.error, Tn = self.startTn, Jt = self.startJt, Nf = self.startNf, Bk = self.startBk)
      M.Tn.range(40, 300)
      M.Jt.range(0.4, 2)
      M.Nf.range(100,500)
      M.Bk.range(.1, 8)
      # self.Jt = self.fixedJt #DELETE ME LATERRRRRRRRRR
      # self.Nf = self.fixedNf

      self.x, dx, chisq, params = OrderMethods.fit(M)
      self.Nf = self.x[2]
      self.Jt = self.x[1]
      self.Bk = self.x[0] 
      self.x = self.x[3]
      # self.x = self.x[0]
      if self.print : print("ORDER PARAM RESULT (X2):", self.x)
      if self.print : print("THE JT RESULT {X0}:", self.Jt)
      if self.print : print("THE NF RESULT {X1}:", self.Nf)
      if self.print : print("THE BK RESULT {X3}:", self.Bk)

      dx = params[0].dx

      if self.print : print("chisqds: ", chisq)
      if self.print : print("BUT ARE WE SURE?:", dx)

      # plt.plot(self.T, self.I, 'ro')
      # plt.show()
      # plt.close()

      # if(action > 0.5):
      #   reward += 1000

      if chisq < 300 and chisq >= 1:
        reward += 100*(1/chisq)
      elif chisq < 1:
        reward += 100
        if dx < 1:
          reward += 300

      self.chisqds.append(chisq)
      self.transTemps.append(self.x)
      self.Jts.append(self.Jt)
      self.Nfs.append(self.Nf)
      self.Bks.append(self.Bk)
      self.acts.append(self.action_options[action])

      if not self.goodTn :
        if abs(self.x - self.fixedTn) < 0.1 :
          #print("ADDINGtn")
          self.convergsTn.append(self.steps)
          self.goodTn = True
      if not self.goodNf : 
        if abs (self.Nf - self.fixedNf) < 0.1:
          #print("ADDINGnf")
          self.convergsNf.append(self.steps)
          self.goodNf = True 
      if not self.goodJt : 
        if abs(self.Jt - self.fixedJt) < 0.1: 
          #print("ADDINGjt")
          self.convergsJt.append(self.steps)
          self.goodJt = True 
      if not self.goodBk : 
        if abs(self.Bk - self.fixedBk) < 0.1: 
          #print("ADDINGbk")
          self.convergsBk.append(self.steps)
          self.goodBk = True 

    if self.print : print("rewaRD:: ", reward)

    self.totReward += reward
    self.state = np.array([self.curTemp])

    if (self.steps > 4 and chisq < 0.05 and dx < 1): #less than or equal to?
      if self.print : print("terminated: excellent conditions")
      terminal = True
      self.log()
    elif (self.curTemp >= 340):
      if self.print : print("terminated: over max temp")
      terminal = True
      self.log()
    elif (self.steps > 100):
      if self.print : print("terminated: too long")
      terminal = True
      self.log()
    else:
      terminal = False
    # print(self.goodTn)
    # print(self.goodJt)
    # print(self.goodNf)
    # print(self.goodBk)
    if terminal:
      if not self.goodTn :
        #print("ADDINGtn")
        self.convergsTn.append(self.steps)
        self.goodTn = True
      if not self.goodNf : 
        #print("ADDINGnf")
        self.convergsNf.append(self.steps)
        self.goodNf = True
      if not self.goodJt : 
        #print("ADDINGjt")
        self.convergsJt.append(self.steps)
        self.goodJt = True
      if not self.goodBk : 
        #print("ADDINGbk")
        self.convergsBk.append(self.steps)
        self.goodBk = True

    return self.state, reward, terminal, self.info


  def reset(self):
    if self.print : print("reset")
    self.steps = 0
    self.T = np.array([])
    self.I = []
    self.chisqds = []
    self.transTemps = []
    self.acts = []
    self.Jts = []
    self.Nfs = []
    self.totReward = 0
    self.curTemp = 3
    self.state = np.array([self.curTemp])
    self.goodTn = False 
    self.goodNf = False 
    self.goodJt = False 
    self.goodBk = False 

    #basic multiple param set training
    # if (self.steps % 3 == 0):
    #   self.setVars(130, 1.2, 150)
    # elif (self.steps % 3 == 1):
    #   self.setVars(180, .9, 270)
    # else:
    #   self.setVars(160, .885, 330)
    newTn = self.fixedTn
    #newTn = random.randrange(60, 280, 10)
    newJt = self.fixedJt
    #newJt = random.randrange(7, 15, 1)/10.0
    newNf = self.fixedNf
    #newNf = random.randrange(120, 480, 10)
    newBk = self.fixedBk
    #newBk = random.randrange(30, 78, 2)/10.0

    self.setVars(newTn, newJt, newNf, newBk)

    self.startTn = random.randrange(int(newTn) - 20, int(newTn) + 20, 2)
    if self.print : print("start Tn: ", self.startTn)
    rangea = int(newJt*100 - 30)
    rangeb  = int(newJt*100 + 30)
    self.startJt = random.randrange(rangea, rangeb, 2)/100.0
    if self.print : print("start Jt: ", self.startJt)
    self.startNf = random.randrange(int(newNf) - 20, int(newNf) + 20, 2)
    if self.print : print("start Nf: ", self.startNf)
    rangea = int(newBk*100 - 20)
    rangeb  = int(newBk*100 + 20)
    self.startBk = random.randrange(rangea, rangeb, 2)/100.0
    if self.print : print("start Bk: ", self.startBk)

    return self.state #starting state

  def log(self):
    self.episodeNum += 1

    logdir = "/wrk/kmm11/orderout/thirdpaperrun/"

    filename = logdir + "chis/chiLog-" + str(self.episodeNum) + ".npy"
    np.savetxt(filename, self.chisqds)

    filename = logdir + "temps/tnLog-" + str(self.episodeNum) + ".npy"
    np.savetxt(filename, self.transTemps)

    filename = logdir + "jt/jtLog-" + str(self.episodeNum) + ".npy"
    np.savetxt(filename, self.Jts)

    filename = logdir + "nf/nfLog-" + str(self.episodeNum) + ".npy"
    np.savetxt(filename, self.Nfs)

    filename = logdir + "bk/bkLog-" + str(self.episodeNum) + ".npy"
    np.savetxt(filename, self.Bks)

    filename = logdir + "acts/actLog-" + str(self.episodeNum) + ".npy"
    np.savetxt(filename, self.acts)

    filename = logdir + "convergsTn.npy"
    np.savetxt(filename, self.convergsTn)

    filename = logdir + "convergsNf.npy"
    np.savetxt(filename, self.convergsNf)
    
    filename = logdir + "convergsJt.npy"
    np.savetxt(filename, self.convergsJt)

    filename = logdir + "convergsBk.npy"
    np.savetxt(filename, self.convergsBk)

    self.rewards.append(self.totReward)
    filename = logdir + "runrewards.npy"
    np.savetxt(filename, self.rewards)

  #
  # def action_space(self):
  #     return spaces.Box(low = np.array([0.5]), high = np.array([340 - self.curTemp]))
  #     # return spaces.Box(low = np.array([self.curTemp + 0.5]), high = np.array([340]))

  @property 
  def observation_space(self):
    return spaces.Box(low=3.0, high=340.0, shape=(self.steps + 1,), dtype=np.float32)

  def round_to(self, n, precision):
    correction = 0.5 if n >= 0 else -0.5
    return int( n/precision+correction ) * precision

  def getData(self, T):
    return OrderMethods.orderparameter(T, self.fixedTn, self.fixedJt, self.fixedNf, self.fixedBk)

  def setVars(self, Tn, Jt, Nf, Bk):
    self.fixedTn = Tn
    if self.print : print("fixed Tn: ", self.fixedTn)
    self.fixedJt = Jt
    if self.print : print("fixed Jt: ", self.fixedJt)
    self.fixedNf = Nf
    if self.print : print("fixed Nf: ", self.fixedNf)
    self.fixedBk = Bk
    if self.print : print("fixed Bk: ", self.fixedBk)

  def getVars(self):
    return self.x, self.Jt, self.Nf, self.Bk

  def getFixedVars(self):
    return self.fixedTn, self.fixedJt, self.fixedNf, self.fixedBk

if __name__ == "__main__":
  # Instantiate the env
  env = OrnlEnv()
  # wrap it
  env = make_vec_env(lambda: env, n_envs=1) #retrieve?

  model = PPO('MlpPolicy', env, verbose=1).learn(10000)

  model.save("/wrk/kmm11/orderout/models/thirdpaperrun")