Skip to content

Commit 9e7fb55

Browse files
TomDonoghueTomDonoghue
committed
io updates, make consistent set of load funcs, and updates test
1 parent c7375a7 commit 9e7fb55

File tree

7 files changed

+142
-40
lines changed

7 files changed

+142
-40
lines changed

‎specparam/tests/core/test_io.py‎

Lines changed: 32 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -49,14 +49,14 @@ def test_save_model_str(tfm):
4949
save_model(tfm, file_name_set, TEST_DATA_PATH, False, False, True, False)
5050
save_model(tfm, file_name_dat, TEST_DATA_PATH, False, False, False, True)
5151

52-
assert os.path.exists(os.path.join(TEST_DATA_PATH, file_name_res + '.json'))
53-
assert os.path.exists(os.path.join(TEST_DATA_PATH, file_name_set + '.json'))
54-
assert os.path.exists(os.path.join(TEST_DATA_PATH, file_name_dat + '.json'))
52+
assert os.path.exists(TEST_DATA_PATH / (file_name_res + '.json'))
53+
assert os.path.exists(TEST_DATA_PATH / (file_name_set + '.json'))
54+
assert os.path.exists(TEST_DATA_PATH / (file_name_dat + '.json'))
5555

5656
# Test saving out all save elements
5757
file_name_all = 'test_all'
5858
save_model(tfm, file_name_all, TEST_DATA_PATH, False, True, True, True)
59-
assert os.path.exists(os.path.join(TEST_DATA_PATH, file_name_all + '.json'))
59+
assert os.path.exists(TEST_DATA_PATH / (file_name_all + '.json'))
6060

6161
def test_save_model_append(tfm):
6262
"""Check saving fm data, appending to a file."""
@@ -66,20 +66,20 @@ def test_save_model_append(tfm):
6666
save_model(tfm, file_name, TEST_DATA_PATH, True, True, True, True)
6767
save_model(tfm, file_name, TEST_DATA_PATH, True, True, True, True)
6868

69-
assert os.path.exists(os.path.join(TEST_DATA_PATH, file_name + '.json'))
69+
assert os.path.exists(TEST_DATA_PATH / (file_name + '.json'))
7070

7171
def test_save_model_fobj(tfm):
7272
"""Check saving fm data, with file object file specifier."""
7373

7474
file_name = 'test_fileobj'
7575

7676
# Save, using file-object: three successive lines with three possible save settings
77-
with open(os.path.join(TEST_DATA_PATH, file_name + '.json'), 'w') as f_obj:
77+
with open(TEST_DATA_PATH / (file_name + '.json'), 'w') as f_obj:
7878
save_model(tfm, f_obj, TEST_DATA_PATH, False, True, False, False)
7979
save_model(tfm, f_obj, TEST_DATA_PATH, False, False, True, False)
8080
save_model(tfm, f_obj, TEST_DATA_PATH, False, False, False, True)
8181

82-
assert os.path.exists(os.path.join(TEST_DATA_PATH, file_name + '.json'))
82+
assert os.path.exists(TEST_DATA_PATH / (file_name + '.json'))
8383

8484
def test_save_group(tfg):
8585
"""Check saving fg data."""
@@ -92,14 +92,14 @@ def test_save_group(tfg):
9292
save_group(tfg, file_name=set_file_name, file_path=TEST_DATA_PATH, save_settings=True)
9393
save_group(tfg, file_name=dat_file_name, file_path=TEST_DATA_PATH, save_data=True)
9494

95-
assert os.path.exists(os.path.join(TEST_DATA_PATH, res_file_name + '.json'))
96-
assert os.path.exists(os.path.join(TEST_DATA_PATH, set_file_name + '.json'))
97-
assert os.path.exists(os.path.join(TEST_DATA_PATH, dat_file_name + '.json'))
95+
assert os.path.exists(TEST_DATA_PATH / (res_file_name + '.json'))
96+
assert os.path.exists(TEST_DATA_PATH / (set_file_name + '.json'))
97+
assert os.path.exists(TEST_DATA_PATH / (dat_file_name + '.json'))
9898

9999
# Test saving out all save elements
100100
file_name_all = 'test_group_all'
101101
save_group(tfg, file_name_all, TEST_DATA_PATH, False, True, True, True)
102-
assert os.path.exists(os.path.join(TEST_DATA_PATH, file_name_all + '.json'))
102+
assert os.path.exists(TEST_DATA_PATH / (file_name_all + '.json'))
103103

104104
def test_save_group_append(tfg):
105105
"""Check saving fg data, appending to file."""
@@ -109,17 +109,17 @@ def test_save_group_append(tfg):
109109
save_group(tfg, file_name, TEST_DATA_PATH, True, save_results=True)
110110
save_group(tfg, file_name, TEST_DATA_PATH, True, save_results=True)
111111

