enable option to calculate geometric invariants using rockit interface

examples/calculate_invariants_position_longtrial.py

@@ -0,0 +1,104 @@
+# Example for calculating invariants from a long trial of position data.
+
+import pandas as pd
+from mpl_toolkits.mplot3d import Axes3D
+from invariants_py.calculate_invariants.rockit_calculate_vector_invariants_position import OCP_calc_pos
+import numpy as np
+import matplotlib.pyplot as plt
+from invariants_py.reparameterization import reparameterize_positiontrajectory_arclength
+from invariants_py.data_handler import find_data_path
+
+# Load the CSV file
+df = pd.read_csv(find_data_path('trajectory_long.csv'))
+
+# Plot xyz coordinates with respect to timestamp
+plt.figure()
+plt.plot(df['timestamp'], df['x'], label='x')
+plt.plot(df['timestamp'], df['y'], label='y')
+plt.plot(df['timestamp'], df['z'], label='z')
+plt.xlabel('Timestamp')
+plt.ylabel('Coordinates')
+plt.legend()
+plt.title('XYZ Coordinates with respect to Timestamp')
+plt.show()
+
+# Plot the trajectory in 3D
+fig = plt.figure()
+ax = fig.add_subplot(111, projection='3d')
+ax.plot(df['x'], df['y'], df['z'])
+ax.set_aspect('equal')
+ax.set_xlabel('X')
+ax.set_ylabel('Y')
+ax.set_zlabel('Z')
+ax.set_title('3D Trajectory')
+plt.show()
+
+# Prepare the trajectory data
+trajectory = np.column_stack((df['x'], df['y'], df['z']))
+timestamps = df['timestamp'].values
+stepsize = np.mean(np.diff(timestamps))
+
+# Downsample the trajectory to 100 samples
+downsampled_indices = np.linspace(0, len(trajectory) - 1, 200, dtype=int)
+trajectory = trajectory[downsampled_indices]/1000 # Convert to meters
+timestamps = timestamps[downsampled_indices]
+stepsize = np.mean(np.diff(timestamps))
+
+# Reparameterize the trajectory based on arclength
+# Note: The reparameterization is not necessary if the data size is within the limit
+trajectory, arclength, arclength_n, nb_samples, stepsize = reparameterize_positiontrajectory_arclength(trajectory)
+
+# Use the standard approach if the data size is within the limit
+ocp = OCP_calc_pos(window_len=len(trajectory),fatrop_solver=True,geometric=True)
+invariants, reconstructed_trajectory, moving_frames = ocp.calculate_invariants(trajectory, stepsize)
+
+# Plot the calculated invariants as subplots
+fig, axs = plt.subplots(3, 1, figsize=(10, 8))
+
+axs[0].plot(timestamps, invariants[:, 0], label='Invariant 1')
+axs[0].plot(0, 0, label='Invariant 1')
+axs[0].set_xlabel('Timestamp')
+axs[0].set_ylabel('Invariant 1')
+axs[0].legend()
+axs[0].set_title('Calculated Geometric Invariant 1')
+
+axs[1].plot(timestamps, invariants[:, 1], label='Invariant 2')
+axs[1].set_xlabel('Timestamp')
+axs[1].set_ylabel('Invariant 2')
+axs[1].legend()
+axs[1].set_title('Calculated Geometric Invariant 2')
+
+axs[2].plot(timestamps, invariants[:, 2], label='Invariant 3')
+axs[2].set_xlabel('Timestamp')
+axs[2].set_ylabel('Invariant 3')
+axs[2].legend()
+axs[2].set_title('Calculated Geometric Invariant 3')
+
+plt.tight_layout()
+plt.show()
+
+# Plot reconstructed trajectory versus original trajectory
+plt.figure()
+plt.plot(timestamps, trajectory[:, 0], label='Original x')
+plt.plot(timestamps, trajectory[:, 1], label='Original y')
+plt.plot(timestamps, trajectory[:, 2], label='Original z')
+plt.plot(timestamps, reconstructed_trajectory[:, 0], label='Reconstructed x')
+plt.plot(timestamps, reconstructed_trajectory[:, 1], label='Reconstructed y')
+plt.plot(timestamps, reconstructed_trajectory[:, 2], label='Reconstructed z')
+plt.xlabel('Timestamp')
+plt.ylabel('Coordinates')
+plt.legend()
+plt.title('Original and Reconstructed Trajectory')
+plt.show()
+
+# Plot the reconstructed trajectory in 3D
+fig = plt.figure()
+ax = fig.add_subplot(111, projection='3d')
+ax.plot(trajectory[:, 0], trajectory[:, 1], trajectory[:, 2])
+ax.plot(reconstructed_trajectory[:, 0], reconstructed_trajectory[:, 1], reconstructed_trajectory[:, 2])
+ax.set_aspect('equal')
+ax.set_xlabel('X')
+ax.set_ylabel('Y')
+ax.set_zlabel('Z')
+ax.set_title('Reconstructed 3D Trajectory')
+plt.show()

invariants_py/calculate_invariants/rockit_calculate_vector_invariants_position.py

@@ -26,7 +26,7 @@
 
 class OCP_calc_pos:
 
