Rastrigin Function Landscape

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.colors import LinearSegmentedColormap

x = np.linspace(-5.12, 5.12, 200)
y = np.linspace(-5.12, 5.12, 200)
X, Y = np.meshgrid(x, y)
A = 10
Z = 2*A + (X**2 - A*np.cos(2*np.pi*X)) + (Y**2 - A*np.cos(2*np.pi*Y))

fig, ax = plt.subplots(figsize=(10, 8), facecolor='#ffffff')
ax.set_facecolor('#ffffff')

# Custom colormap: amber to orange
colors = ['#fffbeb', '#fcd34d', '#F5B027', '#F54927']
cmap = LinearSegmentedColormap.from_list('amber_orange', colors, N=256)

cs = ax.contourf(X, Y, Z, levels=35, cmap=cmap)
ax.contour(X, Y, Z, levels=20, colors='#92400e', linewidths=0.3, alpha=0.5)
cbar = plt.colorbar(cs, ax=ax, pad=0.02)
cbar.set_label('f(x,y)', color='#374151', fontsize=11)
cbar.ax.yaxis.set_tick_params(color='#374151')
plt.setp(plt.getp(cbar.ax.axes, 'yticklabels'), color='#374151')

ax.set_xlabel('X', fontsize=11, color='#374151', fontweight='500')
ax.set_ylabel('Y', fontsize=11, color='#374151', fontweight='500')
ax.set_title('Rastrigin Function Landscape', fontsize=14, color='#1f2937', fontweight='bold', pad=15)

ax.tick_params(colors='#6b7280', labelsize=9)
for spine in ax.spines.values():
    spine.set_color('#d1d5db')

plt.tight_layout()
plt.show()