brocksam / pycollo

General-purpose optimal control, trajectory optimisation and parameter optimisation using direct collocation
https://brocksam.github.io/pycollo/
MIT License
8 stars 3 forks source link

Feature Request: add setting enabling users to save a solution as a json file #20

Open jackirvine97 opened 4 years ago

jackirvine97 commented 4 years ago

This function should provide a starting point:

def parse_solution_as_json(problem, json_filename='solution'):
    """
    Parses optimal_control_problem.solution object as json and saves
    it in working directory.
    """
    try:
        solution = problem.solution
    except Exception as e:
        print(e)
        return

    solution_dict = {}
    num_phases = len(solution.state)

    for phase_index in range(num_phases):
        time_list = solution.time[phase_index].tolist()

        state_symbols = [str(symbol) for symbol in problem.phases[phase_index].state_variables]
        states = solution.state[phase_index].tolist()
        state_dict = dict(zip(state_symbols, states))

        control_symbols = [str(symbol) for symbol in problem.phases[phase_index].control_variables]
        controls = solution.control[phase_index].tolist()
        control_dict = dict(zip(control_symbols, controls))

        phase_dict = {"Times": time_list, "States": state_dict, "Controls": control_dict}

        solution_dict[str(problem.phases[phase_index])] = phase_dict

    solution_dict["Parameter Variables"] = solution.parameter.tolist()

    with open(f'{json_filename}.json', 'w') as json_file:
        json.dump(solution_dict, json_file)

    return