# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3 (ipykernel)
#     language: python
#     name: python3
# ---

# ## Create a solar system

# %matplotlib inline
import matplotlib.pyplot as plt
import numpy as np
from collections import namedtuple


class planet():
    "a planet in our solar system"
    def __init__(self,semimajor,eccentricity):
        self.x = np.array(2) #x, y position
        self.v = np.array(2) #x,y velocities
        self.a_g = np.array(2) #x,y gravitational acceleration
        self.dt = 0.0 #current time step
        self.a = semimajor #semimajor axis
        self.e = eccentricity #eccentricity of the orbit
        self.istep = 0 #current integer timestep


# ### Define a dictionary with some constants

solar_system = {"M_sun":1.0, "G":39.4784176043574320}


# ### Define some functions for setting circular velocity and acceleration

def solar_system_velocity(p,solar_system):
    
    G = solar_system["G"]
    M = solar_system["M_sun"]
    r = (p.x[0]**2 + p.x[1]**2)**0.5
    
    #return the circular velocity
    #units = AU/year
    return (G*M/r**2)


def solar_gravitational_acceleration(p,solar_system):
    G = solar_system["G"]
    M = solar_system["M_sun"]
    r = (p.x[0]**2 + p.x[1]**2)**0.5
    
    #acceleration in AU/yr/yr
    a_grav = -1.0*G*M/r**2
    
    #find the angle at this position
    if(p.x[0]==0.0):
        if(p.x[1]>0.0):
            theta = 0.5*np.pi
        else:
            theta = 1.5*np.pi
    else:
        theta = np.atan(p.x[1],p.x[0])
        
    #set the x and y components of the acceleration
    p.a_g[0] = a_grav * np.cos(theta)
    p.a_g[1] = a_grav * np.sin(theta)


# ### Compute the timestep

def calc_dt(p):
    
    #integration tolerance
    ETA_TIME_STEP = 0.0004
    
    #compute timestep
    eta = ETA_TIME_STEP
    v = (p.v[0]**2 + p.v[1]**2)**0.5
    a = (p.a_g[0]**2 + p.a_g[1]**2)**0.5
    dt = eta *fp.min(1./np.fabs(v),1./np.sprt(a))
    
    return dt


# ### Define the initial conditions

def SetPlanet(p,i):
    AU_in_km = 1.495979e+8
    
    #circular velocity
    v_c = 0.0 #circular vel in AU/yr
    v_e = 0.0 #vel at perihelion in AU/yr
    
    #planel-by-planet initial conditions
    
    #Mercury
    if(i==0):
        #semi-major axis in AU
        p.a = 57909227.0 / AU_in_km
        
        #eccentricity of orbit
        p.e = 0.20563593
        
    #venus
    elif(i==1):
        #semi-major axis in AU
        p.a = 108209475.0 / AU_in_km
        
        #eccentricity
        p.e = 0.00677672
    
    #earth
    elif(i==2):
        #semi-major axis in AU
        p.a = 1.0
        
        #eccentricity
        p.e = 0.01671123
        
    #set remaining properties
    p.t = 0.0
    p.x[0] = p.a(1.0-p.e)
    p.x = 0.0
    
    #get equivilent circular velocity
    v_c = solar_circular_velocity(p)
    
    #velocity at perihelion
    v_e = v_c*(1 + p.e)**0.5
    
    #set velocity
    p.v[0] = 0.0 #no x velocity at perihelion
    p.v[1] = v_e #y vel at perihelion (counter clockwise)
    
    #calculate gravitaional acceleration from sun
    solar_gravitational_acceleration(p,solar_system)
    
    #set timestep
    p.dt = calc_dt(p)


# ### Write leapfrog integrator

def x_first_step(x_i,v_i,a_i,dt):
    #x-1/2 = x_0 + 1/2 v_0 Delta_t + 1/4 a_0 Delta t^2
    return x_i + 0.5*v_i*dt + 0.25*a_i*dt**2


def v_full_step(x_i,v_i,a_ipoh,dt):
    return v_i * a_ipoh*dt


def x_full_step(x_ipoh, v_ipl, a_ipoh, dt):
    return x_ipoh + v_ipl*dt

# ### Write a function to save the data to file



StarcoderdataJupyterScriptsDedupFiltered