-    def __init__(self, window_len=100, rms_error_traj=10**-3, fatrop_solver=False, bool_unsigned_invariants=False, planar_task=False, solver_options = {}):
+    def __init__(self, window_len=100, rms_error_traj=10**-3, fatrop_solver=False, bool_unsigned_invariants=False, planar_task=False, solver_options = {}, geometric=False):
         """
         Initializes an instance of the RockitCalculateVectorInvariantsPosition class.
         It specifies the optimal control problem (OCP) for calculating the invariants of a trajectory in a symbolic way.
@@ -113,8 +113,12 @@ def __init__(self, window_len=100, rms_error_traj=10**-3, fatrop_solver=False, b
             # TODO make normal to plane a parameter instead of hardcoding the Z-axis
             ocp.subject_to( cas.dot(R_t[:,2],np.array([0,0,1])) > 0)
 
-        # TODO can we implement geometric constraint (constant i1) by making i1 a state?
-
+        # Geometric invariants (optional), i.e. enforce constant speed
+        if geometric:
+            L = ocp.state()  # introduce extra state L for the speed
+            ocp.set_next(L, L)  # enforce constant speed
+            ocp.subject_to(invars[0] - L == 0, include_last=False)  # relate to first invariant
+
         """ Objective function """
 
         # Minimize moving frame invariants to deal with singularities and noise
@@ -142,7 +146,7 @@ def __init__(self, window_len=100, rms_error_traj=10**-3, fatrop_solver=False, b
             ocp._method.set_option("linsol_lu_fact_tol",1e-12)
         else:
             ocp.method(rockit.MultipleShooting(N=N-1))
-            ocp.solver('ipopt', {'expand':True, 'print_time':False, 'ipopt.tol':tolerance,'ipopt.print_info_string':'yes', 'ipopt.max_iter':max_iter,'ipopt.print_level':print_level, 'ipopt.ma57_automatic_scaling':'no', 'ipopt.linear_solver':'mumps'})
+            ocp.solver('ipopt', {'expand':True, 'print_time':False, 'ipopt.tol':tolerance, 'ipopt.print_info_string':'yes', 'ipopt.max_iter':max_iter, 'ipopt.print_level':print_level, 'ipopt.ma57_automatic_scaling':'no', 'ipopt.linear_solver':'mumps'})
 
         """ Encapsulate solver in a casadi function so that it can be easily reused """
         ocp.solve() # code generation

invariants_py/calculate_invariants/rockit_calculate_vector_invariants_position_mf.py

@@ -26,7 +26,7 @@
 
 class OCP_calc_pos:
 
-    def __init__(self, window_len=100, rms_error_traj=10**-2, fatrop_solver=False, bool_unsigned_invariants=False, planar_task=False, solver_options = {}):
+    def __init__(self, window_len=100, rms_error_traj=10**-2, fatrop_solver=False, bool_unsigned_invariants=False, planar_task=False, geometric = False, solver_options = {}):
         """
         Initializes an instance of the RockitCalculateVectorInvariantsPosition class.
         It specifies the optimal control problem (OCP) for calculating the invariants of a trajectory in a symbolic way.
@@ -117,8 +117,12 @@ def __init__(self, window_len=100, rms_error_traj=10**-2, fatrop_solver=False, b
             # TODO make normal to plane a parameter instead of hardcoding the Z-axis
             ocp.subject_to( cas.dot(R_t[:,2],np.array([0,0,1])) > 0)
 
-        # TODO can we implement geometric constraint (constant i1) by making i1 a state?
-
+        # Geometric invariants (optional), i.e. enforce constant speed
+        if geometric:
+            L = ocp.state()  # introduce extra state L for the speed
+            ocp.set_next(L, L)  # enforce constant speed
+            ocp.subject_to(invars[0] - L == 0, include_last=False)  # relate to first invariant
+
         """ Objective function """
 
         # Minimize moving frame invariants to deal with singularities and noise

invariants_py/calculate_invariants/rockit_calculate_vector_invariants_position_mj.py

@@ -8,7 +8,7 @@
 
 class OCP_calc_pos:
 
-    def __init__(self, window_len = 100, w_pos = 1, w_regul_jerk = 10**-5 , bool_unsigned_invariants = False, w_regul_invars = 0, fatrop_solver = False, planar_task = False, solver_options = {}):
+    def __init__(self, window_len = 100, w_pos = 1, w_regul_jerk = 10**-5 , bool_unsigned_invariants = False, w_regul_invars = 0, fatrop_solver = False, planar_task = False, geometric = False, solver_options = {}):
         fatrop_solver = check_solver(fatrop_solver)               
 
         # Set solver options
@@ -42,8 +42,8 @@ def __init__(self, window_len = 100, w_pos = 1, w_regul_jerk = 10**-5 , bool_uns
         ''' Constraints '''
 
         # test initial constraint on position
-        self.ocp.subject_to(self.ocp.at_tf(self.p_obj == self.p_obj_m))
-        self.ocp.subject_to(self.ocp.at_t0(self.p_obj == self.p_obj_m))
+        #self.ocp.subject_to(self.ocp.at_tf(self.p_obj == self.p_obj_m))
+        #self.ocp.subject_to(self.ocp.at_t0(self.p_obj == self.p_obj_m))
 
         # Orthonormality of rotation matrix (only needed for one timestep, property is propagated by integrator)  
         self.ocp.subject_to(self.ocp.at_t0(ocp_helper.tril_vec(R_t.T @ R_t - np.eye(3))==0.))
@@ -71,6 +71,12 @@ def __init__(self, window_len = 100, w_pos = 1, w_regul_jerk = 10**-5 , bool_uns
             self.ocp.subject_to(self.i1 >= 0) # velocity always positive
             #ocp.subject_to(invars[1,:]>=0) # curvature rate always positive
 
+        # Geometric invariants (optional), i.e. enforce constant speed
+        if geometric:
+            L = self.ocp.state()  # introduce extra state L for the speed
+            self.ocp.set_next(L, L)  # enforce constant speed
+            self.ocp.subject_to(self.i1 - L == 0)  # relate to first invariant
+
         ''' Objective function '''
         self.N_controls = window_len-1