imt-ring 1.3.7__py3-none-any.whl → 1.3.8__py3-none-any.whl

{imt_ring-1.3.7.dist-info → imt_ring-1.3.8.dist-info}/METADATA

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

@@ -15,8 +15,8 @@ ring/algorithms/custom_joints/rr_joint.py,sha256=jnRtjtOCALMaq2_0bcu2d7qgfQ6etXp
 ring/algorithms/custom_joints/suntay.py,sha256=7-kym1kMDwqYD_2um1roGcBeB8BlTCPe1wljuNGNARA,16676
 ring/algorithms/generator/__init__.py,sha256=p4ucl0zQtp5NwNoXIRjmTzGGRu2WOAWFfNmYRPwQles,912
 ring/algorithms/generator/base.py,sha256=AKH7GXEmRGV1kK8okiqa12uq0Ah9VYlqgdLw-99oFoQ,14840
-ring/algorithms/generator/batch.py,sha256=EOCX0vOxDwVOweArGsUneeeYysdSY2mFB55W052Wd9o,9161
-ring/algorithms/generator/motion_artifacts.py,sha256=aKdkZU5OF4_aKyL4Yo-ftZRwrDCve1LuuREGAUlTqtI,8551
+ring/algorithms/generator/batch.py,sha256=ABmaNguQKVStphHN9h5_oNaNC2Un_YoR0iswdKaLOHs,9178
+ring/algorithms/generator/motion_artifacts.py,sha256=Hk088xPYxnKsqmQEtr_bqRWq5zYBVM7PEG8IvHmMY6g,8711
 ring/algorithms/generator/pd_control.py,sha256=3pOaYig26vmp8gippDfy2KNJRZO3kr0rGd_PBIuEROM,5759
 ring/algorithms/generator/randomize.py,sha256=G_vBIo0OwQkXL2u0djwbaoaeb02C4LQCTNNloOYIU2M,3699
 ring/algorithms/generator/transforms.py,sha256=nvNDvM20tEw9Zd0ra0TxA25uf01L40Y2UKvtQmOrlGo,12782
@@ -50,13 +50,14 @@ 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=-bryExVoKJYSF_G_KYc5hI_GciIhj2xZ8WGi6TdRghw,1836
+ring/ml/__init__.py,sha256=52LpEjni5lG-ov5-3ocodH-vKZxNcFMU7W9XfjDicp0,2113
 ring/ml/base.py,sha256=PQ72VasEqlecBZgWP5HE5rWYyLiLq7nCVLymXo9f0dw,8959
-ring/ml/callbacks.py,sha256=DkSy5c7IRqAAks2dx8acEBExYxUv-xiUFwZn4odPYq4,13253
+ring/ml/callbacks.py,sha256=W19QF6_uvaNCjs8ObsjNXD7mv9gFgJBixdRSbB_BynE,13301
 ring/ml/ml_utils.py,sha256=hQEmeZoahdJyFrz0NZXYi1Yijl7GvPBdqwzZBzlUIUM,7638
 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=uDW6JMdbMcjUKr3wCL2drWzDUd0Pc3BoroUwLcYoUx4,10914
+ring/ml/ringnet.py,sha256=rgje5AKUKpT8K-vbE9_SgZ3IijR8TJEHnaqxsE57Mhc,8617
+ring/ml/rnno_v1.py,sha256=MpuDyd9B32Qq_gRqiZf4x1q0-2HzADK5JhjjzMWQfkE,1026
+ring/ml/train.py,sha256=huUfMK6eotS6BRrQKoZ-AUG0um3jlqpfQFZNJT8LKiE,10854
 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
@@ -77,7 +78,7 @@ 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=mIcKNv5v2de8HrG7bAhl2bNfmwkMZyIIwFkJq2XWMOI,5357
-imt_ring-1.3.7.dist-info/METADATA,sha256=V6Oow_ZZwpBuHuIbyPIoKFtrhFboxMmuIPx1Rilq3-A,3104
-imt_ring-1.3.7.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-imt_ring-1.3.7.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
-imt_ring-1.3.7.dist-info/RECORD,,
+imt_ring-1.3.8.dist-info/METADATA,sha256=V0FWiPAeTUFF7hD-Y4cucpF9HGZETgXgWaZgYIz41ho,3104
+imt_ring-1.3.8.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+imt_ring-1.3.8.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
+imt_ring-1.3.8.dist-info/RECORD,,

ring/algorithms/generator/batch.py

