Multi-Stage Rocket Delta-V Optimizer

#Delta-v Partition Optimizer
#   Ian Mallett, 2016
#   ian [^at^] geometrian.com

#This program takes your rocket configuration and calculates
#   optimal partitioning of your rocket's delta-v into its
#   stages.

#=========== BEGIN USER-CONFIGURABLE SECTION ===========

#Choose the minimum practical number of stages; more stages
#   improves overall mass ratio but adds more dead weight
#   and expense for a given delta-v.

#Stage parameters are:
#   Specific impulse (s)
#   Inert mass fraction (roughly ~0.08 to ~0.7)

stages = [ #Saturn V
    (263.0,130000.0/2290000.0),
    (421.0,40100.0/496200.0),
    (421.0,13500.0/123000.0)
]
M_0 = 48600.0 #Payload mass, kg
d_vm = 17911.9 #Total delta-v for mission

##stages = [ #example I made up (toy rocket)
##    (3000.0/9.80665,0.08)
##]
##M_0 = 1.0 #Payload mass, kg
##d_vm = 4000.0 #Total delta-v for mission

##stages = [ #example I made up
##    (390.0, 0.08),
##    (400.0, 0.10),
##    (800.0, 0.30)
##]
##M_0 = 1000.0 #Payload mass, kg
##d_vm = 21000.0 #Total delta-v for mission

tol = 1e-9 #Solver tolerance for each fraction

#=========== END USER-CONFIGURABLE SECTION ===========


from math import *
import random
import sys


#Basic setup
in_idle = "idlelib.run" in sys.modules
g_0 = -9.80665 #Standard gravity
n = len(stages)
stages = stages[::-1]

#Validate input
if n == 0:
    raise ValueError("Must have at least one stage in rocket!")
for I, f in stages:
    if I<=0: raise ValueError("Specific impulse must be positive for every stage!")
    if f<=0.0 or f>=1.0: raise ValueError("Inert mass fraction must lie in the open interval (0.0,1.0) for every stage!")
if M_0 <= 0.0:
    raise ValueError("Payload must have a positive mass!")
if d_vm <= 0.0:
    raise ValueError("Mission total delta-v must be positive!")

#Helper functions
def get_total_mass(f_dvs):
    total_mass = M_0
    for i in range(n):
        I = stages[i][0] #stage specific impulse
        f = stages[i][1] #stage inert mass fraction
        V_e = -g_0 * I #stage exhaust speed
        assert f_dvs[i] > 0.0
        d_vs = d_vm * f_dvs[i] #stage target delta v
        R = e ** (d_vs / V_e) #stage mass ratio
        if R*f >= 1.0: return False #Stage cannot be built!
        M_i = total_mass * (R-1.0) / (1.0 - f*R)
        assert M_i > 0.0
        total_mass += M_i
    return total_mass
def get_normalized_fractions(f_dvs):
    s = sum([f_dv for f_dv in f_dvs])
    f_dvs = [f_dv/s for f_dv in f_dvs]
    return f_dvs

print("Beginning solve for %d-stage rocket . . ."%n)
#Find a feasible configuration
print("Attempting to find a valid initial configuration . . .")
iterations = 0
best = False
while best == False:
    f_dvs = get_normalized_fractions( [random.random() for i in range(n)] )
    best = get_total_mass(f_dvs)
    iterations += 1
    if iterations == 100000:
        raise ValueError("Unable to find a valid initial configuration!  Rocket may be impossible or nearly impossible!")

#Optimize it
print("Optimizing configuration . . .")
failed = 0
radius = 1.0
while radius > tol:
    f_dvs1 = list(f_dvs)

    rand_ind = random.randint(0,n-1)
    while True:
        rand_frac = f_dvs1[rand_ind] + radius*((random.random()**0.1)-0.5)
        if rand_frac>0.0: break
    f_dvs1[rand_ind] = rand_frac

    f_dvs1 = get_normalized_fractions(f_dvs1)

    proposed = get_total_mass(f_dvs1)
    if proposed != False and proposed<best:
        f_dvs = list(f_dvs1)
        best = proposed
        if failed < 10:
            radius /= 0.8
        failed = 0

        s = "Mass: %f, Fractions:"%best
        for i in range(n): s+=" %f"%f_dvs[n-i-1]
        if in_idle: print(s)
        else: sys.stdout.write("\r"+s)
    else:
        failed += 1

    if failed > 10:
        radius *= 0.5
        failed = 0
