Change interface for nak_torch

.gitignore

@@ -4,3 +4,4 @@ __pycache__
 .vscode
 *.xml
 *.pdf
+.env

README.md

@@ -1 +1,34 @@
 # Kernel-based quantization algorithms
+
+## Installation
+We recommend installing with `uv`. Currently, the way to install this locally would be
+```bash
+$ uv pip install -e git+https://github.com/Nodes-and-Kernels/nak_torch
+```
+
+If you plan on using the examples, make sure that `[examples]` option is installed. Also, make sure that there is no other installation of `pystan`, which is a dependency---we use a fork of the original package to reduce latency for our algorithms when using a stan posterior.
+
+## List of Algorithms
+### MSIP
+We largely focus on _mean-shift interacting particle_ (MSIP) algorithms, and we are working to implement several of these. Currently, we have:
+
+- MSIP
+- MSIPGS
+
+For these algorithms, we have multiple estimators---each of these produces a certain set of dynamics. In particular, we have:
+
+- MSIPFredholm
+- MSIPGradientFree
+- MSIPGradientInformed
+- MSIPGMMGaussianKernel
+
+### Other algorithms
+We also include several other typical interacting-particle sampling algorithms.
+
+- Consensus-based sampler (`CBS`)
+- Deep ensembles (`DeepEnsembles`)
+- Ensemble Kalman Sampler (`EKS`)
+- Gradient-informed affine-invariant Langevin dynamics (`GradALDI`)
+- Gradient-free ALDI (`GradFreeALDI`)
+- Stein variational gradient descent (`SVGD`)
+

examples/aristoff_bangerth.py

@@ -2,14 +2,17 @@
 import math
 import torch
 from functions import aristoff_bangerth as ab, build_aristoff_bangerth
-from nak_torch.algorithms import msip, svgd
+import nak_torch
+from nak_torch.algorithms import MSIP, SVGD
+from nak_torch.algorithms.msip import MSIPFredholm
 from matplotlib import ticker
 import gc
 import matplotlib.pyplot as plt
 from nak_torch.tools.kernel import sqexp_kernel_matrix
 from tqdm import tqdm
 import pandas as pd
-import pyro_tools
+from nak_torch.tools.types import BatchGradLogDensityEvaluator
+from nak_torch.tools import pyro_tools
 from pyro.infer import mcmc
 
 if torch.cuda.is_available():
@@ -18,25 +21,73 @@
     torch.set_default_device("cpu")
 torch.set_default_dtype(torch.float64)
 
+# %%
+def plot_samples(pts, max_side_len = 6):
+    n_particles = pts.shape[0]
+    side_len = min(max_side_len, int(math.floor(math.sqrt(n_particles))))
+    pts = pts[:side_len**2]
+    fig = plt.figure(figsize=(9, 6), layout='constrained')
+    gs = fig.add_gridspec(side_len, side_len + 2)
+    vabs = max(pts.min().abs(), pts.max().abs())
+    plt_kwargs = {'vmin': -vabs, 'vmax': vabs, 'extent': (0, 8, 0, 8)}
+
+    for i in range(side_len):
+        for j in range(side_len):
+            ax = fig.add_subplot(gs[i, j])
+            # ax.set_axis_off()
+            ax.set_aspect('equal')
+            t = ax.matshow(pts[i*side_len + j].reshape(8, 8), **plt_kwargs)
+            # ax.vlines(jnp.arange(1,8), -0.1, 8.1, color='w', lw=0.75)
+            # ax.hlines(jnp.arange(1,8), -0.1, 8.1, color='w', lw=0.75)
+            ax.minorticks_on()
+            ax.set_xticks([])
+            ax.set_yticks([])
+            ax.xaxis.set_minor_locator(ticker.MultipleLocator())
+            ax.yaxis.set_minor_locator(ticker.MultipleLocator())
+            ax.grid(which="both", linewidth=1.5, color="w")
+            ax.tick_params(which="minor", length=0)
+    ax_cb = fig.add_subplot(gs[:-2, -2:])
+    ax_cb.set_title(r"Scale of $\log\theta$", y=0.6)
+    cax_cb = ax_cb.inset_axes((0.1, 0.45, 0.8, 0.1))
+    ax_cb.axis('off')
+    fig.colorbar(t, cax=cax_cb, orientation='horizontal') # type: ignore
+    ax_true = fig.add_subplot(gs[-2:, -2:])
+    ax_true.set_aspect('equal')
+    ax_true.matshow(ab.theta_true.log().reshape(8, 8), **plt_kwargs)
+    ax_true.minorticks_on()
+    ax_true.set_xticks([])
+    ax_true.set_yticks([])
+    ax_true.xaxis.set_minor_locator(ticker.MultipleLocator())
+    ax_true.yaxis.set_minor_locator(ticker.MultipleLocator())
+    ax_true.set_title(r"True $\theta$")
+    ax_true.grid(which="both", linewidth=1.5, color="w")
+    ax_true.tick_params(which="minor", length=0)
+    return fig
+
+
 # %%
 use_compiled = True
 model = build_aristoff_bangerth(use_compiled=use_compiled, dtype=torch.float64)
 log_p = model.to_log_dens(use_compiled=use_compiled)
