pv-current/test_stations.py

# Created by: leo
# Created on: 2018.12.14

import pytest
import pandas as pd
from coef_matrix import classify

@pytest.fixture
def load_data():
    raw = pd.read_csv('../HuBeiShaYang/hbpv10days.csv',
                      usecols=[0, 1, 2, 3, 4, 6],
                      names=['time', 'station', 'lev1', 'lev2', 'strno', 'current'],
                      dtype={'lev1':str, 'lev2':str, 'strno':str},
                      parse_dates = ['time']).dropna()
    inp = raw[(raw['current'] >= 0) & (raw['current'] < 10)]
    pd.options.mode.chained_assignment = None
    inp['day'] = inp['time'].map(lambda x: x.date())
    # 没有 set_index 会导致 grps.apply() 出现日期不在 Index 中错误
    inp2 = inp.set_index(['station', 'lev1', 'lev2', 'day'])
    grps = inp2.groupby(inp2.index)
    return grps

def test_classify(load_data):
    df1 = load_data.apply(classify)
    df2 = df1.reset_index()
    res = df2.rename(columns={'level_0': 'cid'})
    res.to_csv('res_one_process.csv', index=False)