Improve type annotations and DataFrame code

Author: simicd

README.md

@@ -44,26 +44,30 @@ This Python package implements the Smith-Wilson yield curve fitting algorithm. I
 1. Call the Smiwth-Wilson fitting algorithm. This returns the rates as numpy vector with each element corresponding to the maturity in `terms_target`
    ```py
    # Calculate fitted rates based on actual observations and two parametes alpha & UFR
-   rates_ext = sw.fit_smithwilson_rates(rates_obs=rates, t_obs=terms,
-                                        t_target=terms_target, alpha=alpha, ufr=ufr)
+   fitted_rates = sw.fit_smithwilson_rates(rates_obs=rates, t_obs=terms,
+                                           t_target=terms_target, alpha=alpha, ufr=ufr)
    ```
 
 1. To display the results and/or processing them it can be useful to turn them into a table, here using the pandas library:
    ```py
+   # Ensure pandas package is imported
    import pandas as pd
 
-   # Create dictionary with maturity as key and rate as value
-   observed = dict(zip(terms, rates))
-   extrapolated = dict(zip(terms_target, rates_ext.flatten()))
-   # Create and print dataframe
-   print(pd.DataFrame({"Observed": observed, "Extrapolated": extrapolated}))
+   # ...
+
+   # Turn inputs & outputs into dataframe
+   observed_df = pd.DataFrame(data=rates, index=terms, columns=["observed"])
+   extrapolated_df = pd.DataFrame(data=fitted_rates, index=terms_target, columns=["extrapolated"])
+
+   # Combine and print dataframe
+   print(observed_df.join(extrapolated_df, how="outer"))
    ```
 
 A complete example can be found in [main.py](https://github.com/simicd/smith-wilson-py/blob/master/main.py)
 <br /><br />
 
 ## Algorithm
-The algorithm is fully vectorized and uses numpy, making it very performant. All functions are in [core.py](https://github.com/simicd/smith-wilson-py/blob/master/smithwilson/core.py).
+The algorithm is fully vectorized and uses numpy, making it very performant. The code is in [core.py](https://github.com/simicd/smith-wilson-py/blob/master/smithwilson/core.py).
 
 The function `fit_smithwilson_rates()` expects following parameters:
 - Observed rates

main.py

@@ -4,26 +4,28 @@
 # Input - Switzerland EIOPA spot rates with LLP 25 years and extrapolation period of 150 years
 # Source: https://eiopa.europa.eu/Publications/Standards/EIOPA_RFR_20190531.zip
 #         EIOPA_RFR_20190531_Term_Structures.xlsx; Tab: RFR_spot_no_VA
-rates = [-0.00803, -0.00814, -0.00778, -0.00725, -0.00652,
-         -0.00565, -0.0048, -0.00391, -0.00313, -0.00214,
-         -0.0014, -0.00067, -0.00008, 0.00051, 0.00108,
-         0.00157, 0.00197, 0.00228, 0.0025, 0.00264,
-         0.00271, 0.00274, 0.0028, 0.00291, 0.00309]
-terms = [float(y + 1) for y in range(len(rates))] # 1.0, 2.0, ..., 25.0
+rates = [
+    -0.00803, -0.00814, -0.00778, -0.00725, -0.00652, -0.00565, -0.0048, -0.00391, -0.00313, -0.00214, -0.0014, -0.00067,
+    -0.00008, 0.00051, 0.00108, 0.00157, 0.00197, 0.00228, 0.0025, 0.00264, 0.00271, 0.00274, 0.0028, 0.00291, 0.00309
+]
+terms = [float(y + 1) for y in range(len(rates))]     # 1.0, 2.0, ..., 25.0
 ufr = 0.029
 alpha = 0.128562
 
 # Target - Extrapolate to 150 years
-terms_target = [float(y + 1) for y in range(150)] # 1.0, 2.0, ..., 150.0
+terms_target = [float(y + 1) for y in range(150)]     # 1.0, 2.0, ..., 150.0
 
 # Calculate fitted rates based on actual observations and two parametes alpha & UFR
-rates_ext = sw.fit_smithwilson_rates(rates_obs=rates, t_obs=terms,
-                                     t_target=terms_target, alpha=alpha, ufr=ufr)
+fitted_rates = sw.fit_smithwilson_rates(rates_obs=rates, t_obs=terms, t_target=terms_target, alpha=alpha, ufr=ufr)
 
 # Display Outputs
 # Create dictionary with maturity as key and rate as value
-observed = dict(zip(terms, rates))
-extrapolated = dict(zip(terms_target, rates_ext.flatten()))
+extrapolated = dict(zip(terms_target, fitted_rates.flatten()))
+print(extrapolated)
 
-# Create and print dataframe
-print(pd.DataFrame({"Observed": observed, "Extrapolated": extrapolated}))
+# Create dataframe
+observed_df = pd.DataFrame(data=rates, index=terms, columns=["observed"])
+extrapolated_df = pd.DataFrame(data=fitted_rates, index=terms_target, columns=["extrapolated"])
+
+# Combie and print dataframe
+print(observed_df.join(extrapolated_df, how="outer"))

smithwilson/core.py

@@ -1,8 +1,9 @@
-
-from math import exp, log
+from math import log
 import numpy as np
+from typing import Union, List
+
 
-def calculate_prices(rates: np.ndarray, t: np.ndarray) -> np.ndarray:
+def calculate_prices(rates: Union[np.ndarray, List[float]], t: Union[np.ndarray, List[float]]) -> np.ndarray:
     """Calculate prices from zero-coupon rates
 
     Args:
@@ -13,10 +14,14 @@ def calculate_prices(rates: np.ndarray, t: np.ndarray) -> np.ndarray:
         Prices as vector of length n
     """
 
+    # Convert list into numpy array
+    rates = np.array(rates)
+    t = np.array(t)
+
     return np.power(1 + rates, -t)
 
 
-def ufr_discount_factor(ufr: float, t: np.ndarray) -> np.ndarray:
+def ufr_discount_factor(ufr: float, t: Union[np.ndarray, List[float]]) -> np.ndarray:
     """Calculate Ultimate Forward Rate (UFR) discount factors.
 
     Takes the UFR with a vector of maturities and returns for each of the
@@ -37,10 +42,14 @@ def ufr_discount_factor(ufr: float, t: np.ndarray) -> np.ndarray:
 
     # Convert annualized ultimate forward rate to log-return
     ufr = log(1 + ufr)
+
+    # Convert list into numpy array
+    t = np.array(t)
+
     return np.exp(-ufr * t)
 
 
-def wilson_function(t1: np.ndarray, t2: np.ndarray, alpha: float, ufr: float) -> np.ndarray:
+def wilson_function(t1: Union[np.ndarray, List[float]], t2: Union[np.ndarray, List[float]], alpha: float, ufr: float) -> np.ndarray:
     """Calculate matrix of Wilson functions
 
     The Smith-Wilson method requires the calculation of a series of Wilson
@@ -85,7 +94,7 @@ def wilson_function(t1: np.ndarray, t2: np.ndarray, alpha: float, ufr: float) ->
     return W
 
 
-def fit_parameters(rates: np.ndarray, t: np.ndarray, alpha: float, ufr: float) -> np.ndarray:
+def fit_parameters(rates: Union[np.ndarray, List[float]], t: Union[np.ndarray, List[float]], alpha: float, ufr: float) -> np.ndarray:
     """Calculate Smith-Wilson parameter vector ζ
 
     Given the Wilson-matrix, vector of discount factors and prices,
@@ -115,13 +124,13 @@ def fit_parameters(rates: np.ndarray, t: np.ndarray, alpha: float, ufr: float) -
     # Calculate vector of parameters (p. 17)
     # To invert the Wilson-matrix, conversion to type matrix is required
     zeta = np.matrix(W).I * (mu - P)
-    zeta = np.array(zeta) # Convert back to more general array type
+    zeta = np.array(zeta)     # Convert back to more general array type
 
     return zeta
 
 
-def fit_smithwilson_rates(rates_obs: np.ndarray, t_obs: np.ndarray, t_target: np.ndarray,
-                                alpha: float, ufr: float) -> np.ndarray:
+def fit_smithwilson_rates(rates_obs: Union[np.ndarray, List[float]], t_obs: Union[np.ndarray, List[float]],
+                          t_target: Union[np.ndarray, List[float]], alpha: float, ufr: float) -> np.ndarray:
     """Calculate zero-coupon yields with Smith-Wilson method based on observed rates.
 
     This function expects the rates and initial maturity vector to be
@@ -170,7 +179,7 @@ def fit_smithwilson_rates(rates_obs: np.ndarray, t_obs: np.ndarray, t_target: np
 
     # Price vector - equivalent to discounting with zero-coupon yields 1/(1 + r)^t
     # for prices where t_obs = t_target. All other matuirites are interpolated or extrapolated
-    P = ufr_disc - W @ zeta  # '@' in numpy is the dot product of two matrices
+    P = ufr_disc - W @ zeta     # '@' in numpy is the dot product of two matrices
 
     # Transform price vector to zero-coupon rate vector (1/P)^(1/t) - 1
-    return np.power(1/P, 1/t_target) - 1
+    return np.power(1 / P, 1 / t_target) - 1