What can I choose?
What might it become at 62?

This is the complete program used by the notebook, including table display, chart display, and ZIP export.

"""Reproduce the retirement paper's 12 tables and five figures.

Python 3.9+; install dependencies once:
    python -m pip install numpy matplotlib

Run from a terminal:
    python retirement_paper_all_tables_graphs.py --output-dir C:/ai

Or open this file in Spyder/IDLE/VS Code and Run. OUTPUT_DIR below is the
default. The folder is created if necessary; same-named outputs are replaced.
No external data, Stata, JavaScript runtime, or internet access is required
after installing NumPy and Matplotlib. This script uses the original
JavaScript random-number algorithm, not NumPy's default random generator.

All plotted values are calculated. Reference values are used ONLY to test
replication. Baseline 60/40 median rounds to $3,950,073, as in the paper.
Small floating-point differences across systems may remain below rounding.
Plots reproduce values and meaning; fonts/layout can vary by computer.

CSV numbers are unrounded; format currency in Excel for paper-style tables.
PNG figures are 300 dpi; PDFs provide vector versions for publication.
The journal's three numerical tables are also exported.

Model: illustrative assumptions, parametric Monte Carlo, 3,000 paths,
40 years, monthly end-of-period deposits beginning at $1,000 and rising
5% annually. No retirement withdrawals, spending target, or success rate.
No real-world account limits or detailed taxes; no historical calibration.
"""

import argparse
import csv
import json
from pathlib import Path
import sys

import numpy as np
import matplotlib
matplotlib.use("Agg")  # Save figures without needing an interactive window.
import matplotlib.pyplot as plt

OUTPUT_DIR = Path("/content/retirement_results") if Path("/content").exists() else Path.cwd() / "retirement_results"
PATHS = 3000
MONTHS = 480
SEED = 424242
FEE = 0.001
INITIAL_MONTHLY = 1000.0
DEPOSIT_GROWTH = 0.05
TAX_NOW = 0.22
TAX_LATER = 0.20

ASSET_NAMES = [
    "S&P 500 index", "Blue-chip US stocks", "QQQ-style growth",
    "US + international stocks", "High-growth mutual fund", "US Treasuries",
    "Investment-grade corporate bonds", "Inflation-linked Treasuries (TIPS)",
    "Gold", "Crypto / Bitcoin-like exposure",
]
PORTFOLIO_NAMES = ASSET_NAMES + [
    "Treasury + corporate bond blend", "Balanced 60/40", "Conservative 40/60",
    "Aggressive 80/20", "Balanced + 10% gold", "Balanced + 5% crypto",
]
ASSET_MU = np.array([.08, .075, .08, .077, .08, .04, .05, .015, .04, 0.0])
ASSET_SD = np.array([.16, .14, .23, .15, .25, .06, .08, .05, .18, .65])
WEIGHTS = np.vstack((np.eye(10), np.zeros((6, 10))))
WEIGHTS[10, [5, 6]] = [.5, .5]
WEIGHTS[11, [0, 5]] = [.6, .4]
WEIGHTS[12, [0, 5]] = [.4, .6]
WEIGHTS[13, [0, 5]] = [.8, .2]
WEIGHTS[14, [0, 5, 8]] = [.55, .35, .10]
WEIGHTS[15, [0, 5, 9]] = [.6, .35, .05]

LOADINGS = np.zeros((10, 4))
RESIDUAL = np.zeros(10)
LOADINGS[:5, 0] = .9
RESIDUAL[:5] = np.sqrt(.19)
LOADINGS[5:8, 1] = .8
RESIDUAL[5:8] = np.sqrt(.36)
LOADINGS[8, 2] = 1
LOADINGS[9, 0] = .4
LOADINGS[9, 3] = np.sqrt(.84)
WEIGHTED_SD = WEIGHTS * ASSET_SD
SIGMA = np.sqrt(np.sum((WEIGHTED_SD @ LOADINGS)**2, axis=1)
                + np.sum((WEIGHTED_SD * RESIDUAL)**2, axis=1))

SCENARIOS = [
    ("2.5% for 40 years", np.full(40, .025)),
    ("4% for 10 years then 2.5%", np.r_[np.full(10, .04), np.full(30, .025)]),
    ("4% for 40 years", np.full(40, .04)),
    ("6% for 10 years then 3%", np.r_[np.full(10, .06), np.full(30, .03)]),
]


