From 7e49d935ba8e3404fcb2ebc6c4b55ed150e295af Mon Sep 17 00:00:00 2001
From: Alessio Quaglino <quaglino@google.com>
Date: Thu, 4 May 2023 01:49:41 -0700
Subject: [PATCH] `mj_multiRay`: Reset elevation angle if crossing the
 coordinate singularity.

PiperOrigin-RevId: 529341397
Change-Id: I0ca6ed5b97bcc585b0f426d6fb8a6b4cf9367a2b
---
 src/engine/engine_ray.c        | 8 +++++++-
 test/engine/engine_ray_test.cc | 4 ++--
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/src/engine/engine_ray.c b/src/engine/engine_ray.c
index abfe02e1..04ed3408 100644
--- a/src/engine/engine_ray.c
+++ b/src/engine/engine_ray.c
@@ -56,7 +56,11 @@ static mjtNum longitude(const mjtNum vec[3]) {
 
 // map to elevation angle in spherical coordinates
 static mjtNum latitude(const mjtNum vec[3]) {
-  return mju_atan2(mju_sqrt(vec[0]*vec[0] + vec[1]*vec[1]), vec[2]);
+  if (vec[2]>0) {
+    return mju_atan2(mju_sqrt(vec[0]*vec[0] + vec[1]*vec[1]), vec[2]);
+  } else {
+    return mju_atan2(mju_sqrt(vec[0]*vec[0] + vec[1]*vec[1]), vec[2]) - mjPI;
+  }
 }
 
 
@@ -919,7 +923,9 @@ void mju_multiRayPrepare(const mjModel* m, const mjData* d, const mjtNum pnt[3],
 
       if (AABB[2]-AABB[0] > mjPI) {
         AABB[0] = -mjPI;
+        AABB[1] = -mjPI/2;
         AABB[2] =  mjPI;
+        AABB[3] =  mjPI/2;
       }
 
       if (AABB[3]-AABB[1] > mjPI) {  // SHOULD NOT OCCUR
diff --git a/test/engine/engine_ray_test.cc b/test/engine/engine_ray_test.cc
index fa9e8e37..2726967a 100644
--- a/test/engine/engine_ray_test.cc
+++ b/test/engine/engine_ray_test.cc
@@ -234,9 +234,9 @@ TEST_F(RayTest, EdgeCases) {
   m->geom_aabb[3] = m->geom_aabb[4] = m->geom_aabb[5] = 0;
   mju_multiRayPrepare(m, d, pnt4, NULL, NULL, 1, -1, geom_ba, NULL);
   EXPECT_FLOAT_EQ(geom_ba[0], 0);
-  EXPECT_FLOAT_EQ(geom_ba[1], mjPI/2);
+  EXPECT_FLOAT_EQ(geom_ba[1], -mjPI/2);
   EXPECT_FLOAT_EQ(geom_ba[2], 0);
-  EXPECT_FLOAT_EQ(geom_ba[3], mjPI/2);
+  EXPECT_FLOAT_EQ(geom_ba[3], -mjPI/2);
   mjtNum vec4[] = {1, 0, 0};
   mj_multiRay(m, d, pnt4, vec4, NULL, 1, -1, &rgeomid, &dist, 1);
   EXPECT_FLOAT_EQ(dist, 0.9);