112-
assert os.path.exists(os.path.join(TEST_DATA_PATH, file_name + '.json'))
112+
assert os.path.exists(TEST_DATA_PATH / (file_name + '.json'))
113113

114114
def test_save_group_fobj(tfg):
115115
"""Check saving fg data, with file object file specifier."""
116116

117117
file_name = 'test_fileobj'
118118

119-
with open(os.path.join(TEST_DATA_PATH, file_name + '.json'), 'w') as f_obj:
119+
with open(TEST_DATA_PATH / (file_name + '.json'), 'w') as f_obj:
120120
save_group(tfg, f_obj, TEST_DATA_PATH, False, True, False, False)
121121

122-
assert os.path.exists(os.path.join(TEST_DATA_PATH, file_name + '.json'))
122+
assert os.path.exists(TEST_DATA_PATH / (file_name + '.json'))
123123

124124
def test_save_time(tft):
125125
"""Check saving ft data."""
@@ -132,9 +132,14 @@ def test_save_time(tft):
132132
save_group(tft, file_name=set_file_name, file_path=TEST_DATA_PATH, save_settings=True)
133133
save_group(tft, file_name=dat_file_name, file_path=TEST_DATA_PATH, save_data=True)
134134

135-
assert os.path.exists(os.path.join(TEST_DATA_PATH, res_file_name + '.json'))
136-
assert os.path.exists(os.path.join(TEST_DATA_PATH, set_file_name + '.json'))
137-
assert os.path.exists(os.path.join(TEST_DATA_PATH, dat_file_name + '.json'))
135+
assert os.path.exists(TEST_DATA_PATH / (res_file_name + '.json'))
136+
assert os.path.exists(TEST_DATA_PATH / (set_file_name + '.json'))
137+
assert os.path.exists(TEST_DATA_PATH / (dat_file_name + '.json'))
138+
139+
# Test saving out all save elements
140+
file_name_all = 'test_time_all'
141+
save_group(tft, file_name_all, TEST_DATA_PATH, False, True, True, True)
142+
assert os.path.exists(TEST_DATA_PATH / (file_name_all + '.json'))
138143

139144
def test_save_event(tfe):
140145
"""Check saving fe data."""
@@ -147,10 +152,16 @@ def test_save_event(tfe):
147152
save_event(tfe, file_name=set_file_name, file_path=TEST_DATA_PATH, save_settings=True)
148153
save_event(tfe, file_name=dat_file_name, file_path=TEST_DATA_PATH, save_data=True)
149154

150-
assert os.path.exists(os.path.join(TEST_DATA_PATH, set_file_name + '.json'))
155+
assert os.path.exists(TEST_DATA_PATH / (set_file_name + '.json'))
156+
for ind in range(len(tfe)):
157+
assert os.path.exists(TEST_DATA_PATH / (res_file_name + '_' + str(ind) + '.json'))
158+
assert os.path.exists(TEST_DATA_PATH / (dat_file_name + '_' + str(ind) + '.json'))
159+
160+
# Test saving out all save elements
161+
file_name_all = 'test_event_all'
162+
save_event(tfe, file_name_all, TEST_DATA_PATH, False, True, True, True)
151163
for ind in range(len(tfe)):
152-
assert os.path.exists(os.path.join(TEST_DATA_PATH, res_file_name + '_' + str(ind) + '.json'))
153-
assert os.path.exists(os.path.join(TEST_DATA_PATH, dat_file_name + '_' + str(ind) + '.json'))
164+
assert os.path.exists(TEST_DATA_PATH / (file_name_all + '_' + str(ind) + '.json'))
154165