if not in_idle: sys.stdout.write("\r"+" "*80)

#Print results
print("Completed optimization!")
art = ["+---"+n*"-"+"-->"]
stats = [""]
total_mass_dry = 0.0
total_mass_wet = 0.0
for i in range(-1,n,1):
    sp = (n - i)*" "
    if i == -1:
        art.append("|")
        art.append("|   "+sp+"  ^")
        art.append("|   "+sp+" / \\")
        art.append("|   "+sp+"/   \\")
        art.append("|   "+sp+"*---*")
        stats.append("Stage %d (payload):"%(n+1))
        stats.append("  Mass:                  %f kg"%M_0)
        stats.append("")
        stats.append("")
        stats.append("")
        total_mass_dry += M_0
        total_mass_wet += M_0
    else:
        art.append("+----"+n*"-"+"->")
        art.append("|    "+sp+"_____"+"__"*i)
        art.append("|   "+sp+"/     "+"  "*i+"\\")
        art.append("|   "+sp+"|     "+"  "*i+"|")
        art.append("|   "+sp+"|     "+"  "*i+"|")
        art.append("|   "+sp+"|     "+"  "*i+"|")
        art.append("|   "+sp+"|     "+"  "*i+"|")
        art.append("|   "+sp+"|     "+"  "*i+"|")
        art.append("|   "+sp+"|     "+"  "*i+"|")
        art.append("|   "+sp+"|     "+"  "*i+"|")
        if i < n-1:
            art.append("|   "+sp+"|     "+"  "*i+"|")
            art.append("|   "+sp+"^^^^^^"+"^^"*i+"^")
        else:
            art.append("| "+sp+" /|     "+"  "*i+"|\\")
            art.append("| "+sp+"<_^^^^^^"+"^^"*i+"^_>")
        I = stages[i][0]
        f = stages[i][1]
        V_e = -g_0 * I
        d_vs = d_vm * f_dvs[i]
        R = e ** (d_vs / V_e)
        M_i = total_mass_wet * (R-1.0) / (1.0 - f*R)
        stats.append("")
        stats.append("Stage %d:"%(n-i))
        stats.append("  Mass (dry):            %f kg"%(M_i*f))
        stats.append("  Mass (wet):            %f kg"%M_i)
        stats.append("  Mass (propellant):     %f kg"%(M_i*(1.0-f)))
        stats.append("  Specific impulse:      %f s"%I)
        stats.append("  Exhaust velocity:      %f m/s"%V_e)
        stats.append("  Delta-v:               %f m/s"%d_vs)
        stats.append("  Delta-v fraction:      %f"%f_dvs[i])
        stats.append("  Propellant fraction:   %f"%(1.0-f))
        stats.append("  Inert fraction:        %f"%f)
        stats.append("  Mass ratio:            %f"%R)
        total_mass_wet += M_i
        total_mass_dry += M_i * f

art.append(art[0])
stats.append("")

total_mass_prop = total_mass_wet - total_mass_dry
for i in range(8): art.append("")
stats.append("Total:")
stats.append("  Mass (dry):            %f kg"%total_mass_dry)
stats.append("  Mass (wet):            %f kg"%total_mass_wet)
stats.append("  Mass (propellant):     %f kg"%(total_mass_wet-total_mass_dry))
stats.append("  Delta-v:               %f m/s"%d_vm)
stats.append("  Propellant fraction:   %f"%(total_mass_prop/total_mass_wet))
stats.append("  Inert fraction:        %f"%(total_mass_dry/total_mass_wet))
stats.append("  Mass ratio:            %f"%(M_0/total_mass_wet))

assert len(art) == len(stats)
max_len = 0
for stat in stats:
    max_len = max([max_len,len(stat)])
lines = []
for i in range(len(art)):
    s = stats[i]
    if len(art[i])!=0: s += " "*(max_len-len(stats[i])) + "   " + art[i]
    print(s)

if not in_idle: input("ENTER to exit.")