imt-ring 1.6.8__py3-none-any.whl → 1.6.10__py3-none-any.whl

{imt_ring-1.6.8.dist-info → imt_ring-1.6.10.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: imt-ring
-Version: 1.6.8
+Version: 1.6.10
 Summary: RING: Recurrent Inertial Graph-based Estimator
 Author-email: Simon Bachhuber <simon.bachhuber@fau.de>
 Project-URL: Homepage, https://github.com/SimiPixel/ring

{imt_ring-1.6.8.dist-info → imt_ring-1.6.10.dist-info}/RECORD

@@ -1,6 +1,6 @@
 ring/__init__.py,sha256=k7tL-XgggUwWxHCXyv60rQn-OcXHPg82QcIUkKLEd-c,5057
 ring/algebra.py,sha256=F0GwbP8LQP5qGVkoMUYJmkp9Hn2nKAVIkCVYDEjNjGU,3128
-ring/base.py,sha256=BGAJE3PSOUnTHte4UesJc1J7MQraIEiVpStkhrgXhaI,33245
+ring/base.py,sha256=MkkziQx01sdMOpB8MFUDFgFlZUrXCFjpb8hS9yKHUyM,33751
 ring/maths.py,sha256=qPHH6TpHCK3TgExI98gNEySoSRKOwteN9McUlyUFipI,12207
 ring/spatial.py,sha256=nmZ-UhRanhyM34bez8uCS4wMwaKqLkuEbgKGP5XNH60,2351
 ring/algorithms/__init__.py,sha256=IiK9EN5Xgs3dB075-A-H-Yad0Z7vzvKIJF2g6X_-C_8,1224
@@ -14,8 +14,8 @@ ring/algorithms/custom_joints/rr_imp_joint.py,sha256=_YJK0p8_0MHFtr1NuGnNZoxTbwa
 ring/algorithms/custom_joints/rr_joint.py,sha256=jnRtjtOCALMaq2_0bcu2d7qgfQ6etXpoh43MioRaDmY,1000
 ring/algorithms/custom_joints/suntay.py,sha256=tOEGM304XciHO4pmvxr4faA4xXVO4N2HlPdFmXKbcrw,16726
 ring/algorithms/generator/__init__.py,sha256=bF-CW3x2x-o6KWESKy-DuxzZPh3UNSjJb_MaAcSHGsQ,277
-ring/algorithms/generator/base.py,sha256=YKUr9UnjFbWAYjP8V1j0FaKGVpK92W6E1iuoet7qBNg,14522
-ring/algorithms/generator/batch.py,sha256=ylootnXmj-JyuB_f5OCknHst9wFKO3gkjQbMrFNXY2g,2513
+ring/algorithms/generator/base.py,sha256=JHCTbHtmYEdmsHyQnJN9vMP6rzrlBoqzTZ27c3zhCDI,15655
+ring/algorithms/generator/batch.py,sha256=9yFxVv11hij-fJXGPxA3zEh1bE2_jrZk0R7kyGaiM5c,2551
 ring/algorithms/generator/finalize_fns.py,sha256=LUw1Wc2YrmMRRh4RF704ob3bZOXktAZAbbLoBm_p1yw,9131
 ring/algorithms/generator/motion_artifacts.py,sha256=2VJbldVDbI3PSyboshIbtYvSAKzBBwGV7cQfYjqvluM,9167
 ring/algorithms/generator/pd_control.py,sha256=XJ_Gd5AkIRh-jBrMfQyMXjVwhx2gCNHznjzFbmAwhZs,5767
@@ -83,7 +83,7 @@ ring/utils/randomize_sys.py,sha256=G_vBIo0OwQkXL2u0djwbaoaeb02C4LQCTNNloOYIU2M,3
 ring/utils/utils.py,sha256=oGC7kh19s5zvmnUvWy8B3fBl9loVU58ppz91osk2m3w,6550
 ring/utils/register_gym_envs/__init__.py,sha256=PtPIRBQJ16339xZ9G9VpvqrvcGbQ_Pk_SUz4tQPa9nQ,94
 ring/utils/register_gym_envs/saddle.py,sha256=tA5CyW_akSXyDm0xJ83CtOrUMVElH0f9vZtEDDJQalI,4422
-imt_ring-1.6.8.dist-info/METADATA,sha256=FxO8BZbMRGagmYW1cQCHPuUIuiXxsIr8F-zlBkBVHNQ,3820
-imt_ring-1.6.8.dist-info/WHEEL,sha256=HiCZjzuy6Dw0hdX5R3LCFPDmFS4BWl8H-8W39XfmgX4,91
-imt_ring-1.6.8.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
-imt_ring-1.6.8.dist-info/RECORD,,
+imt_ring-1.6.10.dist-info/METADATA,sha256=j5LtBbakAQBMu_XP_642TYIaq0PFkG0S9h49nrnReoc,3821
+imt_ring-1.6.10.dist-info/WHEEL,sha256=Mdi9PDNwEZptOjTlUcAth7XJDFtKrHYaQMPulZeBCiQ,91
+imt_ring-1.6.10.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
+imt_ring-1.6.10.dist-info/RECORD,,

{imt_ring-1.6.8.dist-info → imt_ring-1.6.10.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (72.2.0)
+Generator: setuptools (73.0.1)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

ring/algorithms/generator/base.py

@@ -6,6 +6,7 @@ import jax
 import jax.numpy as jnp
 import numpy as np
 import tree_utils
+from tree_utils import PyTree
 
 from ring import base
 from ring import utils
@@ -143,9 +144,7 @@ class RCMG:
 
         return n_calls
 
-    def to_list(
-        self, sizes: int | list[int] = 1, seed: int = 1
-    ) -> list[tree_utils.PyTree[np.ndarray]]:
+    def _generators_ncalls(self, sizes: int | list[int] = 1):
         "Returns list of unbatched sequences as numpy arrays."
         repeats = self._compute_repeats(sizes)
         sizes = list(jnp.array(repeats) * jnp.array(self._size_of_generators))
@@ -165,7 +164,45 @@ class RCMG:
                 batch.generators_lazy([self.gens[i]], [reduced_repeats[i]], jits[i])
             )
 
-        return batch.generators_eager_to_list(gens, n_calls, seed, self._disable_tqdm)
+        return gens, n_calls
+
+    def to_list(
+        self, sizes: int | list[int] = 1, seed: int = 1
+    ) -> list[tree_utils.PyTree[np.ndarray]]:
+        "Returns list of unbatched sequences as numpy arrays."
+        gens, n_calls = self._generators_ncalls(sizes)
+
+        data = []
+        batch.generators_eager(
+            gens, n_calls, lambda d: data.extend(d), seed, self._disable_tqdm
+        )
+        return data
+
+    def to_folder(
+        self,
+        path: str,
+        sizes: int | list[int] = 1,
+        seed: int = 1,
+        overwrite: bool = True,
+        file_prefix: str = "seq",
+        save_fn: Callable[[PyTree[np.ndarray], str], None] = utils.pickle_save,
+        verbose: bool = True,
+    ):
+
+        i = 0
+
+        def callback(data: list[PyTree[np.ndarray]]) -> None:
+            nonlocal i
+            data = utils.replace_elements_w_nans(data, verbose=verbose)
+            for d in data:
+                file = utils.parse_path(
+                    path, file_prefix + str(i), file_exists_ok=overwrite
+                )
+                save_fn(d, file)
+                i += 1
+
+        gens, n_calls = self._generators_ncalls(sizes)
+        batch.generators_eager(gens, n_calls, callback, seed, self._disable_tqdm)
 
     def to_pickle(
         self,

ring/algorithms/generator/batch.py

@@ -1,8 +1,10 @@
+from typing import Callable
+
 import jax
 import jax.numpy as jnp
 import numpy as np
 from tqdm import tqdm
-import tree_utils
+from tree_utils import PyTree
 
 from ring import utils
 from ring.algorithms.generator import types
@@ -50,15 +52,15 @@ def generators_lazy(
     return generator
 
 
-def generators_eager_to_list(
+def generators_eager(
     generators: list[types.BatchedGenerator],
     n_calls: list[int],
+    callback: Callable[[list[PyTree[np.ndarray]]], None],
     seed: int = 1,
     disable_tqdm: bool = False,
-) -> list[tree_utils.PyTree]:
+) -> None:
 
     key = jax.random.PRNGKey(seed)
-    data = []
     for gen, n_call in tqdm(
         zip(generators, n_calls),
         desc="executing generators",
@@ -81,6 +83,4 @@ def generators_eager_to_list(
 
             sample_flat, _ = jax.tree_util.tree_flatten(sample)
             size = 1 if len(sample_flat) == 0 else sample_flat[0].shape[0]
-            data.extend([jax.tree_map(lambda a: a[i], sample) for i in range(size)])
-
-    return data
+            callback([jax.tree_map(lambda a: a[i], sample) for i in range(size)])

ring/base.py

@@ -590,12 +590,14 @@ class System(_Base):
 
         return sys
 
-    def findall_imus(self) -> list[str]:
-        return [name for name in self.link_names if name[:3] == "imu"]
+    def findall_imus(self, names: bool = True) -> list[str] | list[int]:
+        bodies = [name for name in self.link_names if name[:3] == "imu"]
+        return bodies if names else [self.name_to_idx(n) for n in bodies]
 
-    def findall_segments(self) -> list[str]:
-        imus = self.findall_imus()
-        return [name for name in self.link_names if name not in imus]
+    def findall_segments(self, names: bool = True) -> list[str] | list[int]:
+        imus = self.findall_imus(names=True)
+        bodies = [name for name in self.link_names if name not in imus]
+        return bodies if names else [self.name_to_idx(n) for n in bodies]
 
     def _bodies_indices_to_bodies_name(self, bodies: list[int]) -> list[str]:
         return [self.idx_to_name(i) for i in bodies]
@@ -615,6 +617,11 @@ class System(_Base):
         bodies = [i for i, _typ in enumerate(self.link_types) if _typ == typ]
         return self._bodies_indices_to_bodies_name(bodies) if names else bodies
 
+    def children(self, name: str, names: bool = False) -> list[int] | list[str]:
+        p = self.name_to_idx(name)
+        bodies = [i for i in range(self.num_links()) if self.link_parents[i] == p]
+        return bodies if (not names) else [self.idx_to_name(i) for i in bodies]
+
     def scan(self, f: Callable, in_types: str, *args, reverse: bool = False):
         """Scan `f` along each link in system whilst carrying along state.