"""A direct deep Laplacian, independently checked against a dense Hessian."""
import torch
from omnibias.torch.laplacian import deep_field_laplacian

weight = torch.tensor([[0.2, -0.3], [0.5, 0.1]], dtype=torch.float64,
                      requires_grad=True)
bias = torch.tensor([0.1, -0.1], dtype=torch.float64)
readout = torch.tensor([[0.7, -0.4]], dtype=torch.float64)
layers = [(weight, bias, "tanh"), (readout, None, None)]
point = torch.tensor([[0.3, -0.2]], dtype=torch.float64)
laplacian = deep_field_laplacian(point, layers)

def field(x):
    return (readout @ torch.tanh(weight @ x + bias)).squeeze()

reference = torch.func.hessian(field)(point[0]).trace()
assert laplacian.shape == (1, 1)
torch.testing.assert_close(laplacian.squeeze(), reference, rtol=1e-10, atol=1e-12)
laplacian.square().mean().backward()
assert weight.grad is not None and torch.isfinite(weight.grad).all()
print("Laplacian agrees with Hessian trace; parameter gradients retained.")