A python script to rebalance/allocate investments via contributions

allocate.py

#!/usr/bin/python

import sys, getopt

def total(balance):
    tot = 0
    for key in balance:
        tot += balance[key]
    return tot

def al(balance):
    tot = total(balance)
    al={}
    for key in balance:
        al[key] = balance[key]/tot
    return al

def adjust(balance, contribution):
    newbalance = {}
    for key in balance:
        newbalance[key] = balance[key] + contribution[key]
    return newbalance


def allocation(target_al, balance, amount_new):
    target_d_stock = target_al['d_stock']
    target_i_stock = target_al['i_stock']
    target_d_bond = target_al['d_bond']
    target_i_bond = target_al['i_bond']
    
    bal_d_stock = balance['d_stock']
    bal_i_stock = balance['i_stock']
    bal_d_bond = balance['d_bond']
    bal_i_bond = balance['i_bond']
    
    
    total_invested_dol = bal_d_stock+bal_i_stock+bal_d_bond+bal_i_bond
    
    contrib = {}
    contrib['d_stock'] = round(target_d_stock*(total_invested_dol+amount_new) - bal_d_stock,2)
    contrib['i_stock'] = round(target_i_stock*(total_invested_dol+amount_new) - bal_i_stock, 2)
    contrib['d_bond'] = round(target_d_bond*(total_invested_dol+amount_new) - bal_d_bond, 2)
    contrib['i_bond'] = round(target_i_bond*(total_invested_dol+amount_new) - bal_i_bond, 2)
    
    
    for contrib_key in contrib:
        
        if contrib[contrib_key] < 0:
            
            old_target_al = target_al[contrib_key]
            for al_key in target_al:
                if al_key == contrib_key:
                    target_al[al_key] = 0.0
                    balance[al_key] = 0.0
                    
                else:
                    target_al[al_key] = target_al[al_key]/(1-old_target_al)
                    
            return allocation(target_al, balance, amount_new)
        
    return contrib

def main(argv):
   inputfile = ''
   outputfile = ''
   try:
      opts, args = getopt.getopt(argv,"h")
   except getopt.GetoptError:
      print 'allocate VTSAX_Balance VTIAX_Balance VBTLX_Balance VTABX_Balance Contribution'
      sys.exit(2)
   for opt, arg in opts:
      if opt == '-h':
         print 'allocate: requires five arguments'
         print 'usage: allocate VTSAX_Balance VTIAX_Balance VBTLX_Balance VTABX_Balance Contribution'
         sys.exit()

   balance = {}      
   balance['d_stock'] = float(argv[0])
   balance['i_stock'] = float(argv[1])
   balance['d_bond'] =  float(argv[2])
   balance['i_bond'] =  float(argv[3])
   
   contribution = float(argv[4])

   target_allocation = {}
   target_allocation['d_stock']=0.54
   target_allocation['i_stock']=0.36
   target_allocation['d_bond']=0.07
   target_allocation['i_bond']=0.03
 
   initial_allocation = al(balance)

   allocated_contribution = allocation(target_allocation, balance, contribution)

   final_balance = adjust(balance, allocated_contribution)

   final_allocation = al(final_balance)

   print 'VERIFY TARGET ALLOCATIONS AND ENTERED BALANCES ARE CORRECT!'
   print ''
   print 'Contribution =', contribution
   print '' 
   print 'Initial Balances:'
   print 'Domestic Stock (VTSAX)      =', balance['d_stock']
   print 'International Stock (VTIAX) =', balance['i_stock']
   print 'Domestic Bonds (VBTLX)      =', balance['d_bond']  
   print 'International Bond (VTABX)  =', balance['i_bond'] 
   print ''
   print 'Target Allocation:'
   print 'Domestic Stock (VTSAX)      =',target_allocation['d_stock']
   print 'International Stock (VTIAX) =',target_allocation['i_stock']
   print 'Domestic Bonds (VBTLX)      =',target_allocation['d_bond']
   print 'International Bond (VTABX)  =',target_allocation['i_bond']
   print ''
   print 'Initial Actual Allocation:'
   print 'Domestic Stock (VTSAX)      =',initial_allocation['d_stock']
   print 'International Stock (VTIAX) =',initial_allocation['i_stock']
   print 'Domestic Bonds (VBTLX)      =',initial_allocation['d_bond'] 
   print 'International Bond (VTABX)  =',initial_allocation['i_bond'] 
   print ''
   print 'Allocated Contribution:'
   print 'Domestic Stock (VTSAX)      =',allocated_contribution['d_stock']
   print 'International Stock (VTIAX) =',allocated_contribution['i_stock']
   print 'Domestic Bonds (VBTLX)      =',allocated_contribution['d_bond'] 
   print 'International Bond (VTABX)  =',allocated_contribution['i_bond'] 
   print ''
   print 'Final Balances:'
   print 'Domestic Stock (VTSAX)      =', final_balance['d_stock']
   print 'International Stock (VTIAX) =', final_balance['i_stock']
   print 'Domestic Bonds (VBTLX)      =', final_balance['d_bond'] 
   print 'International Bond (VTABX)  =', final_balance['i_bond'] 
   print ''
   print 'Final Actual Allocation:'
   print 'Domestic Stock (VTSAX)      =', final_allocation['d_stock']
   print 'International Stock (VTIAX) =', final_allocation['i_stock']
   print 'Domestic Bonds (VBTLX)      =', final_allocation['d_bond'] 
   print 'International Bond (VTABX)  =', final_allocation['i_bond'] 
   


if __name__ == "__main__":
   main(sys.argv[1:])