idealpolyhedra / scripts /test_from_bad_init.py
igriv's picture
Major reorganization and feature additions
d7d27f0
import numpy as np, torch
from ideal_poly_volume_toolkit.geometry import (
delaunay_triangulation_indices,
triangle_volume_from_points_torch,
)
def build_Z(thetas: torch.Tensor) -> torch.Tensor:
Z = torch.empty(thetas.numel() + 2, dtype=torch.complex128, device=thetas.device)
Z[0] = 1 + 0j
Z[1] = 0 + 0j
Z[2:] = torch.exp(1j * thetas.to(torch.complex128))
return Z
def torch_sum_volume(Z_t: torch.Tensor, idx, series_terms: int) -> torch.Tensor:
total = torch.zeros((), dtype=torch.float64, device=Z_t.device)
for (i, j, k) in idx:
total = total + triangle_volume_from_points_torch(
Z_t[i], Z_t[j], Z_t[k], series_terms=series_terms
)
return total
# Start with points very close together (bad configuration)
thetas = torch.tensor([0.1, 0.2, 0.3], dtype=torch.float64, requires_grad=True)
print(f"Initial thetas: {thetas}")
# Setup LBFGS
opt = torch.optim.LBFGS([thetas], lr=1.0, max_iter=20, line_search_fn='strong_wolfe')
for it in range(1, 21):
with torch.no_grad():
Z_np = build_Z(thetas).detach().cpu().numpy()
idx = delaunay_triangulation_indices(Z_np)
def closure():
opt.zero_grad(set_to_none=True)
Z_t = build_Z(thetas)
total = torch_sum_volume(Z_t, idx, 96)
loss = -total
loss.backward()
return loss
opt.step(closure)
with torch.no_grad():
Z_post = build_Z(thetas)
val_post = torch_sum_volume(Z_post, idx, 96)
print(f'[{it:02d}] volume = {val_post.item():.8f}, thetas = {thetas.numpy()}')