Adds eval, lags, corr and season commands to the econometrics menu

openbb_terminal/econometrics/econometrics_controller.py

@@ -1,7 +1,7 @@
 """Econometrics Controller Module"""
 __docformat__ = "numpy"
 
-# pylint: disable=too-many-lines, too-many-branches, inconsistent-return-statements
+# pylint: disable=too-many-arguments,too-many-lines,too-many-branches,inconsistent-return-statements,R0904
 
 import argparse
 import logging
@@ -18,6 +18,7 @@
 from openbb_terminal.custom_prompt_toolkit import NestedCompleter
 from openbb_terminal.decorators import log_start_end
 from openbb_terminal.econometrics import (
+    econometrics_helpers,
     econometrics_model,
     econometrics_view,
     regression_model,
@@ -54,12 +55,16 @@ class EconometricsController(BaseController):
         "show",
         "type",
         "desc",
+        "corr",
+        "season",
         "index",
         "clean",
         "add",
+        "eval",
         "delete",
         "combine",
         "rename",
+        "lag",
         "ols",
         "norm",
         "root",
@@ -177,10 +182,13 @@ def __init__(self, queue: Optional[List[str]] = None):
                 "root",
                 "granger",
                 "coint",
+                "corr",
+                "season",
+                "lag",
             ]:
                 choices[feature] = dict()
 
-            # Inititialzie this for regressions to be able to use -h flag
+            # Initialize this for regressions to be able to use -h flag
             choices["regressions"] = {}
             self.choices = choices
 
@@ -202,6 +210,8 @@ def update_runtime_choices(self):
                 "norm",
                 "root",
                 "coint",
+                "season",
+                "lag",
                 "regressions",
                 "ols",
                 "panel",
@@ -216,6 +226,7 @@ def update_runtime_choices(self):
                 "remove",
                 "combine",
                 "rename",
+                "corr",
             ]:
                 self.choices[feature] = {c: {} for c in self.files}
 
@@ -258,12 +269,16 @@ def print_help(self):
         mt.add_cmd("plot", self.files)
         mt.add_cmd("type", self.files)
         mt.add_cmd("desc", self.files)
+        mt.add_cmd("corr", self.files)
+        mt.add_cmd("season", self.files)
         mt.add_cmd("index", self.files)
         mt.add_cmd("clean", self.files)
         mt.add_cmd("add", self.files)
+        mt.add_cmd("eval", self.files)
         mt.add_cmd("delete", self.files)
         mt.add_cmd("combine", self.files)
         mt.add_cmd("rename", self.files)
+        mt.add_cmd("lag", self.files)
         mt.add_cmd("export", self.files)
         mt.add_info("_tests_")
         mt.add_cmd("norm", self.files)
@@ -670,9 +685,123 @@ def call_desc(self, other_args: List[str]):
                 else:
                     console.print("Empty dataset")
 
