From d4f3624bd86564cf3d7468af979c123c11a8dac2 Mon Sep 17 00:00:00 2001
From: robertmartin8 <martin.robertandrew@gmail.com>
Date: Tue, 10 Dec 2019 22:48:25 +0000
Subject: [PATCH] added save_weights_to_file (#54)

---
 pypfopt/base_optimizer.py    | 35 +++++++++++++++++++++++++++++------
 tests/test_base_optimizer.py | 31 +++++++++++++++++++++++++++++++
 2 files changed, 60 insertions(+), 6 deletions(-)

diff --git a/pypfopt/base_optimizer.py b/pypfopt/base_optimizer.py
index 31a9465..7ef8a6b 100644
--- a/pypfopt/base_optimizer.py
+++ b/pypfopt/base_optimizer.py
@@ -7,6 +7,7 @@ Additionally, we define a general utility function ``portfolio_performance`` to
 evaluate return and risk for a given set of portfolio weights.
 """
 
+import json
 import numpy as np
 import pandas as pd
 from . import objective_functions
@@ -20,6 +21,12 @@ class BaseOptimizer:
     - ``n_assets`` - int
     - ``tickers`` - str list
     - ``weights`` - np.ndarray
+
+    Public methods:
+
+    - ``set_weights()`` creates self.weights (np.ndarray) from a weights dict
+    - ``clean_weights()`` rounds the weights and clips near-zeros.
+    - ``save_weights_to_file()`` saves the weights to csv, json, or txt.
     """
 
     def __init__(self, n_assets, tickers=None):
@@ -44,11 +51,7 @@ class BaseOptimizer:
         :param weights: {ticker: weight} dictionary
         :type weights: dict
         """
-        if self.weights is None:
-            self.weights = [0] * self.n_assets
-        for i, ticker in enumerate(self.tickers):
-            if ticker in weights:
-                self.weights[i] = weights[ticker]
+        self.weights = np.array([weights[ticker] for ticker in self.tickers])
 
     def clean_weights(self, cutoff=1e-4, rounding=5):
         """
@@ -63,6 +66,8 @@ class BaseOptimizer:
         :return: asset weights
         :rtype: dict
         """
+        if self.weights is None:
+            raise AttributeError("Weights not yet computed")
         clean_weights = self.weights.copy()
         clean_weights[np.abs(clean_weights) < cutoff] = 0
         if rounding is not None:
@@ -71,6 +76,25 @@ class BaseOptimizer:
             clean_weights = np.round(clean_weights, rounding)
         return dict(zip(self.tickers, clean_weights))
 
+    def save_weights_to_file(self, filename="weights.csv"):
+        """
+        Utility method to save weights to a text file.
+
+        :param filename: name of file. Should be csv, json, or txt.
+        :type filename: str
+        """
+        clean_weights = self.clean_weights()
+
+        ext = filename.split(".")[1]
+        if ext == "csv":
+            pd.Series(clean_weights).to_csv(filename, header=False)
+        elif ext == "json":
+            with open(filename, "w") as fp:
+                json.dump(clean_weights, fp)
+        else:
+            with open(filename, "w") as f:
+                f.write(str(clean_weights))
+
 
 class BaseScipyOptimizer(BaseOptimizer):
 
@@ -113,7 +137,6 @@ class BaseScipyOptimizer(BaseOptimizer):
         :rtype: tuple of tuples
         """
         # If it is a collection with the right length, assume they are all bounds.
-        print(test_bounds)
         if len(test_bounds) == self.n_assets and not isinstance(
             test_bounds[0], (float, int)
         ):
diff --git a/tests/test_base_optimizer.py b/tests/test_base_optimizer.py
index 3ce3ec5..384b767 100644
--- a/tests/test_base_optimizer.py
+++ b/tests/test_base_optimizer.py
@@ -1,3 +1,5 @@
+import json
+import os
 import numpy as np
 import pytest
 from pypfopt.efficient_frontier import EfficientFrontier
@@ -100,6 +102,8 @@ def test_clean_weights_short():
 
 def test_clean_weights_error():
     ef = setup_efficient_frontier()
+    with pytest.raises(AttributeError):
+        ef.clean_weights()
     ef.max_sharpe()
     with pytest.raises(ValueError):
         ef.clean_weights(rounding=1.3)
@@ -128,3 +132,30 @@ def test_efficient_frontier_init_errors():
 
     with pytest.raises(TypeError):
         EfficientFrontier(mean_returns, mean_returns)
+
+
+def test_set_weights():
+    ef = setup_efficient_frontier()
+    w1 = ef.max_sharpe()
+    test_weights = ef.weights
+    ef.min_volatility()
+    ef.set_weights(w1)
+    np.testing.assert_array_almost_equal(test_weights, ef.weights)
+
+
+def test_save_weights_to_file():
+    ef = setup_efficient_frontier()
+    ef.max_sharpe()
+    ef.save_weights_to_file("tests/test.txt")
+    with open("tests/test.txt", "r") as f:
+        file = f.read()
+    parsed = json.loads(file.replace("'", '"'))
+    assert ef.clean_weights() == parsed
+
+    ef.save_weights_to_file("tests/test.json")
+    with open("tests/test.json", "r") as f:
+        parsed = json.load(f)
+    assert ef.clean_weights() == parsed
+
+    os.remove("tests/test.txt")
+    os.remove("tests/test.json")