From 6ed0d59002777b6bae72a295534072405b841c3d Mon Sep 17 00:00:00 2001
From: BuildingAtom <84liadam@gmail.com>
Date: Thu, 15 Feb 2024 01:17:02 -0500
Subject: [PATCH] backport offline jrs loader & add reachablesets documentation

---
 docs/source/conf.py                           |  2 +-
 docs/source/index.rst                         |  1 +
 docs/source/reachablesets.rst                 | 14 ++++++
 zonopyrobots/__init__.py                      |  1 +
 .../joint_reachable_set/offline_jrs.py        | 44 +++++++++++++++++++
 5 files changed, 61 insertions(+), 1 deletion(-)
 create mode 100644 docs/source/reachablesets.rst
 create mode 100644 zonopyrobots/joint_reachable_set/offline_jrs.py

diff --git a/docs/source/conf.py b/docs/source/conf.py
index 95bc8154..1c220cab 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -71,7 +71,7 @@
 html_title = f"{project} v{version}"
 
 # Version switcher code
-json_url = "https://roahmlab.github.io/zonopy/versions.json"
+json_url = "https://roahmlab.github.io/zonopy-robots/versions.json"
 
 html_theme_options = {
     "path_to_docs": version_match,
diff --git a/docs/source/index.rst b/docs/source/index.rst
index f06460b2..1ad5b532 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -12,6 +12,7 @@ This project is still early in development, so much of the API is subject to cha
 
    robots
    trajectories
+   reachablesets
    kinematics
    dynamics
    utils
diff --git a/docs/source/reachablesets.rst b/docs/source/reachablesets.rst
new file mode 100644
index 00000000..47438367
--- /dev/null
+++ b/docs/source/reachablesets.rst
@@ -0,0 +1,14 @@
+Reachable Sets
+==============
+.. automodule:: zonopyrobots
+    :members:
+    :show-inheritance:
+    :noindex:
+.. currentmodule:: zonopyrobots
+
+.. autosummary::
+    :toctree: generated
+    :nosignatures:
+
+    JrsGenerator
+    OfflineJRS
diff --git a/zonopyrobots/__init__.py b/zonopyrobots/__init__.py
index 823627c2..bf3b4cfe 100644
--- a/zonopyrobots/__init__.py
+++ b/zonopyrobots/__init__.py
@@ -6,6 +6,7 @@
 import zonopyrobots.trajectory as trajectory
 
 from zonopyrobots.joint_reachable_set.gen_jrs import JrsGenerator
+from zonopyrobots.joint_reachable_set.offline_jrs import OfflineJRS
 from zonopyrobots.joint_reachable_set.jrs_trig.load_jrs_trig import *
 from zonopyrobots.joint_reachable_set.jrs_trig.process_jrs_trig import *
 
diff --git a/zonopyrobots/joint_reachable_set/offline_jrs.py b/zonopyrobots/joint_reachable_set/offline_jrs.py
new file mode 100644
index 00000000..05afae2a
--- /dev/null
+++ b/zonopyrobots/joint_reachable_set/offline_jrs.py
@@ -0,0 +1,44 @@
+from __future__ import annotations
+import torch
+from .jrs_trig.process_jrs_trig import process_batch_JRS_trig as _process_batch_JRS_trig
+from .jrs_trig.load_jrs_trig import preload_batch_JRS_trig as _preload_batch_JRS_trig
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import Tuple, Union
+    from numpy import ndarray
+    from torch import Tensor
+    Array = Union[Tensor, ndarray]
+
+class OfflineJRS:
+    def __init__(
+        self,
+        device: torch.device = 'cpu',
+        dtype: torch.dtype = torch.float,
+        ):
+        """ Wrapper for preloading and processing JRS tensors """
+        from .jrs_trig.load_jrs_trig import g_ka
+        self.jrs_tensor = _preload_batch_JRS_trig(device=device, dtype=dtype)
+        self.g_ka = g_ka
+        self.device = device
+        self.dtype = dtype
+
+    def __call__(
+        self,
+        qpos: Array,
+        qvel: Array,
+        joint_axes: Array,
+        ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """ Returns the JRS and the corresponding rotatotopes for a given configuration and velocity
+        Args:
+            qpos (torch.Tensor): The configuration of the robot
+            qvel (torch.Tensor): The velocity of the robot
+            joint_axes (torch.Tensor): The joint axes of the robot
+
+        Returns:
+            Tuple[torch.Tensor, torch.Tensor]: The JRS and the corresponding rotatotopes
+        """
+        qpos = torch.as_tensor(qpos, dtype=self.dtype, device=self.device)
+        qvel = torch.as_tensor(qvel, dtype=self.dtype, device=self.device)
+        joint_axes = torch.as_tensor(joint_axes, dtype=self.dtype, device=self.device)
+        return _process_batch_JRS_trig(self.jrs_tensor, qpos, qvel, joint_axes)