diff --git a/shell.nix b/shell.nix index 11d0f30..d5ed8a2 100644 --- a/shell.nix +++ b/shell.nix @@ -9,6 +9,7 @@ stdenv.mkDerivation rec { python311 python311Packages.virtualenv python311Packages.pandas + python311Packages.numpy python311Packages.xlrd python311Packages.python-lsp-server python311Packages.python-lsp-black diff --git a/zacks.py b/zacks.py index f37ccf0..1e3a51b 100644 --- a/zacks.py +++ b/zacks.py @@ -1,4 +1,5 @@ import pandas as pd +import numpy as np import argparse from datetime import datetime @@ -47,6 +48,30 @@ def add_cols(df): df['PE0 Mean'] = df.groupby('Industry')['PE0'].transform('mean').round(2) df['PE1 Mean'] = df.groupby('Industry')['PE1'].transform('mean').round(2) df['PE2 Mean'] = df.groupby('Industry')['PE2'].transform('mean').round(2) + # Long Profile 1 + mask = ((df['EG1'] > df['EG1 Mean']) & (df['EG2'] > df['EG2 Mean']) & (df['EG2'] > df['EG1']) & (df['EG2 Mean'] > df['EG1 Mean'])) + long_1 = np.where(mask, 'Long EG 1', '') + df['Profile'] = long_1 + # Long Profile 2 + mask = ((df['EG1'] > df['EG1 Mean']) & (df['EG2'] > df['EG2 Mean']) & (df['EG2'] > df['EG1']) & (df['EG2 Mean'] == df['EG1 Mean'])) + long_2 = np.where(mask, 'Long EG 2', df['Profile']) + df['Profile'] = long_2 + # Long Profile 3 + mask = ((df['EG1'] > df['EG1 Mean']) & (df['EG2'] > df['EG2 Mean']) & (df['EG2'] > df['EG1']) & (df['EG2 Mean'] < df['EG1 Mean'])) + long_3 = np.where(mask, 'Long EG 3', df['Profile']) + df['Profile'] = long_3 + # Short Profile 1 + mask = ((df['EG1'] < df['EG1 Mean']) & (df['EG2'] < df['EG2 Mean']) & (df['EG2'] < df['EG1']) & (df['EG2 Mean'] > df['EG1 Mean']) & (df['EG1'] < 0) & (df['EG2'] < 0) & (df['EG1 Mean'] > 0) & (df['EG2 Mean'] > 0)) + short_1 = np.where(mask, 'Short EG 1', df['Profile']) + df['Profile'] = short_1 + # Short Profile 2 + mask = ((df['EG1'] < df['EG1 Mean']) & (df['EG2'] < df['EG2 Mean']) & (df['EG2'] > df['EG1']) & (df['EG2 Mean'] > df['EG1 Mean']) & (df['EG1'] < 0) & (df['EG2'] < 0) & (df['EG1 Mean'] > 0) & (df['EG2 Mean'] > 0)) + short_2 = np.where(mask, 'Short EG 2', df['Profile']) + df['Profile'] = short_2 + # Short Profile 3 + mask = ((df['EG1'] < df['EG1 Mean']) & (df['EG2'] < df['EG2 Mean']) & (df['EG2'] < df['EG1']) & (df['EG2 Mean'] > df['EG1 Mean']) & (df['EG1'] > 0) & (df['EG2'] < 0) & (df['EG1 Mean'] > 0) & (df['EG2 Mean'] > 0)) + short_3 = np.where(mask, 'Short EG 3', df['Profile']) + df['Profile'] = short_3 return df