Add argument to remove masking for heatmap in hcorr and print correla…

…tions to command line (#1913) * Remove masking for heatmap * Add a comment * Add a flag to show either half or the full correlation matrix. * Fix printed correlation table cutting off for large data sets * Update the test * Update some formatting issues * Fix another formatting error * Formatting issue * Remove a double blank line * Black formatting changes * Fix a pylint formatting issue * Use rich table and add a raw flag * Fix formatting issue * Capitalize ticker name in table columns * Update a test * Add docs for new arguments * retrigger checks
OpenBB-finance · Jun 8, 2022 · a9a7771 · a9a7771
1 parent 7e52dac
commit a9a7771
Show file tree

Hide file tree

Showing 4 changed files with 31 additions and 4 deletions.
diff --git a/openbb_terminal/stocks/comparison_analysis/ca_controller.py b/openbb_terminal/stocks/comparison_analysis/ca_controller.py
@@ -578,10 +578,15 @@ def call_hcorr(self, other_args: List[str]):
             dest="start",
             help="The starting date (format YYYY-MM-DD) of the stock",
         )
+        parser.add_argument(
+            "--display-full-matrix",
+            action="store_true",
+            help="Display all matrix values, rather than masking off half.",
+        )
         if other_args and "-" not in other_args[0][0]:
             other_args.insert(0, "-t")
         ns_parser = parse_known_args_and_warn(
-            parser, other_args, EXPORT_BOTH_RAW_DATA_AND_FIGURES
+            parser, other_args, EXPORT_BOTH_RAW_DATA_AND_FIGURES, raw=True
         )
         if ns_parser:
             if self.similar and len(self.similar) > 1:
@@ -590,6 +595,8 @@ def call_hcorr(self, other_args: List[str]):
                     start=ns_parser.start.strftime("%Y-%m-%d"),
                     candle_type=ns_parser.type_candle,
                     export=ns_parser.export,
+                    display_full_matrix=ns_parser.display_full_matrix,
+                    raw=ns_parser.raw,
                 )
             else:
                 console.print("Please make sure there are similar tickers selected. \n")

diff --git a/openbb_terminal/stocks/comparison_analysis/yahoo_finance_view.py b/openbb_terminal/stocks/comparison_analysis/yahoo_finance_view.py
@@ -20,6 +20,7 @@
     export_data,
     plot_autoscale,
     is_valid_axes_count,
+    print_rich_table,
 )
 from openbb_terminal.rich_config import console
 from openbb_terminal.stocks.comparison_analysis import yahoo_finance_model
@@ -169,6 +170,8 @@ def display_correlation(
     candle_type: str = "a",
     external_axes: Optional[List[plt.Axes]] = None,
     export: str = "",
+    display_full_matrix: bool = False,
+    raw: bool = False,
 ):
     """
     Correlation heatmap based on historical price comparison
@@ -186,6 +189,8 @@ def display_correlation(
         External axes (1 axis is expected in the list), by default None
     export : str, optional
         Format to export correlation prices, by default ""
+    display_full_matrix : bool, optional
+        Optionally display all values in the matrix, rather than masking off half, by default False
 
     """
     df_similar = yahoo_finance_model.get_historical(similar_tickers, start, candle_type)
@@ -200,8 +205,10 @@ def display_correlation(
 
     df_similar = df_similar.dropna(axis=1, how="all")
 
-    mask = np.zeros((df_similar.shape[1], df_similar.shape[1]), dtype=bool)
-    mask[np.triu_indices(len(mask))] = True
+    mask = None
+    if not display_full_matrix:
+        mask = np.zeros((df_similar.shape[1], df_similar.shape[1]), dtype=bool)
+        mask[np.triu_indices(len(mask))] = True
 
     # This plot has 1 axis
     if not external_axes:
@@ -211,8 +218,17 @@ def display_correlation(
     else:
         return
 
+    # Print correlations to command line as well
+    correlations = df_similar.corr()
+    if raw:
+        print_rich_table(
+            correlations,
+            headers=[x.title().upper() for x in correlations.columns],
+            show_index=True,
+        )
+
     sns.heatmap(
-        df_similar.corr(),
+        correlations,
         cbar_kws={"ticks": [-1.0, -0.5, 0.0, 0.5, 1.0]},
         cmap="RdYlGn",
         linewidths=1,

diff --git a/tests/openbb_terminal/stocks/comparison_analysis/test_ca_controller.py b/tests/openbb_terminal/stocks/comparison_analysis/test_ca_controller.py
@@ -257,6 +257,8 @@ def test_call_func_expect_queue(expected_queue, queue, func):
                 start="2020-12-01",
                 candle_type="h",
                 export="",
+                display_full_matrix=False,
+                raw=False,
             ),
         ),
         (

diff --git a/website/content/terminal/stocks/comparison_analysis/hcorr/_index.md b/website/content/terminal/stocks/comparison_analysis/hcorr/_index.md
@@ -11,6 +11,8 @@ optional arguments:
   -s START, --start START
                         The starting date (format YYYY-MM-DD) of the stock (default: 2021-02-14)
   -h, --help            show this help message (default: False)
+  --display-full-matrix Display all values in the matrix, rather than masking off half (default: False)
+  --raw                 show the raw output as a table in the terminal window (default: False)
 ```
 
 ![hcorr](https://user-images.githubusercontent.com/46355364/154073186-45336f5f-85e1-4cb9-9307-9694295b1f80.png)