BWR_optimization.py

import numpy as np
import matplotlib.pyplot as plt
import os
import subprocess
import random
import shutil
import time

import tensorflow as tf


# INPUTS AND PARAMETERS

### General inputs
base_input_file = "./base_case/ge14_nominal.inp"
base_batch_file = "./base_case/run_scale.sh"
fixed_pos_pins = ["W1"]
num_epochs = 5
num_episodes = 40
num_steps = 40
result_file = "results.txt"

### Hyperparameters


# FUNCTIONS

### Function to open an input file and read pinmap information
def pinmap_read(file):
    f = open(file,'r')
    lines = f.readlines()
    f.close()
    for i in range(len(lines)):
        line = lines[i]
        line_split = line.split()
        if len(line_split) == 0:
            continue
        if len(line_split) > 3:
            if line_split[3] == 'ASSM': 
                pin_dim = int(line_split[4])
        if line_split[0] == 'pinmap':
            index_map  = i
            exit
    pinmap = []
    for i in range(pin_dim):
        line = lines[i+index_map+1]
        line_split = line.split()
        pinmap.append(line_split)
    pinmap = np.array(pinmap)
    return(pinmap)

### Function to shuffle the pinmap
def shuffle(pinmap,fixed_pos_pins,direction):
    tmp_pinmap = np.copy(pinmap)
    upd_pinmap = np.copy(pinmap)
    nrows,ncols = np.shape(pinmap)
    
    if direction == 'left':
        for i in range(nrows):
            for j in range(ncols):
                if j==ncols-1:
                    upd_pinmap[i][j] = tmp_pinmap[i][0]
                else:
                    if tmp_pinmap[i][j] in fixed_pos_pins:
                        continue
                    elif tmp_pinmap[i][j+1] in fixed_pos_pins:
                        if tmp_pinmap[i][j+2] in fixed_pos_pins:
                            upd_pinmap[i][j] = tmp_pinmap[i][j+3]
                        else:
                            upd_pinmap[i][j] = tmp_pinmap[i][j+2]
                    else:
                        upd_pinmap[i][j] = tmp_pinmap[i][j+1]
    elif direction == 'right':
        for i in range(nrows):
            for j in range(ncols):
                if j==0:
                    upd_pinmap[i][j] = tmp_pinmap[i][ncols-1]
                else:
                    if tmp_pinmap[i][j] in fixed_pos_pins:
                        continue
                    elif tmp_pinmap[i][j-1] in fixed_pos_pins:
                        if tmp_pinmap[i][j-2] in fixed_pos_pins:
                            upd_pinmap[i][j] = tmp_pinmap[i][j-3]
                        else:
                            upd_pinmap[i][j] = tmp_pinmap[i][j-2]
                    else:
                        upd_pinmap[i][j] = tmp_pinmap[i][j-1]
    elif direction == 'up':
        for i in range(ncols):
            for j in range(nrows):
                if j==nrows-1:
                    upd_pinmap[i][j] = tmp_pinmap[i][0]
                else:
                    if tmp_pinmap[i][j] in fixed_pos_pins:
                        continue
                    elif tmp_pinmap[i][j+1] in fixed_pos_pins:
                        if tmp_pinmap[i][j+2] in fixed_pos_pins:
                            upd_pinmap[i][j] = tmp_pinmap[i][j+3]
                        else:
                            upd_pinmap[i][j] = tmp_pinmap[i][j+2]
                    else:
                        upd_pinmap[i][j] = tmp_pinmap[i][j+1]
    elif direction == 'down':
        for i in range(ncols):
            for j in range(nrows):
                if j==0:
                    upd_pinmap[i][j] = tmp_pinmap[i][nrows-1]
                else:
                    if tmp_pinmap[i][j] in fixed_pos_pins:
                        continue
                    elif tmp_pinmap[i][j-1] in fixed_pos_pins:
                        if tmp_pinmap[i][j-2] in fixed_pos_pins:
                            upd_pinmap[i][j] = tmp_pinmap[i][j-3]
                        else:
                            upd_pinmap[i][j] = tmp_pinmap[i][j-2]
                    else:
                        upd_pinmap[i][j] = tmp_pinmap[i][j-1]
    return(upd_pinmap)

### Function to generate input file
def gen_inp(file,inp_file,pinmap):
    f1 = open(file,'r')
    f2 = open(inp_file,'w')
    lines = f1.readlines()
    f1.close()
    rows,cols = np.shape(pinmap)
    for line in lines:
        line_split=line.split()
        if len(line_split) == 0: 
            continue
        if line_split[0] == "pinmap":
            break
        f2.write(line)
    f2.write("pinmap"+"\n")
    for i in range (rows):
        str_pinmapi = ""
        for j in range(cols):
            if len(pinmap[i][j]) == 1:
                space = "  "
            else:
                space = " "
            str_pinmapi += space + pinmap[i][j]
        f2.write(str_pinmapi+"\n")
    f2.write("end"+"\n")
    f2.close()
    
    return()
    

