Spaces:
Sleeping
Sleeping
File size: 1,577 Bytes
82a8f4b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 |
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()}') |