momentum/pymomentum/tensor_momentum/tensor_kd_tree.h at a2a967486c02e025fd8678964c24d647871bbdb1 · facebookresearch/momentum · GitHub

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
/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

#pragma once

#include <ATen/ATen.h>

#include <tuple>

namespace pymomentum {

// Returns [closest point, closest index, valid]
std::tuple<at::Tensor, at::Tensor, at::Tensor>
findClosestPoints(at::Tensor points_source, at::Tensor points_target, float maxDist);

std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor> findClosestPointsWithNormals(
    at::Tensor points_source,
    at::Tensor normals_source,
    at::Tensor points_target,
    at::Tensor normals_target,
    float maxDist,
    float maxNormalDot);

std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor> findClosestPointsOnMesh(
    at::Tensor points_source,
    at::Tensor vertices_target,
    at::Tensor faces_target);

} // namespace pymomentum