def paper_uniforms(count, seed):
    """JavaScript Mulberry32, using exact unsigned integer operations.

    Math.imul's low 32 bits are preserved; a sign conversion is unnecessary
    because every subsequent XOR/shift uses the same 32-bit pattern.
    This function does not modify global NumPy random state.
    """
    mask = np.uint64(0xFFFFFFFF)
    state = (np.uint64(seed) + np.arange(1, count + 1, dtype=np.uint64)
             * np.uint64(0x6D2B79F5)) & mask
    t = ((state ^ (state >> np.uint64(15))) * (state | np.uint64(1))) & mask
    product = (((t ^ (t >> np.uint64(7))) * (t | np.uint64(61))) & mask)
    t = t ^ ((t + product) & mask)
    return ((t ^ (t >> np.uint64(14))) & mask).astype(np.float64) / 4294967296.0


def paper_shocks(seed=SEED):
    """Box-Muller normals in the original path-first, then month order."""
    z = np.empty((PATHS, MONTHS))
    state = int(seed)
    for start in range(0, PATHS, 100):
        n = min(100, PATHS - start)
        count = 2 * n * MONTHS
        u = paper_uniforms(count, state)
        state = (state + count * 0x6D2B79F5) % 4294967296
        normal = (np.sqrt(-2 * np.log(np.maximum(u[0::2], 1e-12)))
                  * np.cos(2 * np.pi * u[1::2]))
        z[start:start + n] = normal.reshape(n, MONTHS)
    return z


