portfolio.py

import numpy as np
import copy

###############################
# INFLATION RATE and INIT_AGE #
###############################

INFLATION_RATE = 0.03
INIT_AGE = 31

################
# UTILITY FUNC #
################

def set_INIT_AGE(age):
    global INIT_AGE
    INIT_AGE = age
    return 0

def pretty_num(num):
    neg = num < 0
    num = abs(num)
    if num > 9.9e5:
        snum = round(num / 1e6, 1)
        s = "$" + str(snum) + "M"
    elif num > 9.9e2:
        snum = round(num / 1e3, 1)
        s = "$" + str(snum) + "K"
    else:
        snum = round(num, -1)
        s = "$" + str(snum)

    if neg:
        return "-" + s
    return s


#############
# TAX STUFF #
#############

FEDERAL_INCOME_TAX_BRACKETS = [
    (11000, 0.10),
    (44725, 0.12),
    (95375, 0.22),
    (182_100, 0.24),
    (231_250, 0.32),
    (578_125, 0.35),
    (np.inf, 0.37),
]
NEW_YORK_INCOME_TAX_BRACKETS = [
    (8500, 0.0400),
    (11700, 0.0450),
    (13900, 0.0525),
    (80650, 0.0585),
    (215_400, 0.0625),
    (1_077_550, 0.0685),
    (5_000_000, 0.0965),
    (25_000_000, 0.1030),
    (np.inf, 0.1090),
]

NEW_YORK_CITY_INCOME_TAX_BRACKETS = [
    (14400, 0.03078),
    (30000, 0.03762),
    (60000, 0.03819),
    (np.inf, 0.03876),
]

FEDERAL_GAINS_TAX_BRACKETS = [(41675, 0.00), (492_300, 0.15), (np.inf, 0.2)]


def marginal_tax(pay, brackets):

    left_to_pay = pay
    tax = 0
    for ii,bb in enumerate(brackets):
        bracket, rate = bb
        if left_to_pay < 1:
            break
        if ii > 0:
            tax += (min(pay, bracket) - brackets[ii-1][0]) * rate
        else:
            tax += min(pay, bracket) * rate
        left_to_pay = pay - bracket#min(margin, bracket) 
    return tax


class IRS_agent(object):
    def __init__(
        self,
        year_offset=0.0,
        federal_gains_brackets=FEDERAL_GAINS_TAX_BRACKETS,
        federal_income_brackets=FEDERAL_INCOME_TAX_BRACKETS,
        state_income_brackets=NEW_YORK_INCOME_TAX_BRACKETS,
        city_income_brackets=NEW_YORK_CITY_INCOME_TAX_BRACKETS,
    ):

        self.year_offset = year_offset
        self.federal_gains_brackets = copy.deepcopy(federal_gains_brackets)
        self.federal_income_brackets = copy.deepcopy(federal_income_brackets)
        self.state_income_brackets = copy.deepcopy(state_income_brackets)
        self.city_income_brackets = copy.deepcopy(city_income_brackets)

    def inflate_tax_brackets(self):
        for i, (bracket, rate) in enumerate(self.federal_gains_brackets):
            self.federal_gains_brackets[i] = (bracket * (1.0 + INFLATION_RATE), rate)
        for i, (bracket, rate) in enumerate(self.federal_income_brackets):
            self.federal_income_brackets[i] = (bracket * (1.0 + INFLATION_RATE), rate)
        for i, (bracket, rate) in enumerate(self.state_income_brackets):
            self.state_income_brackets[i] = (bracket * (1.0 + INFLATION_RATE), rate)
        for i, (bracket, rate) in enumerate(self.city_income_brackets):
            self.city_income_brackets[i] = (bracket * (1.0 + INFLATION_RATE), rate)

    def combined_tax(self, pay, gain=0.0):
        # Federal treats gains and pay differently (insert political opinion here).
        tax = marginal_tax(gain, self.federal_gains_brackets)
        tax += marginal_tax(pay, self.federal_income_brackets)
        # But New York treats them the same? Tax is so complicated.
        tax += marginal_tax(pay + gain, self.state_income_brackets)
        tax += marginal_tax(pay + gain, self.city_income_brackets)

        return tax


#################
# BANK ACCOUNTS #
#################


class checkings_accnt(object):
    def __init__(self):
        self.balance = 0.0
        self.last_tax = 0.0

    def add(self, x):
        self.balance += x

    def remove(self, x):
        assert x <= self.balance
        self.balance -= x

    def get_balance(self):
        return self.balance


