Visualization for the objective function in the gadget lattice optimization problem

gadget-lattice-visualization.py

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits import mplot3d
from matplotlib import cm

# 2D
b = 4

def cost_func(z1, z2):
    d1 = 2 * (b**2 + 1) * z1 - 2 * b * z2
    d2 = 2 * b**2 * z2 - 2 * b * z1
    u1 = (d1 < 0).astype(float)
    u2 = (d2 < 0).astype(float)
    return (b**2 + 1) * (z1**2 + 2 * z1 * u1) + b**2 * (z2**2 + 2 * z2 * u2) - 2 * b * (z1 * z2 + u1 * z2 + u2 * z1)

# defining surface and axes
num_pts = 100
x = np.outer(np.linspace(0, 2, num_pts), np.ones(num_pts))
y = x.copy().T
z = cost_func(x, y)
  
fig = plt.figure()
  
# syntax for 3-D plotting
ax = plt.axes(projection ='3d')
  
# syntax for plotting
surf = ax.plot_surface(x, y, z, cmap=cm.coolwarm, linewidth=0)

# Add a color bar which maps values to colors.
fig.colorbar(surf, shrink=0.5, aspect=5)

pt_x = np.array([0, -1, -1, 1, 1])
pt_y = np.array([0, -1, 1, -1, 1])
pt_z = cost_func(pt_x, pt_y)
# ax.scatter(pt_x, pt_y, pt_z, c='black')

ax.view_init(45, 0)

plt.show()

# 3D
b = 2

def cost_func(z1, z2, z3):
    d1 = 2 * (b**2 + 1) * z1 - 2 * b * z2
    d2 = 2 * (b**2 + 1) * z2 - 2 * b * (z1 + z3)
    d3 = 2 * b**2 * z3 - 2 * b * z2
    u1 = (d1 < 0).astype(float)
    u2 = (d2 < 0).astype(float)
    u3 = (d3 < 0).astype(float)
    return (b**2 + 1) * (z1**2 + 2 * z1 * u1 + z2**2 + 2 * z2 * u2) + b**2 * (z3**2 + 2 * z3 * u3) - 2 * b * (z1 * z2 + u1 * z2 + u2 * z1 + z2 * z3 + u2 * z3 + u3 * z2)

# defining surface and axes
num_pts = 100
z1 = np.outer(np.linspace(0, 2, num_pts), np.ones(num_pts))
z2 = np.full_like(z1, 1)#z1.copy().T
z3 = z1.copy().T#np.full_like(z1, 0)
z = cost_func(z1, z2, z3)
  
fig = plt.figure()
  
# syntax for 3-D plotting
ax = plt.axes(projection ='3d')
  
# syntax for plotting
surf = ax.plot_surface(z1, z3, z, cmap=cm.coolwarm, linewidth=0)

# Add a color bar which maps values to colors.
fig.colorbar(surf, shrink=0.5, aspect=5)

#pt_x = np.array([0, -1, -1, 1, 1])
#pt_y = np.array([0, -1, 1, -1, 1])
#pt_z = cost_func(pt_x, pt_y)
# ax.scatter(pt_x, pt_y, pt_z, c='black')

ax.view_init(60, 0)

plt.show()