Different colormap of predicted mean by GP after loading data from csv file

[MJeong19]

I have a problem when drawing colormap of predicted mean by GP.
The colormap of predicted mean by GP with original data is ordinarily drawn.

However, the colormap of predicted mean by GP with data saved and loaded from csv file was different with original colormap.

I would like to know how can I draw the same colormap of predicted mean by GP with data saved and loaded from csv file with that with original data.
Here is the code I used:

import torch
import matplotlib.pyplot as plt
import numpy as np
from botorch.models import SingleTaskGP
from gpytorch.mlls.exact_marginal_log_likelihood import ExactMarginalLogLikelihood
from botorch.optim.fit import fit_gpytorch_mll_torch
from matplotlib.cm import ScalarMappable
import csv

# 1. Define a synthetic 2D objective function (e.g., the Hartmann 2D function)
def objective_function(x):
    # x is expected to be a tensor of shape (*batch_shape, 2)
    # This is a simple function for demonstration; replace with your actual function
    return -torch.sum(x**2, dim=-1) + 2 * x[..., 0] * x[..., 1]

# 2. Generate initial training data
train_X = torch.rand(10, 2,dtype=torch.float64) * 4 - 2 # 10 points between -2 and 2
train_Y = objective_function(train_X).unsqueeze(-1) # Unsqueeze to add output dimension
trainx=np.array(train_X)
trainy=np.array(train_Y)
# 3. Initialize and fit the GP model
model = SingleTaskGP(train_X, train_Y)
mll = ExactMarginalLogLikelihood(model.likelihood, model)
fit_gpytorch_mll_torch(mll)

# 4. Define the 2D grid for plotting
x1_range = np.linspace(-2, 2, 51)
x2_range = np.linspace(-2, 2, 51)
X1, X2 = np.meshgrid(x1_range, x2_range)
# Convert meshgrid to a tensor of test points
test_X = torch.tensor(np.stack([X1.ravel(), X2.ravel()], axis=1), dtype=torch.float64)

# 5. Get the posterior mean from the model
model.eval() # Set model to evaluation mode
with torch.no_grad():
    posterior = model.posterior(test_X)
    mean = posterior.mean.cpu().numpy().reshape(X1.shape) # Reshape to 2D grid

# 6. Plot the 2D mean using matplotlib
fig,ax=plt.subplots()
plt.rc('font', size=20)
ax.pcolormesh(X1, X2, mean,cmap='Blues')
cm = plt.get_cmap("Blues")
norm = plt.Normalize(mean.min(), mean.max())
sm = ScalarMappable(norm=norm, cmap=cm)
plt.xticks(fontsize=20)
plt.yticks(fontsize=20)
sm.set_array([])
fig.subplots_adjust(right=0.87)
cbar_ax1 = fig.add_axes([0.9, 0.15, 0.01, 0.7])
cbar1 = fig.colorbar(sm,cax=cbar_ax1)
cbar1.set_label("Y", rotation=90, labelpad=20)
cbar1.ax.yaxis.set_label_position('left')
ax.set_xlabel('X1',fontsize=20)
ax.set_ylabel('X2',fontsize=20)
plt.show()
# Save the data

np.savetxt('testx1.csv', trainx, delimiter=",")
np.savetxt('testy1.csv', trainy, delimiter=",")
#load the data
def read_csv(file,n:int):
        f = open(file, 'r', encoding='utf-8')
        rdr = csv.reader(f)
        arri=np.zeros((1,n))
        arr=np.zeros((0,n))
        for line in rdr:
                a1=[float(x) for x in line]
                for i in range(n):
                        arri[0][i]=a1[i]
                arr=np.vstack((arr,arri))
        return arr

trainx_fin=read_csv('testx1.csv',n=2)
trainy_fin=read_csv('testx1.csv',n=1)

train_Xfin=torch.tensor(trainx_fin,dtype=torch.float64)
train_Yfin=torch.tensor(trainy_fin,dtype=torch.float64)
model2 = SingleTaskGP(train_Xfin, train_Yfin)
mll2 = ExactMarginalLogLikelihood(model2.likelihood, model2)
fit_gpytorch_mll_torch(mll2)

model2.eval() # Set model to evaluation mode
with torch.no_grad():
    posterior2 = model2.posterior(test_X)
    mean2 = posterior2.mean.cpu().numpy().reshape(X1.shape) # Reshape to 2D grid

# 6. Plot the 2D mean using matplotlib
fig4,ax4=plt.subplots()
plt.rc('font', size=20)
ax4.pcolormesh(X1, X2, mean2,cmap='Blues')
cm = plt.get_cmap("Blues")
norm = plt.Normalize(mean2.min(), mean2.max())
sm = ScalarMappable(norm=norm, cmap=cm)
plt.xticks(fontsize=20)
plt.yticks(fontsize=20)
sm.set_array([])
fig4.subplots_adjust(right=0.87)
cbar_ax1 = fig4.add_axes([0.9, 0.15, 0.01, 0.7])
cbar1 = fig4.colorbar(sm,cax=cbar_ax1)
cbar1.set_label("Y", rotation=90, labelpad=20)
cbar1.ax.yaxis.set_label_position('left')
ax4.set_xlabel('X1',fontsize=20)
ax4.set_ylabel('X2',fontsize=20)
plt.show()