dhb-xr 0.2.1__py3-none-any.whl

dhb_xr/encoder/dhb_ti.py

@@ -0,0 +1,204 @@
+"""
+DHB-TI: Time-invariant reparameterization.
+
+Reparameterize a trajectory by a geometric progress variable (translational arc-length,
+angular, or hybrid) and resample at uniform progress knots so that DHB-DR/DHB-QR
+invariants are approximately independent of execution speed and sampling rate.
+
+Progress variables:
+  - translation: s_{i+1} = s_i + ||Δp_i||
+  - angular:     θ_{i+1} = θ_i + ||Δr_i||
+  - hybrid:      σ_{i+1} = σ_i + α||Δp_i|| + (1-α)||Δr_i||, α in [0,1]
+
+Uniform knots σ_k = k * Σ/(M-1), then interpolate poses at σ_k (position spline, quat SLERP).
+"""
+
+from __future__ import annotations
+
+import numpy as np
+from typing import Dict, Any, Optional, Literal
+from scipy.interpolate import CubicSpline
+
+from dhb_xr.core import geometry as geom
+from dhb_xr.core.types import DHBMethod, EncodingMethod
+
+ProgressKind = Literal["translation", "angular", "hybrid"]
+_EPS = 1e-12
+_MIN_STEP = 1e-10  # minimum progress step to avoid degeneracy in interpolation
+
+
+def compute_progress(
+    positions: np.ndarray,
+    quaternions: np.ndarray,
+    kind: ProgressKind = "hybrid",
+    alpha: float = 0.5,
+    min_step: float = _MIN_STEP,
+) -> np.ndarray:
+    """
+    Compute cumulative progress along the trajectory.
+
+    positions: (N, 3), quaternions: (N, 4) wxyz.
+    kind: 'translation' (arc-length), 'angular' (rotation magnitude), 'hybrid'.
+    alpha: weight for translation in hybrid; (1-alpha) for rotation. Ignored if kind != 'hybrid'.
+    min_step: minimum increment per segment to avoid degenerate progress (clamped).
+
+    Returns progress (N,) with progress[0] = 0.
+    """
+    positions = np.asarray(positions, dtype=np.float64)
+    quaternions = np.asarray(quaternions, dtype=np.float64)
+    n = positions.shape[0]
+    if n < 2:
+        return np.zeros(n)
+
+    delta_p = np.diff(positions, axis=0)
+    step_p = np.linalg.norm(delta_p, axis=1)
+    step_p = np.maximum(step_p, min_step)
+
+    delta_r = np.array([
+        geom.quat_relative_axis_angle(quaternions[i], quaternions[i + 1])
+        for i in range(n - 1)
+    ])
+    step_r = np.linalg.norm(delta_r, axis=1)
+    step_r = np.maximum(step_r, min_step)
+
+    if kind == "translation":
+        steps = step_p
+    elif kind == "angular":
+        steps = step_r
+    else:
+        steps = alpha * step_p + (1.0 - alpha) * step_r
+
+    progress = np.concatenate([[0.0], np.cumsum(steps)])
+    return progress
+
+
+def resample_by_progress(
+    positions: np.ndarray,
+    quaternions: np.ndarray,
+    M: int,
+    progress_kind: ProgressKind = "hybrid",
+    alpha: float = 0.5,
+    progress: Optional[np.ndarray] = None,
+    min_step: float = _MIN_STEP,
+) -> tuple[np.ndarray, np.ndarray]:
+    """
+    Resample trajectory to M poses at uniform progress knots.
+
+    progress_knots: σ_k = k * Σ/(M-1), k = 0,...,M-1.
+    Positions interpolated with cubic spline in progress; orientations with SLERP.
+
+    Returns (positions_M, quaternions_M) (M, 3), (M, 4) wxyz.
+    """
+    positions = np.asarray(positions, dtype=np.float64)
+    quaternions = np.asarray(quaternions, dtype=np.float64)
+    n = positions.shape[0]
+    if n < 2:
+        raise ValueError("resample_by_progress requires at least 2 poses")
+
+    if progress is None:
+        progress = compute_progress(
+            positions, quaternions, kind=progress_kind, alpha=alpha, min_step=min_step
+        )
+    else:
+        progress = np.asarray(progress, dtype=np.float64)
+        if progress.shape[0] != n:
+            raise ValueError("progress length must match positions/quaternions")
+
+    sigma_total = progress[-1]
+    if sigma_total <= 0:
+        sigma_total = 1.0
+    sigma_knots = np.linspace(0, sigma_total, M, dtype=np.float64)
+
+    pos_resample = np.zeros((M, 3))
+    quat_resample = np.zeros((M, 4))
+    for j in range(3):
+        cs = CubicSpline(progress, positions[:, j])
+        pos_resample[:, j] = cs(sigma_knots)
+
+    for k in range(M):
+        s = sigma_knots[k]
+        if s <= progress[0] + _EPS:
+            quat_resample[k] = quaternions[0].copy()
+            continue
+        if s >= progress[-1] - _EPS:
+            quat_resample[k] = quaternions[-1].copy()
+            continue
+        i = np.searchsorted(progress, s, side="right") - 1
+        i = min(max(i, 0), n - 2)
+        p_lo, p_hi = progress[i], progress[i + 1]
+        segment = p_hi - p_lo
+        if segment <= _EPS:
+            t = 0.0
+        else:
+            t = float((s - p_lo) / segment)
+        t = np.clip(t, 0.0, 1.0)
+        quat_resample[k] = geom.quat_slerp(quaternions[i], quaternions[i + 1], t)
+
+    return pos_resample, quat_resample
+
+
+def encode_dhb_dr_ti(
+    positions: np.ndarray,
+    quaternions: np.ndarray,
+    M: int,
+    progress_kind: ProgressKind = "hybrid",
+    alpha: float = 0.5,
+    method: Union[str, EncodingMethod] = EncodingMethod.POSITION,
+    use_default_initial_frames: bool = True,
+    init_pose: Optional[Dict[str, np.ndarray]] = None,
+    dhb_method: DHBMethod = DHBMethod.DOUBLE_REFLECTION,
+    min_step: float = _MIN_STEP,
+    **encode_kw,
+) -> Dict[str, Any]:
+    """
+    Time-invariant encode: reparameterize by progress to M samples, then DHB-DR encode.
+
+    Returns same structure as encode_dhb_dr (linear_motion_invariants, angular_motion_invariants,
+    initial_pose, ...).
+    """
+    from dhb_xr.encoder.dhb_dr import encode_dhb_dr
+
+    pos_m, quat_m = resample_by_progress(
+        positions, quaternions, M,
+        progress_kind=progress_kind, alpha=alpha, min_step=min_step,
+    )
+    return encode_dhb_dr(
+        pos_m, quat_m,
+        method=method,
+        use_default_initial_frames=use_default_initial_frames,
+        init_pose=init_pose,
+        dhb_method=dhb_method,
+        **encode_kw,
+    )
+
+
+def encode_dhb_qr_ti(
+    positions: np.ndarray,
+    quaternions: np.ndarray,
+    M: int,
+    progress_kind: ProgressKind = "hybrid",
+    alpha: float = 0.5,
+    method: EncodingMethod = EncodingMethod.POSITION,
+    use_default_initial_frames: bool = True,
+    init_pose: Optional[Dict[str, np.ndarray]] = None,
+    min_step: float = _MIN_STEP,
+    **encode_kw,
+) -> Dict[str, Any]:
+    """
+    Time-invariant encode: reparameterize by progress to M samples, then DHB-QR encode.
+
+    Returns same structure as encode_dhb_qr.
+    """
+    from dhb_xr.encoder.dhb_qr import encode_dhb_qr
+
+    pos_m, quat_m = resample_by_progress(
+        positions, quaternions, M,
+        progress_kind=progress_kind, alpha=alpha, min_step=min_step,
+    )
+    return encode_dhb_qr(
+        pos_m, quat_m,
+        method=method,
+        use_default_initial_frames=use_default_initial_frames,
+        init_pose=init_pose,
+        **encode_kw,
+    )

