diff --git a/heliopy/data/imp.py b/heliopy/data/imp.py index 735d17bf6..2659aca2d 100644 --- a/heliopy/data/imp.py +++ b/heliopy/data/imp.py @@ -56,13 +56,7 @@ def __init__(self, probe): ('DWp', u.dimensionless_unscaled)]) def intervals(self, starttime, endtime): - ret = [] - while starttime < endtime: - start_month = datetime(starttime.year, starttime.month, 1) - end_month = start_month + relativedelta(months=1) - ret.append(sunpy.time.TimeRange(start_month, end_month)) - starttime += relativedelta(months=1) - return ret + return self.intervals_monthly(starttime, endtime) def fname(self, interval): start = interval.start.to_datetime() diff --git a/heliopy/data/test/test_util.py b/heliopy/data/test/test_util.py new file mode 100644 index 000000000..d502947bd --- /dev/null +++ b/heliopy/data/test/test_util.py @@ -0,0 +1,14 @@ +from heliopy.data import util +from datetime import datetime + + +def test_montly_intervals(): + intervals = util.Downloader.intervals_monthly( + datetime(1992, 11, 1), datetime(1992, 12, 1)) + assert len(intervals) == 2 + + +def test_yearly_intervals(): + intervals = util.Downloader.intervals_yearly( + datetime(1992, 11, 1), datetime(1993, 2, 1)) + assert len(intervals) == 2 diff --git a/heliopy/data/util.py b/heliopy/data/util.py index e29a8752c..449950952 100644 --- a/heliopy/data/util.py +++ b/heliopy/data/util.py @@ -5,6 +5,7 @@ """ import abc import datetime as dt +import dateutil.relativedelta as reldelt import ftplib import io import os @@ -132,19 +133,6 @@ def intervals(self, starttime, endtime): """ pass - @staticmethod - def intervals_yearly(starttime, endtime): - """ - Returns all annual intervals between *starttime* and *endtime*. - """ - out = [] - # Loop through years - for year in range(starttime.year, endtime.year + 1): - out.append(sunpy.time.TimeRange(dt.datetime(year, 1, 1), - dt.datetime(year + 1, 1, 1))) - return out - - @abc.abstractmethod def fname(self, interval): """ Return the filename to which the data is saved for a given interval. @@ -211,6 +199,39 @@ def load_local_file(self, interval): """ pass + @staticmethod + def intervals_yearly(starttime, endtime): + """ + Returns all annual intervals between *starttime* and *endtime*. + """ + out = [] + # Loop through years + for year in range(starttime.year, endtime.year + 1): + out.append(sunpy.time.TimeRange(dt.datetime(year, 1, 1), + dt.datetime(year + 1, 1, 1))) + return out + + @staticmethod + def intervals_monthly(starttime, endtime): + """ + Returns all monthly intervals between *starttime* and *endtime*. + """ + out = [] + while starttime < endtime + reldelt.relativedelta(months=1): + start_month = dt.datetime(starttime.year, starttime.month, 1) + end_month = start_month + reldelt.relativedelta(months=1) + out.append(sunpy.time.TimeRange(start_month, end_month)) + starttime += reldelt.relativedelta(months=1) + return out + + @staticmethod + def intervals_daily(starttime, endtime): + interval = sunpy.time.TimeRange(starttime, endtime) + daylist = interval.get_dates() + intervallist = [sunpy.time.TimeRange(t, t + dt.timedelta(days=1)) for + t in daylist] + return intervallist + def process(dirs, fnames, extension, local_base_dir, remote_base_url, download_func, processing_func, starttime, endtime, diff --git a/heliopy/data/wind.py b/heliopy/data/wind.py index c1ccce13b..844b53219 100644 --- a/heliopy/data/wind.py +++ b/heliopy/data/wind.py @@ -12,12 +12,18 @@ def _docstring(identifier, description): return cdasrest._docstring(identifier, 'W', description) -def _wind(starttime, endtime, identifier, badvalues=None, units=None): +def _wind(starttime, endtime, identifier, badvalues=None, units=None, + intervals='monthly'): """ Generic method for downloading ACE data. """ dl = cdasrest.CDASDwonloader('wi', identifier, 'wind', badvalues=badvalues, units=units) + # Override intervals + if intervals == 'daily': + dl.intervals = dl.intervals_daily + else: + dl.intervals = dl.intervals_monthly return dl.load(starttime, endtime) @@ -37,7 +43,8 @@ def swe_h1(starttime, endtime): def mfi_h0(starttime, endtime): identifier = 'WI_H0_MFI' units = {'BGSEa_0': u.nT, 'BGSEa_1': u.nT, 'BGSEa_2': u.nT} - return _wind(starttime, endtime, identifier, units=units) + return _wind(starttime, endtime, identifier, units=units, + intervals='daily') mfi_h0.__doc__ = _docstring( @@ -46,7 +53,7 @@ def mfi_h0(starttime, endtime): def mfi_h2(starttime, endtime): identifier = 'WI_H2_MFI' - return _wind(starttime, endtime, identifier) + return _wind(starttime, endtime, identifier, intervals='daily') mfi_h2.__doc__ = _docstring( @@ -55,7 +62,7 @@ def mfi_h2(starttime, endtime): def threedp_pm(starttime, endtime): identifier = 'WI_PM_3DP' - return _wind(starttime, endtime, identifier) + return _wind(starttime, endtime, identifier, intervals='daily') threedp_pm.__doc__ = _docstring(