implement support for cuboids in attachment

src/impl/vamp/bindings/environment.cc

@@ -215,6 +215,16 @@ void vamp::binding::init_environment(nanobind::module_ &pymodule)
             "add_spheres",
             [](vc::Attachment<float> &a, std::vector<collision::Sphere<float>> &spheres)
             { a.spheres.insert(a.spheres.end(), spheres.cbegin(), spheres.cend()); })
+        .def("clear_spheres", [](vc::Attachment<float> &a) { a.spheres.clear(); })
+        .def(
+            "add_cuboid",
+            [](vc::Attachment<float> &a, collision::Cuboid<float> &cuboid)
+            { a.cuboids.emplace_back(cuboid); })
+        .def(
+            "add_cuboids",
+            [](vc::Attachment<float> &a, std::vector<collision::Cuboid<float>> &cuboids)
+            { a.cuboids.insert(a.cuboids.end(), cuboids.cbegin(), cuboids.cend()); })
+        .def("clear_cuboids", [](vc::Attachment<float> &a) { a.cuboids.clear(); })
         .def(
             "set_ee_pose",
             [](vc::Attachment<float> &a, Eigen::Matrix4f &tf)
@@ -224,5 +234,6 @@ void vamp::binding::init_environment(nanobind::module_ &pymodule)
                 a.pose(iso);
             },
             "tf"_a)
-        .def_ro("posed_spheres", &vc::Attachment<float>::posed_spheres);
+        .def_ro("posed_spheres", &vc::Attachment<float>::posed_spheres)
+        .def_ro("posed_cuboids", &vc::Attachment<float>::posed_cuboids);
 }

src/impl/vamp/collision/attachments.hh

@@ -32,12 +32,20 @@ namespace vamp::collision
             {
                 spheres.emplace_back(sphere);
             }
+
+            cuboids.reserve(o.cuboids.size());
+            for (const auto &cuboid : o.cuboids)
+            {
+                cuboids.emplace_back(cuboid);
+            }
         }
 
         std::vector<Sphere<DataT>> spheres;
+        std::vector<Cuboid<DataT>> cuboids;
         // HACK: To get around passing the environment as const but needing to re-pose the
         // attachments
         mutable std::vector<Sphere<DataT>> posed_spheres;
+        mutable std::vector<Cuboid<DataT>> posed_cuboids;
         Eigen::Transform<DataT, 3, Eigen::Isometry> tf;
 
         inline void pose(const Eigen::Transform<DataT, 3, Eigen::Isometry> &p_tf) const noexcept
@@ -52,6 +60,39 @@ namespace vamp::collision
                 auto tfs = n_tf * sp;
                 posed_spheres[i] = Sphere<DataT>(tfs[0], tfs[1], tfs[2], s.r);
             }
+
+            posed_cuboids.resize(cuboids.size());
+            for (auto i = 0U; i < cuboids.size(); ++i)
+            {
+                const auto &c = cuboids[i];
+                Eigen::Matrix<DataT, 3, 1> cp(c.x, c.y, c.z);
+                Eigen::Matrix<DataT, 3, 1> ax1(c.axis_1_x, c.axis_1_y, c.axis_1_z);
+                Eigen::Matrix<DataT, 3, 1> ax2(c.axis_2_x, c.axis_2_y, c.axis_2_z);
+                Eigen::Matrix<DataT, 3, 1> ax3(c.axis_3_x, c.axis_3_y, c.axis_3_z);
+
+                const auto t_center = n_tf * cp;
+                const auto rot = n_tf.linear();
+                const auto t_axis_1 = rot * ax1;
+                const auto t_axis_2 = rot * ax2;
+                const auto t_axis_3 = rot * ax3;
+
+                posed_cuboids[i] = Cuboid<DataT>(
+                    t_center[0],
+                    t_center[1],
+                    t_center[2],
+                    t_axis_1[0],
+                    t_axis_1[1],
+                    t_axis_1[2],
+                    t_axis_2[0],
+                    t_axis_2[1],
+                    t_axis_2[2],
+                    t_axis_3[0],
+                    t_axis_3[1],
+                    t_axis_3[2],
+                    c.axis_1_r,
+                    c.axis_2_r,
+                    c.axis_3_r);
+            }
         }
     };
 }  // namespace vamp::collision

src/impl/vamp/collision/cuboid_capsule.hh

