imt-ring 1.3.0__py3-none-any.whl → 1.3.2__py3-none-any.whl

{imt_ring-1.3.0.dist-info → imt_ring-1.3.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: imt-ring
-Version: 1.3.0
+Version: 1.3.2
 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.3.0.dist-info → imt_ring-1.3.2.dist-info}/RECORD

@@ -1,6 +1,6 @@
 ring/__init__.py,sha256=iNvbAZi7Qfa69IbL1z4lB7zHL8WusV5fBrKah2la-Gc,1566
 ring/algebra.py,sha256=F0GwbP8LQP5qGVkoMUYJmkp9Hn2nKAVIkCVYDEjNjGU,3128
-ring/base.py,sha256=lG1SAVi6VpJT20Xvdhv_NrObMb4008leEqPEaQ0anR8,33566
+ring/base.py,sha256=gqdXejZ4E4liB5mZ6gPof3EDYTThlfro2MQs0bc5eOM,33530
 ring/maths.py,sha256=jJr_kr78-XDce8B4tXQ2Li-jBntVQhaS8csxglCsj8A,12193
 ring/spatial.py,sha256=nmZ-UhRanhyM34bez8uCS4wMwaKqLkuEbgKGP5XNH60,2351
 ring/algorithms/__init__.py,sha256=t3YXcgqMJxadUjFiILVD0HlQRPLdrQyc8aKiB36w0vE,1701
@@ -9,13 +9,13 @@ ring/algorithms/dynamics.py,sha256=nqq5I0RYSbHNlGiLMlohz08IfL9Njsrid4upDnwkGbI,1
 ring/algorithms/jcalc.py,sha256=oqSiwz3Be1VfIpmJXEFTNM_9_o3tyuTtyZt2aqttyN4,28213
 ring/algorithms/kinematics.py,sha256=DOboHI517Vx0pRJUFZtZPmK_qFaiKiQe-37B-M0aC-c,7422
 ring/algorithms/sensors.py,sha256=Y3Wo9qj3BWKoIHB0V04QwyD-Z5m4BrAjfBX8Pn6y9Lg,18005
-ring/algorithms/custom_joints/__init__.py,sha256=33WBnaBJMtq3vVcpMm7zmyeMrLY9PyV_8-wk5oSF65g,227
+ring/algorithms/custom_joints/__init__.py,sha256=_kUyC4TbzjngTQrJVtS6JBKPzTMNbH27jVRJYXViepI,270
 ring/algorithms/custom_joints/rr_imp_joint.py,sha256=a3JT0w7pB94kZ95eBR8ZO853eSeyjFoiXmhYlaXoHDE,2392
 ring/algorithms/custom_joints/rr_joint.py,sha256=jnRtjtOCALMaq2_0bcu2d7qgfQ6etXpoh43MioRaDmY,1000
-ring/algorithms/custom_joints/suntay.py,sha256=d0Z54tIXiepMixE40W5H8JKxrT5U6VskPm2L2kKnQPw,13680
+ring/algorithms/custom_joints/suntay.py,sha256=3aFDfqdC2vUAhD30kkuQltgU_WZmYDyVhKPSoEotEYo,15292
 ring/algorithms/generator/__init__.py,sha256=p4ucl0zQtp5NwNoXIRjmTzGGRu2WOAWFfNmYRPwQles,912
-ring/algorithms/generator/base.py,sha256=QDmzMAgtaK5M9WDl39qjXYfBa99d83vCPWEkYYmsplk,14952
-ring/algorithms/generator/batch.py,sha256=MZurZmQDH1vncoNbCspVNGNlfP0R87J6_HC7MMIqQ6A,8478
+ring/algorithms/generator/base.py,sha256=zmrRK_I6BWoo4WbEcEVK7iFKdPfetc6txs7U8iu1xEk,14771
+ring/algorithms/generator/batch.py,sha256=BGzmwH1AItXjPRyHtsYnAfYnoogw8jxhng9oyVw72lw,9019
 ring/algorithms/generator/motion_artifacts.py,sha256=aKdkZU5OF4_aKyL4Yo-ftZRwrDCve1LuuREGAUlTqtI,8551
 ring/algorithms/generator/pd_control.py,sha256=3pOaYig26vmp8gippDfy2KNJRZO3kr0rGd_PBIuEROM,5759
 ring/algorithms/generator/randomize.py,sha256=G_vBIo0OwQkXL2u0djwbaoaeb02C4LQCTNNloOYIU2M,3699