class savings_accnt(checkings_accnt):
    def __init__(self, interest_rate=0.005):  # Gotta get a better savings rate...
        super().__init__()
        self.interest_rate = interest_rate

    def get_interest(self):
        return self.balance * (self.interest_rate)


#####################
# INVESTMENTS STUFF #
#####################

STOCK_GAINS_RATE = 0.06
BONDS_GAINS_RATE = 0.04


def default_stock_alloc_strategy(i):
    age = INIT_AGE + i
    alloc = min(100.0, max(0.0, 100.0 - age))
    return alloc


class investment_accnt(object):
    def __init__(
        self,
        stock_gains_rate=STOCK_GAINS_RATE,
        bonds_gains_rate=STOCK_GAINS_RATE,
        stock_alloc_strategy=default_stock_alloc_strategy,
    ):
        super().__init__()
        self.i = 0
        self.principal = 0.0
        self.gains = 0.0
        self.stock_gains_rate = stock_gains_rate
        self.bonds_gains_rate = bonds_gains_rate
        self.stock_alloc_strategy = stock_alloc_strategy
        self.change_stock_alloc_perc(self.stock_alloc_strategy(self.i))

    def add(self, x):
        self.principal += x

    def remove(self, x):
        total = max(1, self.gains + self.principal)
        assert x <= total
        # Not a correct removal strategy
        remove_from_principal = x * self.principal / total
        remove_from_gains = x * self.gains / total
        self.principal -= remove_from_principal
        self.gains -= remove_from_gains
        # First element is not gains taxed, the second is.
        return remove_from_principal, remove_from_gains

    def change_stock_alloc_perc(self, perc):
        self.stock_alloc = perc / 100.0
        eff_rate = self.stock_gains_rate * self.stock_alloc
        eff_rate += self.bonds_gains_rate * (1.0 - self.stock_alloc)
        self.effective_gains_rate = eff_rate

    def increment_year(self):
        self.i += 1
        self.gains += (self.principal + self.gains) * (self.effective_gains_rate)
        self.change_stock_alloc_perc(self.stock_alloc_strategy(self.i))

    def get_balance(self):
        return self.principal + self.gains


def default_401k_roth_alloc_strategy(i, retired=False):
    age = INIT_AGE + i
    if age < 40 or retired:
        roth_perc = 100.0
    else:
        roth_perc = 4.0 * abs(age - 60.0) + 20.0

    return roth_perc


class i401k_accnt:
    def __init__(
        self,
        pretax_cap=20.5e3,
        total_cap=66e3,  # With the megabackdoor, assuming size correlated with inflation
        roth_stock_alloc_strategy=default_stock_alloc_strategy,
        pretax_stock_alloc_strategy=default_stock_alloc_strategy,
        i401k_roth_alloc_strategy=default_401k_roth_alloc_strategy,
    ):
        self.roth = investment_accnt(stock_alloc_strategy=roth_stock_alloc_strategy)
        self.pretax = investment_accnt(stock_alloc_strategy=pretax_stock_alloc_strategy)
        self.pretax_cap = pretax_cap
        self.total_cap = total_cap
        self.i401k_roth_alloc_strategy = i401k_roth_alloc_strategy
        self.roth_perc = 100.0

    def add(self, x):
        pretax_x = (1.0 - self.roth_perc / 100.0) * x
        pretax_x = max(self.pretax_cap, pretax_x)
        roth_x = x - pretax_x
        assert roth_x <= self.total_cap - pretax_x

        self.pretax.add(pretax_x)
        self.roth.add(roth_x)

        # Remove pretax_x from taxed amount.
        return pretax_x, roth_x

    # Only allowed after retirement
    def remove(self, x):
        assert x <= self.get_balance()
        # Remove from pretax first
        pretax_x = min(self.pretax.get_balance(), x)
        roth_x = min(self.roth.get_balance(), x - pretax_x)
        untaxed_y, taxed_y = self.pretax.remove(pretax_x)
        if roth_x > 0:
            untaxed_y += np.sum(self.roth.remove(roth_x))

        return untaxed_y, taxed_y

    def increment_year(self, retired=False):
        self.roth.increment_year()
        self.pretax.increment_year()
        i = self.roth.i
        self.roth_perc = self.i401k_roth_alloc_strategy(i, retired=retired)
        self.pretax_cap *= 1.0 + INFLATION_RATE
        self.total_cap *= 1.0 + INFLATION_RATE

    def get_balance(self):
        return self.roth.get_balance() + self.pretax.get_balance()


