Code for optimization 3

optimization3.py

def plot_space(model, WOB, SURF_RPM, constant_inputs):
  N_matrix = np.empty((len(WOB), len(SURF_RPM)))
  for i in range(len(WOB)):
    for j in range(len(SURF_RPM)):
      # Unwrap constant variables
      Depth, PHIF, VSH, SW, KLOGH = constant_inputs.values()
      N = predict(model, Depth, PHIF, VSH, SW, KLOGH, WOB=WOB[i], SURF_RPM=SURF_RPM[j])
      N_matrix[i][j] = N  
    
  plt.imshow(N_matrix, origin='lower', aspect='auto', cmap='inferno',
             extent=(min(SURF_RPM), max(SURF_RPM), min(WOB), max(WOB)))  
  plt.xlabel('SURF_RPM')
  plt.ylabel('WOB')
  plt.title('ROP_AVG', pad=10)    
  plt.colorbar()

# Range of WOB and RPM values
WOB = np.linspace(17e3, 90e3, 30)
SURF_RPM = np.linspace(1.5, 2.5, 30)

# Constant input variables
constant_inputs = {'Depth': 3480, 'PHIF': 0.09, 'VSH': 0.1, 
                   'SW': 1, 'KLOGH': 0.001}                   

# Plot prediction space
plot_space(pipe, WOB, SURF_RPM, constant_inputs)