@@ -50,17 +50,17 @@ ring/io/xml/from_xml.py,sha256=8b44sPVWgoY8JGJZLpJ8M_eLfcfu3IsMtBzSytPTPmw,9234
 ring/io/xml/test_from_xml.py,sha256=bckVrVVmEhCwujd_OF9FGYnX3zU3BgztpqGxxmd0htM,1562
 ring/io/xml/test_to_xml.py,sha256=NGn4VSiFdwhYN5YTBduWMiY9B5dwtxZhCQAR_PXeqKU,946
 ring/io/xml/to_xml.py,sha256=fohb-jWMf2cxVdT5dmknsGyrNMseICSbKEz_urbaWbQ,3407
-ring/ml/__init__.py,sha256=hHQUeYAdqRoEFtpKW4zXyNkAdeH2cPh17vc29hXcWWw,1746
+ring/ml/__init__.py,sha256=-bryExVoKJYSF_G_KYc5hI_GciIhj2xZ8WGi6TdRghw,1836
 ring/ml/base.py,sha256=PQ72VasEqlecBZgWP5HE5rWYyLiLq7nCVLymXo9f0dw,8959
 ring/ml/callbacks.py,sha256=DkSy5c7IRqAAks2dx8acEBExYxUv-xiUFwZn4odPYq4,13253
 ring/ml/ml_utils.py,sha256=hQEmeZoahdJyFrz0NZXYi1Yijl7GvPBdqwzZBzlUIUM,7638
-ring/ml/optimizer.py,sha256=OP70P70YcX-2Z-cuoMluFk-L5Vhh_MmqiHdM9OZqyhI,4703
+ring/ml/optimizer.py,sha256=fWyF__ezUltrA16SLfOC1jvS3zBh9NJsMYa6-V0frhs,4709
 ring/ml/ringnet.py,sha256=OWRDu2COmptzbpJWlRLbPIn_ioKZCAd_iu-eiY_aPjk,8521
 ring/ml/train.py,sha256=ftt2MOSSNGCdL7ZoAXcbIgeHW1Wkpgp6XYyLIBUIClI,10872
 ring/ml/training_loop.py,sha256=CEokvPQuuk_WCd-J60ZDodJYcPVvyxLfgXDr_DnbzRI,3359
 ring/ml/params/0x13e3518065c21cd8.pickle,sha256=Zh2k1zK-TNxJl5F7nyTeQ9001qqRE_dfvaq1HWV287A,9355838
 ring/rendering/__init__.py,sha256=Zf7qOdzK3t2hljIrs5P4zFhzHljLSMRyDDZO2YlZk4k,75
-ring/rendering/base_render.py,sha256=L2x83oWC2KkneuA5Lubg2qoDInDFBfH5wYlATymP9-0,8764
+ring/rendering/base_render.py,sha256=s5dF-GVBqjiWkqVuPQMtTLuM7EtA-YrB7RVWFfIaQ1I,8956
 ring/rendering/mujoco_render.py,sha256=aluzQJp3jrDdPfAyNmQuXIHRfgfBTCCZQqxKOx_0D2s,7770
 ring/rendering/vispy_render.py,sha256=QmRyA7Hqk3uS1SKjcncwc4_vd1m4yWryW2X0i4jRvCw,10260
 ring/rendering/vispy_visuals.py,sha256=ooBZqppnebeL0ANe6V6zUgnNTtDcdkOsa4vZuM4sx-I,7873
@@ -70,14 +70,14 @@ ring/sys_composer/__init__.py,sha256=5J_JJJIHfTPcpxh0v4FqiOs81V1REPUd7pgiw2nAN5E
 ring/sys_composer/delete_sys.py,sha256=cIM9KbyLfg7B9121g7yjzuFbjeNu9cil1dPavAYEgzk,3408
 ring/sys_composer/inject_sys.py,sha256=Mj-q-mUjXKwkg-ol6IQAjf9IJfk7pGhez0_WoTKTgm0,3503
 ring/sys_composer/morph_sys.py,sha256=2GpPtS5hT0eZMptdGpt30Hc97OykJNE67lEVRf7sHrc,12700