+    @log_start_end(log=logger)
+    def call_corr(self, other_args: List[str]):
+        """Process correlation command"""
+        parser = argparse.ArgumentParser(
+            add_help=False,
+            formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+            prog="corr",
+            description="Plot correlation coefficients.",
+        )
+        parser.add_argument(
+            "-d",
+            "--dataset",
+            help="The name of the dataset you want to select",
+            dest="target_dataset",
+            type=str,
+            choices=list(self.datasets.keys()),
+        )
+
+        # if user does not put in --dataset
+        if other_args and "-" not in other_args[0][0]:
+            other_args.insert(0, "--dataset")
+
+        ns_parser = self.parse_known_args_and_warn(
+            parser,
+            other_args,
+            EXPORT_ONLY_FIGURES_ALLOWED,
+        )
+
+        if ns_parser:
+            # check proper file name is provided
+            if not ns_parser.target_dataset:
+                console.print("[red]Please enter valid dataset.\n[/red]")
+                return
+
+            data = self.datasets[ns_parser.target_dataset]
+
+            econometrics_view.display_corr(
+                data,
+                ns_parser.export,
+            )
+
+    @log_start_end(log=logger)
+    def call_season(self, other_args: List[str]):
+        """Process season command"""
+        parser = argparse.ArgumentParser(
+            add_help=False,
+            formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+            prog="season",
+            description="The seasonality for a given column",
+        )
+        parser.add_argument(
+            "-v",
+            "--values",
+            help="Dataset.column values to be displayed in a plot",
+            dest="values",
+            choices={
+                f"{dataset}.{column}": {column: None, dataset: None}
+                for dataset, dataframe in self.datasets.items()
+                for column in dataframe.columns
+            },
+            type=str,
+        )
+        parser.add_argument(
+            "-m",
+            help="A time lag to highlight on the plot",
+            dest="m",
+            type=int,
+            default=None,
+        )
+        parser.add_argument(
+            "--max_lag",
+            help="The maximal lag order to consider",
+            dest="max_lag",
+            type=int,
+            default=24,
+        )
+        parser.add_argument(
+            "-a",
+            "--alpha",
+            help="The confidence interval to display",
+            dest="alpha",
+            type=float,
+            default=0.05,
+        )
+        if other_args and "-" not in other_args[0][0]:
+            other_args.insert(0, "-v")
+        ns_parser = self.parse_known_args_and_warn(
+            parser, other_args, export_allowed=EXPORT_ONLY_FIGURES_ALLOWED
+        )
+
+        if not ns_parser:
+            return
+
+        if not ns_parser.values:
+            console.print("[red]Please enter valid dataset.\n[/red]")
+            return
+
+        try:
+            dataset, col = ns_parser.values.split(".")
+            data = self.datasets[dataset]
+            data.name = dataset
+        except ValueError:
+            console.print("[red]Please enter 'dataset'.'column'.[/red]\n")
+            return
+
+        econometrics_view.display_seasonality(
+            data=data,
+            column=col,
+            export=ns_parser.export,
+            m=ns_parser.m,
+            max_lag=ns_parser.max_lag,
+            alpha=ns_parser.alpha,
+        )
+
     @log_start_end(log=logger)
     def call_type(self, other_args: List[str]):