### Function to generate batch file
def gen_sh(file,sh_file,eps,abs_path):
    f1 = open(file,'r')
    f2 = open(sh_file,'w')
    lines = f1.readlines()
    f1.close()
    rows,cols = np.shape(pinmap)
    for line in lines:
        line_split=line.split()
        if len(line_split) == 0: 
            continue
        if line_split[0] == "cd":
            break
        f2.write(line)

    f2.write("cd "+abs_path+"\n")
    f2.write("/package/scale/6.2.3/bin/scalerte "+"eps"+str(eps)+".inp")
    f2.close()
    
    return()
    

def get_keff(out_file):
    f = open(out_file,'r')
    lines = f.readlines()
    f.close()
    for line in lines:
        line_split=line.split()
        if len(line_split) < 2: 
            continue
        if line_split[1] == "k-eff":
            keff  = float(line_split[3])
            break
    return(keff)
    
def get_PPF(out_file,pinmap):
    f = open(out_file,'r')
    lines = f.readlines()
    f.close()
    rows,cols = np.shape(pinmap)
    i = 0
    for line in lines:
        i = i + 1
        line_split=line.split()
        if len(line_split) < 2: 
            continue
        if (line_split[0] == "Fission" and line_split[1] == "Density"):
            break
    pin_power = np.zeros([rows,cols])
    PPF = 0.0
    for j in range(rows):
        line = lines[j+i+1]
        line_split=line.split()
        for k in range(cols):
            pin_power[j][k] = float(line_split[k])
            if pin_power[j][k]>PPF:
                PPF = pin_power[j][k]
    return(PPF)
    
def get_fobj(keff,PPF):
    w = 0.5
    penalty = 1.0
    fobj = keff - w*PPF
    if keff<1.0 or keff>1.1:
        fobj = fobj - penalty
    
    return(fobj)

# INITIALIZATION

### Read the base case pinmap
pinmap = pinmap_read(base_input_file)
rf = open(result_file,'w')
out_file = base_input_file[:-4]+".out"
keff = get_keff(out_file)
PPF = get_PPF(out_file,pinmap)
fobj = get_fobj(keff,PPF)
print("base case: "+"k-eff = "+str(keff)+", PPF = "+str(PPF)+", objective function = "+str(fobj))
rf.write("base case: "+"k-eff = "+str(keff)+", PPF = "+str(PPF)+", objective function = "+str(fobj)+"\n")

# SHUFFLING AND OPTIMIZATION PROCESS
directions = ['left','right','up','down']
for epoch in range(num_epochs):
    directory = './epoch'+str(epoch+1)
    exist = os.path.exists(directory)
    if exist:
        shutil.rmtree(directory)
    os.mkdir(directory)
    
    ### Loop for all episodes in each epoch
    for eps in range(num_episodes):
        ### randomize shuffling direction 
        ### ==> need to be replaced with Reinforcement Learning later
        direction = random.choice(directions)
        
        ### shuffle the pinmap
        pinmap = shuffle(pinmap,fixed_pos_pins,direction)
        
        ### generate input file and batch run file
        inp_file = directory+"/eps"+str(eps)+".inp"
        sh_file = directory+"/eps"+str(eps)+".sh"
        abs_path = os.path.abspath(directory)
        gen_inp(base_input_file,inp_file,pinmap)
        gen_sh(base_batch_file,sh_file,eps,abs_path)
        
        ### run the file in Polaris
        print("Running Polaris for input file"+inp_file)
        subprocess.call('cd '+abs_path,shell=True)
        subprocess.call('sbatch '+sh_file,shell=True)
    
    ### check if the running finished
    while(True):
        time.sleep(10)
        lfinish = True
        for fname in os.listdir(directory):
            if fname.endswith('.path'):
                lfinish = False
        if lfinish:
            break

    ### read the Polaris output
    for eps in range(num_episodes):
    
        ### get k-eff from Polaris output
        out_file = directory+"/eps"+str(eps)+".out"
        keff = get_keff(out_file)
        
        ### get PPF from Polaris output
        PPF = get_PPF(out_file,pinmap)
        
        ### determine objective function
        fobj = get_fobj(keff,PPF)
        print("epoch "+str(epoch)+" episode "+str(eps)+": k-eff = "+str(keff)+", PPF = "+str(PPF)+", objective function = "+str(fobj))
        rf.write("epoch "+str(epoch)+" episode "+str(eps)+":k-eff = "+str(keff)+", PPF = "+str(PPF)+", objective function = "+str(fobj)+"\n")
        
rf.close()