@@ -147,7 +147,7 @@ def _data_fn_from_paths(
     paths = [utils.parse_path(p, mkdir=False) for p in paths]
 
     extensions = list(set([Path(p).suffix for p in paths]))
-    assert len(extensions) == 1
+    assert len(extensions) == 1, f"{extensions}"
 
     if extensions[0] == ".h5":
         N = sum([utils.hdf5_load_length(p) for p in paths])

ring/algorithms/generator/motion_artifacts.py

@@ -180,9 +180,13 @@ def setup_fn_randomize_damping_stiffness_factory(
             link_spring_stiffness = link_spring_stiffness.at[slice].set(stif)
             link_damping = link_damping.at[slice].set(damp)
 
-        assert len(imus_surely_rigid) == len(triggered_surely_rigid)
+        assert len(imus_surely_rigid) == len(
+            triggered_surely_rigid
+        ), f"{imus_surely_rigid}, {triggered_surely_rigid}"
         for imu_surely_rigid in imus_surely_rigid:
-            assert imu_surely_rigid in triggered_surely_rigid
+            assert (
+                imu_surely_rigid in triggered_surely_rigid
+            ), f"{imus_surely_rigid} not in {triggered_surely_rigid}"
 
         return sys.replace(
             link_damping=link_damping, link_spring_stiffness=link_spring_stiffness

ring/ml/__init__.py

@@ -3,6 +3,7 @@ from . import callbacks
 from . import ml_utils
 from . import optimizer
 from . import ringnet
+from . import rnno_v1
 from . import train
 from . import training_loop
 from .base import AbstractFilter
@@ -42,17 +43,28 @@ def RNNO(
     params=None,
     eval: bool = True,
     samp_freq: float | None = None,
+    v1: bool = False,
     **kwargs,
 ):
     assert "message_dim" not in kwargs
     assert "link_output_normalize" not in kwargs
     assert "link_output_dim" not in kwargs
 
+    if v1:
+        kwargs.update(
+            dict(forward_factory=rnno_v1.rnno_v1_forward_factory, output_dim=output_dim)
+        )
+    else:
+        kwargs.update(
+            dict(
+                message_dim=0,
+                link_output_normalize=False,
+                link_output_dim=output_dim,
+            )
+        )
+
     ringnet = RING(  # noqa: F811
         params=params,
-        message_dim=0,
-        link_output_normalize=False,
-        link_output_dim=output_dim,
         **kwargs,
     )
     ringnet = base.NoGraph_FilterWrapper(ringnet, quat_normalize=return_quats)

ring/ml/callbacks.py

@@ -245,7 +245,8 @@ class SaveParamsTrainingLoopCallback(training_loop.TrainingLoopCallback):
             else:
                 value = "{:.2f}".format(ele.value).replace(".", ",")
                 filename = parse_path(
-                    self.path_to_file + f"_episode={ele.episode}_value={value}",
+                    str(Path(self.path_to_file).with_suffix(""))
+                    + f"_episode={ele.episode}_value={value}",
                     extension="pickle",
                 )
 
@@ -404,7 +405,7 @@ class CheckpointCallback(training_loop.TrainingLoopCallback):
         # only checkpoint if run has been killed
         if training_loop.recv_kill_run_signal():
             path = parse_path(
-                "~/.xxy_checkpoints", ml_utils.unique_id(), extension="pickle"
+                "~/.ring_checkpoints", ml_utils.unique_id(), extension="pickle"
             )
             data = {"params": self.params, "opt_state": self.opt_state}
             pickle_save(

ring/ml/ringnet.py

@@ -191,8 +191,16 @@ class LSTM(hk.RNNCore):
 
 
 class RING(ml_base.AbstractFilter):
-    def __init__(self, params=None, lam=None, jit: bool = True, name=None, **kwargs):
-        self.forward_lam_factory = partial(make_ring, **kwargs)
+    def __init__(
+        self,
+        params=None,
+        lam=None,
+        jit: bool = True,
+        name=None,
+        forward_factory=make_ring,
+        **kwargs,
+    ):
+        self.forward_lam_factory = partial(forward_factory, **kwargs)
         self.params = self._load_params(params)
         self.lam = lam
         self._name = name

ring/ml/rnno_v1.py

@@ -0,0 +1,37 @@
+from typing import Sequence
+
+import haiku as hk
+import jax
+import jax.numpy as jnp
+
+
+def rnno_v1_forward_factory(
+    output_dim: int,
+    rnn_layers: Sequence[int] = (400, 300),
+    linear_layers: Sequence[int] = (200, 100, 50, 50, 25, 25),
+    layernorm: bool = True,
+    act_fn_linear=jax.nn.relu,
+    act_fn_rnn=jax.nn.elu,
+):
+    @hk.without_apply_rng
+    @hk.transform_with_state
+    def forward_fn(X):
+        assert X.shape[-2] == 1
+
+        for i, n_units in enumerate(rnn_layers):
+            state = hk.get_state(f"rnn_{i}", shape=[n_units], init=jnp.zeros)
+            X, state = hk.dynamic_unroll(hk.GRU(n_units), X, state)
+            hk.set_state(f"rnn_{i}", state)
+
+            if layernorm:
+                X = hk.LayerNorm(axis=-1, create_scale=False, create_offset=False)(X)
+            X = act_fn_rnn(X)
+
+        for n_units in linear_layers:
+            X = hk.Linear(n_units)(X)
+            X = act_fn_linear(X)
+
+        y = hk.Linear(output_dim)(X)
+        return y[..., None, :]
+
+    return forward_fn

ring/ml/train.py

@@ -15,7 +15,6 @@ from ring.ml import ml_utils
 from ring.ml import training_loop
 from ring.utils import distribute_batchsize
 from ring.utils import expand_batchsize
-from ring.utils import parse_path
 from ring.utils import pickle_load
 import wandb
 
@@ -217,7 +216,7 @@ def train_fn(
 
         callbacks_all.append(
             ml_callbacks.SaveParamsTrainingLoopCallback(
-                path_to_file=parse_path(callback_save_params, extension=""),
+                path_to_file=callback_save_params,
                 last_n_params=3,
                 track_metrices=callback_save_params_track_metrices,
                 cleanup=False,