Add mj_multiRayNormal for multi-ray casting with normal computation (not exposed in public header)

PiperOrigin-RevId: 847821078
Change-Id: I616441d3460406a92a10731d1a086af6163e96ad

Author: yuvaltassa

src/engine/engine_ray.c

@@ -955,14 +955,6 @@ static mjtNum mj_raySdfNormal(const mjModel* m, const mjData* d, int g,
   return -1;
 }
 
-
-// intersect ray with signed distance field
-static mjtNum ray_sdf(const mjModel* m, const mjData* d, int g,
-                      const mjtNum pnt[3], const mjtNum vec[3]) {
-  return mj_raySdfNormal(m, d, g, pnt, vec, NULL);
-}
-
-
 // intersect ray with mesh, compute normal if given
 static mjtNum mj_rayMeshNormal(const mjModel* m, const mjData* d, int id, const mjtNum pnt[3],
                                const mjtNum vec[3], mjtNum normal[3]) {
@@ -1488,14 +1480,18 @@ void mju_multiRayPrepare(const mjModel* m, const mjData* d, const mjtNum pnt[3],
 }
 
 
-// Performs single ray intersection
+// Performs single ray intersection, compute normal if given
 static mjtNum mju_singleRay(const mjModel* m, mjData* d, const mjtNum pnt[3], const mjtNum vec[3],
-                            int* ray_eliminate, mjtNum* geom_ba, int geomid[1]) {
+                            int* ray_eliminate, mjtNum* geom_ba, int geomid[1],
+                            mjtNum normal[3]) {
   mjtNum dist, newdist;
+  mjtNum normal_local[3];
+  mjtNum* p_normal = normal ? normal_local : NULL;
 
   // clear result
   dist = -1;
   *geomid = -1;
+  if (normal) mju_zero3(normal);
 
   // get ray spherical coordinates
   mjtNum azimuth = longitude(vec);
@@ -1530,25 +1526,24 @@ static mjtNum mju_singleRay(const mjModel* m, mjData* d, const mjtNum pnt[3], co
         }
       }
 
-      // handle mesh and hfield separately
-      if (m->geom_type[i] == mjGEOM_MESH) {
-        newdist = mj_rayMesh(m, d, i, pnt, vec);
-      } else if (m->geom_type[i] == mjGEOM_HFIELD) {
-        newdist = mj_rayHfield(m, d, i, pnt, vec);
-      } else if (m->geom_type[i] == mjGEOM_SDF) {
-        newdist = ray_sdf(m, d, i, pnt, vec);
-      }
-
-      // otherwise general dispatch
-      else {
-        newdist = mju_rayGeom(d->geom_xpos+3*i, d->geom_xmat+9*i,
-                              m->geom_size+3*i, pnt, vec, m->geom_type[i]);
+      // dispatch to type-specific ray function
+      int type = m->geom_type[i];
+      if (type == mjGEOM_MESH) {
+        newdist = mj_rayMeshNormal(m, d, i, pnt, vec, p_normal);
+      } else if (type == mjGEOM_HFIELD) {
+        newdist = mj_rayHfieldNormal(m, d, i, pnt, vec, p_normal);
+      } else if (type == mjGEOM_SDF) {
+        newdist = mj_raySdfNormal(m, d, i, pnt, vec, p_normal);
+      } else {
+        newdist = mju_rayGeomNormal(d->geom_xpos+3*i, d->geom_xmat+9*i,
+                                    m->geom_size+3*i, pnt, vec, type, p_normal);
       }
 
       // update if closer intersection found
       if (newdist >= 0 && (newdist < dist || dist < 0)) {
         dist = newdist;
         *geomid = i;
+        if (normal) mju_copy3(normal, normal_local);
       }
     }
   }
@@ -1557,10 +1552,10 @@ static mjtNum mju_singleRay(const mjModel* m, mjData* d, const mjtNum pnt[3], co
 }
 
 