def simulate(z, portfolio_id, inflation, fee=FEE, stock_return=.08):
    """Return nominal Roth-style terminal percentiles P10, P50 and P90.

    portfolio_id is zero-based. Sigma is an annualized log-return diffusion,
    not the exact SD of annual simple returns. We simulate the aggregate
    fixed-weight portfolio approximation, not individual traded holdings.
    """
    w = WEIGHTS[portfolio_id]
    mu = ASSET_MU.copy()
    mu[0] = stock_return
    annual_means = []
    for pi in inflation:
        mu[7] = (1 + ASSET_MU[7]) * (1 + pi) - 1
        annual_means.append(float(w @ mu))
    sigma = SIGMA[portfolio_id]
    drift = (np.log1p(annual_means) - .5 * sigma**2 + np.log1p(-fee)) / 12
    b = np.zeros(z.shape[0])
    deposit = INITIAL_MONTHLY
    for month in range(MONTHS):
        b = b * np.exp(drift[month // 12] + sigma / np.sqrt(12) * z[:, month]) + deposit
        if (month + 1) % 12 == 0:
            deposit *= 1 + DEPOSIT_GROWTH
    return np.quantile(b, [.1, .5, .9], method="linear")


def save_csv(folder, filename, columns, rows):
    with (folder / filename).open("w", newline="", encoding="utf-8-sig") as f:
        writer = csv.writer(f)
        writer.writerow(columns)
        writer.writerows(rows)


def save_figure(fig, folder, name):
    fig.tight_layout()
    for extension in ("png", "pdf"):
        fig.savefig(folder / f"{name}.{extension}", dpi=300, bbox_inches="tight")
    plt.close(fig)


def make_figures(folder, base, real, fees, returns):
    teal, dark, gold = "#146f70", "#214b50", "#ae812b"
    plt.rcParams.update({"font.size": 10, "axes.spines.top": False,
                         "axes.spines.right": False, "axes.titleweight": "bold",
                         "axes.labelcolor": dark, "text.color": dark})
    for ids, title, filename in [
        ([0, 1, 2, 3, 4], "Stock strategies: 10th to 90th percentile ranges", "figure1_stock_strategies"),
        ([12, 11, 13, 14, 15], "Mixed portfolios: lower outcomes and upper outcomes", "figure2_mixed_portfolios"),
    ]:
        fig, ax = plt.subplots(figsize=(9, 4.4))
        q = base[ids] / 1e6
        y = np.arange(len(ids))
        ax.hlines(y, q[:, 0], q[:, 2], color="#79a99f", linewidth=4)
        ax.scatter(q[:, 0], y, color=dark, s=18, zorder=3)
        ax.scatter(q[:, 2], y, color=dark, s=18, zorder=3)
        ax.scatter(q[:, 1], y, color=teal, s=42, zorder=4)
        ax.set_yticks(y)
        ax.set_yticklabels([PORTFOLIO_NAMES[i] for i in ids])
        ax.invert_yaxis()
        ax.set_xlabel("Ending balance at age 62 ($ millions)")
        ax.set_title(title)
        ax.grid(axis="x", alpha=.15)
        ax.set_axisbelow(True)
        save_figure(fig, folder, filename)

    fig, ax = plt.subplots(figsize=(8.5, 4.5))
    heights = np.array([real[s, 11, 1] for s in range(3)]) / 1e6
    bars = ax.bar(["2.5% throughout", "4% for 10 years\nthen 2.5%", "4% throughout"],
                  heights, color=[teal, "#60998e", "#b8ceca"])
    for bar, value in zip(bars, heights):
        ax.text(bar.get_x() + bar.get_width()/2, value + .025,
                f"${value:.3f}m", ha="center", weight="bold")
    ax.set_ylim(0, max(heights) * 1.22)
    ax.set_title("The same 60/40 nominal median buys less")
    ax.set_ylabel("Median buying power ($ millions)")
    save_figure(fig, folder, "figure3_inflation")

    fig, ax = plt.subplots(figsize=(8.5, 4.5))
    for j, (label, color) in enumerate([(ASSET_NAMES[0], dark), (PORTFOLIO_NAMES[11], teal)]):
        ax.plot([.1, .5, 1], fees[:, j, 1]/1e6, "o-", label=label, color=color)
    ax.set_xlabel("Annual fee (%)")
    ax.set_ylabel("Median ending balance ($ millions)")
    ax.set_title("Annual fees compound over the savings horizon")
    ax.grid(alpha=.15)
    ax.legend(frameon=False)
    save_figure(fig, folder, "figure4_fee_sensitivity")

    fig, ax = plt.subplots(figsize=(8.5, 4.5))
    for j, pid in enumerate([0, 11, 13]):
        ax.plot([6, 8, 10], returns[:, j, 1]/1e6, "o-",
                label=PORTFOLIO_NAMES[pid], color=[dark, teal, gold][j])
    ax.set_xlabel("Assumed annual S&P 500 simple return (%)")
    ax.set_ylabel("Median ending balance ($ millions)")
    ax.set_title("Changing expected returns changes the conclusion")
    ax.grid(alpha=.15)
    ax.legend(frameon=False)
    save_figure(fig, folder, "figure5_return_sensitivity")


def main(output_dir=OUTPUT_DIR):
    folder = Path(output_dir).expanduser()
    folder.mkdir(parents=True, exist_ok=True)
    print("Generating the original seeded shocks...", flush=True)
    z = paper_shocks()
    factors = np.array([np.prod(1 + p) for _, p in SCENARIOS])
    nominal = np.array([[simulate(z, j, p) for j in range(16)] for _, p in SCENARIOS])
    real = nominal / factors[:, None, None]
    base = nominal[0]
    fee_values = [.001, .005, .01]
    fees = np.array([[simulate(z, j, SCENARIOS[0][1], fee=f) for j in [0, 11]] for f in fee_values])
    return_values = [.06, .08, .10]
    returns = np.array([[simulate(z, j, SCENARIOS[0][1], stock_return=r)
                         for j in [0, 11, 13]] for r in return_values])
    seed_values = [424242, 12345, 67890, 13579, 24680]
    seeds = np.array([simulate(z if seed == SEED else paper_shocks(seed), 11, SCENARIOS[0][1])
                      for seed in seed_values])

    # Replication gates are reference checks, not inputs to the simulation.
    np.testing.assert_allclose(paper_uniforms(3, SEED),
        [.1287079993635416, .23631658218801022, .18376533640548587], rtol=0, atol=0)
    np.testing.assert_allclose(base[11],
        [2475478.4676766302, 3950073.4406768735, 6629328.1098863855], rtol=0, atol=.02)
    np.testing.assert_allclose(fees[-1, 1, 1], 3278581.89431493, rtol=0, atol=.02)
    np.testing.assert_allclose(returns[-1, 1, 1], 5040026.35460735, rtol=0, atol=.02)
    np.testing.assert_allclose(seeds[:, 1],
        [3950073.4406768735, 3933419.5919726267, 3934369.398448094,
         3961692.221452865, 3905098.581126632], rtol=0, atol=.02)
    np.testing.assert_allclose(WEIGHTS.sum(axis=1), 1)
    assert np.all(np.diff(nominal, axis=2) >= 0)
    non_tips = [i for i in range(16) if i != 7]
    np.testing.assert_array_equal(nominal[0, non_tips], nominal[2, non_tips])
    print("Replication checks passed; exporting computed tables and figures.", flush=True)

    total = 12000 * ((1.05**40 - 1) / .05)
    save_csv(folder, "table01_savings_schedule.csv", ["Saving year", "Monthly deposit", "Annual deposits"],
             [[y, 1000*1.05**(y-1), 12000*1.05**(y-1)] for y in [1, 5, 10, 20, 30, 40]])
    asset_rows = [[ASSET_NAMES[j], ASSET_MU[j], "real" if j == 7 else "nominal", ASSET_SD[j]] for j in range(10)]
    save_csv(folder, "table02_asset_assumptions.csv", ["Asset", "Annual return input", "Return units", "Annual log diffusion"], asset_rows)
    weight_text = ["50% Treasury / 50% corporate", "60% SP500 / 40% Treasury", "40% SP500 / 60% Treasury",
                   "80% SP500 / 20% Treasury", "55% SP500 / 35% Treasury / 10% gold", "60% SP500 / 35% Treasury / 5% crypto"]
    save_csv(folder, "table03_mixed_portfolio_assumptions.csv", ["Strategy", "Weights", "Annual return input", "Annual log diffusion"],
             [[PORTFOLIO_NAMES[j], weight_text[j-10], WEIGHTS[j] @ ASSET_MU, SIGMA[j]] for j in range(10, 16)])
    save_csv(folder, "table04_inflation_price_factors.csv", ["Inflation scenario", "Price factor", "Buying power of 1 million"],
             [[SCENARIOS[s][0], factors[s], 1e6/factors[s]] for s in range(3)])
    qcols = ["Strategy", "10th percentile", "Median", "90th percentile"]
    save_csv(folder, "table05_baseline_nominal.csv", qcols, [[PORTFOLIO_NAMES[j], *base[j]] for j in range(16)])
    save_csv(folder, "table06_baseline_real.csv", qcols, [[PORTFOLIO_NAMES[j], *real[0, j]] for j in range(16)])
    save_csv(folder, "table07_bonds_TIPS_gold.csv", ["Choice", "Real median", "Nominal P90 minus P10"],
             [[PORTFOLIO_NAMES[j], real[0, j, 1], base[j, 2]-base[j, 0]] for j in [5, 6, 10, 7, 8]])
    save_csv(folder, "table08_crypto_terminal_stress.csv", ["Crypto component loss", "Nominal median after stress", "Real median after stress"],
             [[loss, base[15, 1]*(1-.05*loss), real[0, 15, 1]*(1-.05*loss)] for loss in [0, .8, 1]])
    save_csv(folder, "table09_fee_sensitivity.csv", ["Annual fee percent", "SP500 median", "60/40 median"],
             [[100*f, *fees[i, :, 1]] for i, f in enumerate(fee_values)])
    save_csv(folder, "table10_return_sensitivity.csv", ["SP500 return percent", "All-stock median", "60/40 median", "80/20 median"],
             [[100*r, *returns[i, :, 1]] for i, r in enumerate(return_values)])
    save_csv(folder, "table11_savings_scale.csv", ["Initial monthly deposit", "Total nominal deposits", "60/40 median", "Real median"],
             [[d, total*d/1000, base[11, 1]*d/1000, real[0, 11, 1]*d/1000] for d in [250, 500, 1000]])
    save_csv(folder, "table12_seed_sensitivity.csv", ["Seed", "10th percentile", "Median", "90th percentile"],
             [[seed, *seeds[i]] for i, seed in enumerate(seed_values)])

    # The journal manuscript's three numerical tables.
    settings = {"paths": PATHS, "seed": SEED, "years": 40, "months": MONTHS,
                "starting_age": 22, "retirement_age": 62, "initial_balance": 0,
                "initial_monthly": INITIAL_MONTHLY, "annual_deposit_growth": DEPOSIT_GROWTH,
                "annual_fee": FEE, "current_tax_rate": TAX_NOW, "withdrawal_tax_rate": TAX_LATER,
                "equal_cost": 0, "total_nominal_deposits": total,
                "rng": "JavaScript Mulberry32 with Box-Muller", "draw_order": "path then month",
                "deposit_timing": "after monthly return", "quantiles": "linear interpolation"}
    save_csv(folder, "journal_table01_core_assumptions.csv", ["Input", "Setting"], list(settings.items()))
    save_csv(folder, "journal_table02_baseline.csv", ["Alternative", "Nominal P10 millions", "Nominal P50 millions", "Nominal P90 millions", "Real P50 millions"],
             [[PORTFOLIO_NAMES[j], *(base[j]/1e6), real[0, j, 1]/1e6] for j in [0, 5, 7, 11, 13, 14, 15]])
    save_csv(folder, "journal_table03_inflation.csv", ["Inflation path", "Price factor", "60/40 nominal median millions", "60/40 real median millions", "TIPS real median millions"],
             [[SCENARIOS[s][0], factors[s], nominal[s, 11, 1]/1e6, real[s, 11, 1]/1e6, real[s, 7, 1]/1e6] for s in range(3)])
    save_csv(folder, "paper_run_settings.csv", list(settings.keys()), [list(settings.values())])
    save_csv(folder, "paper_portfolio_weights.csv", ["Strategy", *ASSET_NAMES],
             [[PORTFOLIO_NAMES[j], *WEIGHTS[j]] for j in range(16)])

    # One combined results file, including simplified account comparisons.
    # With equal deposits, pre-tax after-tax values are 80% of Roth values.
    # At equal household cost, the multiplier is .80/.78 instead.
    records = []
    def record(experiment, scenario, pid, q, fac, fee, stock, seed):
        records.append([experiment, scenario, PORTFOLIO_NAMES[pid], pid+1, fee, stock, seed,
                        fac, *q, *(q/fac), *(q*(1-TAX_LATER)),
                        *(q*(1-TAX_LATER)/(1-TAX_NOW))])
    for s, (name, _) in enumerate(SCENARIOS):
        for j in range(16): record("baseline/inflation", name, j, nominal[s, j], factors[s], FEE, .08, SEED)
    for i, fee in enumerate(fee_values):
        for k, j in enumerate([0, 11]): record("fee", SCENARIOS[0][0], j, fees[i, k], factors[0], fee, .08, SEED)
    for i, r in enumerate(return_values):
        for k, j in enumerate([0, 11, 13]): record("return", SCENARIOS[0][0], j, returns[i, k], factors[0], FEE, r, SEED)
    for i, seed in enumerate(seed_values): record("seed", SCENARIOS[0][0], 11, seeds[i], factors[0], FEE, .08, seed)
    columns = ["experiment", "scenario", "portfolio", "portfolio_id", "annual_fee", "stock_return", "seed", "price_factor",
               "roth_p10", "roth_p50", "roth_p90", "real_roth_p10", "real_roth_p50", "real_roth_p90",
               "pretax_equal_deposits_p10", "pretax_equal_deposits_p50", "pretax_equal_deposits_p90",
               "pretax_equal_cost_p10", "pretax_equal_cost_p50", "pretax_equal_cost_p90"]
    assert len(records) == 84
    save_csv(folder, "paper_simulation_results.csv", columns, records)
    make_figures(folder, base, real, fees, returns)
    (folder / "replication_check.json").write_text(json.dumps({
        "checks_passed": True, "baseline_6040_nominal_percentiles": base[11].tolist(),
        "baseline_6040_real_median": float(real[0, 11, 1]),
        "numpy_version": np.__version__, "matplotlib_version": matplotlib.__version__,
        "python_version": sys.version.split()[0], "settings": settings,
    }, indent=2), encoding="utf-8")
    print(f"Baseline 60/40 median: ${base[11, 1]:,.0f}")
    print(f"Baseline 60/40 real median: ${real[0, 11, 1]:,.0f}")
    print(f"Saved 12 study tables, 3 journal tables, 5 PNG/PDF figures, and supporting CSV files to {folder.resolve()}")



# Run the complete simulation, then show every exported table and figure.
main(OUTPUT_DIR)

import pandas as pd
from IPython.display import display, Image, Markdown, FileLink
import shutil

print("\nALL SAVED TABLES (including supporting results)")
for file in sorted(OUTPUT_DIR.glob("*.csv")):
    display(Markdown("### " + file.stem.replace("_", " ")))
    with pd.option_context("display.max_rows", None,
                           "display.max_columns", None,
                           "display.max_colwidth", None):
        display(pd.read_csv(file))

print("\nALL FIVE FIGURES")
for file in sorted(OUTPUT_DIR.glob("figure*.png")):
    display(Markdown("### " + file.stem.replace("_", " ")))
    display(Image(filename=str(file), width=1000))

archive = shutil.make_archive(str(OUTPUT_DIR.parent / "retirement_results_all"), "zip", OUTPUT_DIR)
print("\nDownload all CSV tables, PNG/PDF figures, and replication checks:")
try:
    from google.colab import files
    files.download(archive)
except ImportError:
    display(FileLink(archive))

"""Standalone assumption experiment. Paste ALL code into one new Colab cell."""
# CHANGE ONLY THE VALUES IN THIS BOX, THEN RUN THE ENTIRE CELL AGAIN.
# Use decimals: 0.06 means 6%; 0.005 means 0.50%.
NEW_STOCK_RETURN = 0.08       # Try 0.06: lower expected S&P 500 return
NEW_STOCK_VOLATILITY = 0.16   # Try 0.20: greater modeled S&P 500 volatility
NEW_MONTHLY_DEPOSIT = 1000.0  # Try 500.0: smaller initial monthly saving
NEW_DEPOSIT_GROWTH = 0.05     # Try 0.03: deposits rise 3% per year
NEW_ANNUAL_FEE = 0.001       # Try 0.01: a 1% annual fee
NEW_INFLATION = 0.025        # Try 0.04: 4% inflation every year
NEW_WITHDRAWAL_TAX = 0.20    # Try 0.25: simplified pre-tax withdrawal tax
NEW_STOCK_WEIGHT = 0.60      # Try 0.40: custom mix becomes 40% stock / 60% Treasury
NEW_SEED = 424242            # Try 12345: different random draws
# Keep 3,000 paths and ages 22 to 62 for comparisons with the study.

import argparse
import csv
import json
from pathlib import Path
import sys

import numpy as np
import matplotlib
matplotlib.use("Agg")  # Save figures without needing an interactive window.
import matplotlib.pyplot as plt

OUTPUT_DIR = Path("/content/retirement_results") if Path("/content").exists() else Path.cwd() / "retirement_results"
PATHS = 3000
MONTHS = 480
SEED = 424242
FEE = 0.001
INITIAL_MONTHLY = 1000.0
DEPOSIT_GROWTH = 0.05
TAX_NOW = 0.22
TAX_LATER = 0.20

ASSET_NAMES = [
    "S&P 500 index", "Blue-chip US stocks", "QQQ-style growth",
    "US + international stocks", "High-growth mutual fund", "US Treasuries",
    "Investment-grade corporate bonds", "Inflation-linked Treasuries (TIPS)",
    "Gold", "Crypto / Bitcoin-like exposure",
]
PORTFOLIO_NAMES = ASSET_NAMES + [
    "Treasury + corporate bond blend", "Balanced 60/40", "Conservative 40/60",
    "Aggressive 80/20", "Balanced + 10% gold", "Balanced + 5% crypto",
]
ASSET_MU = np.array([.08, .075, .08, .077, .08, .04, .05, .015, .04, 0.0])
ASSET_SD = np.array([.16, .14, .23, .15, .25, .06, .08, .05, .18, .65])
WEIGHTS = np.vstack((np.eye(10), np.zeros((6, 10))))
WEIGHTS[10, [5, 6]] = [.5, .5]
WEIGHTS[11, [0, 5]] = [.6, .4]
WEIGHTS[12, [0, 5]] = [.4, .6]
WEIGHTS[13, [0, 5]] = [.8, .2]
WEIGHTS[14, [0, 5, 8]] = [.55, .35, .10]
WEIGHTS[15, [0, 5, 9]] = [.6, .35, .05]

LOADINGS = np.zeros((10, 4))
RESIDUAL = np.zeros(10)
LOADINGS[:5, 0] = .9
RESIDUAL[:5] = np.sqrt(.19)
LOADINGS[5:8, 1] = .8
RESIDUAL[5:8] = np.sqrt(.36)
LOADINGS[8, 2] = 1
LOADINGS[9, 0] = .4
LOADINGS[9, 3] = np.sqrt(.84)
WEIGHTED_SD = WEIGHTS * ASSET_SD
SIGMA = np.sqrt(np.sum((WEIGHTED_SD @ LOADINGS)**2, axis=1)
                + np.sum((WEIGHTED_SD * RESIDUAL)**2, axis=1))

SCENARIOS = [
    ("2.5% for 40 years", np.full(40, .025)),
    ("4% for 10 years then 2.5%", np.r_[np.full(10, .04), np.full(30, .025)]),
    ("4% for 40 years", np.full(40, .04)),
    ("6% for 10 years then 3%", np.r_[np.full(10, .06), np.full(30, .03)]),
]


def paper_uniforms(count, seed):
    """JavaScript Mulberry32, using exact unsigned integer operations.

    Math.imul's low 32 bits are preserved; a sign conversion is unnecessary
    because every subsequent XOR/shift uses the same 32-bit pattern.
    This function does not modify global NumPy random state.
    """
    mask = np.uint64(0xFFFFFFFF)
    state = (np.uint64(seed) + np.arange(1, count + 1, dtype=np.uint64)
             * np.uint64(0x6D2B79F5)) & mask
    t = ((state ^ (state >> np.uint64(15))) * (state | np.uint64(1))) & mask
    product = (((t ^ (t >> np.uint64(7))) * (t | np.uint64(61))) & mask)
    t = t ^ ((t + product) & mask)
    return ((t ^ (t >> np.uint64(14))) & mask).astype(np.float64) / 4294967296.0


def paper_shocks(seed=SEED):
    """Box-Muller normals in the original path-first, then month order."""
    z = np.empty((PATHS, MONTHS))
    state = int(seed)
    for start in range(0, PATHS, 100):
        n = min(100, PATHS - start)
        count = 2 * n * MONTHS
        u = paper_uniforms(count, state)
        state = (state + count * 0x6D2B79F5) % 4294967296
        normal = (np.sqrt(-2 * np.log(np.maximum(u[0::2], 1e-12)))
                  * np.cos(2 * np.pi * u[1::2]))
        z[start:start + n] = normal.reshape(n, MONTHS)
    return z


def simulate(z, portfolio_id, inflation, fee=FEE, stock_return=.08):
    """Return nominal Roth-style terminal percentiles P10, P50 and P90.

    portfolio_id is zero-based. Sigma is an annualized log-return diffusion,
    not the exact SD of annual simple returns. We simulate the aggregate
    fixed-weight portfolio approximation, not individual traded holdings.
    """
    w = WEIGHTS[portfolio_id]
    mu = ASSET_MU.copy()
    mu[0] = stock_return
    annual_means = []
    for pi in inflation:
        mu[7] = (1 + ASSET_MU[7]) * (1 + pi) - 1
        annual_means.append(float(w @ mu))
    sigma = SIGMA[portfolio_id]
    drift = (np.log1p(annual_means) - .5 * sigma**2 + np.log1p(-fee)) / 12
    b = np.zeros(z.shape[0])
    deposit = INITIAL_MONTHLY
    for month in range(MONTHS):
        b = b * np.exp(drift[month // 12] + sigma / np.sqrt(12) * z[:, month]) + deposit
        if (month + 1) % 12 == 0:
            deposit *= 1 + DEPOSIT_GROWTH
    return np.quantile(b, [.1, .5, .9], method="linear")


import pandas as pd
from IPython.display import display, Image

for value in [NEW_STOCK_RETURN, NEW_DEPOSIT_GROWTH]:
    assert np.isfinite(value) and value > -1
assert NEW_STOCK_VOLATILITY >= 0 and np.isfinite(NEW_STOCK_VOLATILITY)
assert NEW_MONTHLY_DEPOSIT >= 0 and np.isfinite(NEW_MONTHLY_DEPOSIT)
assert 0 <= NEW_ANNUAL_FEE < 1
assert NEW_INFLATION > -1 and np.isfinite(NEW_INFLATION)
assert 0 <= NEW_WITHDRAWAL_TAX <= 1
assert 0 <= NEW_STOCK_WEIGHT <= 1
assert isinstance(NEW_SEED, int) and 0 <= NEW_SEED < 2**32

folder = Path('/content/retirement_custom') if Path('/content').exists() else Path.cwd() / 'retirement_custom'
folder.mkdir(parents=True, exist_ok=True)
print('Calculating original and changed assumptions...')
original_z = paper_shocks(424242)
baseline = np.array([simulate(original_z, j, np.full(40, .025), fee=.001, stock_return=.08) for j in range(16)])
base_factor = 1.025**40

# Apply your changes, then recompute risk and all 16 portfolio distributions.
ASSET_MU[0] = NEW_STOCK_RETURN
ASSET_SD[0] = NEW_STOCK_VOLATILITY
INITIAL_MONTHLY = NEW_MONTHLY_DEPOSIT
DEPOSIT_GROWTH = NEW_DEPOSIT_GROWTH
WEIGHTS[11] = 0
WEIGHTS[11, [0, 5]] = [NEW_STOCK_WEIGHT, 1-NEW_STOCK_WEIGHT]
weighted_sd = WEIGHTS * ASSET_SD
SIGMA = np.sqrt(np.sum((weighted_sd @ LOADINGS)**2, axis=1)
                + np.sum((weighted_sd * RESIDUAL)**2, axis=1))
z = original_z if NEW_SEED == 424242 else paper_shocks(NEW_SEED)
inflation = np.full(40, NEW_INFLATION)
changed = np.array([simulate(z, j, inflation, fee=NEW_ANNUAL_FEE,
                            stock_return=NEW_STOCK_RETURN) for j in range(16)])
factor = (1 + NEW_INFLATION)**40
assert np.isfinite(changed).all()
assert (np.diff(changed, axis=1) >= 0).all()
assert np.allclose(WEIGHTS.sum(axis=1), 1)
# Changed assumptions are intentionally NOT tested against the original dollar results.
labels = PORTFOLIO_NAMES.copy()
labels[11] = f'Custom mix: {NEW_STOCK_WEIGHT:.0%} S&P 500 / {1-NEW_STOCK_WEIGHT:.0%} Treasury'
deposits = sum(12 * INITIAL_MONTHLY * (1 + DEPOSIT_GROWTH)**y for y in range(40))
rows=[]
for j, label in enumerate(labels):
    rows.append({
        'Strategy':label,
        'Original nominal P10':baseline[j,0], 'Original nominal median':baseline[j,1],
        'Original nominal P90':baseline[j,2], 'Original real median':baseline[j,1]/base_factor,
        'New nominal P10':changed[j,0], 'New nominal median':changed[j,1],
        'New nominal P90':changed[j,2], 'New real median':changed[j,1]/factor,
        'Change in real median':changed[j,1]/factor-baseline[j,1]/base_factor,
        'New pre-tax after-tax nominal median':changed[j,1]*(1-NEW_WITHDRAWAL_TAX),
        'New pre-tax after-tax real median':changed[j,1]*(1-NEW_WITHDRAWAL_TAX)/factor,
        'New total nominal deposits':deposits,
    })
results=pd.DataFrame(rows)
print('For the custom mix, the original comparison is the study 60/40 portfolio.')
with pd.option_context('display.max_rows',None,'display.max_columns',None):
    display(results.round(2))
results.to_csv(folder/'original_vs_changed_assumptions.csv',index=False)
settings={k:v for k,v in list(globals().items()) if k.startswith('NEW_')}
(folder/'changed_assumptions.json').write_text(json.dumps(settings,indent=2))

for title, values, fname in [
    ('Nominal median at age 62', (baseline[:,1], changed[:,1]), 'custom_nominal'),
    ('Median purchasing power in starting-year dollars', (baseline[:,1]/base_factor,changed[:,1]/factor), 'custom_real')]:
    fig,ax=plt.subplots(figsize=(11,9))
    positions=np.arange(16)
    ax.barh(positions-.18,values[0]/1e6,height=.36,label='Original assumptions',color='#b8ceca')
    ax.barh(positions+.18,values[1]/1e6,height=.36,label='Your changed assumptions',color='#126b68')
    ax.set_yticks(positions);ax.set_yticklabels(labels);ax.invert_yaxis()
    ax.set_xlabel('Millions of dollars');ax.set_title(title);ax.legend()
    fig.tight_layout();fig.savefig(folder/(fname+'.png'),dpi=160);plt.close(fig)
    display(Image(filename=str(folder/(fname+'.png'))))
print('New custom-mix nominal median: $'+format(changed[11,1],',.0f'))
print('New custom-mix real median: $'+format(changed[11,1]/factor,',.0f'))
print('Saved comparison CSV, assumptions JSON, and two charts to',folder)
print('To try another assumption: edit the values at the TOP and rerun THIS ENTIRE CELL.')