@@ -0,0 +1,75 @@
+#pragma once
+
+#include <array>
+#include <limits>
+#include <type_traits>
+
+#include <vamp/collision/shapes.hh>
+#include <vamp/collision/cuboid_cuboid.hh>
+#include <vamp/collision/math.hh>
+
+namespace vamp::collision
+{
+
+    template <typename DataT>
+    inline constexpr auto cuboid_capsule(const Cuboid<DataT> &box, const Capsule<DataT> &capsule) noexcept
+        -> DataT
+    {
+        const DataT zero = 0.;
+        const DataT one = 1.;
+        const DataT eps = 1e-6F;
+        const DataT inf = std::numeric_limits<float>::infinity();
+
+        const auto axes = vamp::collision::cuboid_helpers::cuboid_axes(box);
+        const auto extents = vamp::collision::cuboid_helpers::cuboid_extents(box);
+
+        const auto sx = capsule.x1 - box.x;
+        const auto sy = capsule.y1 - box.y;
+        const auto sz = capsule.z1 - box.z;
+
+        std::array<DataT, 3> p0 = {
+            vamp::collision::cuboid_helpers::project_onto_axis(axes, 0, sx, sy, sz),
+            vamp::collision::cuboid_helpers::project_onto_axis(axes, 1, sx, sy, sz),
+            vamp::collision::cuboid_helpers::project_onto_axis(axes, 2, sx, sy, sz)};
+
+        std::array<DataT, 3> dir = {
+            vamp::collision::cuboid_helpers::project_onto_axis(axes, 0, capsule.xv, capsule.yv, capsule.zv),
+            vamp::collision::cuboid_helpers::project_onto_axis(axes, 1, capsule.xv, capsule.yv, capsule.zv),
+            vamp::collision::cuboid_helpers::project_onto_axis(axes, 2, capsule.xv, capsule.yv, capsule.zv)};
+
+        const auto eval_at = [&](const DataT &t)
+        {
+            const auto px = p0[0] + dir[0] * t;
+            const auto py = p0[1] + dir[1] * t;
+            const auto pz = p0[2] + dir[2] * t;
+
+            const auto dx = (px.abs() - extents[0]).max(zero);
+            const auto dy = (py.abs() - extents[1]).max(zero);
+            const auto dz = (pz.abs() - extents[2]).max(zero);
+
+            return dot_3(dx, dy, dz, dx, dy, dz);
+        };
+
+        auto best = eval_at(zero);
+        best = vamp::collision::cuboid_helpers::min_value(best, eval_at(one));
+
+        const std::array<DataT, 2> signs = {one, DataT(-1.)};
+
+        for (auto axis = 0U; axis < 3; ++axis)
+        {
+            const auto non_zero_mask = dir[axis].abs() > eps;
+            const auto denom = vamp::collision::cuboid_helpers::mask_select(non_zero_mask, dir[axis], one);
+
+            for (const auto &sign : signs)
+            {
+                const auto numer = sign * extents[axis] - p0[axis];
+                const auto t = (numer / denom).clamp(0.F, 1.F);
+                const auto candidate =
+                    vamp::collision::cuboid_helpers::mask_select(non_zero_mask, eval_at(t), inf);
+                best = vamp::collision::cuboid_helpers::min_value(best, candidate);
+            }
+        }
+
+        return best - (capsule.r * capsule.r);
+    }
+}  // namespace vamp::collision

src/impl/vamp/collision/cuboid_cuboid.hh