-// performs multiple ray intersections with the precomputed bv and flags
-void mj_multiRay(const mjModel* m, mjData* d, const mjtNum pnt[3], const mjtNum* vec,
-                 const mjtByte* geomgroup, mjtByte flg_static, int bodyexclude,
-                 int* geomid, mjtNum* dist, int nray, mjtNum cutoff) {
+// performs multiple ray intersections, compute normals if given
+void mj_multiRayNormal(const mjModel* m, mjData* d, const mjtNum pnt[3], const mjtNum* vec,
+                       const mjtByte* geomgroup, mjtByte flg_static, int bodyexclude,
+                       int* geomid, mjtNum* dist, mjtNum* normal, int nray, mjtNum cutoff) {
   mj_markStack(d);
 
   // allocate source
@@ -1576,9 +1571,20 @@ void mj_multiRay(const mjModel* m, mjData* d, const mjtNum pnt[3], const mjtNum*
     if (mju_dot3(vec+3*i, vec+3*i) < mjMINVAL) {
       dist[i] = -1;
     } else {
-      dist[i] = mju_singleRay(m, d, pnt, vec+3*i, geom_eliminate, geom_ba, geomid+i);
+      dist[i] = mju_singleRay(m, d, pnt, vec+3*i, geom_eliminate, geom_ba, geomid+i,
+                            normal ? normal+3*i : NULL);
     }
   }
 
   mj_freeStack(d);
 }
+
+
+// performs multiple ray intersections with the precomputed bv and flags
+void mj_multiRay(const mjModel* m, mjData* d, const mjtNum pnt[3], const mjtNum vec[3],
+                 const mjtByte* geomgroup, mjtByte flg_static, int bodyexclude,
+                 int* geomid, mjtNum* dist, int nray, mjtNum cutoff) {
+  mj_multiRayNormal(m, d, pnt, vec, geomgroup, flg_static, bodyexclude,
+                    geomid, dist, NULL, nray, cutoff);
+}
+

src/engine/engine_ray.h

@@ -30,11 +30,18 @@ MJAPI void mju_multiRayPrepare(const mjModel* m, const mjData* d,
                                int* geom_eliminate);
 
 // intersect multiple rays emanating from a single source
-// similar semantics to mj_ray, but vec is an array of (nray x 3) directions.
+// similar semantics to mj_ray, but vec is (nray x 3) and dist is (nray).
 MJAPI void mj_multiRay(const mjModel* m, mjData* d, const mjtNum pnt[3], const mjtNum* vec,
                        const mjtByte* geomgroup, mjtByte flg_static, int bodyexclude,
                        int* geomid, mjtNum* dist, int nray, mjtNum cutoff);
 
+// intersect multiple rays, compute normals if given
+// similar semantics to mj_rayNormal, but vec, normal and dist are arrays.
+MJAPI void mj_multiRayNormal(const mjModel* m, mjData* d, const mjtNum pnt[3], const mjtNum* vec,
+                             const mjtByte* geomgroup, mjtByte flg_static, int bodyexclude,
+                             int* geomid, mjtNum* dist, mjtNum* normal, int nray, mjtNum cutoff);
+
+
 // intersect ray (pnt+x*vec, x>=0) with visible geoms, except geoms on bodyexclude
 //  return geomid and distance (x) to nearest surface, or -1 if no intersection
 //  geomgroup, flg_static are as in mjvOption; geomgroup==NULL skips group exclusion

test/engine/engine_ray_test.cc

@@ -187,8 +187,8 @@ TEST_F(RayTest, MultiRayEqualsSingleRay) {
   constexpr int N = 80;
   constexpr int M = 60;
   mjtNum vec[3*N*M];
-  mjtNum pnt[3] = {1, 2, 3};
-  mjtNum cone[4][3] = {{1, 1, -1}, {1, 1, 1}, {1, -1, -1}, {1, -1, 1}};
+  mjtNum pnt[3] = {-1, 0, 0};
+  mjtNum cone[4][3] = {{1, .2, -.2}, {1, .2, .2}, {1, -.2, -.2}, {1, -.2, .2}};
   memset(vec, 0, 3*N*M*sizeof(mjtNum));
 
   for (int i = 0; i < N; ++i) {
@@ -211,15 +211,75 @@ TEST_F(RayTest, MultiRayEqualsSingleRay) {
   // compare results with single ray function
   mjtNum dist;
   int rgeomid;
-
+  int nhits = 0;
   for (int i = 0; i < N; ++i) {
     for (int j = 0; j < M; ++j) {
       int idx = i * M + j;
       dist = mj_ray(m, d, pnt, vec + 3 * idx, NULL, 1, -1, &rgeomid);
       EXPECT_FLOAT_EQ(dist, dist_multiray[idx]);
       EXPECT_EQ(rgeomid, rgeomid_multiray[idx]);
+      nhits += dist >= 0;
+    }
+  }
+  EXPECT_GT(nhits, 10);
+
+  mj_deleteData(d);
+  mj_deleteModel(m);
+}
+
+TEST_F(RayTest, MultiRayNormalEqualsSingleRayNormal) {
+  char error[1024];
+  mjModel* m = LoadModelFromString(kRayCastingModel, error, sizeof(error));
+  ASSERT_THAT(m, NotNull()) << error;
+  mjData* d = mj_makeData(m);
+  ASSERT_THAT(d, NotNull());
+  mj_forward(m, d);
+
+  // create ray array
+  constexpr int N = 80;
+  constexpr int M = 60;
+  mjtNum vec[3*N*M];
+  mjtNum pnt[3] = {-1, 0, 0};
+  mjtNum cone[4][3] = {{1, .2, -.2}, {1, .2, .2}, {1, -.2, -.2}, {1, -.2, .2}};
+  memset(vec, 0, 3*N*M*sizeof(mjtNum));
+
+  for (int i = 0; i < N; ++i) {
+    for (int j = 0; j < M; ++j) {
+      for (int k = 0; k < 3; ++k) {
+        vec[3 * (i * M + j) + k] =           i * cone[0][k] / (N - 1) +
+                                             j * cone[1][1] / (M - 1) +
+                                   (N - i - 1) * cone[2][k] / (N - 1) +
+                                   (M - j - 1) * cone[3][k] / (M - 1);
+      }
+    }
+  }
+
+  // compute intersections with multiray normal function
+  mjtNum dist_multiray[N*M];
+  int rgeomid_multiray[N*M];
+  mjtNum normal_multiray[3*N*M];
+  mj_multiRayNormal(m, d, pnt, vec, NULL, 1, -1, rgeomid_multiray,
+                    dist_multiray, normal_multiray, N * M, mjMAXVAL);
+
+  // compare results with single ray normal function
+  mjtNum dist;
+  int rgeomid;
+  mjtNum normal[3];
+  int nhits = 0;
+  for (int i = 0; i < N; ++i) {
+    for (int j = 0; j < M; ++j) {
+      int idx = i * M + j;
+      dist = mj_rayNormal(m, d, pnt, vec + 3 * idx, NULL, 1, -1, &rgeomid,
+                          normal);
+      EXPECT_FLOAT_EQ(dist, dist_multiray[idx]);
+      EXPECT_EQ(rgeomid, rgeomid_multiray[idx]);
+      EXPECT_FLOAT_EQ(normal[0], normal_multiray[3*idx]);
+      EXPECT_FLOAT_EQ(normal[1], normal_multiray[3*idx + 1]);
+      EXPECT_FLOAT_EQ(normal[2], normal_multiray[3*idx + 2]);
+      nhits += dist >= 0;
     }
   }
+  EXPECT_GT(nhits, 10);
 
   mj_deleteData(d);
   mj_deleteModel(m);