Skip to content

Entropy estimation with non-invertible generative models #60

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
austin-hoover wants to merge 31 commits into
base: main
Choose a base branch
from
Open

Entropy estimation with non-invertible generative models #60

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
31 commits
Select commit Hold shift + click to select a range
a9220d6
Test jacobian calculation
austin-hoover Aug 28, 2025
fec9b20
Add 6D example
austin-hoover Aug 28, 2025
c0d38ac
Modify example
austin-hoover Aug 28, 2025
488793d
Merge branch 'roussel-ryan:main' into jacobian
austin-hoover Sep 30, 2025
47eee9b
Change name from NSF to NSFDist
austin-hoover Sep 30, 2025
75663a2
Update generative model classes
austin-hoover Sep 30, 2025
9374005
Change example to use gpsr.beams module
austin-hoover Sep 30, 2025
e7f8c57
Format
austin-hoover Sep 30, 2025
f181398
Format
austin-hoover Sep 30, 2025
1c85029
Format
austin-hoover Sep 30, 2025
3be021b
Change name from test_kl_nn to test_kl.ipynb
austin-hoover Oct 2, 2025
bb314ca
Add test test_jacobian_sin
austin-hoover Oct 2, 2025
bb43bf1
Delete test_jacobian.ipynb
austin-hoover Oct 2, 2025
913c247
Add test test_nn_transform_jacobian
austin-hoover Oct 2, 2025
e50ef45
Format
austin-hoover Oct 2, 2025
05e74b3
pylint
austin-hoover Oct 3, 2025
9d9c988
Rename arguments in NSFDist to match NNDist
austin-hoover Oct 9, 2025
d7c0ef0
Update Jacobian examples
austin-hoover Oct 9, 2025
7e379bd
Format
austin-hoover Oct 9, 2025
47023ba
Rename examples/jacobian to examples/kl
austin-hoover Oct 9, 2025
c3872dc
Format
austin-hoover Oct 9, 2025
8a3ea6a
Fix log_prob
austin-hoover Oct 13, 2025
4a9c015
Reverts 8a3ea6aa0391ad761b40ebf01e38aba3727bb787
austin-hoover Oct 13, 2025
5c99d53
Rename variables
austin-hoover Oct 13, 2025
61b6a2e
Add reverse KL minimization examples with ring distribution
austin-hoover Oct 13, 2025
493ea6a
Merge branch 'roussel-ryan:main' into jacobian
austin-hoover Oct 13, 2025
c61528c
Merge branch 'roussel-ryan:main' into jacobian
austin-hoover Apr 30, 2026
5be340b
Shorten comment
austin-hoover Apr 30, 2026
7ca1261
Add test of log_prob calculation using NSF
austin-hoover Apr 30, 2026
1c97186
Merge branch 'jacobian' of https://github.com/austin-hoover/gpsr into…
austin-hoover Apr 30, 2026
892e245
Update kl examples
austin-hoover Apr 30, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
243 changes: 243 additions & 0 deletions docs/examples/kl/kl_nn_gauss.ipynb
View file
Open in desktop

@roussel-ryan roussel-ryan Oct 2, 2025

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this is closer to a usable example, so we should rename to kl_nn.ipynb. A separate unit test would be useful for validating changes to beams.py

@austin-hoover austin-hoover Oct 2, 2025

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Renamed kl_nn.ipynb

