imt-ring 1.2.1__py3-none-any.whl

ring/algorithms/custom_joints/suntay.py

@@ -0,0 +1,424 @@
+from dataclasses import dataclass
+from typing import Callable, NamedTuple, Optional
+
+import haiku as hk
+import jax
+import jax.numpy as jnp
+from tree_utils import PyTree
+
+import ring
+from ring import maths
+from ring.algorithms._random import random_angle_over_time
+
+Params = PyTree
+
+
+class DrawnFnPair(NamedTuple):
+    # (key) -> tree
+    init: Callable[[jax.Array], Params]
+    # (params, q) -> (1,)
+    apply: Callable[[Params, jax.Array], jax.Array]
+
+
+# (flexions, min, max) -> DrawnFnPair
+DrawnFnPairFactory = Callable[[jax.Array, float, float], DrawnFnPair]
+
+
+def deg2rad(deg: float):
+    return (deg / 180.0) * 3.1415926535
+
+
+def GP_DrawFnPair(
+    length_scale: float = 1.4, large_abs_values_of_gps: float = 0.25
+) -> DrawnFnPairFactory:
+
+    def factory(xs, mn, mx):
+        def init(key):
+            return {
+                "xs": xs,
+                "ys": _gp_draw_and_rom(
+                    key=key,
+                    xs=xs,
+                    ys=None,
+                    length_scale=length_scale,
+                    mn=mn,
+                    mx=mx,
+                    amin=-large_abs_values_of_gps,
+                    amax=large_abs_values_of_gps,
+                ),
+            }
+
+        def apply(params, q):
+            return jnp.interp(q, params["xs"], params["ys"])
+
+        return DrawnFnPair(init, apply)
+
+    return factory
+
+
+@dataclass
+class SuntayConfig:
+    flexion_rot_min: float = -deg2rad(5.0)
+    flexion_rot_max: float = deg2rad(95.0)
+    flexion_rot_restrict_method: str = "minmax"
+    ###
+    flexion_pos_min: float = -0.015
+    flexion_pos_max: float = 0.015
+    flexion_pos_factory: DrawnFnPairFactory = GP_DrawFnPair()
+    ###
+    abduction_rot_min: float = deg2rad(-4)
+    abduction_rot_max: float = deg2rad(4)
+    abduction_rot_factory: DrawnFnPairFactory = GP_DrawFnPair()
+    ###
+    abduction_pos_min: float = -0.015
+    abduction_pos_max: float = 0.015
+    abduction_pos_factory: DrawnFnPairFactory = GP_DrawFnPair()
+    ###
+    external_rot_min: float = deg2rad(-10)
+    external_rot_max: float = deg2rad(10)
+    external_rot_factory: DrawnFnPairFactory = GP_DrawFnPair()
+    ###
+    external_pos_min: float = -0.06
+    external_pos_max: float = 0.0
+    external_pos_factory: DrawnFnPairFactory = GP_DrawFnPair()
+    ###
+    num_points: int = 50
+    mconfig: Optional[ring.MotionConfig] = None
+
+
+def register_suntay(sconfig: SuntayConfig, name: str = "suntay"):
+    """Ref to 'E.S. Grood and W.J. Suntay' paper"""
+
+    flexion_xs = jnp.linspace(
+        sconfig.flexion_rot_min, sconfig.flexion_rot_max, num=sconfig.num_points
+    )
+
+    draw_fn_pairs = {}
+    for config_name, params_name in zip(
+        [
+            "flexion_pos",
+            "abduction_rot",
+            "abduction_pos",
+            "external_rot",
+            "external_pos",
+        ],
+        ["ys_S1", "ys_beta", "ys_S2", "ys_gamma", "ys_S3"],
+    ):
+        get = lambda key: getattr(sconfig, config_name + "_" + key)
+        factory = get("factory")
+        draw_fn_pairs[params_name] = factory(flexion_xs, get("min"), get("max"))
+
+    def _suntay_rotation_matrix_R_transpose_eq26(alpha, beta, gamma):
+        sin_alp, sin_bet, sin_gam = jnp.sin(alpha), jnp.sin(beta), jnp.sin(gamma)
+        cos_alp, cos_bet, cos_gam = jnp.cos(alpha), jnp.cos(beta), jnp.cos(gamma)
+        return jnp.array(
+            [
+                [cos_gam * sin_bet, sin_gam * sin_bet, cos_bet],
+                [
+                    -cos_alp * sin_gam - cos_gam * sin_alp * cos_bet,
+                    cos_alp * cos_gam - sin_gam * sin_alp * cos_bet,
+                    sin_bet * sin_alp,
+                ],
+                [
+                    sin_alp * sin_gam - cos_gam * cos_alp * cos_bet,
+                    -cos_gam * sin_alp - cos_alp * cos_bet * sin_gam,
+                    cos_alp * sin_bet,
+                ],
+            ]
+        ).T
+
+    def _suntay_translation_vector_H_eq9(alpha, beta, S):
+        sin_alp, sin_bet = jnp.sin(alpha), jnp.sin(beta)
+        cos_alp, cos_bet = jnp.cos(alpha), jnp.cos(beta)
+        # eq (10)
+        U = jnp.array(
+            [
+                [1, 0, cos_bet],
+                [0, cos_alp, sin_alp * sin_bet],
+                [0, -sin_alp, cos_alp * sin_bet],
+            ]
+        )
+        return U @ S
+
+    def _alpha_beta_gamma_S(q_flexion, params):
+        assert q_flexion.shape == (1,)
+
+        # (1,) -> (,)
+        q_flexion = q_flexion[0]
+
+        S_123 = []
+        for i in range(1, 4):
+            key = f"ys_S{i}"
+            S_123.append(draw_fn_pairs[key].apply(params[key], q_flexion))
+        S = jnp.stack(S_123)
+        # table 2 of suntay paper
+        alpha = q_flexion
+        # note the minus sign, because in config we specify `abduction` not `adduction`
+        adduction = -draw_fn_pairs["ys_beta"].apply(params["ys_beta"], q_flexion)
+        beta = jnp.pi / 2 + adduction
+        gamma = draw_fn_pairs["ys_gamma"].apply(params["ys_gamma"], q_flexion)
+        return alpha, beta, gamma, S
+
+    def _utils_find_suntay_joint(sys: ring.System) -> str:
+        suntay_link_name = None
+        for link_name, link_type in zip(sys.link_names, sys.link_types):
+            if link_type == name:
+                if suntay_link_name is not None:
+                    raise Exception(
+                        f"multiple links of type `{name}` found, link_names "
+                        f"are [{suntay_link_name}, {link_name}]"
+                    )
+                suntay_link_name = link_name
+
+        if suntay_link_name is None:
+            raise Exception(
+                f"no link with type `{name}` found, link_types are {sys.link_types}"
+            )
+        return suntay_link_name
+
+    def _utils_Q_S_H_alpha_beta_gamma(sys: ring.System, qs: jax.Array):
+        # qs.shape = (timesteps, q_size)
+        assert qs.ndim == 2
+        assert qs.shape[-1] == sys.q_size()
+
+        suntay_link_name = _utils_find_suntay_joint(sys)
+
+        params = jax.tree_map(
+            lambda arr: arr[sys.idx_map("l")[suntay_link_name]],
+            sys.links.joint_params[name],
+        )
+        # shape = (timesteps, 1)
+        q_flexion = qs[:, sys.idx_map("q")[suntay_link_name]]
+
+        @jax.vmap
+        def _Q_S_H_alpha_beta_gamma_from_q_flexion(q_flexion):
+            alpha, beta, gamma, S = _alpha_beta_gamma_S(q_flexion, params)
+            cos_bet = jnp.cos(beta)
+            Q = jnp.array([S[0] + S[2] * cos_bet, S[1], -S[2] - S[0] * cos_bet])
+            # translation from femur to tibia
+            H = _suntay_translation_vector_H_eq9(alpha, beta, S)
+            return Q, S, H, alpha, beta, gamma
+
+        return _Q_S_H_alpha_beta_gamma_from_q_flexion(q_flexion)
+
+    def _transform_suntay(q_flexion, params):
+        alpha, beta, gamma, S = _alpha_beta_gamma_S(q_flexion, params)
+
+        # rotation from femur to tibia
+        R_T = _suntay_rotation_matrix_R_transpose_eq26(alpha, beta, gamma)
+        q_fem_tib = maths.quat_from_3x3(R_T)
+        # translation from femur to tibia
+        H = _suntay_translation_vector_H_eq9(alpha, beta, S)
+
+        return ring.Transform.create(pos=H, rot=q_fem_tib)
+
+    def _init_joint_params_suntay(key):
+        params = dict()
+        for params_name, draw_fn_pair in draw_fn_pairs.items():
+            key, consume = jax.random.split(key)
+            params[params_name] = draw_fn_pair.init(consume)
+
+        return params
+
+    def _draw_flexion_angle(
+        mconfig: ring.MotionConfig,
+        key_t: jax.random.PRNGKey,
+        key_value: jax.random.PRNGKey,
+        dt: float,
+        _: jax.Array,
+    ) -> jax.Array:
+        key_value, consume = jax.random.split(key_value)
+
+        if sconfig.mconfig is not None:
+            mconfig = sconfig.mconfig
+
+        ANG_0 = jax.random.uniform(
+            consume, minval=mconfig.ang0_min, maxval=mconfig.ang0_max
+        )
+        # `random_angle_over_time` always returns wrapped angles, thus it would be
+        # inconsistent to allow an initial value that is not wrapped
+        ANG_0 = maths.wrap_to_pi(ANG_0)
+        qs_flexion = random_angle_over_time(
+            key_t,
+            key_value,
+            ANG_0,
+            mconfig.dang_min,
+            mconfig.dang_max,
+            mconfig.delta_ang_min,
+            mconfig.delta_ang_max,
+            mconfig.t_min,
+            mconfig.t_max,
+            mconfig.T,
+            dt,
+            5,
+            mconfig.randomized_interpolation_angle,
+            mconfig.range_of_motion_hinge,
+            mconfig.range_of_motion_hinge_method,
+            mconfig.cdf_bins_min,
+            mconfig.cdf_bins_max,
+            mconfig.interpolation_method,
+        )
+        return restrict(
+            qs_flexion,
+            sconfig.flexion_rot_min,
+            sconfig.flexion_rot_max,
+            -jnp.pi,
+            jnp.pi,
+            method=sconfig.flexion_rot_restrict_method,
+        )
+
+    joint_model = ring.JointModel(
+        transform=_transform_suntay,
+        rcmg_draw_fn=_draw_flexion_angle,
+        init_joint_params=_init_joint_params_suntay,
+        utilities=dict(
+            Q_S_H_alpha_beta_gamma=_utils_Q_S_H_alpha_beta_gamma,
+            find_suntay_joint=_utils_find_suntay_joint,
+        ),
+    )
+    ring.register_new_joint_type(name, joint_model, 1, qd_width=0, overwrite=True)
+
+
+def MLP_DrawnFnPair(
+    center: bool = False, flexion_center: Optional[float] = None
+) -> DrawnFnPairFactory:
+
+    def factory(xs, mn, mx):
+        nonlocal flexion_center
+
+        flexion_mn = jnp.min(xs)
+        flexion_mx = jnp.max(xs)
+
+        if flexion_center is None:
+            flexion_center = (flexion_mn + flexion_mx) / 2
+
+        @hk.without_apply_rng
+        @hk.transform
+        def mlp(x):
+            # normalize the x input; [0, 1]
+            x = _shift(x, flexion_mn, flexion_mx)
+            # center the x input; [-0.5, 0.5]
+            x = x - 0.5
+            net = hk.nets.MLP(
+                [10, 5, 1],
+                activation=jnp.tanh,
+                w_init=hk.initializers.RandomNormal(),
+            )
+            return net(x)
+
+        example_q = jnp.zeros((1,))
+
+        def init(key):
+            return mlp.init(key, example_q)
+
+        def _apply(params, q):
+            q = q[None]
+            return jnp.squeeze(_shift_inv(jax.nn.sigmoid(mlp.apply(params, q)), mn, mx))
+
+        if center:
+
+            def apply(params, q):
+                return _apply(params, q) - _apply(params, flexion_center)
+
+        else:
+            apply = _apply
+
+        return DrawnFnPair(init, apply)
+
+    return factory
+
+
+def _gp_draw_and_rom(key, xs, ys, length_scale, mn, mx, amin, amax):
+    randomized_ys = _gp_draw(key, xs, ys, length_scale)
+    if ys is not None:
+        amin += jnp.min(ys)
+        amax += jnp.max(ys)
+    return restrict(randomized_ys, mn, mx, amin, amax)
+
+
+def _gp_draw(key, xs, ys=None, length: float = 1.0, noise=0.0, method="svd", **kwargs):
+    if ys is None:
+        ys = jnp.zeros_like(xs)
+    cov = _gp_K(lambda *args: _rbf_kernel(*args, length=length), xs, noise)
+    return jax.random.multivariate_normal(
+        key=key, mean=ys, cov=cov, method=method, **kwargs
+    )
+
+
+def _gp_K(kernel, xs, noise: float):
+    assert xs.ndim == 1
+    N = len(xs)
+    xs = xs[:, None]
+
+    K = jax.vmap(lambda x1: jax.vmap(lambda x2: kernel(x1, x2))(xs))(xs)
+    assert K.shape == (N, N, 1)
+    return K[..., 0] + jnp.eye(N) * noise
+
+
+def _rbf_kernel(x1: float, x2: float, length: float):
+    return jnp.exp(-((x1 - x2) ** 2) / (2 * length**2))
+
+
+def _shift(ys, min, max):
+    return (ys - min) / (max - min)
+
+
+def _shift_inv(ys, min, max):
+    return (ys * (max - min)) + min
+
+
+def _normalize(ys, amin=None, amax=None):
+    if amin is None:
+        amin = jnp.min(ys)
+    else:
+        amin = jnp.min(jnp.array([amin, jnp.min(ys)]))
+    if amax is None:
+        amax = jnp.max(ys)
+    else:
+        amax = jnp.max(jnp.array([amax, jnp.max(ys)]))
+    return _shift(ys, amin, amax)
+
+
+def _smoothclamp(x, mi, mx):
+    return mi + (mx - mi) * (
+        lambda t: jnp.where(t < 0, 0, jnp.where(t <= 1, 3 * t**2 - 2 * t**3, 1))
+    )((x - mi) / (mx - mi))
+
+
+def _sigmoidclamp(x, mi, mx):
+    return mi + (mx - mi) * (lambda t: (1 + 200 ** (-t + 0.5)) ** (-1))(
+        (x - mi) / (mx - mi)
+    )
+
+
+def restrict(
+    ys,
+    min: float,
+    max: float,
+    actual_min=None,
+    actual_max=None,
+    method: str = "minmax",
+    method_kwargs=dict(),
+):
+    if method == "minmax":
+        # scale to [0, 1]
+        ys = _normalize(ys, actual_min, actual_max)
+        # scale to [min, max]
+        return _shift_inv(ys, min, max)
+    elif method == "clip":
+        return jnp.clip(ys, min, max)
+    elif method == "smoothclamp":
+        return _smoothclamp(ys, min, max)
+    elif method == "sigmoidclamp":
+        return _sigmoidclamp(ys, min, max)
+    elif method == "sigmoid":
+        # scale to [0, 1]
+        ys = _normalize(ys, actual_min, actual_max)
+        # scale to [-stepness, stepness]
+        stepness = method_kwargs.get("stepness", 3.0)
+        ys = _shift_inv(ys, -stepness, stepness)
+        # scale to [0, 1]
+        ys = jax.nn.sigmoid(ys)
+        return _shift_inv(ys, min, max)
+    else:
+        raise NotImplementedError()