From ba9ae88c8b3d6e9a080f6c648de3678b745a2284 Mon Sep 17 00:00:00 2001
From: Victor Reijgwart <victorreijgwart@gmail.com>
Date: Thu, 12 Dec 2024 19:43:32 +0100
Subject: [PATCH] Expose map transformation method

---
 examples/python/edit/transform_map.py | 20 ++++++++++++++++++++
 library/python/src/edit.cc            | 16 ++++++++++++++++
 library/python/src/measurements.cc    |  3 +++
 3 files changed, 39 insertions(+)
 create mode 100644 examples/python/edit/transform_map.py

diff --git a/examples/python/edit/transform_map.py b/examples/python/edit/transform_map.py
new file mode 100644
index 000000000..4d70ce5c6
--- /dev/null
+++ b/examples/python/edit/transform_map.py
@@ -0,0 +1,20 @@
+import os
+import numpy as np
+import pywavemap as wave
+
+# Load the map
+user_home = os.path.expanduser('~')
+input_map_path = os.path.join(user_home, "your_map.wvmp")
+your_map = wave.Map.load(input_map_path)
+
+# Define a transformation that flips the map upside down, for illustration
+translation = np.array([0.0, 0.0, 0.0])
+rotation = wave.Rotation(w=0.0, x=1.0, y=0.0, z=0.0)
+transformation = wave.Pose(rotation, translation)
+
+# Transform the map
+your_map = wave.edit.transform(your_map, transformation)
+
+# Save the map
+output_map_path = os.path.join(user_home, "your_map_transformed.wvmp")
+your_map.store(output_map_path)
diff --git a/library/python/src/edit.cc b/library/python/src/edit.cc
index 00ddcb98e..df69e465f 100644
--- a/library/python/src/edit.cc
+++ b/library/python/src/edit.cc
@@ -3,10 +3,12 @@
 #include <memory>
 
 #include <nanobind/eigen/dense.h>
+#include <nanobind/stl/unique_ptr.h>
 #include <wavemap/core/map/hashed_chunked_wavelet_octree.h>
 #include <wavemap/core/map/hashed_wavelet_octree.h>
 #include <wavemap/core/utils/edit/crop.h>
 #include <wavemap/core/utils/edit/multiply.h>
+#include <wavemap/core/utils/edit/transform.h>
 
 using namespace nb::literals;  // NOLINT
 
@@ -45,5 +47,19 @@ void add_edit_module(nb::module_& m_edit) {
         edit::multiply(map, multiplier, std::make_shared<ThreadPool>());
       },
       "map"_a, "multiplier"_a);
+
+  // Map transformation methods
+  m_edit.def(
+      "transform",
+      [](HashedWaveletOctree& B_map, const Transformation3D& T_AB) {
+        return edit::transform(B_map, T_AB, std::make_shared<ThreadPool>());
+      },
+      "map"_a, "transformation"_a);
+  m_edit.def(
+      "transform",
+      [](HashedChunkedWaveletOctree& B_map, const Transformation3D& T_AB) {
+        return edit::transform(B_map, T_AB, std::make_shared<ThreadPool>());
+      },
+      "map"_a, "transformation"_a);
 }
 }  // namespace wavemap
diff --git a/library/python/src/measurements.cc b/library/python/src/measurements.cc
index 1c51f7052..5f5d39c4f 100644
--- a/library/python/src/measurements.cc
+++ b/library/python/src/measurements.cc
@@ -13,6 +13,9 @@ void add_measurement_bindings(nb::module_& m) {
   nb::class_<Rotation3D>(m, "Rotation",
                          "A class representing rotations in 3D space.")
       .def(nb::init<Rotation3D::RotationMatrix>(), "rotation_matrix"_a)
+      .def(nb::init<FloatingPoint, FloatingPoint, FloatingPoint,
+                    FloatingPoint>(),
+           "w"_a, "x"_a, "y"_a, "z"_a)
       .def("inverse", &Rotation3D::inverse, "Compute the rotation's inverse.");
   nb::class_<Transformation3D>(m, "Pose",
                                "A class representing poses in 3D space.")