from math import sqrt

import torch
from torch_geometric.utils import geodesic_distance


def test_geodesic_distance():
    pos = torch.Tensor([[0, 0, 0], [2, 0, 0], [0, 2, 0], [2, 2, 0]])
    face = torch.tensor([[0, 1, 3], [0, 2, 3]]).t()

    out = geodesic_distance(pos, face)
    expected = [
        [0, 1, 1, sqrt(2)],
        [1, 0, sqrt(2), 1],
        [1, sqrt(2), 0, 1],
        [sqrt(2), 1, 1, 0],
    ]
    assert torch.allclose(out, torch.tensor(expected))
    assert torch.allclose(out, geodesic_distance(pos, face, num_workers=-1))

    out = geodesic_distance(pos, face, norm=False)
    expected = [
        [0, 2, 2, 2 * sqrt(2)],
        [2, 0, 2 * sqrt(2), 2],
        [2, 2 * sqrt(2), 0, 2],
        [2 * sqrt(2), 2, 2, 0],
    ]
    assert torch.allclose(out, torch.tensor(expected))

    src = torch.tensor([0, 0, 0, 0])
    dest = torch.tensor([0, 1, 2, 3])
    out = geodesic_distance(pos, face, src=src, dest=dest)
    expected = [0, 1, 1, sqrt(2)]
    assert torch.allclose(out, torch.tensor(expected))

    out = geodesic_distance(pos, face, src=src[0:1])
    expected = [0, 1, 1, sqrt(2)]
    assert torch.allclose(out, torch.tensor(expected))

    out = geodesic_distance(pos, face, dest=dest)
    expected = [0, 0, 0, 0]
    assert torch.allclose(out, torch.Tensor(expected))

[((130, 188), 'torch.Tensor', 'torch.Tensor', (['[[0, 0, 0], [2, 0, 0], [0, 2, 0], [2, 2, 0]]'], {}), '([[0, 0, 0], [2, 0, 0], [0, 2, 0], [2, 2, 0]])\n', (142, 188), False, 'import torch\n'), ((252, 280), 'torch_geometric.utils.geodesic_distance', 'geodesic_distance', (['pos', 'face'], {}), '(pos, face)\n', (269, 280), False, 'from torch_geometric.utils import geodesic_distance\n'), ((559, 599), 'torch_geometric.utils.geodesic_distance', 'geodesic_distance', (['pos', 'face'], {'norm': '(False)'}), '(pos, face, norm=False)\n', (576, 599), False, 'from torch_geometric.utils import geodesic_distance\n'), ((817, 843), 'torch.tensor', 'torch.tensor', (['[0, 0, 0, 0]'], {}), '([0, 0, 0, 0])\n', (829, 843), False, 'import torch\n'), ((855, 881), 'torch.tensor', 'torch.tensor', (['[0, 1, 2, 3]'], {}), '([0, 1, 2, 3])\n', (867, 881), False, 'import torch\n'), ((892, 940), 'torch_geometric.utils.geodesic_distance', 'geodesic_distance', (['pos', 'face'], {'src': 'src', 'dest': 'dest'}), '(pos, face, src=src, dest=dest)\n', (909, 940), False, 'from torch_geometric.utils import geodesic_distance\n'), ((1041, 1083), 'torch_geometric.utils.geodesic_distance', 'geodesic_distance', (['pos', 'face'], {'src': 'src[0:1]'}), '(pos, face, src=src[0:1])\n', (1058, 1083), False, 'from torch_geometric.utils import geodesic_distance\n'), ((1184, 1223), 'torch_geometric.utils.geodesic_distance', 'geodesic_distance', (['pos', 'face'], {'dest': 'dest'}), '(pos, face, dest=dest)\n', (1201, 1223), False, 'from torch_geometric.utils import geodesic_distance\n'), ((447, 469), 'torch.tensor', 'torch.tensor', (['expected'], {}), '(expected)\n', (459, 469), False, 'import torch\n'), ((502, 546), 'torch_geometric.utils.geodesic_distance', 'geodesic_distance', (['pos', 'face'], {'num_workers': '(-1)'}), '(pos, face, num_workers=-1)\n', (519, 546), False, 'from torch_geometric.utils import geodesic_distance\n'), ((782, 804), 'torch.tensor', 'torch.tensor', (['expected'], {}), '(expected)\n', (794, 804), False, 'import torch\n'), ((966, 973), 'math.sqrt', 'sqrt', (['(2)'], {}), '(2)\n', (970, 973), False, 'from math import sqrt\n'), ((1006, 1028), 'torch.tensor', 'torch.tensor', (['expected'], {}), '(expected)\n', (1018, 1028), False, 'import torch\n'), ((1109, 1116), 'math.sqrt', 'sqrt', (['(2)'], {}), '(2)\n', (1113, 1116), False, 'from math import sqrt\n'), ((1149, 1171), 'torch.tensor', 'torch.tensor', (['expected'], {}), '(expected)\n', (1161, 1171), False, 'import torch\n'), ((1283, 1305), 'torch.Tensor', 'torch.Tensor', (['expected'], {}), '(expected)\n', (1295, 1305), False, 'import torch\n'), ((200, 236), 'torch.tensor', 'torch.tensor', (['[[0, 1, 3], [0, 2, 3]]'], {}), '([[0, 1, 3], [0, 2, 3]])\n', (212, 236), False, 'import torch\n'), ((316, 323), 'math.sqrt', 'sqrt', (['(2)'], {}), '(2)\n', (320, 323), False, 'from math import sqrt\n'), ((341, 348), 'math.sqrt', 'sqrt', (['(2)'], {}), '(2)\n', (345, 348), False, 'from math import sqrt\n'), ((366, 373), 'math.sqrt', 'sqrt', (['(2)'], {}), '(2)\n', (370, 373), False, 'from math import sqrt\n'), ((391, 398), 'math.sqrt', 'sqrt', (['(2)'], {}), '(2)\n', (395, 398), False, 'from math import sqrt\n'), ((639, 646), 'math.sqrt', 'sqrt', (['(2)'], {}), '(2)\n', (643, 646), False, 'from math import sqrt\n'), ((668, 675), 'math.sqrt', 'sqrt', (['(2)'], {}), '(2)\n', (672, 675), False, 'from math import sqrt\n'), ((697, 704), 'math.sqrt', 'sqrt', (['(2)'], {}), '(2)\n', (701, 704), False, 'from math import sqrt\n'), ((726, 733), 'math.sqrt', 'sqrt', (['(2)'], {}), '(2)\n', (730, 733), False, 'from math import sqrt\n')]