def default_IRA_roth_alloc_strategy(i, pay):
    roth_perc = 100.0 if pay < 115e3 else 0.0
    return roth_perc


class IRA_accnt(i401k_accnt):
    def __init__(
        self,
        cap=6e3,
        roth_stock_alloc_strategy=default_stock_alloc_strategy,
        pretax_stock_alloc_strategy=default_stock_alloc_strategy,
        IRA_roth_alloc_strategy=default_IRA_roth_alloc_strategy,
    ):
        super().__init__(
            pretax_cap=cap,
            total_cap=cap,
            roth_stock_alloc_strategy=roth_stock_alloc_strategy,
            pretax_stock_alloc_strategy=pretax_stock_alloc_strategy,
            i401k_roth_alloc_strategy=IRA_roth_alloc_strategy,
        )
        self.cap = cap

    def increment_year(self, pay=0.0):
        self.roth.increment_year()
        self.pretax.increment_year()
        i = self.roth.i
        self.roth_perc = self.i401k_roth_alloc_strategy(i, pay=pay)
        self.pretax_cap *= 1.0 + INFLATION_RATE
        self.total_cap *= 1.0 + INFLATION_RATE
        # After retirement, cap raises to from $6k to $7k
        if i - INIT_AGE == 65:
            self.pretax_cap *= 6.0 / 7.0
            self.total_cap *= 6.0 / 7.0


###########################
# PUTTING IT ALL TOGETHER #
###########################

EMERGENCY_TO_SPENDING_FACTOR = 0.7