dhb_xr/encoder/encoder_torch.py

@@ -0,0 +1,54 @@
+"""Batched DHB encoder in PyTorch: wrapper over numpy encode for GPU-friendly batch API."""
+
+try:
+    import torch
+    import torch.nn as nn
+except ImportError:
+    torch = None
+    nn = None
+
+if torch is not None:
+
+    import numpy as np
+    from dhb_xr.encoder.dhb_dr import encode_dhb_dr
+    from dhb_xr.core.types import DHBMethod, EncodingMethod
+
+    class DHBEncoderTorch(nn.Module):
+        """
+        Batched DHB-DR encoder.
+        positions (B, N, 3), quaternions (B, N, 4) wxyz -> invariants (B, N-2, 8).
+        Uses numpy encode_dhb_dr per batch item and stacks; supports .to(device) for moving data to GPU.
+        """
+
+        def __init__(self, dhb_method: str = "double_reflection"):
+            super().__init__()
+            self.dhb_method = DHBMethod.DOUBLE_REFLECTION if dhb_method == "double_reflection" else DHBMethod.ORIGINAL
+            self.k = 4 if self.dhb_method == DHBMethod.DOUBLE_REFLECTION else 3
+
+        def forward(
+            self,
+            positions: torch.Tensor,
+            quaternions: torch.Tensor,
+        ) -> torch.Tensor:
+            B, N, _ = positions.shape
+            assert quaternions.shape == (B, N, 4)
+            device = positions.device
+            dtype = positions.dtype
+            inv_list = []
+            for b in range(B):
+                pos = positions[b].detach().cpu().numpy()
+                quat = quaternions[b].detach().cpu().numpy()
+                out = encode_dhb_dr(
+                    pos, quat,
+                    method=EncodingMethod.POSITION,
+                    use_default_initial_frames=True,
+                    dhb_method=self.dhb_method,
+                )
+                lin = out["linear_motion_invariants"]
+                ang = out["angular_motion_invariants"]
+                inv_list.append(torch.from_numpy(np.concatenate([lin, ang], axis=1)))
+            out_t = torch.stack(inv_list, dim=0).to(device=device, dtype=dtype)
+            return out_t
+
+else:
+    DHBEncoderTorch = None

