Passed
Push — master ( 8ed98a...de2059 )
by Stefan
03:31
created

spaceweather.core._dl_file()   A

Complexity

Conditions 4

Size

Total Lines 5
Code Lines 5

Duplication

Lines 0
Ratio 0 %

Importance

Changes 0
Metric Value
cc 4
eloc 5
nop 1
dl 0
loc 5
rs 10
c 0
b 0
f 0
1
# Copyright (c) 2020 Stefan Bender
2
#
3
# This module is part of pyspaceweather.
4
# pyspaceweather is free software: you can redistribute it or modify
5
# it under the terms of the GNU General Public License as published
6
# by the Free Software Foundation, version 2.
7
# See accompanying COPYING.GPLv2 file or http://www.gnu.org/licenses/gpl-2.0.html.
8
"""Python interface for space weather indices
9
10
"""
11
import os
12
from pkg_resources import resource_filename
13
import requests
14
15
import numpy as np
16
import pandas as pd
17
18
__all__ = ["sw_daily", "ap_kp_3h"]
19
20
DL_URL = "https://celestrak.com/SpaceData/SW-All.txt"
21
SW_FILE = "SW-All.txt"
22
23
24
def _dl_file(swfile):
25
	with requests.get(DL_URL, stream=True) as r:
26
		with open(swfile, 'wb') as fd:
27
			for chunk in r.iter_content(chunk_size=1024):
28
				fd.write(chunk)
29
30
31
def _get_last_update(swpath):
32
	for line in open(swpath):
33
		if line.startswith("UPDATED"):
34
			# closes the file automatically
35
			break
36
	return pd.to_datetime(line.lstrip("UPDATED"), utc=True)
0 ignored issues
show
introduced by
The variable line does not seem to be defined in case the for loop on line 32 is not entered. Are you sure this can never be the case?
Loading history...
37
38
39
def check_for_update(swpath, max_age="30days"):
40
	last_update = _get_last_update(swpath)
41
	file_age = pd.Timestamp.utcnow() - last_update
42
	return file_age > pd.Timedelta(max_age)
43
44
45
def _read_sw(swpath):
46
	kpns = ["Kp{0}".format(i) for i in range(0, 23, 3)] + ["Kpsum"]
47
	sw = np.genfromtxt(
48
		swpath,
49
		skip_header=3,
50
				# yy mm dd br rd kp kp kp kp kp kp kp kp Kp ap ap ap ap ap ap ap ap Ap cp c9 is f1  q f2 f3 f4 f5 f6
51
		delimiter=[4, 3, 3, 5, 3, 3, 3, 3, 3, 3, 3, 3, 3, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 2, 4, 6, 2, 6, 6, 6, 6, 6],
52
		dtype=   "i4,i4,i4,i4,i4,i4,i4,i4,i4,i4,i4,i4,i4,i4,i4,i4,i4,i4,i4,i4,i4,i4,i4,f8,i4,i4,f8,i4,f8,f8,f8,f8,f8",
53
		names=[
54
			"year", "month", "day", "bsrn", "rotd",
55
			"Kp0", "Kp3", "Kp6", "Kp9", "Kp12", "Kp15", "Kp18", "Kp21", "Kpsum",
56
			"Ap0", "Ap3", "Ap6", "Ap9", "Ap12", "Ap15", "Ap18", "Ap21", "Apavg",
57
			"Cp", "C9", "isn", "f107_adj", "Q", "f107_81ctr_adj", "f107_81lst_adj",
58
			"f107_obs", "f107_81ctr_obs", "f107_81lst_obs"
59
		]
60
	)[2:-1]
61
	sw = sw[sw["year"] != -1]
62
	ts = pd.to_datetime([
63
		"{0:04d}-{1:02d}-{2:02d}".format(yy, mm, dd)
64
		for yy, mm, dd in sw[["year", "month", "day"]]
65
	])
66
	sw_df = pd.DataFrame(sw, index=ts)
67
	sw_df[kpns] = 0.1 * sw_df[kpns]
68
	return sw_df
69
70
71
def sw_daily(swfile=SW_FILE, update_interval="30days"):
72
	"""Daily Ap, Kp, and f10.7 index values
73
	"""
74
	swpath = resource_filename(__name__, os.path.join("data", swfile))
75
	# ensure that the file exists and is up to date
76
	if (
77
		not os.path.exists(swpath)
78
		or check_for_update(swpath, max_age=update_interval)
79
	):
80
		_dl_file(swpath)
81
	return _read_sw(swpath)
82
83
84
def ap_kp_3h(swfile=SW_FILE):
85
	"""3h Ap and Kp index values
86
	"""
87
	daily_df = sw_daily(swfile)
88
	ret = daily_df.copy()
89
	apns = ["Ap{0}".format(i) for i in range(0, 23, 3)]
90
	kpns = ["Kp{0}".format(i) for i in range(0, 23, 3)]
91
	for i, (ap, kp) in enumerate(zip(apns, kpns)):
92
		ret[ap].index = daily_df[ap].index + pd.Timedelta((i * 3 + 1.5), unit="h")
93
		ret[kp].index = daily_df[kp].index + pd.Timedelta((i * 3 + 1.5), unit="h")
94
	sw_ap = pd.concat([ret[ap] for ap in apns])
95
	sw_kp = pd.concat([ret[kp] for kp in kpns])
96
	return pd.DataFrame({"Ap": sw_ap, "Kp": sw_kp})
97