Skip to content
Snippets Groups Projects
optimizer.py 3.21 KiB
Newer Older
  • Learn to ignore specific revisions
  • class Optimizer:
        def __init__(self, problem):
            self.problem = problem
    
    
    rachelmoan's avatar
    rachelmoan committed
        def solve_optimization_problem(self, initial_guesses=None, solver_options=None, lam_g=None):
    
            opti = self.problem['opti']
    
            
            if initial_guesses:
                for param, value in initial_guesses.items():
                    opti.set_initial(self.problem[param], value)
    
    
    rachelmoan's avatar
    rachelmoan committed
            if lam_g is not None:
                opti.set_initial(opti.lam_g, lam_g)
    
    
            # Set numerical backend, with options if provided
            if solver_options:
                opti.solver('ipopt', solver_options)
            else:
                opti.solver('ipopt')
    
    
            def print_intermediates_callback(i):
                # print("Iteration:", i, "Current solution:", opti.debug.value(X), opti.debug.value(U))
                X_debug = opti.debug.value(X)
                U_debug = opti.debug.value(U)
    
                # plot the state and the control 
                # split a figure in half. The left side will show the positions, the right side will plot the controls
                # X[i*3, :] is the ith robot's x position, X[i*3+1, :] is the y position, X[i*3+2, :] is the heading
                # U[i*2, :] is the ith robot's linear velocity, U[i*2+1, :] is the ith robot's angular velocity
                import matplotlib.pyplot as plt
                fig, axs = plt.subplots(1, 2, figsize=(12, 6))
                for j in range(X_debug.shape[0]//3):
                    axs[0].plot(X_debug[j*3, :], X_debug[j*3+1, :], label=f"Robot {j}")
                    axs[0].scatter(X_debug[j*3, 0], X_debug[j*3+1, 0], color='green')
                    axs[0].scatter(X_debug[j*3, -1], X_debug[j*3+1, -1], color='red')
                    axs[0].set_title("Robot Positions")
                    axs[0].set_xlabel("X")
                    axs[0].set_ylabel("Y")
                    axs[0].legend()
    
                    axs[1].plot(U_debug[j*2, :], label=f"Robot {j}")
                    axs[1].plot(U_debug[j*2+1, :], label=f"Robot {j}")
                    axs[1].set_title("Robot Controls")
                    axs[1].set_xlabel("Time")
                    axs[1].set_ylabel("Control")
                    axs[1].legend()
    
                plt.show()
                
    
        
            # opti.callback(print_intermediates_callback)
    
    
            try:
                sol = opti.solve()   # actual solve
                status = 'succeeded'
            except:
                sol = None
                status = 'failed'
    
            results = {
                'status' : status,
                'solution' : sol,
            }
    
            # print(f"Final total = {sol.value(self.problem['cost'])}")
            # print(f"robot costs = {sol.value(self.problem['robot_cost'])}")
            # print(f"obstacle costs = {sol.value(self.problem['obs_cost'])}")
            # print(f"control costs = {sol.value(self.problem['control_cost'])}")
            # print(f"time costs = {sol.value(self.problem['time_cost'])}")
            # print(f"goal costs = {sol.value(self.problem['goal_cost'])}")
    
    
            if sol:
                for var_name, var in self.problem.items():
                    if var_name != 'opti':
                        results[var_name] = sol.value(var)
    
    
    rachelmoan's avatar
    rachelmoan committed
            opti = self.problem['opti']
            lam_g = sol.value(opti.lam_g)
            results['lam_g'] = lam_g
    
            return results,sol