main.py

import numpy as np
import streamlit as st
import pandas as pd
from scipy.integrate import cumulative_trapezoid

BPS = 0.0001


def nelson_siegel(
    t: np.ndarray, beta0: float, beta1: float, beta2: float, tau: float
) -> np.ndarray:
    """
    Calculate yield curve using Nelson-Siegel model.

    Args:
        t: Time to maturity (years)
        beta0: Long-term interest rate level
        beta1: Short-term component
        beta2: Medium-term component
        tau: Time decay parameter

    Returns:
        Yield curve values
    """
    exp_term = np.exp(-t / tau)
    factor1 = (1 - exp_term) / (t / tau)
    factor2 = factor1 - exp_term
    return beta0 + beta1 * factor1 + beta2 * factor2


def nelson_siegel_forward(
    t: np.ndarray, beta0: float, beta1: float, beta2: float, tau: float
) -> np.ndarray:
    """
    Calculate instantaneous forward rate using Nelson-Siegel model.

    Args:
        t: Time to maturity (years)
        beta0: Long-term interest rate level
        beta1: Short-term component
        beta2: Medium-term component
        tau: Time decay parameter

    Returns:
        Instantaneous forward rate curve values
    """
    exp_term = np.exp(-t / tau)
    return beta0 + beta1 * exp_term + beta2 * (t / tau) * exp_term


def forward_rate_shift(t: np.ndarray, t1: float, t2: float, shift: float) -> np.ndarray:
    """Shift the forward rate in by a constant in a bucket."""
    return np.where(
        (t >= t1) & (t <= t2),
        shift,
        0,
    )


def zero_rate_shift(t: np.ndarray, t1: float, t2: float, shift: float) -> np.ndarray:
    """The zero rate shift implied by the forward rate shift."""
    # The zero rate shift is the integral of the forward rate shift
    # divided by the time to maturity
    return shift / t * np.minimum(np.maximum(t - t1, 0), t2 - t1)


def main():
    st.title("Nelson-Siegel Yield Curve Calculator")

    col1, col2 = st.columns(2)
    with col1:
        beta0 = st.number_input(
            r"Long-term interest rate level ($$\beta_0$$)",
            value=0.03,
            step=0.01,
            min_value=0.0,
            max_value=0.1,
        )
        beta1 = st.number_input(
            r"Short-term component ($$\beta_1$$)",
            value=-0.02,
            min_value=-0.05,
            max_value=0.05,
            step=0.01,
        )
        plot_points = st.number_input(
            "Number of plot points",
            value=1000,
            step=1,
            min_value=1,
            max_value=10000,
        )
    with col2:
        beta2 = st.number_input(
            r"Medium-term component ($$\beta_2$$)",
            value=0.02,
            min_value=-0.05,
            max_value=0.05,
            step=0.01,
        )
        tau = st.number_input(
            r"Time decay parameter ($$\tau$$)",
            min_value=0.0,
            value=2.0,
            max_value=5.0,
            step=0.1,
        )
        ytm = st.number_input("Years to maturity", value=30, step=1, min_value=1)

    col1, col2, col3 = st.columns(3)
    with col1:
        bps_shift = st.number_input(
            "Basis point shift",
            value=0,
            step=1,
            min_value=-100,
            max_value=100,
        )
    with col2:
        shift_t1 = st.number_input(
            "Start of shift (years)",
            value=1,
            step=1,
            min_value=0,
            max_value=30,
        )
    with col3:
        shift_len = st.number_input(
            "Length of shift (years)",
            value=1,
            step=1,
            min_value=0,
            max_value=30,
        )

    # Calculate yield curve
    time_grid = np.linspace(0.0, ytm, plot_points)

    # Base zero and forward rates
    zero_rate = nelson_siegel(time_grid, beta0, beta1, beta2, tau)
    forward_rate = nelson_siegel_forward(time_grid, beta0, beta1, beta2, tau)

    # Apply basis point shift
    shifted_forward_rate = nelson_siegel_forward(
        time_grid, beta0, beta1, beta2, tau
    ) + forward_rate_shift(time_grid, shift_t1, shift_t1 + shift_len, bps_shift * BPS)
    shifted_zero_rate = nelson_siegel(
        time_grid, beta0, beta1, beta2, tau
    ) + zero_rate_shift(time_grid, shift_t1, shift_t1 + shift_len, bps_shift * BPS)

    # Plot yield curve
    st.write("### Yield Curve")
    st.write(
        "This plot shows the yield curve calculated using the Nelson-Siegel model."
    )
    df = pd.DataFrame(
        {
            "time": time_grid,
            "zero": zero_rate,
            "forward": forward_rate,
            "shifted_forward": shifted_forward_rate,
            "shifted_zero": shifted_zero_rate,
        }
    )
    st.line_chart(df, x="time", use_container_width=True)

    if bps_shift != 0:
        st.write("### Implied zero rate shift")
        number_of_shifts = st.number_input(
            "Number of shifts", value=3, step=1, min_value=1, max_value=100
        )
        data = {"time": time_grid}
        for i in range(0, number_of_shifts):
            # Shift the forward rate in by a constant in a bucket
            t1 = shift_t1 + i * shift_len
            t2 = t1 + shift_len
            data[f"shifted_forward_{i}"] = (
                zero_rate_shift(time_grid, t1, t2, bps_shift * BPS) / BPS
            )

        df_s = pd.DataFrame(data)
        df_s = df_s.set_index("time")
        st.line_chart(
            df_s,
            y_label="Shift",
            use_container_width=True,
        )

        st.write("### Cumulative shift")
        st.write("Summing the individual shifts gives the cumulative shift.")
        cs = df_s.sum(axis=1)
        cs.name = "Cumulative Shift"
        st.line_chart(cs, y_label="Shift", use_container_width=True)


if __name__ == "__main__":
    main()