From 327868b86e0d69a3af01299bc5371a56b9ef2ed9 Mon Sep 17 00:00:00 2001
From: Patrick Labatut <plabatut@fb.com>
Date: Fri, 13 Mar 2020 04:28:30 -0700
Subject: [PATCH] Add utility function to tesselate a torus

Summary: Add utility function to tesselate a torus, to be used in more complex mesh I/O benchmarks

Reviewed By: bottler

Differential Revision: D20390724

fbshipit-source-id: 882bbbe9cac81cf340a34495b9aa66e3c1ddeebc
---
 pytorch3d/utils/__init__.py |  1 +
 pytorch3d/utils/torus.py    | 71 +++++++++++++++++++++++++++++++++++++
 2 files changed, 72 insertions(+)
 create mode 100644 pytorch3d/utils/torus.py

diff --git a/pytorch3d/utils/__init__.py b/pytorch3d/utils/__init__.py
index c0c156e9..bcf5f27f 100644
--- a/pytorch3d/utils/__init__.py
+++ b/pytorch3d/utils/__init__.py
@@ -1,5 +1,6 @@
 # Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
 
 from .ico_sphere import ico_sphere
+from .torus import torus
 
 __all__ = [k for k in globals().keys() if not k.startswith("_")]
diff --git a/pytorch3d/utils/torus.py b/pytorch3d/utils/torus.py
new file mode 100644
index 00000000..9280e640
--- /dev/null
+++ b/pytorch3d/utils/torus.py
@@ -0,0 +1,71 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+from itertools import tee
+from math import cos, pi, sin
+from typing import Iterator, Optional, Tuple
+import torch
+
+from pytorch3d.structures.meshes import Meshes
+
+
+# Make an iterator over the adjacent pairs: (-1, 0), (0, 1), ..., (N - 2, N - 1)
+def _make_pair_range(N: int) -> Iterator[Tuple[int, int]]:
+    i, j = tee(range(-1, N))
+    next(j, None)
+    return zip(i, j)
+
+
+def torus(
+    r: float,
+    R: float,
+    sides: int,
+    rings: int,
+    device: Optional[torch.device] = None,
+) -> Meshes:
+    """
+    Create vertices and faces for a torus.
+
+    Args:
+        r: Inner radius of the torus.
+        R: Outer radius of the torus.
+        sides: Number of inner divisions.
+        rings: Number of outer divisions.
+        device: Device on which the outputs will be allocated.
+
+    Returns:
+        Meshes object with the generated vertices and faces.
+    """
+    if not (sides > 0):
+        raise ValueError("sides must be > 0.")
+    if not (rings > 0):
+        raise ValueError("rings must be > 0.")
+    device = device if device else torch.device("cpu")
+
+    verts = []
+    for i in range(rings):
+        # phi ranges from 0 to 2 pi (rings - 1) / rings
+        phi = 2 * pi * i / rings
+        for j in range(sides):
+            # theta ranges from 0 to 2 pi (sides - 1) / sides
+            theta = 2 * pi * j / sides
+            x = (R + r * cos(theta)) * cos(phi)
+            y = (R + r * cos(theta)) * sin(phi)
+            z = r * sin(theta)
+            # This vertex has index i * sides + j
+            verts.append([x, y, z])
+
+    faces = []
+    for i0, i1 in _make_pair_range(rings):
+        index0 = (i0 % rings) * sides
+        index1 = (i1 % rings) * sides
+        for j0, j1 in _make_pair_range(sides):
+            index00 = index0 + (j0 % sides)
+            index01 = index0 + (j1 % sides)
+            index10 = index1 + (j0 % sides)
+            index11 = index1 + (j1 % sides)
+            faces.append([index00, index10, index11])
+            faces.append([index11, index01, index00])
+
+    verts_list = [torch.tensor(verts, dtype=torch.float32, device=device)]
+    faces_list = [torch.tensor(faces, dtype=torch.int64, device=device)]
+    return Meshes(verts_list, faces_list)