From fa59413221728f9dc4c8aebe15a3d65ee4cccc11 Mon Sep 17 00:00:00 2001
From: Gavin Chan <gavincyi@gmail.com>
Date: Thu, 22 Jun 2023 23:14:02 +0100
Subject: [PATCH] feat: add covariance shrinkage

---
 src/fpm_risk_model/cov_estimator.py | 54 ++++++++++++++++++++++++++---
 1 file changed, 50 insertions(+), 4 deletions(-)

diff --git a/src/fpm_risk_model/cov_estimator.py b/src/fpm_risk_model/cov_estimator.py
index db2b670..2b929f7 100644
--- a/src/fpm_risk_model/cov_estimator.py
+++ b/src/fpm_risk_model/cov_estimator.py
@@ -1,5 +1,6 @@
 from typing import Optional
 
+from numpy import diag_indices_from, trace
 from pandas import DataFrame, Series
 
 from .risk_model import RiskModel
@@ -11,7 +12,12 @@ class CovarianceEstimator:
     Covariance estimator.
     """
 
-    def __init__(self, risk_model: RiskModel):
+    def __init__(
+        self,
+        risk_model: RiskModel,
+        shrinkage_method: Optional[str] = None,
+        delta: Optional[float] = None,
+    ):
         """
         Constructor.
 
@@ -19,8 +25,14 @@ def __init__(self, risk_model: RiskModel):
         ----------
         risk_model : RiskModel
           Risk model object.
+        shrinkage_method : Optional[str]
+          Shrinkage method. Options are "constant" and "ledoit_wolf_constant_variance".
+        delta : Optional[float]
+          Delta, only used in constant shrinkage.
         """
         self._risk_model = risk_model
+        self.shrinkage_method = shrinkage_method
+        self.delta = delta
 
     def corr(self) -> DataFrame:
         """
@@ -63,6 +75,25 @@ def cov(
             .mul(volatility, axis=0)
             .mul(volatility, axis=1)
         )
+
+        if self.shrinkage_method == "constant":
+            cov = self.constant_shrinkage(cov, self.delta)
+        elif self.shrinkage_method is not None:
+            raise ValueError(
+                f"Cannot recognize shrinkage method {self.shrinkage_method}"
+            )
+
+        return cov
+
+    @staticmethod
+    def constant_shrinkage(cov: DataFrame, delta: float):
+        """
+        Constant shrinkage.
+        """
+        N = len(cov)
+        avg_var = trace(cov) / N * delta
+        cov *= 1 - delta
+        cov.values[diag_indices_from(cov)] += avg_var
         return cov
 
 
@@ -71,7 +102,12 @@ class RollingCovarianceEstimator:
     Rolling covariance estimator.
     """
 
-    def __init__(self, rolling_risk_model: RollingFactorRiskModel):
+    def __init__(
+        self,
+        rolling_risk_model: RollingFactorRiskModel,
+        shrinkage_method: Optional[str] = None,
+        delta: Optional[float] = None,
+    ):
         """
         Constructor.
 
@@ -79,8 +115,14 @@ def __init__(self, rolling_risk_model: RollingFactorRiskModel):
         ----------
         rolling_risk_model : RollingFactorRiskModel
           Rolling risk model object.
+        shrinkage_method : Optional[str]
+          Shrinkage method. Options are "constant" and "ledoit_wolf_constant_variance".
+        delta : Optional[float]
+          Delta, only used in constant shrinkage.
         """
         self._rolling_risk_model = rolling_risk_model
+        self.shrinkage_method = shrinkage_method
+        self.delta = delta
 
     def cov(self, volatility: Optional[DataFrame] = None):
         """
@@ -92,8 +134,12 @@ def cov(self, volatility: Optional[DataFrame] = None):
           Volaility series to convert from correlation to covariance. Optional.
         """
         return {
-            date: CovarianceEstimator(risk_model).cov(
-                volatility.loc[date], strict=False
+            date: (
+                CovarianceEstimator(
+                    risk_model,
+                    shrinkage_method=self.shrinkage_method,
+                    delta=self.delta,
+                ).cov(volatility.loc[date], strict=False)
             )
             for date, risk_model in self._rolling_risk_model.items()
         }