Original file line number Diff line number Diff line change
@@ -0,0 +1,243 @@
{
"cells": [
{
"cell_type": "markdown",
"id": "ecf3fd1a",
"metadata": {},
"source": [
"# Test KL minimization with non-invertible transformation"
]
},
{
"cell_type": "code",
"id": "8f899f1d",
"metadata": {},
"source": [
"import matplotlib.pyplot as plt\n",
"import numpy as np\n",
"import torch\n",
"\n",
"from gpsr.beams import NNDist"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "code",
"id": "ed6f9f69",
"metadata": {},
"source": [
"ndim = 2\n",
"nsamp = 10_000\n",
"\n",
"model_dist = NNDist(width=32, depth=2, ndim=ndim)\n",
"targ_dist = torch.distributions.MultivariateNormal(\n",
" torch.zeros(ndim),\n",
" torch.eye(ndim),\n",
")\n",
"optimizer = torch.optim.Adam(model_dist.parameters(), lr=0.005)\n",
"\n",
"# Warmup\n",
"for i in range(300):\n",
" x = model_dist.sample(nsamp)\n",
" cov_matrix = torch.cov(x.T)\n",
" loss = torch.mean(torch.abs(cov_matrix - torch.eye(ndim)))\n",
" loss.backward()\n",
" optimizer.step()\n",
" optimizer.zero_grad()\n",
"\n",
"history = {\"loss\": []}\n",
"\n",
"for iteration in range(201):\n",
" loss = -model_dist.entropy(nsamp, prior=targ_dist)\n",
" loss.backward()\n",
" optimizer.step()\n",
" optimizer.zero_grad()\n",
"\n",
" history[\"loss\"].append(loss.item())\n",
"\n",
" if iteration % 50 == 0:\n",
" print(iteration, loss)"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "code",
"id": "664c404d",
"metadata": {},
"source": [
"fig, ax = plt.subplots(figsize=(3.0, 2.0))\n",
"ax.plot(history[\"loss\"])\n",
"ax.set_xlabel(\"Iteration\")\n",
"ax.set_ylabel(\"Loss\")\n",
"plt.show()"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "code",
"id": "2203c1b5",
"metadata": {},
"source": [
"with torch.no_grad():\n",
" # Sample and evaluate density at each point\n",
" x, log_p = model_dist.sample_and_log_prob(100_000)\n",
"\n",
" x = x.numpy()\n",
" log_p = log_p.numpy()\n",
"\n",
" x_targ = targ_dist.rsample((x.shape[0],))\n",
" x_targ = x_targ.numpy()\n",
"\n",
" # Sort by density\n",
" idx = np.argsort(log_p)\n",
" x = x[idx]\n",
" log_p = log_p[idx]\n",
"\n",
" # Plot histogram and scatter plot with points colored by density\n",
" plot_xmax = 4.0\n",
" plot_limits = 2 * [(-plot_xmax, plot_xmax)]\n",
"\n",
" fig, axs = plt.subplots(ncols=3, sharex=True, sharey=True, figsize=(8.5, 2.5))\n",
"\n",
" grid_values, grid_edges = np.histogramdd(\n",
" x, bins=64, range=plot_limits, density=True\n",
" )\n",
" axs[0].hist2d(x[:, 0], x[:, 1], bins=64, range=plot_limits, density=True)\n",
" axs[1].scatter(x[:, 0], x[:, 1], c=np.exp(log_p), s=1)\n",
" axs[2].hist2d(x_targ[:, 0], x_targ[:, 1], bins=64, range=plot_limits, density=True)\n",
" axs[0].set_title(\"Model samples\", fontsize=\"medium\")\n",
" axs[1].set_title(\"Model density\", fontsize=\"medium\")\n",
" axs[2].set_title(\"Target samples\", fontsize=\"medium\")\n",
" plt.show()"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"id": "bd48e24a",
"metadata": {},
"source": [
"### 6D"
]
},
{
"cell_type": "code",
"id": "c7d87086",
"metadata": {},
"source": [
"ndim = 6\n",
"nsamp = 10_000\n",
"\n",
"model_dist = NNDist(width=32, depth=2, ndim=ndim)\n",
"targ_dist = torch.distributions.MultivariateNormal(\n",
" torch.zeros(ndim),\n",
" torch.eye(ndim),\n",
")\n",
"optimizer = torch.optim.Adam(model_dist.parameters(), lr=0.001)\n",
"\n",
"# Warmup\n",
"for i in range(500):\n",
" x = model_dist.sample(nsamp)\n",
" cov_matrix = torch.cov(x.T)\n",
" loss = torch.mean(torch.abs(cov_matrix - torch.eye(ndim)))\n",
" loss.backward()\n",
" optimizer.step()\n",
" optimizer.zero_grad()\n",
"\n",
"history = {\"loss\": []}\n",
"\n",
"for iteration in range(501):\n",
" loss = -model_dist.entropy(nsamp, prior=targ_dist)\n",
" loss.backward()\n",
" optimizer.step()\n",
" optimizer.zero_grad()\n",
"\n",
" history[\"loss\"].append(loss.item())\n",
"\n",
" if iteration % 50 == 0:\n",
" print(iteration, loss)"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "code",
"id": "edd26498",
"metadata": {},
"source": [
"fig, ax = plt.subplots(figsize=(3.0, 2.0))\n",
"ax.plot(history[\"loss\"])\n",
"ax.set_xlabel(\"Iteration\")\n",
"ax.set_ylabel(\"Loss\")\n",
"plt.show()"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"id": "57f67080",
"metadata": {},
"source": [
"This time we won't make the scatter plot: the density at each point is in 6D space, not the 2D space in the plot."
]
},
{
"cell_type": "code",
"id": "cf1efa6e",
"metadata": {},
"source": [
"with torch.no_grad():\n",
" # Sample particles.\n",
" x_model = model_dist.sample(100_000)\n",
" x_targ = targ_dist.rsample((x_model.shape[0],))\n",
"\n",
" # Plot histogram and scatter plot with points colored by density.\n",
" plot_xmax = 4.0\n",
" plot_limits = 2 * [(-plot_xmax, plot_xmax)]\n",
"\n",
" fig, axs = plt.subplots(ncols=2, sharex=True, sharey=True, figsize=(5.5, 2.5))\n",
" for ax, x in zip(axs, [x_model, x_targ]):\n",
" ax.hist2d(x[:, 0], x[:, 1], bins=64, range=plot_limits, density=True)\n",
" axs[0].set_title(\"Model\", fontsize=\"medium\")\n",
" axs[1].set_title(\"Target\", fontsize=\"medium\")\n",
" plt.show()"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "code",
"id": "5bad2980",
"metadata": {},
"source": [],
"outputs": [],
"execution_count": null
}
],
"metadata": {
"kernelspec": {
"display_name": "gpsr",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.12.11"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
153 changes: 153 additions & 0 deletions docs/examples/kl/kl_nn_rings.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,153 @@
{
"cells": [
{
"cell_type": "markdown",
"id": "ecf3fd1a",
"metadata": {},
"source": [
"# Test KL minimization with non-invertible transformation"
]
},
{
"cell_type": "code",
"id": "8f899f1d",
"metadata": {},
"source": [
"import matplotlib.pyplot as plt\n",
"import torch\n",
"\n",
"from gpsr.beams import NNDist"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "code",
"id": "5dea9acd",
"metadata": {},
"source": [
"class TargetDist:\n",
" def log_prob(self, x: torch.Tensor) -> torch.Tensor:\n",
" x1 = x[:, 0]\n",
" x2 = x[:, 1]\n",
" return torch.sin(torch.pi * x1) - 2 * (x1**2 + x2**2 - 2.0) ** 2"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "code",
"id": "ed6f9f69",
"metadata": {},
"source": [
"ndim = 2\n",
"nsamp = 10_000\n",
"\n",
"targ_dist = TargetDist()\n",
"\n",
"activation = torch.nn.Tanh()\n",
"# activation = torch.nn.SiLU()\n",
"# activation = torch.nn.GELU()\n",
"# activation = torch.nn.LeakyReLU(0.1)\n",
"# activation = torch.nn.ELU()\n",
"# activation = torch.nn.GLU()\n",
"\n",
"model_dist = NNDist(width=20, depth=2, ndim=ndim, activation=activation)\n",
"optimizer = torch.optim.Adam(model_dist.parameters(), lr=0.001)\n",
"\n",
"# Warmup\n",
"for i in range(300):\n",
" x = model_dist.sample(nsamp)\n",
" cov_matrix = torch.cov(x.T)\n",
" loss = torch.mean(torch.abs(cov_matrix - torch.eye(ndim)))\n",
" loss.backward()\n",
" optimizer.step()\n",
" optimizer.zero_grad()\n",
"\n",
"history = {\"loss\": []}\n",
"\n",
"for iteration in range(201):\n",
" loss = -model_dist.entropy(nsamp, prior=targ_dist)\n",
" loss.backward()\n",
" optimizer.step()\n",
" optimizer.zero_grad()\n",
"\n",
" history[\"loss\"].append(loss.item())\n",
"\n",
" if iteration % 50 == 0:\n",
" print(iteration, loss)"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "code",
"id": "664c404d",
"metadata": {},
"source": [
"fig, ax = plt.subplots(figsize=(3.0, 2.0))\n",
"ax.plot(history[\"loss\"])\n",
"ax.set_xlabel(\"Iteration\")\n",
"ax.set_ylabel(\"Loss\")\n",
"plt.show()"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "code",
"id": "c96ede97",
"metadata": {},
"source": [
"with torch.no_grad():\n",
" grid_shape = (128, 128)\n",
" grid_edges = [torch.linspace(-4.0, 4.0, n + 1) for n in grid_shape]\n",
" grid_coords = [0.5 * (e[:-1] + e[1:]) for e in grid_edges]\n",
" grid_points = torch.stack(\n",
" [c.ravel() for c in torch.meshgrid(*grid_coords, indexing=\"ij\")], axis=-1\n",
" )\n",
" grid_values = torch.exp(targ_dist.log_prob(grid_points)).reshape(grid_shape)\n",
"\n",
" x = model_dist.sample(256_000)\n",
"\n",
" fig, axs = plt.subplots(ncols=2, sharex=True, sharey=True, figsize=(5.5, 2.5))\n",
" axs[0].hist2d(x[:, 0], x[:, 1], bins=grid_edges)\n",
" axs[1].pcolormesh(grid_coords[0], grid_coords[1], grid_values.T)\n",
" axs[0].set_title(\"Model\", fontsize=\"medium\")\n",
" axs[1].set_title(\"Target\", fontsize=\"medium\")\n",
" plt.show()"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "code",
"id": "5bad2980",
"metadata": {},
"source": [],
"outputs": [],
"execution_count": null
}
],
"metadata": {
"kernelspec": {
"display_name": "gpsr",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.12.11"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
Loading
Loading