Apply suggestions from code review to centroidal expressions

Co-authored-by: Diego Ferigo <[email protected]>

src/hippopt/robot_planning/expressions/centroidal.py

@@ -14,10 +14,10 @@ def centroidal_dynamics_with_point_forces(
 ) -> cs.Function:
     options = {} if options is None else options
 
-    if point_position_names is None:
-        point_position_names = []
-        for i in range(number_of_points):
-            point_position_names.append(f"p{i}")
+    point_position_names = (
+        point_position_names if point_position_names is not None else
+        [f"p{i}" for i in range(number_of_points)]
+    )
 
     assert len(point_position_names) == number_of_points
 
@@ -44,20 +44,15 @@ def centroidal_dynamics_with_point_forces(
 
     p = []
     f = []
-    for i in range(number_of_points):
-        p.append(cs.MX.sym(point_position_names[i], 3))
-        input_vars.append(p[i])
-        f.append(cs.MX.sym(point_force_names[i], 3))
-        input_vars.append(f[i])
-
-    input_names = []
-    for var in input_vars:
-        input_names.append(var.name())
+    for point_position_name, point_force_name in zip(point_position_names, point_force_names):
+        p.append(cs.MX.sym(point_position_name, 3))
+        input_vars.append(p[-1])
+        f.append(cs.MX.sym(point_force_name, 3))
+        input_vars.append(f[-1])
 
-    h_g = m @ g
+    input_names = [var.name() for var in input_vars]
 
-    for i in range(number_of_points):
-        h_g = h_g + cs.vertcat(f[i], cs.cross(p[i] - x, f[i]))
+    h_g = cs.sum([m * g] + [cs.vertcat(fi, cs.cross(pi - x, fi)) for fi, pi in zip(f, p)])
 
     return cs.Function(
         "centroidal_dynamics_with_point_forces",