implements generate_reference_data_from_talib.py (#8196)

Tests/TestData/generate_reference_data_from_talib.py

@@ -0,0 +1,84 @@
+import pandas as pd
+import talib.abstract
+
+
+def format_number(x):
+    if pd.isnull(x):
+        return ""
+    else:
+        return f"{x:.6f}"
+
+
+def write_dataframe(df, fname_stem, format_dict):
+    # format index
+    datetime_output_fmt = "%m/%d/%Y %I:%M:%S %p"
+    df.index = df.index.strftime(datetime_output_fmt)
+    # format columns
+    for col in df.columns:
+        if format_dict is not None and col in format_dict:
+            format_function = format_dict[col]
+            df[col] = df[col].map(format_function)
+    df.to_csv(f"{fname_stem}.csv", sep=",")
+
+
+def generate_reference_data_for_single_output_indicator(
+    df, indicator_type, parameters, fname_stem, output_name, format_dict=None
+):
+    print("* Processing %s" % indicator_type.info)
+    series_output = indicator_type(df, **parameters)
+    series_output.name = output_name
+    df_output = series_output.to_frame()
+    df.columns = df.columns.str.capitalize()
+    df_all = pd.concat([df, df_output], axis=1)
+    write_dataframe(df_all, fname_stem, format_dict)
+
+
+def generate_reference_data_for_multi_output_indicator(
+    df, indicator_type, parameters, fname_stem, output_names=None, format_dict=None
+):
+    print("* Processing %s" % indicator_type.info)
+    df_output = indicator_type(df, **parameters)
+    if output_names is not None:
+        df_output.columns = output_names
+    df.columns = df.columns.str.capitalize()
+    df_all = pd.concat([df, df_output], axis=1)
+    write_dataframe(df_all, fname_stem, format_dict)
+
+
+def main():
+    fname = "spy_daily_klines_2013-01-16_2015-12-01_no_volume.csv"
+    datetime_input_fmt = "%m/%d/%Y %I:%M:%S %p"
+    df = pd.read_csv(fname)
+    df["Date"] = pd.to_datetime(df["Date"], format=datetime_input_fmt)
+    df = df.set_index("Date")
+    df.columns = df.columns.str.lower()
+
+    generate_reference_data_for_single_output_indicator(
+        df.copy(),
+        talib.abstract.ATR,
+        {"timeperiod": 14},
+        "spy_atr",
+        "Average True Range 14",
+        {"Average True Range 14": format_number},
+    )
+
+    generate_reference_data_for_multi_output_indicator(
+        df.copy(),
+        talib.abstract.BBANDS,
+        {"timeperiod": 20, "nbdevup": 2.0, "nbdevdn": 2.0},
+        "spy_bollinger_bands",
+        [
+            "Bollinger Bands® 20 2 Top",
+            "Moving Average 20",
+            "Bollinger Bands® 20 2 Bottom",
+        ],
+        {
+            "Bollinger Bands® 20 2 Top": format_number,
+            "Moving Average 20": format_number,
+            "Bollinger Bands® 20 2 Bottom": format_number,
+        },
+    )
+
+
+if __name__ == "__main__":
+    main()