def test_dump_single_gather(segy_path, tmp_path, name, retain_parent_segy_headers, header_index, header_cols,
dump_index):
survey = Survey(segy_path, header_index=header_index, header_cols=header_cols, name=name)
expected_gather = survey.get_gather(dump_index)
expected_gather.dump(path=tmp_path, name=name, retain_parent_segy_headers=retain_parent_segy_headers)
files = glob.glob(os.path.join(tmp_path, '*'))
assert len(files) == 1, "Dump creates more than one file"
dumped_survey = Survey(files[0], header_index=header_index, header_cols=header_cols)
ix = 1 if header_index == 'TRACE_SEQUENCE_FILE' else dump_index
dumped_gather = dumped_survey.get_gather(index=ix)
drop_columns = ["TRACE_SEQUENCE_FILE"] + list({"TRACE_SAMPLE_INTERVAL"}
& set(tuple(expected_gather.headers.columns)))
compare_gathers(expected_gather, dumped_gather, drop_cols=drop_columns, check_types=True,
same_survey=False)
if retain_parent_segy_headers:
expected_survey = Survey(segy_path, header_index=header_index, header_cols='all')
full_exp_headers = expected_survey.headers
full_exp_headers = full_exp_headers.loc[dump_index:dump_index].reset_index()
full_dump_headers = Survey(files[0], header_index=header_index, header_cols='all').headers
full_dump_headers = full_dump_headers.reset_index()
sample_rates = full_dump_headers['TRACE_SAMPLE_INTERVAL']
assert np.unique(sample_rates) > 1
assert np.allclose(sample_rates[0] / 1000, expected_survey.sample_rate)
full_exp_headers.drop(columns=["TRACE_SEQUENCE_FILE", "TRACE_SAMPLE_INTERVAL"], inplace=True)
full_dump_headers.drop(columns=["TRACE_SEQUENCE_FILE", "TRACE_SAMPLE_INTERVAL"], inplace=True)
assert full_exp_headers.equals(full_dump_headers)
def test_aggregate_segys(segy_path, tmp_path, mode, indices):
expected_survey = Survey(segy_path, header_index='FieldRecord', header_cols='all', name='raw')
indices = expected_survey.headers.index.drop_duplicates() if indices == 'all' else indices
if mode == 'split':
paths = [f'folder/folder_{i}' for i in range(len(indices))]
else:
paths = [''] * len(indices)
for num, (ix, path) in enumerate(zip(indices, paths)):
g = expected_survey.get_gather(ix)
g.dump(os.path.join(tmp_path, path), name=f'{num}_{ix}', retain_parent_segy_headers=True)
aggregate_segys(os.path.join(tmp_path, './**/*.sgy'), os.path.join(tmp_path, 'aggr.sgy'), recursive=True)
dumped_survey = Survey(os.path.join(tmp_path, 'aggr.sgy'), header_index='FieldRecord', header_cols='all')
assert np.allclose(expected_survey.samples, dumped_survey.samples),"Samples don't match"
assert np.allclose(expected_survey.sample_rate, dumped_survey.sample_rate), "Sample rate doesn't match"
assert np.allclose(expected_survey.n_samples, dumped_survey.n_samples), "length of samples doesn't match"
#TODO: optimize
drop_columns = ["TRACE_SEQUENCE_FILE"] + list({"TRACE_SAMPLE_INTERVAL"}
& set(tuple(expected_survey.headers.columns)))
expected_survey_headers = (expected_survey.headers.loc[indices].reset_index()
.sort_values(['FieldRecord', 'TraceNumber'])
.drop(columns=drop_columns)
.reset_index(drop=True))
dumped_survey_headers = (dumped_survey.headers.reset_index()
.sort_values(['FieldRecord', 'TraceNumber'])
.drop(columns=drop_columns)
.reset_index(drop=True))
assert len(expected_survey_headers) == len(dumped_survey_headers), "Length of surveys' headers don't match"
assert expected_survey_headers.equals(dumped_survey_headers), "The headers don't match"
for ix in indices:
expected_gather = expected_survey.get_gather(ix)
expected_gather.sort(by='TraceNumber')
dumped_gather = dumped_survey.get_gather(ix)
dumped_gather.sort(by='TraceNumber')
compare_gathers(expected_gather, dumped_gather, drop_cols=drop_columns, check_types=True,
same_survey=False)
def test_get_gather(self, load_segy, init_limits, load_limits):
"""Test gather loading by its index."""
path, trace_data = load_segy
survey = Survey(path,
header_index="FieldRecord",
limits=init_limits,
bar=False)
index = np.random.choice(survey.indices)
gather = survey.get_gather(index, limits=load_limits)
assert gather.index == index
gather_headers = survey.get_headers_by_indices([index])
assert gather.headers.equals(gather_headers)
# load_limits take priority over init_limits
limits = init_limits if load_limits is None else load_limits
traces_pos = gather["TRACE_SEQUENCE_FILE"].ravel() - 1
gather_data = trace_data[traces_pos, limits]
assert np.allclose(gather.data, gather_data)
assert np.allclose(gather.samples, survey.file_samples[limits])
def test_dump_single_gather_with_invalid_kwargs(segy_path, dump_kwargs, error):
survey = Survey(segy_path, header_index='FieldRecord')
gather = survey.get_gather(1)
with pytest.raises(error):
gather.dump(**dump_kwargs)
def test_gather_get_central_gather(segy_path):
"""test_gather_get_central_gather"""
survey = Survey(segy_path, header_index=['INLINE_3D', 'CROSSLINE_3D'], header_cols=['offset', 'FieldRecord'])
survey = survey.generate_supergathers()
gather = survey.get_gather((0, 0))
gather.get_central_gather()