155166
def test_load_json_str():
156167
"""Test loading JSON file, with str file specifier.
@@ -170,7 +181,7 @@ def test_load_json_fobj():
170181

171182
file_name = 'test_all'
172183

173-
with open(os.path.join(TEST_DATA_PATH, file_name + '.json'), 'r') as f_obj:
184+
with open(TEST_DATA_PATH / (file_name + '.json'), 'r') as f_obj:
174185
data = load_json(f_obj, '')
175186

176187
assert data

‎specparam/tests/core/test_reports.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -15,28 +15,28 @@ def test_save_model_report(tfm, skip_if_no_mpl):
1515

1616
save_model_report(tfm, file_name, TEST_REPORTS_PATH)
1717

18-
assert os.path.exists(os.path.join(TEST_REPORTS_PATH, file_name + '.pdf'))
18+
assert os.path.exists(TEST_REPORTS_PATH / (file_name + '.pdf'))
1919

2020
def test_save_group_report(tfg, skip_if_no_mpl):
2121

2222
file_name = 'test_group_report'
2323

2424
save_group_report(tfg, file_name, TEST_REPORTS_PATH)
2525

26-
assert os.path.exists(os.path.join(TEST_REPORTS_PATH, file_name + '.pdf'))
26+
assert os.path.exists(TEST_REPORTS_PATH / (file_name + '.pdf'))
2727

2828
def test_save_time_report(tft, skip_if_no_mpl):
2929

3030
file_name = 'test_time_report'
3131

3232
save_time_report(tft, file_name, TEST_REPORTS_PATH)
3333

34-
assert os.path.exists(os.path.join(TEST_REPORTS_PATH, file_name + '.pdf'))
34+
assert os.path.exists(TEST_REPORTS_PATH / (file_name + '.pdf'))
3535

3636
def test_save_event_report(tfe, skip_if_no_mpl):
3737

3838
file_name = 'test_event_report'
3939

4040
save_event_report(tfe, file_name, TEST_REPORTS_PATH)
4141

42-
assert os.path.exists(os.path.join(TEST_REPORTS_PATH, file_name + '.pdf'))
42+
assert os.path.exists(TEST_REPORTS_PATH / (file_name + '.pdf'))

‎specparam/tests/objs/test_group.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -219,7 +219,7 @@ def test_save_model_report(tfg):
219219
file_name = 'test_group_model_report'
220220
tfg.save_model_report(0, file_name, TEST_REPORTS_PATH)
221221

222-
assert os.path.exists(os.path.join(TEST_REPORTS_PATH, file_name + '.pdf'))
222+
assert os.path.exists(TEST_REPORTS_PATH / (file_name + '.pdf'))
223223

224224
def test_get_results(tfg):
225225
"""Check get results method."""

‎specparam/tests/plts/test_utils.py‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -81,23 +81,23 @@ def example_plot():
8181

8282
# Test defaults to saving given file path & name
8383
example_plot(file_path=TEST_PLOTS_PATH, file_name='test_savefig1.pdf')
84-
assert os.path.exists(os.path.join(TEST_PLOTS_PATH, 'test_savefig1.pdf'))
84+
assert os.path.exists(TEST_PLOTS_PATH / 'test_savefig1.pdf')
8585

8686
# Test works the same when explicitly given `save_fig`
8787
example_plot(save_fig=True, file_path=TEST_PLOTS_PATH, file_name='test_savefig2.pdf')
88-
assert os.path.exists(os.path.join(TEST_PLOTS_PATH, 'test_savefig2.pdf'))
88+
assert os.path.exists(TEST_PLOTS_PATH / 'test_savefig2.pdf')
8989

9090
# Test giving additional save kwargs
9191
example_plot(file_path=TEST_PLOTS_PATH, file_name='test_savefig3.pdf',
9292
save_kwargs={'facecolor' : 'red'})
93-
assert os.path.exists(os.path.join(TEST_PLOTS_PATH, 'test_savefig3.pdf'))
93+
assert os.path.exists(TEST_PLOTS_PATH / 'test_savefig3.pdf')
9494

9595
# Test does not save when `save_fig` set to False
9696
example_plot(save_fig=False, file_path=TEST_PLOTS_PATH, file_name='test_savefig_nope.pdf')
97-
assert not os.path.exists(os.path.join(TEST_PLOTS_PATH, 'test_savefig_nope.pdf'))
97+
assert not os.path.exists(TEST_PLOTS_PATH / 'test_savefig_nope.pdf')
9898

9999
def test_save_figure():
100100

101101
plt.plot([1, 2], [3, 4])
102102
save_figure(file_name='test_save_figure.pdf', file_path=TEST_PLOTS_PATH)
103-
assert os.path.exists(os.path.join(TEST_PLOTS_PATH, 'test_save_figure.pdf'))
103+
assert os.path.exists(TEST_PLOTS_PATH / 'test_save_figure.pdf')

‎specparam/tests/utils/test_download.py‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import os
44
import shutil
5+
from pathlib import Path
56

67
import numpy as np
78

@@ -10,7 +11,7 @@
1011
###################################################################################################
1112
###################################################################################################
1213

13-
TEST_FOLDER = 'test_data'
14+
TEST_FOLDER = Path('test_data')
1415

1516
def clean_up_downloads():
1617

@@ -29,14 +30,14 @@ def test_check_data_file():
2930
filename = 'freqs.npy'
3031

3132
check_data_file(filename, TEST_FOLDER)
32-
assert os.path.isfile(os.path.join(TEST_FOLDER, filename))
33+
assert os.path.isfile(TEST_FOLDER / filename)
3334

3435
def test_fetch_example_data():
3536

3637
filename = 'spectrum.npy'
3738

3839
fetch_example_data(filename, folder=TEST_FOLDER)
39-
assert os.path.isfile(os.path.join(TEST_FOLDER, filename))
40+
assert os.path.isfile(TEST_FOLDER / filename)
4041

4142
clean_up_downloads()
4243

‎specparam/tests/utils/test_io.py‎

Lines changed: 32 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,8 @@
33
import numpy as np
44

55
from specparam.core.items import OBJ_DESC
6-
from specparam.objs import SpectralModel, SpectralGroupModel
6+
from specparam.objs import (SpectralModel, SpectralGroupModel,
7+
SpectralTimeModel, SpectralTimeEventModel)
78

89
from specparam.tests.settings import TEST_DATA_PATH
910

@@ -30,10 +31,10 @@ def test_load_model():
3031
for meta_dat in OBJ_DESC['meta_data']:
3132
assert getattr(tfm, meta_dat) is not None
3233

33-
def test_load_group():
34+
def test_load_group_model():
3435

3536
file_name = 'test_group_all'
36-
tfg = load_group(file_name, TEST_DATA_PATH)
37+
tfg = load_group_model(file_name, TEST_DATA_PATH)
3738

3839
assert isinstance(tfg, SpectralGroupModel)
3940

@@ -44,3 +45,31 @@ def test_load_group():
4445
assert tfg.power_spectra is not None
4546
for meta_dat in OBJ_DESC['meta_data']:
4647
assert getattr(tfg, meta_dat) is not None
48+
49+
def test_load_time_model(tbands):
50+
51+
file_name = 'test_time_all'
52+
53+
# Load without bands definition
54+
tft = load_time_model(file_name, TEST_DATA_PATH)
55+
assert isinstance(tft, SpectralTimeModel)
56+
57+
# Load with bands definition
58+
tft2 = load_time_model(file_name, TEST_DATA_PATH, tbands)
59+
assert isinstance(tft2, SpectralTimeModel)
60+
assert tft2.time_results
61+
62+
def test_load_event_model(tbands):
63+
64+
file_name = 'test_event_all'
65+
66+
# Load without bands definition
67+
tfe = load_event_model(file_name, TEST_DATA_PATH)
68+
assert isinstance(tfe, SpectralTimeEventModel)
69+
assert len(tfe) > 1
70+
71+
# Load with bands definition
72+
tfe2 = load_event_model(file_name, TEST_DATA_PATH, tbands)
73+
assert isinstance(tfe2, SpectralTimeEventModel)
74+
assert tfe2.event_time_results
75+
assert len(tfe2) > 1

‎specparam/utils/io.py‎

Lines changed: 64 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
###################################################################################################
55

66
def load_model(file_name, file_path=None, regenerate=True):
7-
"""Load a model file.
7+
"""Load a model file into a model object.
88
99
Parameters
1010
----------
@@ -31,8 +31,8 @@ def load_model(file_name, file_path=None, regenerate=True):
3131
return model
3232

3333

34-
def load_group(file_name, file_path=None):
35-
"""Load a group file.
34+
def load_group_model(file_name, file_path=None):
35+
"""Load a group file into a group model object.
3636
3737
Parameters
3838
----------
@@ -55,3 +55,64 @@ def load_group(file_name, file_path=None):
5555
group.load(file_name, file_path)
5656

5757
return group
58+
59+
60+
def load_time_model(file_name, file_path=None, peak_org=None):
61+
"""Load a time file into a time model object.
62+
63+
64+
Parameters
65+
----------
66+
file_name : str
67+
File to load data data.
68+
file_path : Path or str, optional
69+
Path to directory to load from. If None, loads from current directory.
70+
peak_org : int or Bands, optional
71+
How to organize peaks.
72+
If int, extracts the first n peaks.
73+
If Bands, extracts peaks based on band definitions.
74+
75+
Returns
76+
-------
77+
time : SpectralTimeModel
78+
Object with the loaded data.
79+
"""
80+
81+
# Initialize a time object (imported locally to avoid circular imports)
82+
from specparam.objs import SpectralTimeModel
83+
time = SpectralTimeModel()
84+
85+
# Load data into object
86+
time.load(file_name, file_path, peak_org)
87+
88+
return time
89+
90+
91+
def load_event_model(file_name, file_path=None, peak_org=None):
92+
"""Load an event file into an event model object.
93+
94+
Parameters
95+
----------
96+
file_name : str
97+
File to load data data.
98+
file_path : Path or str, optional
99+
Path to directory to load from. If None, loads from current directory.
100+
peak_org : int or Bands, optional
101+
How to organize peaks.
102+
If int, extracts the first n peaks.
103+
If Bands, extracts peaks based on band definitions.
104+
105+
Returns
106+
-------
107+
event : SpectralTimeEventModel
108+
Object with the loaded data.
109+
"""
110+
111+
# Initialize an event object (imported locally to avoid circular imports)
112+
from specparam.objs import SpectralTimeEventModel
113+
event = SpectralTimeEventModel()
114+
115+
# Load data into object
116+
event.load(file_name, file_path, peak_org)
117+
118+
return event

0 commit comments

Comments
 (0)