Allow setting input prefix name when converting to function

ami-iit · Apr 17, 2024 · 9bbca1e · 9bbca1e
1 parent 4d15127
commit 9bbca1e
Show file tree

Hide file tree

Showing 5 changed files with 44 additions and 16 deletions.
diff --git a/src/hippopt/base/opti_solver.py b/src/hippopt/base/opti_solver.py
@@ -591,13 +591,16 @@ def get_free_parameters_names(self) -> list[str]:
         return self._free_parameters
 
     def to_function(
-        self, name: str = "opti_function", options: dict = None
+        self,
+        input_name_prefix: str,
+        function_name: str = "opti_function",
+        options: dict = None,
     ) -> cs.Function:
         self._cost = self._cost if self._cost is not None else cs.MX(0)
         self._solver.minimize(self._cost)
 
-        # Prepend guess to the variable names
-        guess_names = ["guess." + name for name in self._objects_dict]
+        # Prepend input_name_prefix to the variable names
+        guess_names = [input_name_prefix + name for name in self._objects_dict]
         all_variables_values = list(self._objects_dict.values())
 
         # Workaround for https://github.com/casadi/casadi/issues/3655
@@ -622,7 +625,7 @@ def to_function(
 
         options = {} if options is None else options
         return self._solver.to_function(
-            name,
+            function_name,
             all_variables_values,
             output_variables,
             guess_names,

diff --git a/src/hippopt/turnkey_planners/humanoid_kinodynamic/main_periodic_step_sensitivity.py b/src/hippopt/turnkey_planners/humanoid_kinodynamic/main_periodic_step_sensitivity.py
@@ -267,9 +267,11 @@ def compute_state(
             pf_parametric_link_length_multipliers
         )
     pf_function = pf_input.to_function(
-        name="pose_finder", options={"error_on_fail": True}
+        input_name_prefix="pf_in.",
+        function_name="pose_finder",
+        options={"error_on_fail": True},
     )
-    pf_guess_dict = pf_guess.to_dict(prefix="guess.")
+    pf_guess_dict = pf_guess.to_dict(prefix="pf_in.")
     output_pf_dict = pf_function(**pf_guess_dict)
 
     for k in output_pf_dict:
@@ -549,10 +551,12 @@ def get_references(
         )
 
     planner_function = planner.to_function(
-        name="kinodynamic_walking", options={"error_on_fail": True}
+        input_name_prefix="in.",
+        function_name="kinodynamic_walking",
+        options={"error_on_fail": True},
     )
 
-    initial_guess_dict = planner_guess.to_dict(prefix="guess.")
+    initial_guess_dict = planner_guess.to_dict(prefix="in.")
     initial_guess_dict_pruned = {}
     for key in initial_guess_dict:
         if isinstance(initial_guess_dict[key], np.ndarray) or isinstance(

diff --git a/src/hippopt/turnkey_planners/humanoid_kinodynamic/planner.py b/src/hippopt/turnkey_planners/humanoid_kinodynamic/planner.py
@@ -1071,11 +1071,16 @@ def solve(self) -> hp.Output[Variables]:
         return output
 
     def to_function(
-        self, name: str = "opti_function", options: dict = None
+        self,
+        input_name_prefix: str,
+        function_name: str = "opti_function",
+        options: dict = None,
     ) -> cs.Function:
 
         inner_function = self.optimization_solver.to_function(
-            name=name + "_internal", options=options
+            input_name_prefix=input_name_prefix,
+            function_name=function_name + "_internal",
+            options=options,
         )
         variable_names = inner_function.name_in()
         variables_list = []
@@ -1087,9 +1092,13 @@ def to_function(
         optimization_structure = copy.deepcopy(
             self.optimization_solver.get_optimization_objects()
         )
-        optimization_structure.from_dict(input_dict=variables_sym, prefix="guess.")
+        optimization_structure.from_dict(
+            input_dict=variables_sym, prefix=input_name_prefix
+        )
         mass_regularized_vars = self._apply_mass_regularization(optimization_structure)
-        output_dict = inner_function(**mass_regularized_vars.to_dict(prefix="guess."))
+        output_dict = inner_function(
+            **mass_regularized_vars.to_dict(prefix=input_name_prefix)
+        )
         optimization_structure = copy.deepcopy(
             self.optimization_solver.get_optimization_objects()
         )
@@ -1104,7 +1113,12 @@ def to_function(
             output_names.append(n)
 
         return cs.Function(
-            name, variables_list, output_values, variable_names, output_names, options
+            function_name,
+            variables_list,
+            output_values,
+            variable_names,
+            output_names,
+            options,
         )
 
     def get_adam_model(self) -> adam.model.Model:

diff --git a/src/hippopt/turnkey_planners/humanoid_pose_finder/planner.py b/src/hippopt/turnkey_planners/humanoid_pose_finder/planner.py
@@ -652,9 +652,16 @@ def solve(self) -> hp.Output[Variables]:
         return self.op.problem.solve()
 
     def to_function(
-        self, name: str = "opti_function", options: dict = None
+        self,
+        input_name_prefix: str,
+        function_name: str = "opti_function",
+        options: dict = None,
     ) -> cs.Function:
-        return self.optimization_solver.to_function(name=name, options=options)
+        return self.optimization_solver.to_function(
+            input_name_prefix=input_name_prefix,
+            function_name=function_name,
+            options=options,
+        )
 
     def get_adam_model(self) -> adam.model.Model:
         if self.parametric_model:

diff --git a/test/test_optimization_problem.py b/test/test_optimization_problem.py
@@ -312,7 +312,7 @@ def test_opti_to_function():
         ),
     )
 
-    opti_function = opti_solver.to_function()
+    opti_function = opti_solver.to_function(input_name_prefix="guess.")
     output_dict = opti_function(**initial_guess.to_dict(prefix="guess."))
     output = MyTestVarAndPar()
     output.from_dict(output_dict)