-        """Process type"""
+        """Process type command"""
         parser = argparse.ArgumentParser(
             add_help=False,
             formatter_class=argparse.ArgumentDefaultsHelpFormatter,
@@ -746,7 +875,7 @@ def call_type(self, other_args: List[str]):
 
     @log_start_end(log=logger)
     def call_index(self, other_args: List[str]):
-        """Process index"""
+        """Process index command"""
         parser = argparse.ArgumentParser(
             add_help=False,
             formatter_class=argparse.ArgumentDefaultsHelpFormatter,
@@ -864,7 +993,7 @@ def call_index(self, other_args: List[str]):
 
     @log_start_end(log=logger)
     def call_clean(self, other_args: List[str]):
-        """Process clean"""
+        """Process clean command"""
         parser = argparse.ArgumentParser(
             add_help=False,
             formatter_class=argparse.ArgumentDefaultsHelpFormatter,
@@ -914,7 +1043,7 @@ def call_clean(self, other_args: List[str]):
 
     @log_start_end(log=logger)
     def call_add(self, other_args: List[str]):
-        """Process add"""
+        """Process add command"""
         parser = argparse.ArgumentParser(
             add_help=False,
             formatter_class=argparse.ArgumentDefaultsHelpFormatter,
@@ -1017,9 +1146,109 @@ def call_add(self, other_args: List[str]):
             self.update_loaded()
         console.print()
 
+    @log_start_end(log=logger)
+    def call_lag(self, other_args: List[str]):
+        """Process lag command"""
+        parser = argparse.ArgumentParser(
+            add_help=False,
+            formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+            prog="lag",
+            description="Add lag to a variable by shifting a column.",
+        )
+        parser.add_argument(
+            "-v",
+            "--values",
+            help="Dataset.column values to add lag to.",
+            dest="values",
+            choices={
+                f"{dataset}.{column}": {column: None, dataset: None}
+                for dataset, dataframe in self.datasets.items()
+                for column in dataframe.columns
+            },
+            type=str,
+            required="-h" not in other_args,
+        )
+        parser.add_argument(
+            "-l",
+            "--lags",
+            action="store",
+            dest="lags",
+            type=check_positive,
+            default=5,
+            help="How many periods to lag the selected column.",
+            required="-h" not in other_args,
+        )
+        parser.add_argument(
+            "-f",
+            "--fill-value",
+            action="store",
+            dest="fill_value",
+            help="The value used for filling the newly introduced missing values.",
+        )
+
+        if other_args and "-" not in other_args[0][0]:
+            other_args.insert(0, "-v")
+        ns_parser = self.parse_known_args_and_warn(
+            parser, other_args, export_allowed=NO_EXPORT
+        )
+
+        if not ns_parser:
+            return
+
+        try:
+            dataset, col = ns_parser.values.split(".")
+            data = self.datasets[dataset]
+        except ValueError:
+            console.print("[red]Please enter 'dataset'.'column'.[/red]\n")
+            return
+
+        data[col + "_with_" + str(ns_parser.lags) + "_lags"] = data[col].shift(
+            ns_parser.lags, fill_value=ns_parser.fill_value
+        )
+        self.datasets[dataset] = data
+
+        self.update_runtime_choices()
+
+    @log_start_end(log=logger)
+    def call_eval(self, other_args):
+        """Process eval command"""
+        parser = argparse.ArgumentParser(
+            add_help=False,
+            formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+            prog="eval",
+            description="""Create custom data column from loaded datasets.  Can be mathematical expressions supported
+                by pandas.eval() function.
+
+                Example.  If I have loaded `fred DGS2,DGS5` and I want to create a new column that is the difference
+                between these two, I can create a new column by doing `eval spread = DGS2 - DGS5`.
+                Notice that the command is case sensitive, i.e., `DGS2` is not the same as `dgs2`.
+                """,
+        )
+        parser.add_argument(
+            "-q",
+            "--query",
+            type=str,
+            nargs="+",
+            dest="query",
+            required="-h" not in other_args,
+            help="Query to evaluate on loaded datasets",
+        )
+        if other_args and "-" not in other_args[0][0]:
+            other_args.insert(0, "-q")
+
+        ns_parser = self.parse_known_args_and_warn(
+            parser, other_args, export_allowed=EXPORT_ONLY_RAW_DATA_ALLOWED
+        )
+        if ns_parser:
+            self.datasets = econometrics_helpers.create_new_entry(
+                self.datasets, " ".join(ns_parser.query)
+            )
+            self.update_runtime_choices()
+            self.update_loaded()
+
     @log_start_end(log=logger)
     def call_delete(self, other_args: List[str]):
-        """Process add"""
+        """Process delete command"""
         parser = argparse.ArgumentParser(
             add_help=False,
             formatter_class=argparse.ArgumentDefaultsHelpFormatter,
@@ -1060,7 +1289,7 @@ def call_delete(self, other_args: List[str]):
 
     @log_start_end(log=logger)
     def call_combine(self, other_args: List[str]):
-        """Process combine"""
+        """Process combine command"""
         parser = argparse.ArgumentParser(
             add_help=False,
             formatter_class=argparse.ArgumentDefaultsHelpFormatter,
@@ -1118,7 +1347,7 @@ def call_combine(self, other_args: List[str]):
 
     @log_start_end(log=logger)
     def call_rename(self, other_args: List[str]):
-        """Process rename"""
+        """Process rename command"""
         parser = argparse.ArgumentParser(
             add_help=False,
             formatter_class=argparse.ArgumentDefaultsHelpFormatter,