diff --git a/src/KSP_tools/krpc_explore.py b/src/KSP_tools/krpc_explore.py index 1712081..bb3a708 100644 --- a/src/KSP_tools/krpc_explore.py +++ b/src/KSP_tools/krpc_explore.py @@ -1,289 +1,29 @@ -# import krpc -# -# conn = krpc.connect(name='KSP Explore') -# vessel = conn.space_center.active_vessel -# orbit = vessel.orbit -# surface = vessel.surface_velocity_reference_frame -# print(f'Vessel Name: {vessel.name}') -# print(f'Orbit Apoapsis: {orbit.apoapsis} m') -# print(f'Orbit Periapsis: {orbit.periapsis} m') -# print(f'Surface Speed: {vessel.surface_velocity(surface).magnitude} m/s') - import time - -import cvxpy as cp import krpc -import numpy as np -from krpc.services.spacecenter import ReferenceFrame, Vessel -from scipy.spatial.transform import Rotation as R - -from src.optimization import optimize_fuel - -landing_platform = np.array([159780.5, -1018.1, -578410.4]) -origin = np.array([-0.0021359772614235606, -0.004701372718752435, 0.152749662369653]) - -def optimize_fuel( - p_target: np.ndarray, - g: float, - m: float, - p0: np.ndarray, - v0: np.ndarray, - K: int, - h: float, - F_max: float, - alpha: float, - gamma: float, - **kwargs: dict, -) -> Tuple[np.ndarray, np.ndarray, float, cp.Problem]: - """ - - Minimize fuel consumption for a rocket to land on a target - - :param p_target: landing target in m - :param g: gravitational acceleration in m/s^2 - :param m: mass in kg - :param p0: position in m - :param v0: velocity in m/s - :param K: Number of discretization steps - :param h: discretization step in s - :param F_max: maximum thrust of engine in kg*m/s^2 (Newton) - :param alpha: Glide path angle in radian - :param gamma: converts fuel consumption to liters of fuel consumption - :return: position, thrust, fuel consumption, problem - """ - - P_min = p_target[2] - - # Variables - V = cp.Variable((K + 1, 3)) # velocity - P = cp.Variable((K + 1, 3)) # position - F = cp.Variable((K, 3)) # thrust - - # Constraints - # Match initial position and initial velocity - constraints = [ - V[0] == v0, - P[0] == p0, - P[:, 2] >= P_min, - # P[:, 2] <= P_max, - F[:, 2] >= 0, - V[:, 2] <= 0, - ] - - # Match final position and 0 velocity - constraints += [ - V[K] == [0, 0, 0], - P[K] == p_target, - ] - - # Physics dynamics for velocity - constraints += [V[1:, :2] == V[:-1, :2] + h * (F[:, :2] / m)] - constraints += [V[1:, 2] == V[:-1, 2] + h * (F[:, 2] / m - g)] - - # Physics dynamics for position - constraints += [P[1:] == P[:-1] + (h / 2) * (V[:-1] + V[1:])] - - # Maximum thrust constraint - constraints += [cp.norm(F, 2, axis=1) <= F_max] - - fuel_consumption = gamma * cp.sum(cp.norm(F, axis=1)) - - # Regularization - height_regularization = cp.sum(cp.abs(P[:, 2] - P_min)) - xy_regularization = cp.sum(cp.abs(P[:, :2] - p_target[:2].reshape((1, -1)))) - glide_violation = cp.pos(np.tan(alpha) * cp.norm(P[:, :2], axis=1) - P[:, 2]) - glide_violation = cp.sum(glide_violation) - gamma_height = 1 - gamma_xy = 1 - gamma_glide = 1000 - reg = ( - gamma_height * height_regularization - + gamma_xy * xy_regularization - + gamma_glide * glide_violation - ) - - problem = cp.Problem(cp.Minimize(fuel_consumption + reg), constraints) - problem.solve(**kwargs) - return P.value, F.value, fuel_consumption.value, problem -def run() -> None: - """ - The main function containing the MPC loop - """ - conn = krpc.connect(name="Landing") +def get_surface_speed(): + conn = krpc.connect(name='Vessel speed') vessel = conn.space_center.active_vessel - print(f"the vessel's name is: {vessel.name}") - - ref_frame = get_reference_frame(conn, vessel) - - position = conn.add_stream(vessel.position, ref_frame) - velocity = conn.add_stream(vessel.velocity, ref_frame) - - relative_frame = ref_frame.create_relative( - ref_frame, rotation=R.from_euler("yz", [-90, -90], degrees=True).as_quat() - ) - vessel.auto_pilot.reference_frame = relative_frame - vessel.auto_pilot.target_pitch_and_heading(60, 0) - - launch_rocket(vessel) + obt_frame = vessel.orbit.body.non_rotating_reference_frame + srf_frame = vessel.orbit.body.reference_frame while True: + obt_speed = vessel.flight(obt_frame).speed + srf_speed = vessel.flight(srf_frame).speed + print('Orbital speed = %.1f m/s, Surface speed = %.1f m/s' % + (obt_speed, srf_speed)) + time.sleep(1) - time.sleep(0.1) +def get_surface_velocity(): + conn = krpc.connect(name='Orbital speed') + vessel = conn.space_center.active_vessel + ref_frame = conn.space_center.ReferenceFrame.create_hybrid( + position=vessel.orbit.body.reference_frame, + rotation=vessel.surface_reference_frame) - p = np.array(position()) - v = np.array(velocity()) - max_thrust = vessel.max_thrust - - if vessel.situation.name == "landed": - vessel.control.throttle = 0 - break - - # Update the glide angle - alpha = np.arctan(p[2] * 1.05 / np.linalg.norm(p[:2])) - - try: - P, F, _, _ = optimize_fuel( - origin, - 9.81, - vessel.mass, - p, - v, - 100, - 1, - max_thrust, - alpha, - 1, - ) - except (cp.SolverError, AssertionError) as e: - print(e) - if v[2] > 0: - vessel.control.throttle = 0 - continue - - if F is None: - print("no solution found") - vessel.control.throttle = 0 - continue - else: - target_F = F[0] - thrust_level = np.linalg.norm(target_F) / max_thrust - - direction = target_F / np.linalg.norm(target_F) - xy_len = np.linalg.norm(direction[[0, 1]]) - - # pitch between -90 and 90 - target_pitch = np.arctan2(direction[2], xy_len) * 180 / np.pi - - # heading between 0 and 360 - target_heading = (np.arctan2(direction[1], direction[0]) * 180 / np.pi) % 360 - - vessel.auto_pilot.target_pitch_and_heading(target_pitch, target_heading) - vessel.auto_pilot.engage() - - if thrust_level < 0.01: - thrust_level = 0 - vessel.control.throttle = thrust_level - - conn.drawing.clear() - draw_trajectory(conn, P, ref_frame) - draw_F(conn, target_F, ref_frame, vessel.reference_frame) - draw_autopilot_direction( - conn, vessel.auto_pilot.target_direction, relative_frame, vessel.reference_frame - ) - - print("thrust: ", round(thrust_level, 2)) - - conn.close() - - -def draw_trajectory(conn: krpc.Client, P: np.ndarray, ref_frame: ReferenceFrame) -> None: - """ - Draw the next 10 points of the trajectory - """ - for i in range(1, P.shape[0]): - p = P[i] - p0 = P[i - 1] - conn.drawing.add_line(p0, p, ref_frame) - if i > 10: - break - - -def draw_F( - conn: krpc.Client, - target_F: np.ndarray, - ref_frame: ReferenceFrame, - vessel_ref_frame: ReferenceFrame, -) -> None: - """ - Draw the target force vector in blue - """ - transformed_F = conn.space_center.transform_direction(target_F, ref_frame, vessel_ref_frame) - line = conn.drawing.add_line([0, 0, 0], transformed_F, vessel_ref_frame) - line.color = (0, 0, 1) - - -def draw_autopilot_direction( - conn: krpc.Client, - direction: np.ndarray, - ref_frame: ReferenceFrame, - vessel_ref_frame: ReferenceFrame, -) -> None: - """ - Draw the autopilot direction in red - """ - transformed_direction = conn.space_center.transform_direction( - direction, ref_frame, vessel_ref_frame - ) - line = conn.drawing.add_line([0, 0, 0], np.array(transformed_direction) * 100, vessel_ref_frame) - line.color = (1, 0, 0) - - -def draw_coordinate_system(conn: krpc.Client, ref_frame: ReferenceFrame) -> None: - """ - Draw the coordinate system of the reference frame - """ - line = conn.drawing.add_line([0, 0, 0], [10, 0, 0], ref_frame) - line.color = (1, 0, 0) - line = conn.drawing.add_line([0, 0, 0], [0, 10, 0], ref_frame) - line.color = (0, 1, 0) - line = conn.drawing.add_line([0, 0, 0], [0, 0, 10], ref_frame) - line.color = (0, 0, 1) - - -def launch_rocket(vessel: Vessel) -> None: - """ - Launch the rocket on a hard-coded trajectory - """ - vessel.auto_pilot.engage() - vessel.control.throttle = 1 - vessel.control.activate_next_stage() - time.sleep(2) - vessel.auto_pilot.target_pitch_and_heading(60, 180) - time.sleep(6) - vessel.control.throttle = 0 - - -def get_reference_frame(conn: krpc.Client, vessel: Vessel) -> ReferenceFrame: - """ - Get the reference frame of the landing platform - """ - - z = landing_platform - z = z / np.linalg.norm(z) - x = np.array([1, 0, 0]) - y = np.cross(z, x) - y = y / np.linalg.norm(y) - x = np.cross(y, z) - x = x / np.linalg.norm(x) - q = R.from_matrix(np.array([x, y, z]).T).as_quat() - ref_frame = conn.space_center.ReferenceFrame.create_relative( - vessel.orbit.body.reference_frame, position=landing_platform, rotation=q - ) - return ref_frame - - -if __name__ == "__main__": - run() \ No newline at end of file + while True: + velocity = vessel.flight(ref_frame).velocity + print('Surface velocity = (%.1f, %.1f, %.1f)' % velocity) + time.sleep(1) \ No newline at end of file