1
|
|
|
# -*- coding: utf-8 -*- |
2
|
|
|
# vim:fileencoding=utf-8 |
3
|
|
|
# |
4
|
|
|
# Copyright (c) 2022 Stefan Bender |
5
|
|
|
# |
6
|
|
|
# This module is part of pyspaceweather. |
7
|
|
|
# pyspaceweather is free software: you can redistribute it or modify |
8
|
|
|
# it under the terms of the GNU General Public License as published |
9
|
|
|
# by the Free Software Foundation, version 2. |
10
|
|
|
# See accompanying COPYING.GPLv2 file or http://www.gnu.org/licenses/gpl-2.0.html. |
11
|
|
|
"""OMNI data read tests |
12
|
|
|
""" |
13
|
|
|
import os |
14
|
|
|
import requests |
15
|
|
|
|
16
|
|
|
import numpy as np |
17
|
|
|
import pandas as pd |
18
|
|
|
|
19
|
|
|
import pytest |
20
|
|
|
|
21
|
|
|
from spaceweather import cache_omnie, omnie_hourly, omnie_mask_missing, sw_daily |
22
|
|
|
from spaceweather.omni import OMNI_URL_BASE, OMNI_PREFIX, OMNI_EXT |
23
|
|
|
|
24
|
|
|
_TEST_YEAR = 2012 |
25
|
|
|
_TEST_FILE = "{0}_{1:04d}.{2}".format(OMNI_PREFIX, _TEST_YEAR, OMNI_EXT) |
26
|
|
|
_TEST_URL = os.path.join(OMNI_URL_BASE, _TEST_FILE) |
27
|
|
|
_TEST_PATH = os.path.join(".", "tests") |
28
|
|
|
|
29
|
|
|
|
30
|
|
|
@pytest.fixture(scope="module") |
31
|
|
|
def df_d(): |
32
|
|
|
return sw_daily() |
33
|
|
|
|
34
|
|
|
|
35
|
|
|
@pytest.fixture(scope="module") |
36
|
|
|
def df_o(): |
37
|
|
|
return omnie_hourly(2000, local_path=_TEST_PATH, prefix="omni2t") |
38
|
|
|
|
39
|
|
|
|
40
|
|
|
def test_cache(mocker, tmpdir): |
41
|
|
|
mocker.patch("requests.get") |
42
|
|
|
tmpdir = str(tmpdir) |
43
|
|
|
# Check non-existent file |
44
|
|
|
cache_omnie(year=_TEST_YEAR, local_path=tmpdir) |
45
|
|
|
requests.get.assert_called_once_with(_TEST_URL, stream=True) |
46
|
|
|
# Check non-existent (sub)dir |
47
|
|
|
tmppath = os.path.join(tmpdir, "data") |
48
|
|
|
cache_omnie(year=_TEST_YEAR, local_path=tmppath) |
49
|
|
|
requests.get.assert_called_with(_TEST_URL, stream=True) |
50
|
|
|
|
51
|
|
|
|
52
|
|
|
def test_auto_update(mocker, tmpdir): |
53
|
|
|
# test with non-existent file |
54
|
|
|
mocker.patch("requests.get") |
55
|
|
|
tmpdir = str(tmpdir) |
56
|
|
|
with pytest.raises(IOError): |
57
|
|
|
with pytest.warns(UserWarning): |
58
|
|
|
omnie_hourly(year=_TEST_YEAR, cache=True, local_path=tmpdir) |
59
|
|
|
requests.get.assert_called_once_with(_TEST_URL, stream=True) |
60
|
|
|
|
61
|
|
|
|
62
|
|
|
def test_not_avail(mocker, tmpdir): |
63
|
|
|
# test with non-existent file |
64
|
|
|
tmpdir = str(tmpdir) |
65
|
|
|
with pytest.raises(IOError): |
66
|
|
|
with pytest.warns(UserWarning): |
67
|
|
|
omnie_hourly(year=_TEST_YEAR, cache=False, local_path=tmpdir) |
68
|
|
|
|
69
|
|
|
|
70
|
|
|
@pytest.mark.parametrize("hour", range(0, 24, 3)) |
71
|
|
|
@pytest.mark.parametrize("index", ["Ap", "Kp"]) |
72
|
|
|
def test_hourly(hour, index, df_d, df_o): |
73
|
|
|
df1 = df_o |
74
|
|
|
# The last row is for the missing value test. |
75
|
|
|
df1 = df1.iloc[:-1] |
76
|
|
|
df2 = df_d |
77
|
|
|
ind_name = "{0}{1}".format(index, hour) |
78
|
|
|
df1_ind = df1[df1["hour"] == hour][index] |
79
|
|
|
df2_ind = df2[ind_name].loc[df1_ind.index.date].rename(index) |
80
|
|
|
df2_ind.index = df1_ind.index |
81
|
|
|
pd.testing.assert_series_equal(df1_ind, df2_ind) |
82
|
|
|
|
83
|
|
|
|
84
|
|
View Code Duplication |
@pytest.mark.parametrize( |
|
|
|
|
85
|
|
|
"name, result", |
86
|
|
|
[ |
87
|
|
|
("Ap", np.array([56, 39, 27, 18, 32, 15, 32, 22])), |
88
|
|
|
("Kp", np.array([5.3, 4.7, 4.0, 3.3, 4.3, 3.0, 4.3, 3.7])), |
89
|
|
|
] |
90
|
|
|
) |
91
|
|
|
def test_3hourly_index(name, result, df_o): |
92
|
|
|
df = df_o |
93
|
|
|
np.testing.assert_allclose( |
94
|
|
|
df.loc[ |
95
|
|
|
pd.date_range( |
96
|
|
|
"2000-01-01 00:00", "2000-01-01 23:00", freq="3h" |
97
|
|
|
) |
98
|
|
|
][name].values, |
99
|
|
|
result, |
100
|
|
|
rtol=1e-12, |
101
|
|
|
) |
102
|
|
|
|
103
|
|
|
|
104
|
|
|
def test_mask_missing(df_o): |
105
|
|
|
df = df_o |
106
|
|
|
dfp = omnie_mask_missing(df) |
107
|
|
|
# The last row should contain all NaNs. |
108
|
|
|
dfp = dfp.iloc[-1] |
109
|
|
|
for v in filter( |
110
|
|
|
lambda n: n not in ["year", "doy", "hour"], |
111
|
|
|
dfp.index, |
112
|
|
|
): |
113
|
|
|
assert np.isnan(dfp[v]) |
114
|
|
|
|