dhb_xr/encoder/padding.py

@@ -0,0 +1,82 @@
+"""
+Boundary extrapolation for length preservation.
+Prepends 2 poses and appends 1 pose so invariant length aligns with N-1 steps.
+"""
+
+import numpy as np
+from typing import Dict, Any
+
+from dhb_xr.core import geometry as geom
+
+
+def extrapolate_boundary_poses(
+    positions: np.ndarray,
+    quaternions: np.ndarray,
+) -> tuple:
+    """
+    Extrapolate boundary poses for length-preserved encoding.
+
+    Prepends g_{-2}, g_{-1} and appends g_N using first-order extrapolation:
+    - Position: p_{-1} = 2*p_0 - p_1, p_{-2} = 3*p_0 - 2*p_1, p_N = 2*p_{N-1} - p_{N-2}
+    - Rotation: Lie-algebra consistent; R_{-1} = R_0 * exp(-[Δr_0]_×), etc.
+
+    positions: (N, 3)
+    quaternions: (N, 4) wxyz
+
+    Returns (positions_ext, quaternions_ext) with shape (N+3, 3) and (N+3, 4).
+    """
+    positions = np.asarray(positions, dtype=np.float64)
+    quaternions = np.asarray(quaternions, dtype=np.float64)
+    n = positions.shape[0]
+    assert n >= 2 and quaternions.shape[0] == n
+
+    # Rotation differentials for extrapolation
+    rotation_diff = np.zeros((n - 1, 3))
+    for i in range(1, n):
+        R_prev = geom.quat_to_rot(quaternions[i - 1]).T
+        R_curr = geom.quat_to_rot(quaternions[i]).T
+        R_rel = R_curr @ R_prev.T
+        rotation_diff[i - 1] = geom.rot_to_axis_angle(R_rel)
+
+    # Prepend: p_{-1} = 2*p_0 - p_1, p_{-2} = 3*p_0 - 2*p_1
+    dp0 = positions[1] - positions[0]
+    p_minus1 = positions[0] - dp0  # 2*p_0 - p_1
+    p_minus2 = p_minus1 - dp0  # 3*p_0 - 2*p_1
+
+    # R_{-1} = R_0 * exp(-[Δr_0]_×), Δr_{-1} = Δr_0 for first-order
+    dr0 = rotation_diff[0]
+    R0 = geom.quat_to_rot(quaternions[0])
+    R_minus1 = R0 @ geom.axis_angle_to_rot(-dr0).T
+    q_minus1 = geom.rot_to_quat(R_minus1)
+    dr_minus1 = dr0  # first-order
+    R_minus2_mat = geom.quat_to_rot(q_minus1) @ geom.axis_angle_to_rot(-dr_minus1).T
+    q_minus2 = geom.rot_to_quat(R_minus2_mat)
+
+    # Append: p_N = 2*p_{N-1} - p_{N-2}
+    p_last = 2 * positions[-1] - positions[-2]
+    dr_last = rotation_diff[-1]
+    R_last = geom.quat_to_rot(quaternions[-1]) @ geom.axis_angle_to_rot(dr_last).T
+    q_last = geom.rot_to_quat(R_last)
+
+    positions_ext = np.vstack(
+        (p_minus2.reshape(1, 3), p_minus1.reshape(1, 3), positions, p_last.reshape(1, 3))
+    )
+    quaternions_ext = np.vstack(
+        (
+            q_minus2.reshape(1, 4),
+            q_minus1.reshape(1, 4),
+            quaternions,
+            q_last.reshape(1, 4),
+        )
+    )
+    return positions_ext, quaternions_ext
+
+
+def apply_length_preserving_padding(
+    positions: np.ndarray,
+    quaternions: np.ndarray,
+) -> tuple:
+    """
+    Same as extrapolate_boundary_poses. Kept for API compatibility.
+    """
+    return extrapolate_boundary_poses(positions, quaternions)

