rocket.py

from physics_constants import landa, alpha, EARTH_RADIUS, EARTH_ANGULAR_VELOCITY, G, EARTH_MASS
from numpy import sin, cos, sqrt
from parameters import dimensions, MIN_HEIGHT, dt, time
from coords_math import norm, radial_to_cartesian
import parameters
import physics
import copy

R = EARTH_RADIUS

class Rocket():
    def __init__(self, boosters, stage2, stage3, capsule):
        self.stages = [boosters, stage2, stage3, capsule]
        self.active_stages = [stage for stage in self.stages if stage.Is_Active]

        #we are using cartesian coordinates for the position, velocity and accel
        self.initial_position = radial_to_cartesian(R, landa, alpha)
        self.positions = [self.initial_position]

        self.initial_velocity = [R * cos(landa) * -sin(alpha) * EARTH_ANGULAR_VELOCITY, R * cos(landa) * cos(alpha) * EARTH_ANGULAR_VELOCITY, 0]
        self.velocities = [self.initial_velocity]

        self.initial_acceleration = physics.get_acceleration(self, True)
        self.accelerations = [self.initial_acceleration]

        self.stop_get_time_for_max_height_simulation = False
        self.orbiting = False

        #helpers
        self.heights = [0]

    #for when it starts orbiting the earth at 400km
    def get_orbital_velocity(self):
        position = copy.deepcopy(self.positions[-1])

        radial_distance = self.get_radial_distance()

        direction_vector = [-position[1]/radial_distance, position[0]/radial_distance, 0]

        velocity = sqrt(G * EARTH_MASS / radial_distance)
        velocity_vector = [0, 0, 0]
        for i in range(dimensions):
            velocity_vector[i] = direction_vector[i] * velocity

        return velocity_vector

    #time needed until the engines shut down to reach 400km
    def get_time_for_max_height(self):
        if self.stop_get_time_for_max_height_simulation:
            return
        
        test_time = copy.deepcopy(time)
        test_velocities = copy.deepcopy(self.velocities.copy())
        test_positions = copy.deepcopy(self.positions.copy())
        test_accelerations = copy.deepcopy(self.accelerations.copy())
        test_heights = copy.deepcopy(self.heights)

        def tick_no_thrust():
            test_acceleration = physics.get_acceleration(self, False)
            test_velocity = test_velocities[-1]
            test_position = test_positions[-1]

            for i in range(dimensions):
                test_velocity[i] = test_velocity[i] + test_acceleration[i] * dt 
                test_position[i] = test_position[i] + test_velocity[i] * dt 

            test_accelerations.append(test_acceleration)
            test_velocities.append(test_velocity)
            test_positions.append(test_position)

            height =  norm(test_positions[-1]) - EARTH_RADIUS
            test_heights.append(height)

        if len(test_heights) < 2:
            return
        
        while test_heights[-1] > test_heights[-2] and test_heights[-1] > MIN_HEIGHT:
            tick_no_thrust()
            if(test_heights[-1] > 400000):
                parameters.time_using_thrust = test_time[-1]
                self.stop_get_time_for_max_height_simulation = True

    def get_radial_distance(self):
        return norm(self.positions[-1])
    
    def get_height(self):
        return self.get_radial_distance() - EARTH_RADIUS
    
    def get_mass(self):
        mass = 0
        for stage in self.stages:
            if stage.is_detached:
                continue
            mass += stage.get_mass()
        return mass

    def get_thrust(self):
        thrust = 0
        for stage in self.active_stages:
            thrust += stage.get_thrust()
        return thrust
    
    def detach(self, stage):
        self.active_stages.remove(stage)
        stage.is_detached = True

    def tick(self, stop_at_max_height):
        if self.get_height() > parameters.MAX_HEIGHT and not self.orbiting:
            self.orbiting = True

        for stage in self.stages:
            if stage in self.active_stages:
                stage.burn_fuel()

            if stage.Is_Active and not stage in self.active_stages: #set new active classes
                self.active_stages.append(stage)
            elif not stage.Is_Active and stage in self.active_stages: #remove deactive classes
                self.detach(stage)

        acceleration = physics.get_acceleration(self, True)
        if stop_at_max_height and time[-1] >= parameters.time_using_thrust:
            acceleration = physics.get_acceleration(self, False)
        velocity = copy.deepcopy(self.velocities[-1])
        position = copy.deepcopy(self.positions[-1])
            
        if self.orbiting:
            velocity = self.get_orbital_velocity()
        for i in range(dimensions):
            if not self.orbiting:
                velocity[i] = velocity[i] + acceleration[i] * dt
            position[i] = position[i] + velocity[i] * dt

        self.accelerations.append(acceleration)
        self.velocities.append(velocity)
        self.positions.append(position)
        
        self.heights.append(self.get_height())

class RocketStage:
    def __init__(self, name, num_of_boosters, M_empty, M_fuel, thrust, Drag_Area, Burn_Rate, Is_Active, Next_Stage):
        self.name = name
        self.num_of_boosters = num_of_boosters
        self.M_empty = M_empty
        self.M_fuel = M_fuel
        self.thrust = thrust
        self.Drag_Area = Drag_Area
        self.Burn_Rate = Burn_Rate
        self.Is_Active = Is_Active
        self.Next_Stage = Next_Stage
        self.is_detached = False

    def get_mass(self):
        return (self.M_empty + self.M_fuel) * self.num_of_boosters
    
    def get_thrust(self):
        return self.thrust * self.num_of_boosters

    def burn_fuel(self): #we treat this like it was for each booster because we multiply mass and thrust by num of boosters
        fuel_burned = self.Burn_Rate * dt
        if self.M_fuel > fuel_burned:
            self.M_fuel -= fuel_burned
            return
        
        #shut down stage
        self.M_fuel = 0
        self.Is_Active = False
        if self.Next_Stage != None:
            self.Next_Stage.Is_Active = True