공매도 데이터 분석하기

import pandas as pd
from quantylab.common.db.mongodb import MongoDBClient
coll = MongoDBClient.get_coll('short')
cursor = coll.find({'date': {'$gte': '20210601', '$lte': '20220114'}}, {'_id': 0, 'dt_update': 0, 'date_update': 0}, sort=[('date', -1)])
df_short = pd.DataFrame(list(cursor))
coll = MongoDBClient.get_coll('shortbalance')
cursor = coll.find({'date': {'$gte': '20210601', '$lte': '20220114'}}, {'_id': 0, 'dt_update': 0, 'date_update': 0}, sort=[('date', -1)])
df_shortbalance = pd.DataFrame(list(cursor))
coll = MongoDBClient.get_coll('stockcodes')
cursor = coll.find({}, {'_id': 0, 'code': 1, 'name': 1})
df_stockcodes = pd.DataFrame(list(cursor))

df = pd.merge(df_short, df_shortbalance[['code', 'date', 'bal_rto', 'bal_amt', 'bal_qty']], on=['code', 'date'], how='left')
df = pd.merge(df, df_stockcodes, on='code', how='left')
df = df[['name'] + df.columns.tolisZt()[:-1]]
df.to_csv('20220113_analysis_short.csv', index=False)

import pandas as pd
str_columns = ['code', 'date']
df = pd.read_csv('20220113_analysis_short.csv', converters={k: lambda x: str(x) for k in str_columns})

df = df.sort_values(by='date')
# % 단위 떼기
df['short_ratio'] /= 100
df['bal_rto'] /= 100
df = df.dropna().copy()

from tqdm import tqdm
for k, g in tqdm(df.groupby('code')):
    # 상환 규모 확인
    # 당일 short_volume - (당일 bal_qty - 전일 bal_qty)
    df.loc[df['code'] == k, 'bal_qty_diff'] = g['bal_qty'].diff()
    df.loc[df['code'] == k, 'repay_volume'] = (df.loc[df['code'] == k, 'short_volume'] - df.loc[df['code'] == k, 'bal_qty_diff']).clip(0)
    df.loc[df['code'] == k, 'repay_ratio'] = df.loc[df['code'] == k, 'repay_volume'] / df.loc[df['code'] == k, 'volume']
    # 이동평균
    df.loc[df['code'] == k, 'short_ratio_ma5'] = g['short_ratio'].fillna(0).rolling(5).mean()
    df.loc[df['code'] == k, 'bal_rto_ma5'] = g['bal_rto'].fillna(0).rolling(5).mean()
    df.loc[df['code'] == k, 'repay_ratio_ma5'] = df.loc[df['code'] == k, 'repay_ratio'] .fillna(0).rolling(5).mean()

_df = df[df['date'] == '20220114'].sort_values(by='repay_ratio_ma5', ascending=False)
_df