From 69455d6a46b8e9c7ad8af67083abc9de058818d2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Romain=20BR=C3=89GIER?= <romain.bregier@naverlabs.com>
Date: Thu, 21 Mar 2024 11:47:43 +0100
Subject: [PATCH] v1.4.4

---
 docsource/source/index.rst |  2 ++
 roma/utils.py              | 20 +++++++++++++++++++-
 setup.py                   |  2 +-
 test/test_utils.py         | 13 +++++++++++++
 4 files changed, 35 insertions(+), 2 deletions(-)

diff --git a/docsource/source/index.rst b/docsource/source/index.rst
index c08cedd..b3fd4c5 100644
--- a/docsource/source/index.rst
+++ b/docsource/source/index.rst
@@ -231,6 +231,8 @@ Bits of code were adapted from SciPy. Documentation is generated, distributed an
 
 Changelog
 ==========
+Version 1.4.4:
+    - Added :func:`~roma.utils.identity_quat()`.
 Version 1.4.3: 
     - Fix normalization bug in :func:`~roma.utils.quat_composition()` (thanks jamiesalter for reporting).
 Version 1.4.2:
diff --git a/roma/utils.py b/roma/utils.py
index 69917eb..c31093b 100644
--- a/roma/utils.py
+++ b/roma/utils.py
@@ -97,6 +97,25 @@ def random_rotvec(size = tuple(), dtype=torch.float, device=None):
     quat = random_unitquat(size, dtype=dtype, device=device)
     return roma.mappings.unitquat_to_rotvec(quat)
 
+def identity_quat(size = tuple(), dtype=torch.float, device=None):
+    """
+    Return a batch of identity unit quaternions.
+
+    Args:
+        size (tuple or int): batch size. Use for example ``tuple()`` to generate a single element, and ``(5,2)`` to generate a 5x2 batch.
+    Returns:
+        batch of identity quaternions (size x 4 tensor, XYZW convention).
+    Note:
+        All returned batch quaternions refer to the same memory location.
+        Consider cloning the output tensor prior performing any in-place operations.
+    """
+    if type(size) == int:
+        size = (size,)
+    quat = torch.zeros(4, dtype=dtype, device=device)
+    quat[...,-1] = 1.
+    quat = quat.expand(list(size) + [-1])
+    return quat
+
 def rotmat_cosine_angle(R):
     """
     Returns the cosine angle of the input 3x3 rotation matrix R.
@@ -110,7 +129,6 @@ def rotmat_cosine_angle(R):
     assert R.shape[-2:] == (3,3), "Expecting a ...x3x3 batch of rotation matrices"
     return  0.5 * (R[...,0,0] + R[...,1,1] + R[...,2,2] - 1.0)
 
-
 _ONE_OVER_2SQRT2 = 1.0 / (2 * np.sqrt(2))
 def rotmat_geodesic_distance(R1, R2, clamping=1.0):
     """
diff --git a/setup.py b/setup.py
index 0fb36cb..a07036a 100644
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@
 
 setuptools.setup(
     name="roma",
-    version="1.4.3",
+    version="1.4.4",
     author="Romain Brégier",
     author_email="romain.bregier@naverlabs.com",
     description="A lightweight library to deal with 3D rotations in PyTorch.",
diff --git a/test/test_utils.py b/test/test_utils.py
index 2d299eb..b282c02 100644
--- a/test/test_utils.py
+++ b/test/test_utils.py
@@ -64,6 +64,19 @@ def test_other_geodesic_distance(self):
             alpha_rotvec = roma.rotvec_geodesic_distance(rotvec1, rotvec2)
             self.assertTrue(is_close(alpha_rotvec, alpha_q))
 
+    def test_identity_quat(self):
+        q = roma.identity_quat()
+        self.assertTrue(q.shape == (4,))
+        self.assertTrue(is_close(q, roma.quat_inverse(q)))
+
+        q = roma.identity_quat(5)
+        self.assertTrue(q.shape == (5,4))
+        self.assertTrue(is_close(q, roma.quat_inverse(q)))
+
+        q = roma.identity_quat((3,2))
+        self.assertTrue(q.shape == (3,2,4))
+        self.assertTrue(is_close(q, roma.quat_inverse(q)))
+
     def test_random_unitquat(self):
         q = roma.random_unitquat((3,5))
         self.assertTrue(q.shape == (3,5,4))