enter image heat transferthere is my code but I would like to change just the blue color in the background, put one or two colors in the background

it's the heat transfert between beam, earth in the air, in 2 dimension write in python 2D animation model simulates the process of heat transfert during the contact interaction of a steel beam heated to a temperature of 1500 °C with the earth's surface at a temperature of 20 ° C.

[import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from matplotlib.colors import Normalize

# --------------------------------------------------------
# 1) DOMAIN DEFINITION AND GRID
# --------------------------------------------------------
x_min, x_max = -0.5, 0.5   # meters
y_min, y_max = -0.5, 0.2   # meters

Nx = 201  # number of grid points in x
Ny = 151  # number of grid points in y

x = np.linspace(x_min, x_max, Nx)
y = np.linspace(y_min, y_max, Ny)

dx = x\[1\] - x\[0\]
dy = y\[1\] - y\[0\]

X, Y = np.meshgrid(x, y, indexing='xy') 
# X.shape, Y.shape -> (Ny, Nx)

# --------------------------------------------------------
# 2) MATERIAL PROPERTIES
# --------------------------------------------------------
# Steel:
rho_steel = 7800.0       # kg/m3
cp_steel  = 600.0        # J/(kg·K)
k_steel   = 50.0         # W/(m·K)
alpha_steel = k_steel / (rho_steel * cp_steel)

# Earth:
rho_soil = 1800.0
cp_soil  = 900.0
k_soil   = 2.0
alpha_soil = k_soil / (rho_soil * cp_soil)

# Air:
rho_air = 1.225
cp_air  = 1005.0
k_air   = 0.026
alpha_air = k_air / (rho_air * cp_air)

# Convective heat transfer coefficients (example values):
alpha_air1 = 10.0  # W/(m2·K) - between earth and air
alpha_air2 = 20.0  # W/(m2·K) - between steel and air

T_air = 20.0       # Ambient air temperature (°C)

# --------------------------------------------------------
# 3) DEFINE REGION MASKS
# --------------------------------------------------------
# Steel beam: 0 <= y <= 0.05, |x| <= 0.2
steel_mask = (np.abs(X) <= 0.2) & (Y >= 0.0) & (Y <= 0.05)

# Earth: y < 0, or the ground extension for 0 <= y <= 0.05 but |x|>0.2
earth_mask = ((Y < 0.0) 
              | ((Y >= 0.0) & (Y <= 0.05) & (np.abs(X) > 0.2)))

# Air: y > 0.05
air_mask = (Y > 0.05)

# --------------------------------------------------------
# 4) INITIAL TEMPERATURE FIELD
# --------------------------------------------------------
T = np.zeros((Ny, Nx)) + 20.0  # everything starts at 20 °C
T\[steel_mask\] = 1500.0         # steel at 1500 °C initially

# --------------------------------------------------------
# 5) THERMAL DIFFUSIVITY FIELD
# --------------------------------------------------------
alpha_field = np.zeros_like(T)
alpha_field\[steel_mask\] = alpha_steel
alpha_field\[earth_mask\] = alpha_soil
alpha_field\[air_mask\]   = alpha_air

# --------------------------------------------------------
# 6) TIME-STEPPING PARAMETERS
# --------------------------------------------------------
dt = 0.05     # time step (s)
t_end = 30.0    # total simulation time (s)
nsteps = int(t_end / dt)

print(f"dx={dx:.4f} m, dy={dy:.4f} m, dt={dt:.6f} s, steps={nsteps}")

# --------------------------------------------------------
# 7) HELPER FOR CONVECTIVE BOUNDARY (ROBIN) AT Y = y_cst
# --------------------------------------------------------
def convective_bc(T_upper, T_lower, k_mat, h_conv, T_env):
    """
    1D formula for a convective boundary condition:
      -k * (T\[N\] - T\[N-1\]) / dy = h_conv * (T\[N\] - T_env)
    => T\[N\] = \[ T\[N-1\] + (h_conv*dy/k)*T_env \] / \[1 + (h_conv*dy/k)\]

    T_upper: The temperature at boundary node (y = boundary)
    T_lower: The temperature just inside the domain (one step away)
    k_mat:   Thermal conductivity of the material at boundary
    h_conv:  Convective heat transfer coefficient
    T_env:   Ambient temperature outside
    """
    # factor:
    c = (h_conv * dy) / k_mat
    return ( T_lower + c * T_env ) / (1 + c)

# --------------------------------------------------------
# 8) PREPARE PLOTTING / ANIMATION
# --------------------------------------------------------
fig, ax = plt.subplots(figsize=(7,5))
ax.set_xlim(x_min, x_max),
ax.set_ylim(y_min, y_max)
ax.set_xlabel("x (m)")
ax.set_ylabel("y (m)")
plt.plot(0.1, -1.0, color = 'r')
earth_line = ax.axhline(y=0.0, color='black', linestyle='-', linewidth=2)
ax.set_title("2D Heat Conduction with Convective BC")

norm = Normalize(vmin=20, vmax=1500)
cmap = plt.cm.jet

im = ax.imshow(T, origin='lower', extent=\[x_min, x_max, y_min, y_max\],
               cmap=cmap, norm=norm, aspect='auto')

cbar = plt.colorbar(im, ax=ax)
cbar.set_label("Temperature (°C)")

# --------------------------------------------------------
# 9) EXPLICIT FINITE-DIFFERENCE UPDATE FUNCTION
# --------------------------------------------------------
def step_temperature(T):
    # Make a copy for T_{n+1}
    T_new = T.copy()

    # Compute second derivatives in x and y (central difference)
    # We'll slice for the interior \[1:-1, 1:-1\]
    d2Tdx2 = ( T\[:,2:\] - 2*T\[:,1:-1\] + T\[:,:-2\] ) / dx**2
    d2Tdy2 = ( T\[2:,:\] - 2*T\[1:-1,:\] + T\[:-2,:\] ) / dy**2

    # Update interior
    T_new\[1:-1,1:-1\] = (
        T\[1:-1,1:-1\] +
        alpha_field\[1:-1,1:-1\]*dt*( d2Tdx2\[1:-1,:\] + d2Tdy2\[:,1:-1\] )
    )

    # -------------------------------------------
    # Dirichlet at outer edges: T=20°C
    # -------------------------------------------
    # Bottom (y= y_min)
    T_new\[0,:\] = 20.0
    # Top (y= y_max)
    T_new\[-1,:\] = 20.0
    # Left (x= x_min)
    T_new\[:,0\] = 20.0
    # Right (x= x_max)
    T_new\[:,-1\] = 20.0

    # -------------------------------------------
    # CONVECTIVE BC ON EARTH TOP (y=0) OUTSIDE STEEL
    #    +k_soil * dT/dy = alpha_air1 ( T - T_air )
    # => -k_soil * (T\[1, j\] - T\[0, j\])/dy = alpha_air1 (T\[0,j\]-T_air)
    # but we'll solve it using convective_bc() function
    #
    # The top node for earth is where earth_mask is true at y=0,
    # and the next node above is T\[1,j\]. We'll do it only for x > 0.2 or x < -0.2
    # -------------------------------------------
    # We find those j indices for which y=0 is in the domain:
    # Actually, let's do a simpler approach: loop over j from 0..Nx-1, check if
    # (y=0) belongs to earth_mask. Then apply convective formula.
    j_indices = range(Nx)
    # find the i index corresponding to y=0 (within rounding):
    # since y_min<0, we find integer i0 s.t. y\[i0\] is 0 or the closest to 0
    i0 = np.argmin( np.abs(y) )  # index where y is closest to 0
    for j in j_indices:
        if earth_mask\[i0,j\]:
            # the material here is soil => k_soil
            # T_upper = T_new\[i0, j\], T_lower = T_new\[i0+1, j\]
            T_upper = T_new\[i0, j\]
            T_lower = T_new\[i0+1, j\]
            # we use k_soil for conduction
            bc_val = convective_bc(T_upper, T_lower, k_soil, alpha_air1, T_air)
            T_new\[i0, j\] = bc_val

    # -------------------------------------------
    # CONVECTIVE BC ON STEEL TOP (y=0.05)
    #   -k_steel * dT/dy = alpha_air2 (T - T_air)
    # We'll assume that row i_top satisfies y\[i_top\]=0.05
    # Then T\[i_top, j\] is the top boundary cell for steel
    # We'll approximate T\[i_top+1, j\] as the next cell in air (if any).
    # -------------------------------------------
    # find i_top where y\[i_top\] ~ 0.05
    i_top = np.argmin( np.abs(y - 0.05) )
    for j in range(Nx):
        if steel_mask\[i_top, j\]:
            # T_upper is T_new\[i_top, j\]
            # T_lower is T_new\[i_top-1, j\] in standard conduction sign,
            # but for convenience let's define T_upper=T(i_top) and T_lower=T(i_top-1).
            T_upper = T_new\[i_top, j\]
            T_lower = T_new\[i_top-1, j\]
            # BC: -k_steel * (T_upper - T_lower)/dy = alpha_air2 (T_upper - T_air)
            # We'll adapt the convective_bc helper with a negative sign:
            # if we think of "convective_bc" as +k*(T_upper - T_lower)/dy = h(T - T_env),
            # we can just flip the sign for steel:
            # ( T_lower - T_upper ) / dy = ...
            # For simplicity, define a small helper if steel_mask:
            k_mat = k_steel
            h_conv = alpha_air2
            # The rearranged formula for T_upper is:
            #
            #   -k_steel*( T_upper - T_lower )/dy = h_conv*( T_upper - T_air )
            #   => (T_upper - T_lower ) = -( h_conv*dy / k_steel ) ( T_upper - T_air )
            #   => T_upper\[ 1 + (h_conv*dy / k_steel) \] = T_lower + (h_conv*dy / k_steel)*T_air
            #   => T_upper = \[ T_lower + (h_conv*dy/k_steel)*T_air \] / \[1 + (h_conv*dy/k_steel)\]
            c = (h_conv*dy)/k_mat
            T_new\[i_top, j\] = ( T_lower + c*T_air ) / (1 + c)

    return T_new

# --------------------------------------------------------
# 10) ANIMATION UPDATE
# --------------------------------------------------------
def update(frame):
    global T
    T = step_temperature(T)
    im.set_data(T)
    ax.set_title(f"2D Heat Transfer, t = {frame*dt:.3f} s")
    return \[im\]

# --------------------------------------------------------
# 11) RUN ANIMATION
# --------------------------------------------------------
ani = FuncAnimation(fig, update, frames=nsteps, interval=30, blit=False)
plt.tight_layout()
plt.show()][1]