dhb_xr/generative/__init__.py

@@ -0,0 +1,78 @@
+"""
+Generative models for DHB-Token trajectory generation.
+
+This module provides variational flow matching (VFM) and rectified flow matching (V-RFM)
+for multi-modal trajectory generation in the DHB invariant latent space.
+
+Key components:
+- FlowMatcher: Base flow matching model for continuous latent generation
+- VariationalFlowMatcher: VFM with latent conditioning for multi-modal generation
+- VFMTokenGenerator: End-to-end integration with DHB tokenizers
+
+Example usage:
+    >>> from dhb_xr.generative import VariationalFlowMatcher, VFMTokenGenerator
+    >>> from dhb_xr.tokenization import DHBTokenizer
+    >>>
+    >>> # Create tokenizer and flow matcher
+    >>> tokenizer = DHBTokenizer(invariant_dim=8, latent_dim=32, codebook_size=512)
+    >>> flow_matcher = VariationalFlowMatcher(latent_dim=32, hidden_dim=128)
+    >>>
+    >>> # Create end-to-end generator
+    >>> generator = VFMTokenGenerator(tokenizer, flow_matcher)
+    >>>
+    >>> # Generate multi-modal trajectories
+    >>> invariants = generator.generate_multimodal(prefix_invariants, num_modes=4)
+"""
+
+from .flow_matching import (
+    SinusoidalTimeEmbedding,
+    VelocityNetwork,
+    FlowMatcher,
+    VariationalFlowMatcher,
+)
+from .sampling import euler_solve, rk4_solve, ode_solve
+from .latent_encoder import (
+    LatentEncoder,
+    CategoricalLatentEncoder,
+    HybridLatentEncoder,
+)
+from .vfm_tokenizer import (
+    VFMTokenGenerator,
+    ConditionalVFMGenerator,
+)
+from .training import (
+    InvariantDataset,
+    train_vfm_tokenizer,
+    evaluate_model,
+    compute_reconstruction_error,
+    compute_generation_diversity,
+    linear_kl_schedule,
+    cyclical_kl_schedule,
+)
+
+__all__ = [
+    # Flow matching
+    "SinusoidalTimeEmbedding",
+    "VelocityNetwork",
+    "FlowMatcher",
+    "VariationalFlowMatcher",
+    # Latent encoders
+    "LatentEncoder",
+    "CategoricalLatentEncoder",
+    "HybridLatentEncoder",
+    # VFM Token generators
+    "VFMTokenGenerator",
+    "ConditionalVFMGenerator",
+    # Training utilities
+    "InvariantDataset",
+    "train_vfm_tokenizer",
+    "evaluate_model",
+    "compute_reconstruction_error",
+    "compute_generation_diversity",
+    "linear_kl_schedule",
+    "cyclical_kl_schedule",
+    # ODE solvers
+    "euler_solve",
+    "rk4_solve",
+    "ode_solve",
+]