-log_th = torch.randn(500, 64, dtype=torch.float64)
-test_out = log_p(log_th)
+log_th = torch.randn(25, 64, dtype=torch.float64)
+test_out = log_p(log_th, None)
 
 # %%
-grad_log_p = torch.func.grad(lambda t: log_p(t).sum())
-test_eval = grad_log_p(log_th)
+def _tmp_log_p(log_theta, arg: None):
+    ret = log_p(log_theta, arg)
+    return ret.sum(), ret
+
+grad_log_p = torch.func.grad(lambda t,a: log_p(t, a).sum())
+grad_val_log_p = torch.func.grad(_tmp_log_p, has_aux=True)
+test_grad = grad_log_p(log_th, None)
+test_grad_2, test_out_2 = grad_val_log_p(log_th, None)
 
 # %%
-del log_th
-del test_out
-# del test_eval
+del log_th, test_out, test_grad, test_grad_2, test_out_2
 gc.collect()
 
 # %%
-n_particles, n_steps, dim = 500, 25, 64
+n_particles, n_steps, dim = 25, 25, 64
 kernel_bandwidth = 0.75
 
 torch.manual_seed(1)
@@ -45,84 +96,73 @@
     dtype=torch.float64,
 )  # Sample from prior
 
-msip_args = {
+default_kwargs = {
+    "dim": dim,
+    "bounds": (-8, 8),
+    "n_steps": n_steps,
     "n_particles": n_particles,
-    "n_steps": n_steps,  # "epochs" (passes over all particles)
-    "dim": 64,
-    "bounds": (-8, 8),   # [a,b]^d
-    "gradient_informed": True,
+    "keep_all": False,
     "lr": 1e-1,
-    "noise": 0.05,          # currently unused
+    "kernel_lengthscale": 0.1,
     "init_particles": init_particles,
-    "kernel_bandwidth": kernel_bandwidth,
-    "bandwidth_factor": 0.25,
-    "seed": 0,
-    "kernel_diag_infl": 1e-10,
-    "keep_all": False,
-    "device": None
+    "gradient_decay": 0.95,
+    "kernel_diag_infl": 1e-6,
 }
 
 # %%
-trajectories_msip = msip(
-    log_p,
-    **msip_args
-)
+msip_kwargs = default_kwargs.copy()
+msip_kwargs["lr"] = 1e-2
+msip_kwargs["kernel_lengthscale_quantile"] = 0.05
+msip = MSIP(**msip_kwargs)
+target_msip_fr = MSIPFredholm(log_dens_grad_val=grad_val_log_p, **msip_kwargs)
+
+# %%
+trajectories_pts_msip_fr, trajectories_wts_msip_fr = nak_torch.nak(target_msip_fr, msip, **msip_kwargs)
 
 # %%
-n_steps_hmc = 1000
+n_steps_hmc = 100
 pyro_model = pyro_tools.PyroModel(model, dim)
 hmc_kernel = mcmc.NUTS(pyro_model)
-mcmc_setup = mcmc.MCMC(hmc_kernel, num_samples=n_steps_hmc, warmup_steps=100)
+mcmc_setup = mcmc.MCMC(hmc_kernel, num_samples=n_steps_hmc, warmup_steps=10)
 mcmc_setup.run(model.true_obs)
 
 hmc_samples = mcmc_setup.get_samples()["theta"]
 
 # %%
-trajectories_svgd = svgd(
-    log_p,
-    is_log_density_batched=True,
-    **msip_args
+target_svgd = BatchGradLogDensityEvaluator(
+    log_p, is_grad=False, is_batched=True
 )
 
+svgd = SVGD(
+    kernel_lengthscale_quantile=0.5,  # Median heuristic
+    **msip_kwargs
+)
+svgd_kwargs = msip_kwargs.copy()
+svgd_kwargs["lr"] = 1e-1
+svgd_kwargs["n_steps"] = 100
+
+trajectories_pts_svgd = nak_torch.nak(
+    target_svgd,
+    svgd,
+    **svgd_kwargs
+)
+
+# %%
+pts_msip = trajectories_pts_msip_fr[-1] - init_particles
+fig = plot_samples(pts_msip)
+fig.suptitle("MSIP Samples")
+plt.show()
+
+# %%
+pts_hmc = hmc_samples[10::3]
+fig = plot_samples(pts_hmc)
+fig.suptitle("HMC Samples")
+plt.show()
+
 # %%
-side_len = min(6, int(math.floor(math.sqrt(n_particles))))
-pts = trajectories_msip[-1][:side_len**2].detach().cpu()# - init_particles[:side_len**2]
-fig = plt.figure(figsize=(9, 6), layout='constrained')
-gs = fig.add_gridspec(side_len, side_len + 2)
-vabs = max(pts.min().abs(), pts.max().abs())
-plt_kwargs = {'vmin': -vabs, 'vmax': vabs, 'extent': (0, 8, 0, 8)}
-
-for i in range(side_len):
-    for j in range(side_len):
-        ax = fig.add_subplot(gs[i, j])
-        # ax.set_axis_off()
-        ax.set_aspect('equal')
-        t = ax.matshow(pts[i*side_len + j].reshape(8, 8), **plt_kwargs)
-        # ax.vlines(jnp.arange(1,8), -0.1, 8.1, color='w', lw=0.75)
-        # ax.hlines(jnp.arange(1,8), -0.1, 8.1, color='w', lw=0.75)
-        ax.minorticks_on()
-        ax.set_xticks([])
-        ax.set_yticks([])
-        ax.xaxis.set_minor_locator(ticker.MultipleLocator())
-        ax.yaxis.set_minor_locator(ticker.MultipleLocator())
-        ax.grid(which="both", linewidth=1.5, color="w")
-        ax.tick_params(which="minor", length=0)
-ax_cb = fig.add_subplot(gs[:-2, -2:])
-ax_cb.set_title(r"Scale of $\log\theta$", y=0.6)
-cax_cb = ax_cb.inset_axes((0.1, 0.45, 0.8, 0.1))
-ax_cb.axis('off')
-fig.colorbar(t, cax=cax_cb, orientation='horizontal')
-ax_true = fig.add_subplot(gs[-2:, -2:])
-ax_true.set_aspect('equal')
-ax_true.matshow(ab.theta_true.log().reshape(8, 8), **plt_kwargs)
-ax_true.minorticks_on()
-ax_true.set_xticks([])
-ax_true.set_yticks([])
-ax_true.xaxis.set_minor_locator(ticker.MultipleLocator())
-ax_true.yaxis.set_minor_locator(ticker.MultipleLocator())
-ax_true.set_title(r"True $\theta$")
-ax_true.grid(which="both", linewidth=1.5, color="w")
-ax_true.tick_params(which="minor", length=0)
+pts_svgd = trajectories_pts_svgd[-1]
+fig = plot_samples(pts_svgd)
+fig.suptitle("SVGD Samples")
 plt.show()
 
 # %%