| from __future__ import print_function |
| import sys |
| import numpy as np |
| import y2023.control_loops.python.graph_paths as graph_paths |
| |
| |
| def index_function_name(name): |
| return "%sIndex" % name |
| |
| |
| def path_function_name(name): |
| return "Make%sPath" % name |
| |
| |
| def add_edge(cc_file, name, segment, index, reverse): |
| cc_file.append(" // Adding edge %d" % index) |
| vmax = "vmax" |
| if segment.vmax: |
| vmax = "::std::min(vmax, %f)" % segment.vmax |
| |
| alpha_unitizer = "alpha_unitizer" |
| if segment.alpha_unitizer is not None: |
| alpha_unitizer = "(::Eigen::Matrix<double, 3, 3>() << %f, %f, %f, %f, %f, %f, %f, %f, %f).finished()" % ( |
| segment.alpha_unitizer[0, 0], |
| segment.alpha_unitizer[0, 1], |
| segment.alpha_unitizer[0, 2], |
| segment.alpha_unitizer[1, 0], |
| segment.alpha_unitizer[1, 1], |
| segment.alpha_unitizer[1, 2], |
| segment.alpha_unitizer[2, 0], |
| segment.alpha_unitizer[2, 1], |
| segment.alpha_unitizer[2, 2], |
| ) |
| cc_file.append(" trajectories->emplace_back(%s," % (vmax)) |
| cc_file.append(" %s," % (alpha_unitizer)) |
| if reverse: |
| cc_file.append( |
| " Trajectory(dynamics, &hybrid_roll_joint_loop->plant(), Path::Reversed(%s()), 0.005));" |
| % (path_function_name(str(name)))) |
| else: |
| cc_file.append( |
| " Trajectory(dynamics, &hybrid_roll_joint_loop->plant(), %s(), 0.005));" |
| % (path_function_name(str(name)))) |
| |
| start_index = None |
| end_index = None |
| for point, name in graph_paths.points: |
| if (point[:2] == segment.start |
| ).all() and point[2] == segment.alpha_rolls[0][1]: |
| start_index = name |
| if (point[:2] == segment.end |
| ).all() and point[2] == segment.alpha_rolls[-1][1]: |
| end_index = name |
| |
| if reverse: |
| start_index, end_index = end_index, start_index |
| |
| cc_file.append( |
| " edges.push_back(SearchGraph::Edge{%s(), %s()," % |
| (index_function_name(start_index), index_function_name(end_index))) |
| cc_file.append( |
| " (trajectories->back().trajectory.path().length() + 0.2)});" |
| ) |
| |
| # TODO(austin): Allow different vmaxes for different paths. |
| cc_file.append(" trajectories->back().trajectory.OptimizeTrajectory(") |
| cc_file.append(" trajectories->back().alpha_unitizer,") |
| cc_file.append(" trajectories->back().vmax);") |
| cc_file.append("") |
| |
| |
| def main(argv): |
| cc_file = [] |
| cc_file.append("") |
| cc_file.append("#include <memory>") |
| cc_file.append("") |
| cc_file.append( |
| "#include \"frc971/control_loops/double_jointed_arm/graph.h\"") |
| cc_file.append( |
| "#include \"y2023/control_loops/superstructure/arm/generated_graph.h\"" |
| ) |
| cc_file.append( |
| "#include \"y2023/control_loops/superstructure/arm/trajectory.h\"") |
| cc_file.append( |
| "#include \"y2023/control_loops/superstructure/roll/integral_hybrid_roll_plant.h\"" |
| ) |
| |
| cc_file.append("using frc971::control_loops::arm::SearchGraph;") |
| cc_file.append( |
| "using y2023::control_loops::superstructure::arm::Trajectory;") |
| cc_file.append("using y2023::control_loops::superstructure::arm::Path;") |
| cc_file.append("using y2023::control_loops::superstructure::arm::NSpline;") |
| cc_file.append( |
| "using y2023::control_loops::superstructure::arm::CosSpline;") |
| |
| cc_file.append("") |
| cc_file.append("namespace y2023 {") |
| cc_file.append("namespace control_loops {") |
| cc_file.append("namespace superstructure {") |
| cc_file.append("namespace arm {") |
| |
| h_file = [] |
| h_file.append( |
| "#ifndef Y2023_CONTROL_LOOPS_SUPERSTRUCTURE_ARM_GENERATED_GRAPH_H_") |
| h_file.append( |
| "#define Y2023_CONTROL_LOOPS_SUPERSTRUCTURE_ARM_GENERATED_GRAPH_H_") |
| h_file.append("") |
| h_file.append("#include <memory>") |
| h_file.append("") |
| h_file.append( |
| "#include \"frc971/control_loops/double_jointed_arm/graph.h\"") |
| h_file.append( |
| "#include \"y2023/control_loops/superstructure/arm/arm_constants.h\"") |
| h_file.append( |
| "#include \"y2023/control_loops/superstructure/arm/trajectory.h\"") |
| |
| h_file.append("") |
| h_file.append("namespace y2023 {") |
| h_file.append("namespace control_loops {") |
| h_file.append("namespace superstructure {") |
| h_file.append("namespace arm {") |
| |
| h_file.append("using frc971::control_loops::arm::SearchGraph;") |
| h_file.append( |
| "using y2023::control_loops::superstructure::arm::Trajectory;") |
| h_file.append("using y2023::control_loops::superstructure::arm::Path;") |
| h_file.append( |
| "using y2023::control_loops::superstructure::arm::kArmConstants;") |
| |
| h_file.append("") |
| h_file.append("struct TrajectoryAndParams {") |
| h_file.append(" TrajectoryAndParams(double new_vmax,") |
| h_file.append( |
| " const ::Eigen::Matrix<double, 3, 3> &new_alpha_unitizer," |
| ) |
| h_file.append(" Trajectory &&new_trajectory)") |
| h_file.append(" : vmax(new_vmax),") |
| h_file.append(" alpha_unitizer(new_alpha_unitizer),") |
| h_file.append(" trajectory(::std::move(new_trajectory)) {}") |
| h_file.append(" double vmax;") |
| h_file.append(" ::Eigen::Matrix<double, 3, 3> alpha_unitizer;") |
| h_file.append(" Trajectory trajectory;") |
| h_file.append("};") |
| h_file.append("") |
| |
| # Now dump out the vertices and associated constexpr vertex name functions. |
| for index, point in enumerate(graph_paths.points): |
| h_file.append("") |
| h_file.append("constexpr uint32_t %s() { return %d; }" % |
| (index_function_name(point[1]), index)) |
| h_file.append("inline ::Eigen::Matrix<double, 3, 1> %sPoint() {" % |
| point[1]) |
| h_file.append( |
| " return (::Eigen::Matrix<double, 3, 1>() << %f, %f, %f).finished();" |
| % (np.pi / 2.0 - point[0][0], np.pi / 2.0 - point[0][1], |
| np.pi / 2.0 - point[0][2])) |
| h_file.append("}") |
| |
| front_points = [ |
| index_function_name(point[1]) + "()" |
| for point in graph_paths.front_points |
| ] |
| h_file.append("") |
| h_file.append("constexpr ::std::array<uint32_t, %d> FrontPoints() {" % |
| len(front_points)) |
| h_file.append(" return ::std::array<uint32_t, %d>{{%s}};" % |
| (len(front_points), ", ".join(front_points))) |
| h_file.append("}") |
| |
| back_points = [ |
| index_function_name(point[1]) + "()" |
| for point in graph_paths.back_points |
| ] |
| h_file.append("") |
| h_file.append("constexpr ::std::array<uint32_t, %d> BackPoints() {" % |
| len(back_points)) |
| h_file.append(" return ::std::array<uint32_t, %d>{{%s}};" % |
| (len(back_points), ", ".join(back_points))) |
| h_file.append("}") |
| |
| # Add the Make*Path functions. |
| h_file.append("") |
| cc_file.append("") |
| for name, segment in list(enumerate(graph_paths.unnamed_segments)) + [ |
| (x.name, x) for x in graph_paths.named_segments |
| ]: |
| h_file.append("::std::unique_ptr<Path> %s();" % |
| path_function_name(name)) |
| cc_file.append("::std::unique_ptr<Path> %s() {" % |
| path_function_name(name)) |
| cc_file.append( |
| " return ::std::unique_ptr<Path>(new Path(CosSpline{NSpline<4, 2>((Eigen::Matrix<double, 2, 4>() << " |
| ) |
| points = [ |
| segment.start, segment.control1, segment.control2, segment.end |
| ] |
| for i in range(len(points)): |
| cc_file.append("%.12f," % (np.pi / 2.0 - points[i][0])) |
| for i in range(len(points)): |
| cc_file.append("%.12f%s" % (np.pi / 2.0 - points[i][1], |
| ", " if i != len(points) - 1 else "")) |
| |
| cc_file.append(").finished()), {") |
| for alpha, roll in segment.alpha_rolls: |
| cc_file.append( |
| "CosSpline::AlphaTheta{.alpha = %.12f, .theta = %.12f}" % |
| (alpha, np.pi / 2.0 - roll)) |
| if alpha != segment.alpha_rolls[-1][0]: |
| cc_file.append(", ") |
| cc_file.append(" }}));") |
| cc_file.append("}") |
| |
| # Matrix of nodes |
| h_file.append("::std::vector<::Eigen::Matrix<double, 3, 1>> PointList();") |
| |
| cc_file.append( |
| "::std::vector<::Eigen::Matrix<double, 3, 1>> PointList() {") |
| cc_file.append(" ::std::vector<::Eigen::Matrix<double, 3, 1>> points;") |
| for point in graph_paths.points: |
| cc_file.append( |
| " points.push_back((::Eigen::Matrix<double, 3, 1>() << %.12s, %.12s, %.12s).finished());" |
| % (np.pi / 2.0 - point[0][0], np.pi / 2.0 - point[0][1], |
| np.pi / 2.0 - point[0][2])) |
| cc_file.append(" return points;") |
| cc_file.append("}") |
| |
| # Now create the MakeSearchGraph function. |
| h_file.append("") |
| h_file.append("// Builds a search graph.") |
| h_file.append("SearchGraph MakeSearchGraph(" |
| "const frc971::control_loops::arm::Dynamics *dynamics, " |
| "::std::vector<TrajectoryAndParams> *trajectories,") |
| h_file.append(" " |
| "const ::Eigen::Matrix<double, 3, 3> &alpha_unitizer,") |
| h_file.append(" " |
| "double vmax,") |
| h_file.append( |
| "const StateFeedbackLoop<3, 1, 1, double, StateFeedbackHybridPlant<3, 1, 1>, HybridKalman<3, 1, 1>> *hybrid_roll_joint_loop);" |
| ) |
| cc_file.append("SearchGraph MakeSearchGraph(" |
| "const frc971::control_loops::arm::Dynamics *dynamics, " |
| "::std::vector<TrajectoryAndParams> *trajectories,") |
| cc_file.append(" " |
| "const ::Eigen::Matrix<double, 3, 3> &alpha_unitizer,") |
| cc_file.append(" " |
| "double vmax,") |
| cc_file.append( |
| "const StateFeedbackLoop<3, 1, 1, double, StateFeedbackHybridPlant<3, 1, 1>, HybridKalman<3, 1, 1>> *hybrid_roll_joint_loop) {" |
| ) |
| cc_file.append(" ::std::vector<SearchGraph::Edge> edges;") |
| |
| index = 0 |
| segments_and_names = list(enumerate(graph_paths.unnamed_segments)) + [ |
| (x.name, x) for x in graph_paths.named_segments |
| ] |
| |
| for name, segment in segments_and_names: |
| add_edge(cc_file, name, segment, index, False) |
| index += 1 |
| add_edge(cc_file, name, segment, index, True) |
| index += 1 |
| |
| cc_file.append(" return SearchGraph(%d, ::std::move(edges));" % |
| len(graph_paths.points)) |
| cc_file.append("}") |
| |
| h_file.append("") |
| h_file.append("} // namespace arm") |
| h_file.append("} // namespace superstructure") |
| h_file.append("} // namespace control_loops") |
| h_file.append("} // namespace y2023") |
| h_file.append("") |
| h_file.append( |
| "#endif // Y2023_CONTROL_LOOPS_SUPERSTRUCTURE_ARM_GENERATED_GRAPH_H_") |
| |
| cc_file.append("} // namespace arm") |
| cc_file.append("} // namespace superstructure") |
| cc_file.append("} // namespace control_loops") |
| cc_file.append("} // namespace y2023") |
| |
| if len(argv) == 3: |
| with open(argv[1], "w") as hfd: |
| hfd.write("\n".join(h_file)) |
| |
| with open(argv[2], "w") as ccfd: |
| ccfd.write("\n".join(cc_file)) |
| else: |
| print("\n".join(h_file)) |
| print("\n".join(cc_file)) |
| |
| |
| if __name__ == '__main__': |
| main(sys.argv) |