Optimization visualization with pyvista

plot_optim.py

import pyvista as pv
import numpy as np


make_gif = True

# increase n_points for a higher resolution
n_points = 100
xmin, xmax = -1.2, 1.2
bounds = 1.25 * np.array([xmin, xmax, xmin, xmax, 0., 0.])
x = np.linspace(xmin, xmax, n_points)
y = np.linspace(xmin, xmax, n_points)
x, y = np.meshgrid(x, y)
coords = np.array(list(zip(x.flatten(), y.flatten())))

g = x ** 4 + y ** 4
g = g.flatten()
constraint_mask = g <= 1


def func(x, y):
    return x ** 3 - y ** 3 + 2


f = func(x, y)
f = f.flatten()

domain_coords = np.zeros((n_points ** 2, 3))
domain_coords[:, :2] = coords
domain = pv.PolyData(domain_coords)
domain = domain.delaunay_2d()
domain_in, _ = domain.remove_points(~constraint_mask)
domain_out, _ = domain.remove_points(constraint_mask)

minimizer_array = np.array([- 0.5 ** 0.25, 0.5 ** 0.25, 0.])
value_array = minimizer_array + np.array([0., 0., func(*minimizer_array[:2])])
minimizer = pv.PolyData(minimizer_array)
dashed = pv.Spline(np.vstack((minimizer_array, value_array)), n_points=20)

cone_direction = np.array([-1., 1., 0.])
cone = pv.Cone(minimizer_array, cone_direction, angle=60, height=20)
cone.points[:, 2] = 0
cone.points *= 1 / (4 * 20) ** 0.5

surface_data = np.zeros((n_points ** 2, 3))
surface_data[:, :2] = coords
surface_data[:, 2] = f
surface = pv.PolyData(surface_data)
surface = pv.PolyData(surface)
surface = surface.delaunay_2d()
surface_in, _ = surface.remove_points(~constraint_mask)
surface_out, _ = surface.remove_points(constraint_mask)

constraint_title = pv.Text3D("Constraint set K", depth=0.2)
constraint_title.points -= constraint_title.points.mean(0)[None, :]
constraint_title.points *= \
    1.75 / (constraint_title.points.max() - constraint_title.points.min())
constraint_title.rotate_z(90)

text3d = pv.Text3D("Minimizer x", depth=0.2)
text3d.points -= text3d.points.mean(0)[None, :]
text3d.points /= text3d.points.max() - text3d.points.min()
text3d.rotate_z(90)
text3d.rotate_y(90)
text3d.points += 1.75 * minimizer_array

tangeant = pv.Text3D("Tangeant cone at x", depth=0.2)
tangeant.points -= tangeant.points.mean(0)[None, :]
tangeant.points *= 1.75 / (tangeant.points.max() - tangeant.points.min())
tangeant.rotate_z(90)
tangeant.rotate_x(180)
tangeant.rotate_y(90)
tangeant.points += np.array([-1.5, -1.5, 0])

plotter = pv.Plotter()
plotter.add_mesh(domain_out, color="gray", opacity=0.2)
plotter.add_mesh(domain_in, color="black")
plotter.add_mesh(surface_in, scalars=f[constraint_mask], cmap="hot")
plotter.add_mesh(surface_out, cmap="Greys", opacity=0.2)
plotter.add_mesh(cone, color="blue", opacity=0.3)
plotter.add_mesh(minimizer, color="red", render_points_as_spheres=True,
                 point_size=15)
plotter.add_mesh(dashed, color="red")
plotter.add_mesh(text3d, color="red")
plotter.add_mesh(tangeant, color="black")
plotter.add_mesh(constraint_title, color="white")
plotter.background_color = "white"
plotter.show_bounds(grid='front', location='outer',
                    show_zaxis=False, color="black",
                    bounds=bounds)

if make_gif:
    # when the window shows up, close it by pressing the q-Key NOT the quit
    # button
    plotter.show(auto_close=False)

    path = plotter.generate_orbital_path(3., n_points=200,
                                         shift=1.75 * domain_out.length)
    plotter.open_movie('orbit.mp4')
    plotter.orbit_on_path(path, write_frames=True)
    plotter.close()
else:
    plotter.show()