Using SVGD + Torch algorithms¶

In [1]:

import matplotlib.pyplot as plt
import numpy as np
import torch  # ty: ignore[unresolved-import] -- optional extra

import simplesvgd

import matplotlib.pyplot as plt import numpy as np import torch # ty: ignore[unresolved-import] -- optional extra import simplesvgd

This notebook is exactly the same as the first, only with the optimization algorithm replaced by Torch implementations.

In [2]:

smoothing = 30

def himmelblau(input_array: np.ndarray) -> np.ndarray:

    # As this is a 2-dimensional function, assert that the passed input_array
    # is correct.
    assert input_array.shape[1] == 2

    # To simplify reading this function, we do this step in between. It is not
    # the most optimal way to program this.
    x = input_array[:, 0, None]
    y = input_array[:, 1, None]

    output_array = (x ** 2 + y - 11) ** 2 + (x + y ** 2 - 7) ** 2

    # As the output should be a scalar function, assert that the
    # output is also length 1 in dim 2 of the array.
    assert output_array.shape == (input_array.shape[0], 1)

    return output_array / smoothing


def himmelblau_grad(input_array: np.ndarray) -> np.ndarray:

    # As this is a 2-dimensional function, assert that the passed input_array
    # is correct.
    assert input_array.shape[1] == 2

    # To simplify reading this function, we do this step in between. It is not
    # the most optimal way to program this.
    x = input_array[:, 0, None]
    y = input_array[:, 1, None]

    # Compute partial derivatives and combine them
    output_array_dx = 2 * (x ** 2 + y - 11) * (2 * x) + 2 * (x + y ** 2 - 7)
    output_array_dy = 2 * (x ** 2 + y - 11) + 2 * (x + y ** 2 - 7) * (2 * y)
    output_array = np.hstack((output_array_dx, output_array_dy))

    # Check if the output shape is correct
    assert output_array.shape == input_array.shape

    return output_array / smoothing

smoothing = 30 def himmelblau(input_array: np.ndarray) -> np.ndarray: # As this is a 2-dimensional function, assert that the passed input_array # is correct. assert input_array.shape[1] == 2 # To simplify reading this function, we do this step in between. It is not # the most optimal way to program this. x = input_array[:, 0, None] y = input_array[:, 1, None] output_array = (x ** 2 + y - 11) ** 2 + (x + y ** 2 - 7) ** 2 # As the output should be a scalar function, assert that the # output is also length 1 in dim 2 of the array. assert output_array.shape == (input_array.shape[0], 1) return output_array / smoothing def himmelblau_grad(input_array: np.ndarray) -> np.ndarray: # As this is a 2-dimensional function, assert that the passed input_array # is correct. assert input_array.shape[1] == 2 # To simplify reading this function, we do this step in between. It is not # the most optimal way to program this. x = input_array[:, 0, None] y = input_array[:, 1, None] # Compute partial derivatives and combine them output_array_dx = 2 * (x ** 2 + y - 11) * (2 * x) + 2 * (x + y ** 2 - 7) output_array_dy = 2 * (x ** 2 + y - 11) + 2 * (x + y ** 2 - 7) * (2 * y) output_array = np.hstack((output_array_dx, output_array_dy)) # Check if the output shape is correct assert output_array.shape == input_array.shape return output_array / smoothing

In [3]:

domain = [-6, 6, -6, 6]

x1s = np.linspace(domain[0], domain[1], 100)
x2s = np.linspace(domain[2], domain[3], 100)

grid = np.stack(np.meshgrid(x1s, x2s))

background_image = np.empty_like(grid[0, :, :])

for ix1 in range(100):
    background_image[ix1, :] = himmelblau(grid[:, ix1, :].T)[:, 0]


background = (x1s, x2s, background_image)

domain = [-6, 6, -6, 6] x1s = np.linspace(domain[0], domain[1], 100) x2s = np.linspace(domain[2], domain[3], 100) grid = np.stack(np.meshgrid(x1s, x2s)) background_image = np.empty_like(grid[0, :, :]) for ix1 in range(100): background_image[ix1, :] = himmelblau(grid[:, ix1, :].T)[:, 0] background = (x1s, x2s, background_image)

Now we optimize the starting samples using SVGD + Torch. Note the differences are minimal.

In [8]:

%matplotlib notebook

figure = plt.figure(figsize=(6, 6))
plt.xlabel("Parameter 0")
plt.ylabel("Parameter 1")
plt.title("SVGD animation on the himmelblau function")


rng = np.random.default_rng()
initial_samples = 3 * rng.standard_normal((1000, 2))
optimizer_class = torch.optim.Adam
optimizer_parameters = {"lr":1e-1}


final_samples = simplesvgd.update_torch(
    initial_samples,
    himmelblau_grad,
    n_iter=160,
    optimizer_class=optimizer_class,
    optimizer_parameters=optimizer_parameters,
    animate=True,
    background=background,
    figure=figure,
)

%matplotlib notebook figure = plt.figure(figsize=(6, 6)) plt.xlabel("Parameter 0") plt.ylabel("Parameter 1") plt.title("SVGD animation on the himmelblau function") rng = np.random.default_rng() initial_samples = 3 * rng.standard_normal((1000, 2)) optimizer_class = torch.optim.Adam optimizer_parameters = {"lr":1e-1} final_samples = simplesvgd.update_torch( initial_samples, himmelblau_grad, n_iter=160, optimizer_class=optimizer_class, optimizer_parameters=optimizer_parameters, animate=True, background=background, figure=figure, )

  0%|          | 0/160 [00:00<?, ?it/s]

Let's compare the initial and final results:

In [9]:

plt.figure(figsize=(8, 4))
plt.subplot(121)
_, bins_0, bins_1, _ = plt.hist2d(
    initial_samples[:, 0],
    initial_samples[:, 1],
    range=[[-6, 6], [-6, 6]],
    bins=50,
    cmap=plt.get_cmap("Greys"),
)
plt.gca().set_aspect(1)
plt.xlabel("Parameter 0")
plt.ylabel("Parameter 1")
plt.title("Initial Gaussian samples")
plt.xlim(-6, 6)
plt.ylim(-6, 6)

plt.subplot(122)
_ = plt.hist2d(
    final_samples[:, 0],
    final_samples[:, 1],
    bins=(bins_0, bins_1),
    range=[[-6, 6], [-6, 6]],
    cmap=plt.get_cmap("Greys"),
)
plt.gca().set_aspect(1)
plt.xlabel("Parameter 0")
plt.ylabel("Parameter 1")
plt.title("Final SVGD samples")
plt.xlim(-6, 6)
plt.ylim(-6, 6)

plt.show()

plt.figure(figsize=(8, 4)) plt.subplot(121) _, bins_0, bins_1, _ = plt.hist2d( initial_samples[:, 0], initial_samples[:, 1], range=[[-6, 6], [-6, 6]], bins=50, cmap=plt.get_cmap("Greys"), ) plt.gca().set_aspect(1) plt.xlabel("Parameter 0") plt.ylabel("Parameter 1") plt.title("Initial Gaussian samples") plt.xlim(-6, 6) plt.ylim(-6, 6) plt.subplot(122) _ = plt.hist2d( final_samples[:, 0], final_samples[:, 1], bins=(bins_0, bins_1), range=[[-6, 6], [-6, 6]], cmap=plt.get_cmap("Greys"), ) plt.gca().set_aspect(1) plt.xlabel("Parameter 0") plt.ylabel("Parameter 1") plt.title("Final SVGD samples") plt.xlim(-6, 6) plt.ylim(-6, 6) plt.show()

Looks wonderful!