class portfolio(object):
    def __init__(
        self,
        retirement_age=65,
        init_checkings=1e3,
        init_savings=1e3,
        init_brokerage=0.0,
        init_roth_401k=1e3,
        init_pretax_401k=0.0,
        init_roth_IRA=1e3,
        init_pretax_IRA=0.0,
    ):
        self.init_age = INIT_AGE
        self.age = INIT_AGE
        self.retirement_age = retirement_age
        self.tax_agent = IRS_agent()

        self.checkings = checkings_accnt()
        self.savings = savings_accnt()
        self.brokerage = investment_accnt()
        self.i401k = i401k_accnt()
        self.IRA = IRA_accnt()
        self.retired = False

        self.target_emergency_balance = 10e3
        self.checkings.balance = init_checkings
        self.savings.balance = init_savings
        self.brokerage.principal = init_brokerage
        self.i401k.roth.principal = init_roth_401k
        self.i401k.pretax.principal = init_pretax_401k
        self.IRA.roth.principal = init_roth_IRA
        self.IRA.pretax.principal = init_pretax_IRA

        self.effective_self_pay_tax_rate = 0.2
        self.effective_tax_rate = 0.2

    def set_target_emergency_balance(self, x):
        self.target_emergency_balance = x

    def get_interest(self):
        return self.savings.interest()

    def get_emergency_contributions(self, x):
        # Build up emergency savings slowly because we gamble against fat tails.
        delta_target_emergency_checkings = self.target_emergency_balance * 0.1
        delta_target_emergency_checkings -= self.checkings.balance
        delta_target_emergency_checkings = min(
            delta_target_emergency_checkings, max(1e3, 0.05 * x)
        )
        delta_target_emergency_savings = self.target_emergency_balance * 0.9
        delta_target_emergency_savings -= self.savings.balance
        delta_target_emergency_savings = min(
            delta_target_emergency_savings, max(9e3, 0.2 * x)
        )
        delta_target_emergency_checkings = max(0.0, delta_target_emergency_checkings)
        delta_target_emergency_savings = max(0.0, delta_target_emergency_savings)
        return delta_target_emergency_checkings, delta_target_emergency_savings

    def contribute(self, x, verbose=False):

        untax = 0
        delta_checkings, delta_savings = self.get_emergency_contributions(x)
        if verbose:
            print(f"  ADD TO CASH: {pretty_num(delta_checkings+delta_savings)}")
        self.checkings.add(delta_checkings)
        self.savings.add(delta_savings)
        x -= delta_checkings + delta_savings
        if x > 0:
            i401k_contribution = min(x, self.i401k.total_cap)
            if verbose:
                print(f"  401k CONTRIBUTION: {pretty_num(i401k_contribution)}")
            untax += self.i401k.add(i401k_contribution)[0]
            x -= i401k_contribution
        if x > 0:
            IRA_contribution = min(x, self.IRA.cap)
            if verbose:
                print(f"  IRA CONTRIBUTION: {pretty_num(IRA_contribution)}")
            untax += self.IRA.add(IRA_contribution)[0]
            x -= IRA_contribution
        if x > 0:
            if verbose:
                print(f"  BROKERAGE CONTRIBUTION: {pretty_num(x)}")
            self.brokerage.add(x)

        return untax

    def cash_disburse(self, x):
        # Removing in ratio to preserve emergency contribution strategy.
        total_emergency = self.checkings.balance + self.savings.balance
        assert x <= total_emergency
        checkings_x = x * self.checkings.balance / total_emergency
        savings_x = x * self.savings.balance / total_emergency
        self.checkings.remove(checkings_x)
        self.savings.remove(savings_x)
        untaxed_y = checkings_x + savings_x
        return untaxed_y

    def disburse(self, x, verbose=False):
        untaxed_y = 0.0
        taxed_y = 0.0

        x = abs(x)
        if x > 0 and self.brokerage.get_balance() > 0:
            brokerage_x = min(x, self.brokerage.get_balance())
            if verbose:
                print(f"  BROKERAGE WITHDRAWL: {pretty_num(brokerage_x)}")
            ut_y, t_y = self.brokerage.remove(brokerage_x)
            x -= brokerage_x
            untaxed_y += ut_y
            taxed_y += t_y
        if (x > 0) and self.retired and self.i401k.get_balance() > 0:
            i401k_x = min(x, self.i401k.get_balance())
            if verbose:
                print(f"  401k WITHDRAWL: {pretty_num(i401k_x)}")
            ut_y, t_y = self.i401k.remove(i401k_x)
            x -= i401k_x
            untaxed_y += ut_y
            taxed_y += t_y
        if (x > 0) and self.retired and self.IRA.get_balance() > 0:
            IRA_x = min(x, self.IRA.get_balance())
            if verbose:
                print(f"  IRA WITHDRAWL: {pretty_num(IRA_x)}")
            ut_y, t_y = self.IRA.remove(IRA_x)
            x -= IRA_x
            untaxed_y += ut_y
            taxed_y += t_y
        if x > 0:
            if verbose:
                print(f"  CASH WITHDRAWL: {pretty_num(x)}")
            untaxed_y += self.cash_disburse(x)

        return untaxed_y, taxed_y

    def increment_year(self, pay=1e3, spending=1e3, verbose=True):
        if verbose:
            print(f"PAY: {pretty_num(pay)}")
        if verbose:
            print(f"SPENDING: {pretty_num(spending)}")
        self.age += 1
        year = self.age - self.init_age
        deflate = (1.0 + INFLATION_RATE) ** (-year)
        interest = self.savings.get_interest()
        if verbose:
            print(f"INTEREST: {pretty_num(interest)}")
        msg = f"age:{self.age}\n"
        msg += (
            f"Pay:       {pretty_num(pay)} \t(t0 {pretty_num(pay*deflate)})\n"
        )
        msg += (
            f"Spending:  {pretty_num(spending)} \t(t0 {pretty_num(spending*deflate)})\n"
        )
        msg += (
            f"Interest:  {pretty_num(interest)} \t(t0 {pretty_num(interest*deflate)})\n"
        )
        take_home = pay + interest

        if self.age >= self.retirement_age:
            self.retired = True
        self.set_target_emergency_balance(spending * EMERGENCY_TO_SPENDING_FACTOR)
        if self.checkings.balance > self.target_emergency_balance * 0.1:
            checkings_x = self.checkings.balance - self.target_emergency_balance * 0.1
            self.checkings.remove(checkings_x)
            take_home += checkings_x
            if verbose:
                print(f"EXTRA IN CHECKINGS: {pretty_num(checkings_x)}")
        if self.savings.balance > self.target_emergency_balance * 0.9:
            savings_x = self.savings.balance - self.target_emergency_balance * 0.9
            self.savings.remove(savings_x)
            take_home += savings_x
            if verbose:
                print(f"EXTRA IN SAVINGS: {pretty_num(savings_x)}")

        self.brokerage.increment_year()
        self.i401k.increment_year(retired=self.retired)
        self.IRA.increment_year(pay=pay)
        self.tax_agent.inflate_tax_brackets()

        estimated_tax = self.effective_tax_rate * max(spending, take_home)
        estimated_tax = max(estimated_tax, self.tax_agent.combined_tax(take_home))
        estimated_tax *= 1.1  # fudge_factor
        if verbose:
            print(f"ESTIMATED TAX BILL: {pretty_num(estimated_tax)}")

        net_self_pay = 0.0
        self_pay = 0.0
        if take_home - estimated_tax < spending:
            # Overshoot to account for taxes:
            withdrawl = 1.0 + self.effective_self_pay_tax_rate
            withdrawl *= spending - take_home - estimated_tax
            withdrawl *= 1.03  # fudge factor
            if verbose:
                print(f"WITHDRAWING: {pretty_num(withdrawl)}")
            untaxable_disbursement, taxable_disbursement = self.disburse(
                withdrawl, verbose
            )
            untaxable_contribution = 0.0
            self_pay = taxable_disbursement + untaxable_disbursement
            net_self_pay = self_pay
        else:
            withdrawl = 0.0
            untaxable_disbursement = 0.0
            taxable_disbursement = 0.0
            if verbose:
                print(
                    f"CONTRIBUTING: {pretty_num(take_home - spending - estimated_tax)}"
                )
            untaxable_contribution = self.contribute(
                take_home - spending - estimated_tax, verbose
            )

        tax = self.tax_agent.combined_tax(
            pay + interest - untaxable_contribution, taxable_disbursement
        )
        income_tax = self.tax_agent.combined_tax(
            pay + interest - untaxable_contribution, 0.0
        )
        disbursement_tax = tax - income_tax
        net_self_pay = self_pay - disbursement_tax
        if disbursement_tax > 0:
            if verbose:
                print(f"DISBURSEMENT TAX: {pretty_num(disbursement_tax)}")
        if verbose:
            print(f"NET SELF PAY: {pretty_num(net_self_pay)}")

        tax_return = estimated_tax - tax
        if verbose:
            print(f"FULL TAX BILL: {pretty_num(tax)}")
        if tax_return < 0:
            tax_return = abs(tax_return)
            self.cash_disburse(tax_return)
            if verbose:
                print(f"TAX OWED: {pretty_num(tax_return)}")
        else:
            self.savings.add(0.9 * tax_return)
            self.checkings.add(0.1 * tax_return)
            if verbose:
                print(f"TAX RETURNED: {pretty_num(tax_return)}")

        if self_pay > 0:
            self.effective_self_pay_tax_rate = disbursement_tax / (withdrawl + 1e-3)
        effective_tax_rate = tax / (pay + interest + self_pay)
        self.effective_tax_rate = effective_tax_rate

        msg += f"Taxed:     {pretty_num(tax)} \t(t0 {pretty_num(tax*deflate)})"
        msg += f"\t[{round(100*effective_tax_rate, 1)}%]\n"
        msg += f"Self-paid: {pretty_num(net_self_pay)} "
        msg += f"\t(t0 {pretty_num(net_self_pay*deflate)})\n"
        msg += f"Checkings: {pretty_num(self.checkings.balance)}"
        msg += f" \t(t0 {pretty_num(self.checkings.balance*deflate)})\n"
        msg += f"Savings:   {pretty_num(self.savings.balance)}"
        msg += f" \t(t0 {pretty_num(self.savings.balance*deflate)})\n"
        msg += f"Brokerage: {pretty_num(self.brokerage.get_balance())}"
        msg += f" \t(t0 {pretty_num(self.brokerage.get_balance()*deflate)})\n"
        msg += f"401k:      {pretty_num(self.i401k.get_balance())}"
        msg += f" \t(t0 {pretty_num(self.i401k.get_balance()*deflate)})\n"
        msg += f"IRA:       {pretty_num(self.IRA.get_balance())}"
        msg += f" \t(t0 {pretty_num(self.IRA.get_balance()*deflate)})\n"

        msg += f"NET WORTH: {pretty_num(self.net_worth())} "
        msg += f"\t(t0 {pretty_num(self.net_worth()*deflate)})\n\n"
        return msg

    def net_worth(self):
        net_worth = self.checkings.balance
        net_worth += self.savings.balance
        net_worth += self.brokerage.get_balance()
        net_worth += self.i401k.get_balance()
        net_worth += self.IRA.get_balance()
        return net_worth


if __name__ == "__main__":

    p = portfolio()
    p.increment_year(pay=1e6, spending=1e5)