Spaces:
Sleeping
Sleeping
File size: 2,830 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 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 |
import torch
import numpy as np
from ideal_poly_volume_toolkit.geometry import triangle_volume_from_points_torch, _LobFn
# Test if the issue is with the custom autograd function _LobFn
print("Testing if custom autograd function is the issue...\n")
# Test 1: Simple function with standard PyTorch ops
print("Test 1: Standard PyTorch function")
theta1 = torch.tensor(0.5, dtype=torch.float64, requires_grad=True)
opt1 = torch.optim.LBFGS([theta1], lr=1.0)
def closure1():
opt1.zero_grad()
loss = -(torch.sin(theta1) + 0.5 * torch.cos(2*theta1)) # some smooth function
loss.backward()
print(f" theta={theta1.item():.4f}, loss={loss.item():.4f}, grad={theta1.grad.item():.4f}")
return loss
print("Before:", theta1.item())
opt1.step(closure1)
print("After:", theta1.item())
# Test 2: Using the custom _LobFn directly
print("\n\nTest 2: Custom _LobFn autograd function")
theta2 = torch.tensor(0.5, dtype=torch.float64, requires_grad=True)
opt2 = torch.optim.LBFGS([theta2], lr=1.0)
def closure2():
opt2.zero_grad()
# Use the Lobachevsky function directly
lob_val = _LobFn.apply(theta2, 96)
loss = -lob_val
loss.backward()
print(f" theta={theta2.item():.4f}, loss={loss.item():.4f}, grad={theta2.grad.item():.4f}")
return loss
print("Before:", theta2.item())
opt2.step(closure2)
print("After:", theta2.item())
# Test 3: Full triangle volume computation
print("\n\nTest 3: Full triangle volume (our actual case)")
theta3 = torch.tensor(0.5, dtype=torch.float64, requires_grad=True)
opt3 = torch.optim.LBFGS([theta3], lr=1.0)
def closure3():
opt3.zero_grad()
z0 = torch.tensor(0+0j, dtype=torch.complex128)
z1 = torch.tensor(1+0j, dtype=torch.complex128)
z2 = torch.exp(1j * theta3.to(torch.complex128))
volume = triangle_volume_from_points_torch(z0, z1, z2, series_terms=96)
loss = -volume
loss.backward()
print(f" theta={theta3.item():.4f}, loss={loss.item():.4f}, grad={theta3.grad.item():.4f}")
return loss
print("Before:", theta3.item())
opt3.step(closure3)
print("After:", theta3.item())
# Test 4: Try with Adam optimizer instead
print("\n\nTest 4: Same problem with Adam optimizer")
theta4 = torch.tensor(0.5, dtype=torch.float64, requires_grad=True)
opt4 = torch.optim.Adam([theta4], lr=0.1)
print("Initial theta:", theta4.item())
for i in range(10):
opt4.zero_grad()
z0 = torch.tensor(0+0j, dtype=torch.complex128)
z1 = torch.tensor(1+0j, dtype=torch.complex128)
z2 = torch.exp(1j * theta4.to(torch.complex128))
volume = triangle_volume_from_points_torch(z0, z1, z2, series_terms=96)
loss = -volume
loss.backward()
opt4.step()
if i % 3 == 0:
print(f" Step {i}: theta={theta4.item():.4f}, volume={volume.item():.4f}, grad={theta4.grad.item():.4f}")
print("Final theta:", theta4.item()) |