-ring/utils/__init__.py,sha256=6BWUMmbQ4E-Qwd-SNfRlpbzJ0UJ1DpEclstrgbLdDvk,773
+ring/utils/__init__.py,sha256=rTvSA4RiJAVCY_A64FUMd8IJTv94LgoSA3Ps5X63_jA,799
 ring/utils/batchsize.py,sha256=mPFGD7AedFMycHtyIuZtNWCaAvKLLWSWaB7X6u54xvM,1358
 ring/utils/colab.py,sha256=ZLHwP0jNQUsmZJU4l68a5djULPi6T-jYNNHevjIoMn8,1631
 ring/utils/hdf5.py,sha256=BzXwVypZmEZeHVgeGZ78YYdi10NEQtnPhdrb8dQAXo0,5856
 ring/utils/normalizer.py,sha256=67L2BU1MRsMT4pD41ta3JJMppLN0ozFmnwrmXDtnqrQ,1698
 ring/utils/path.py,sha256=hAfSlqRi-ew536RnjDDM7IKapdMJc-EvhrR0Y-BCFWc,1265
-ring/utils/utils.py,sha256=AzOzR95oOyfdtJhjt5iIb35u611NlTb1Ds4QDKrGMOM,4967
-imt_ring-1.3.0.dist-info/METADATA,sha256=fYndGMXxYVbuBmGAqpzGezBcD6-ruTt3VOAMdLxTwDE,3104
-imt_ring-1.3.0.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-imt_ring-1.3.0.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
-imt_ring-1.3.0.dist-info/RECORD,,
+ring/utils/utils.py,sha256=I2f6-DMBrrgy5tpLzPLlezifQgkO2fERZWyX3cfb4sI,5303
+imt_ring-1.3.2.dist-info/METADATA,sha256=OQmB5-CEpy-JWv2K9vYo9fUNQDW0Jg-fFU4kBkiVRGQ,3104
+imt_ring-1.3.2.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+imt_ring-1.3.2.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
+imt_ring-1.3.2.dist-info/RECORD,,

ring/algorithms/custom_joints/__init__.py

@@ -2,5 +2,6 @@ from .rr_imp_joint import register_rr_imp_joint
 from .rr_joint import register_rr_joint
 from .suntay import GP_DrawFnPair
 from .suntay import MLP_DrawnFnPair
+from .suntay import Polynomial_DrawnFnPair
 from .suntay import register_suntay
 from .suntay import SuntayConfig

ring/algorithms/custom_joints/suntay.py

@@ -293,6 +293,63 @@ def register_suntay(sconfig: SuntayConfig, name: str = "suntay"):
     ring.register_new_joint_type(name, joint_model, 1, overwrite=True)
 
 
+def Polynomial_DrawnFnPair(
+    order: int = 2,
+    val: float = 2.0,
+    center: bool = False,
+    flexion_center: Optional[float] = None,
+) -> DrawnFnPairFactory:
+    assert val >= 0.0
+
+    # because 0-th order is also counted
+    order += 1
+
+    def factory(xs, mn, mx):
+        nonlocal flexion_center
+
+        flexion_mn = jnp.min(xs)
+        flexion_mx = jnp.max(xs)
+
+        def _apply_poly_factors(poly_factors, q):
+            return poly_factors @ jnp.power(q, jnp.arange(order))
+
+        if flexion_center is None:
+            flexion_center = (flexion_mn + flexion_mx) / 2
+        else:
+            flexion_center = jnp.array(flexion_center)
+
+        def init(key):
+            c1, c2 = jax.random.split(key)
+            poly_factors = jax.random.uniform(
+                c1, shape=(order,), minval=-val, maxval=val
+            )
+            q0 = jax.random.uniform(c2, minval=flexion_mn, maxval=flexion_mx)
+            values = jax.vmap(_apply_poly_factors, in_axes=(None, 0))(
+                poly_factors, xs - q0
+            )
+            amax = jnp.max(values)
+            amin = jnp.min(values)
+            return amin, amax, poly_factors, q0
+
+        def _apply(params, q):
+            amin, amax, poly_factors, q0 = params
+            q = q - q0
+            value = _apply_poly_factors(poly_factors, q)
+            return restrict(value, mn, mx, amin, amax)
+
+        if center:
+
+            def apply(params, q):
+                return _apply(params, q) - _apply(params, flexion_center)
+
+        else:
+            apply = _apply
+
+        return DrawnFnPair(init, apply)
+
+    return factory
+
+
 def MLP_DrawnFnPair(
     center: bool = False, flexion_center: Optional[float] = None
 ) -> DrawnFnPairFactory:

ring/algorithms/generator/base.py

@@ -4,6 +4,8 @@ import warnings
 
 import jax
 import jax.numpy as jnp
+import tree_utils
+
 from ring import base
 from ring import utils
 from ring.algorithms import jcalc
@@ -13,7 +15,6 @@ from ring.algorithms.generator import motion_artifacts
 from ring.algorithms.generator import randomize
 from ring.algorithms.generator import transforms
 from ring.algorithms.generator import types
-import tree_utils
 
 
 class RCMG:
@@ -108,23 +109,20 @@ class RCMG:
                         partial_build_gen(sys=_sys, config=_config, sys_ml=sys_ml)
                     )
 
-    def _to_data(self, sizes, seed, jit):
-        return batch.batch_generators_eager_to_list(
-            self.gens, sizes, seed=seed, jit=jit
-        )
+    def _to_data(self, sizes, seed):
+        return batch.batch_generators_eager_to_list(self.gens, sizes, seed=seed)
 
-    def to_list(self, sizes: int | list[int] = 1, seed: int = 1, jit: bool = False):
-        return self._to_data(sizes, seed, jit)
+    def to_list(self, sizes: int | list[int] = 1, seed: int = 1):
+        return self._to_data(sizes, seed)
 
     def to_pickle(
         self,
         path: str,
         sizes: int | list[int] = 1,
         seed: int = 1,
-        jit: bool = False,
         overwrite: bool = True,
     ) -> None:
-        data = tree_utils.tree_batch(self._to_data(sizes, seed, jit))
+        data = tree_utils.tree_batch(self._to_data(sizes, seed))
         utils.pickle_save(data, path, overwrite=overwrite)
 
     def to_hdf5(
@@ -132,10 +130,9 @@ class RCMG:
         path: str,
         sizes: int | list[int] = 1,
         seed: int = 1,
-        jit: bool = False,
         overwrite: bool = True,
     ) -> None:
-        data = tree_utils.tree_batch(self._to_data(sizes, seed, jit))
+        data = tree_utils.tree_batch(self._to_data(sizes, seed))
         utils.hdf5_save(path, data, overwrite=overwrite)
 
     def to_eager_gen(
@@ -143,11 +140,8 @@ class RCMG:
         batchsize: int = 1,
         sizes: int | list[int] = 1,
         seed: int = 1,
-        jit: bool = False,
     ) -> types.BatchedGenerator:
-        return batch.batch_generators_eager(
-            self.gens, sizes, batchsize, seed=seed, jit=jit
-        )
+        return batch.batch_generators_eager(self.gens, sizes, batchsize, seed=seed)
 
     def to_lazy_gen(
         self, sizes: int | list[int] = 1, jit: bool = True

ring/algorithms/generator/batch.py

@@ -6,12 +6,13 @@ import warnings
 import jax
 import jax.numpy as jnp
 import numpy as np
-from ring import utils
-from ring.algorithms.generator import types
 from tqdm import tqdm
 import tree_utils
 from tree_utils import tree_batch
 
+from ring import utils
+from ring.algorithms.generator import types
+
 
 def _build_batch_matrix(batchsizes: list[int]) -> jax.Array:
     arr = []
@@ -61,11 +62,24 @@ def batch_generators_lazy(
     return generator
 
 
+def _number_of_executions_required(size: int) -> int:
+    vmap_threshold = 128
+    _, vmap = utils.distribute_batchsize(size)
+
+    primes = iter(utils.primes(vmap))
+    n_calls = 1
+    while vmap > vmap_threshold:
+        prime = next(primes)
+        n_calls *= prime
+        vmap /= prime
+
+    return n_calls
+
+
 def batch_generators_eager_to_list(
     generators: types.Generator | list[types.Generator],
     sizes: int | list[int],
     seed: int = 1,
-    jit: bool = True,
 ) -> list[tree_utils.PyTree]:
     "Returns list of unbatched sequences as numpy arrays."
     generators, sizes = _process_sizes_batchsizes_generators(generators, sizes)
@@ -73,11 +87,20 @@ def batch_generators_eager_to_list(
     key = jax.random.PRNGKey(seed)
     data = []
     for gen, size in tqdm(zip(generators, sizes), desc="eager data generation"):
-        key, consume = jax.random.split(key)
-        sample = batch_generators_lazy(gen, size, jit=jit)(consume)
-        # converts also to numpy
-        sample = jax.device_get(sample)
-        data.extend([jax.tree_map(lambda a: a[i], sample) for i in range(size)])
+
+        n_calls = _number_of_executions_required(size)
+        # decrease size by n_calls times
+        size = int(size / n_calls)
+        jit = True if n_calls > 1 else False
+        gen_jit = batch_generators_lazy(gen, size, jit=jit)
+
+        for _ in range(n_calls):
+            key, consume = jax.random.split(key)
+            sample = gen_jit(consume)
+            # converts also to numpy
+            sample = jax.device_get(sample)
+            data.extend([jax.tree_map(lambda a: a[i], sample) for i in range(size)])
+
     return data
 
 
@@ -243,12 +266,11 @@ def batch_generators_eager(
     shuffle: bool = True,
     drop_last: bool = True,
     seed: int = 1,
-    jit: bool = True,
 ) -> types.BatchedGenerator:
     """Eagerly create a large precomputed generator by calling multiple generators
     and stacking their output."""
 
-    data = batch_generators_eager_to_list(generators, sizes, seed=seed, jit=jit)
+    data = batch_generators_eager_to_list(generators, sizes, seed=seed)
     return batched_generator_from_list(data, batchsize, shuffle, drop_last)
 
 
@@ -270,7 +292,7 @@ def _process_sizes_batchsizes_generators(
 
     assert len(generators) == len(list_sizes)
 
-    _WARN_SIZE = 4096
+    _WARN_SIZE = 1e6  # disable this warning
     for size in list_sizes:
         if size >= _WARN_SIZE:
             warnings.warn(

ring/base.py

@@ -99,15 +99,15 @@ class _Base:
     def ndim(self):
         return tu.tree_ndim(self)
 
-    def shape(self, axis=0) -> int:
-        return tu.tree_shape(self, axis)
-
-    def __len__(self) -> int:
-        Bs = tree_map(lambda arr: arr.shape[0], self)
+    def shape(self, axis: int = 0) -> int:
+        Bs = tree_map(lambda arr: arr.shape[axis], self)
         Bs = set(jax.tree_util.tree_flatten(Bs)[0])
         assert len(Bs) == 1
         return list(Bs)[0]
 
+    def __len__(self) -> int:
+        return self.shape(axis=0)
+
 
 @struct.dataclass
 class Transform(_Base):
@@ -685,14 +685,13 @@ class System(_Base):
         self,
         xs: Transform | list[Transform],
         yhat: dict | jax.Array | np.ndarray,
-        stepframe: int = 1,
         # by default we don't predict the global rotation
         transparent_segment_to_root: bool = True,
         **kwargs,
     ):
         "`xs` matches `sys`. `yhat` matches `sys_noimu`. `yhat` are child-to-parent."
         return ring.rendering.render_prediction(
-            self, xs, yhat, stepframe, transparent_segment_to_root, **kwargs
+            self, xs, yhat, transparent_segment_to_root, **kwargs
         )
 
     def delete_system(self, link_name: str | list[str], strict: bool = True):

ring/ml/__init__.py

@@ -12,6 +12,8 @@ from .optimizer import make_optimizer
 from .ringnet import RING
 from .train import train_fn
 
+_lpf_cutoff_freq = 10.0
+
 
 def RING_ICML24(params=None, eval: bool = True, **kwargs):
     """Create the RING network used in the icml24 paper.
@@ -29,7 +31,7 @@ def RING_ICML24(params=None, eval: bool = True, **kwargs):
     ringnet = RING(params=params, **kwargs)  # noqa: F811
     ringnet = base.ScaleX_FilterWrapper(ringnet)
     if eval:
-        ringnet = base.LPF_FilterWrapper(ringnet, 10.0, samp_freq=None)
+        ringnet = base.LPF_FilterWrapper(ringnet, _lpf_cutoff_freq, samp_freq=None)
     ringnet = base.GroundTruthHeading_FilterWrapper(ringnet)
     return ringnet
 
@@ -39,6 +41,7 @@ def RNNO(
     return_quats: bool = False,
     params=None,
     eval: bool = True,
+    samp_freq: float | None = None,
     **kwargs,
 ):
     assert "message_dim" not in kwargs
@@ -55,7 +58,7 @@ def RNNO(
     ringnet = base.NoGraph_FilterWrapper(ringnet, quat_normalize=return_quats)
     ringnet = base.ScaleX_FilterWrapper(ringnet)
     if eval and return_quats:
-        ringnet = base.LPF_FilterWrapper(ringnet, 10.0, samp_freq=None)
+        ringnet = base.LPF_FilterWrapper(ringnet, _lpf_cutoff_freq, samp_freq=samp_freq)
     if return_quats:
         ringnet = base.GroundTruthHeading_FilterWrapper(ringnet)
     return ringnet

ring/ml/optimizer.py

@@ -14,10 +14,10 @@ from optax._src.transform import AddNoiseState
 def make_optimizer(
     lr: float,
     n_episodes: int,
-    n_steps_per_episode: int,
+    n_steps_per_episode: int = 6,
     adap_clip: Optional[float] = 0.1,
     glob_clip: Optional[float] = 0.2,
-    skip_large_update_max_normsq: float = 5.0,
+    skip_large_update_max_normsq: float = 100.0,
     skip_large_update_warmup: int = 300,
     inner_opt=optax.lamb,
     cos_decay_twice: bool = False,

ring/rendering/base_render.py

@@ -136,12 +136,15 @@ def render_prediction(
     sys: base.System,
     xs: base.Transform | list[base.Transform],
     yhat: dict | jax.Array | np.ndarray,
-    stepframe: int = 1,
     # by default we don't predict the global rotation
     transparent_segment_to_root: bool = True,
     **kwargs,
 ):
     "`xs` matches `sys`. `yhat` matches `sys_noimu`. `yhat` are child-to-parent."
+
+    offset_truth = kwargs.pop("offset_truth", [0, 0, 0])
+    offset_pred = kwargs.pop("offset_pred", [0, 0, 0])
+
     if isinstance(xs, list):
         # list -> batched Transform
         xs = xs[0].batch(*xs[1:])
@@ -180,18 +183,23 @@ def render_prediction(
 
     # swap time axis, and link axis
     xs, xshat = xs.transpose((1, 0, 2)), xshat.transpose((1, 0, 2))
+
+    add_offset = lambda x, offset: algebra.transform_mul(
+        x, base.Transform.create(pos=jnp.array(offset, dtype=jnp.float32))
+    )
+
     # create mapping from `name` -> Transform
     xs_dict = dict(
         zip(
             ["hat_" + name for name in sys_noimu.link_names],
-            [xshat[i] for i in range(sys_noimu.num_links())],
+            [add_offset(xshat[i], offset_pred) for i in range(sys_noimu.num_links())],
         )
     )
     xs_dict.update(
         dict(
             zip(
                 sys.link_names,
-                [xs[i] for i in range(sys.num_links())],
+                [add_offset(xs[i], offset_truth) for i in range(sys.num_links())],
             )
         )
     )
@@ -202,11 +210,8 @@ def render_prediction(
         xs_render.append(xs_dict[name])
     xs_render = xs_render[0].batch(*xs_render[1:])
     xs_render = xs_render.transpose((1, 0, 2))
-    N = xs_render.shape()
-    xs_render = [xs_render[t] for t in range(0, N, stepframe)]
 
     frames = render(sys_render, xs_render, **kwargs)
-
     return frames
 
 

ring/utils/__init__.py

@@ -15,6 +15,7 @@ from .utils import dict_union
 from .utils import import_lib
 from .utils import pickle_load
 from .utils import pickle_save
+from .utils import primes
 from .utils import pytree_deepcopy
 from .utils import sys_compare
 from .utils import to_list

ring/utils/utils.py

@@ -159,3 +159,17 @@ def pickle_load(
     with open(path, "rb") as file:
         obj = pickle.load(file)
     return obj
+
+
+def primes(n: int) -> list[int]:
+    "Primefactor decomposition in ascending order."
+    primfac = []
+    d = 2
+    while d * d <= n:
+        while (n % d) == 0:
+            primfac.append(d)  # supposing you want multiple factors repeated
+            n //= d
+        d += 1
+    if n > 1:
+        primfac.append(n)
+    return primfac