@@ -0,0 +1,165 @@
+#pragma once
+
+#include <array>
+
+#include <vamp/collision/shapes.hh>
+#include <vamp/collision/math.hh>
+
+namespace vamp::collision::cuboid_helpers
+{
+    template <typename DataT>
+    inline constexpr auto cuboid_axes(const Cuboid<DataT> &c) noexcept -> std::array<std::array<DataT, 3>, 3>
+    {
+        return {{
+            {c.axis_1_x, c.axis_1_y, c.axis_1_z},
+            {c.axis_2_x, c.axis_2_y, c.axis_2_z},
+            {c.axis_3_x, c.axis_3_y, c.axis_3_z},
+        }};
+    }
+
+    template <typename DataT>
+    inline constexpr auto cuboid_radii(const Cuboid<DataT> &c) noexcept -> std::array<DataT, 3>
+    {
+        return {c.axis_1_r, c.axis_2_r, c.axis_3_r};
+    }
+
+    template <typename DataT>
+    inline constexpr auto axis_dot(
+        const std::array<std::array<DataT, 3>, 3> &axes,
+        std::size_t idx,
+        const DataT &x,
+        const DataT &y,
+        const DataT &z) noexcept -> DataT
+    {
+        return dot_3(axes[idx][0], axes[idx][1], axes[idx][2], x, y, z);
+    }
+
+    template <typename MaskT, typename ValueT>
+    inline constexpr auto mask_select(const MaskT &mask, const ValueT &truthy, const ValueT &falsy) -> ValueT
+    {
+        if constexpr (std::is_same_v<MaskT, bool>)
+        {
+            return mask ? truthy : falsy;
+        }
+        else
+        {
+            return falsy.blend(truthy, mask);
+        }
+    }
+
+    template <typename DataT>
+    inline constexpr auto min_value(const DataT &a, const DataT &b) -> DataT
+    {
+        const auto mask = (a <= b);
+        return mask_select(mask, a, b);
+    }
+
+    template <typename DataT>
+    inline constexpr auto cuboid_extents(const Cuboid<DataT> &c) noexcept -> std::array<DataT, 3>
+    {
+        return {c.axis_1_r, c.axis_2_r, c.axis_3_r};
+    }
+
+    template <typename DataT>
+    inline constexpr auto project_onto_axis(
+        const std::array<std::array<DataT, 3>, 3> &axes,
+        std::size_t idx,
+        const DataT &x,
+        const DataT &y,
+        const DataT &z) noexcept -> DataT
+    {
+        return dot_3(axes[idx][0], axes[idx][1], axes[idx][2], x, y, z);
+    }
+}  // namespace vamp::collision::cuboid_helpers
+
+namespace vamp::collision
+{
+
+    /**
+     * Separating axis theorem test between two arbitrarily oriented cuboids.
+     *
+     * The returned value is the maximum separation measured across the 15 candidate axes
+     * (three face normals from each cuboid and the nine cross products between them). A
+     * strictly positive result indicates that the cuboids are separated by that margin,
+     * while a non-positive result indicates an overlap (zero meaning they are exactly
+     * touching).
+     */
+    template <typename DataT>
+    inline constexpr auto cuboid_cuboid(const Cuboid<DataT> &a, const Cuboid<DataT> &b) noexcept -> DataT
+    {
+        const auto axes_a = vamp::collision::cuboid_helpers::cuboid_axes(a);
+        const auto axes_b = vamp::collision::cuboid_helpers::cuboid_axes(b);
+        const auto radii_a = vamp::collision::cuboid_helpers::cuboid_radii(a);
+        const auto radii_b = vamp::collision::cuboid_helpers::cuboid_radii(b);
+
+        // Translation from A to B, expressed in A's frame.
+        const auto tx = b.x - a.x;
+        const auto ty = b.y - a.y;
+        const auto tz = b.z - a.z;
+
+        std::array<DataT, 3> t = {
+            vamp::collision::cuboid_helpers::axis_dot(axes_a, 0, tx, ty, tz),
+            vamp::collision::cuboid_helpers::axis_dot(axes_a, 1, tx, ty, tz),
+            vamp::collision::cuboid_helpers::axis_dot(axes_a, 2, tx, ty, tz)};
+
+        std::array<std::array<DataT, 3>, 3> r{};
+        std::array<std::array<DataT, 3>, 3> abs_r{};
+        const DataT eps = DataT(1e-6F);
+
+        for (auto i = 0U; i < 3; ++i)
+        {
+            for (auto j = 0U; j < 3; ++j)
+            {
+                const auto dot =
+                    dot_3(axes_a[i][0], axes_a[i][1], axes_a[i][2], axes_b[j][0], axes_b[j][1], axes_b[j][2]);
+                r[i][j] = dot;
+                abs_r[i][j] = dot.abs() + eps;
+            }
+        }
+
+        const auto axis_test_a = [&](std::size_t i)
+        {
+            const auto proj = t[i].abs();
+            const auto rb = radii_b[0] * abs_r[i][0] + radii_b[1] * abs_r[i][1] + radii_b[2] * abs_r[i][2];
+            return proj - (radii_a[i] + rb);
+        };
+
+        const auto axis_test_b = [&](std::size_t j)
+        {
+            const auto proj = (t[0] * r[0][j] + t[1] * r[1][j] + t[2] * r[2][j]).abs();
+            const auto ra = radii_a[0] * abs_r[0][j] + radii_a[1] * abs_r[1][j] + radii_a[2] * abs_r[2][j];
+            return proj - (ra + radii_b[j]);
+        };
+
+        const auto axis_test_cross = [&](std::size_t i, std::size_t j)
+        {
+            const auto i1 = (i + 1) % 3;
+            const auto i2 = (i + 2) % 3;
+            const auto j1 = (j + 1) % 3;
+            const auto j2 = (j + 2) % 3;
+
+            const auto ra = radii_a[i1] * abs_r[i2][j] + radii_a[i2] * abs_r[i1][j];
+            const auto rb = radii_b[j1] * abs_r[i][j2] + radii_b[j2] * abs_r[i][j1];
+            const auto proj = (t[i2] * r[i1][j] - t[i1] * r[i2][j]).abs();
+            return proj - (ra + rb);
+        };
+
+        auto max_sep = axis_test_a(0);
+        max_sep = max_sep.max(axis_test_a(1));
+        max_sep = max_sep.max(axis_test_a(2));
+
+        max_sep = max_sep.max(axis_test_b(0));
+        max_sep = max_sep.max(axis_test_b(1));
+        max_sep = max_sep.max(axis_test_b(2));
+
+        for (auto i = 0U; i < 3; ++i)
+        {
+            for (auto j = 0U; j < 3; ++j)
+            {
+                max_sep = max_sep.max(axis_test_cross(i, j));
+            }
+        }
+
+        return max_sep;
+    }
+}  // namespace vamp::collision