diff --git a/setup.py b/setup.py index c8204eb8..21bf263b 100644 --- a/setup.py +++ b/setup.py @@ -92,5 +92,5 @@ 'obspy>=1.2.0', 'pyproj', 'tzlocal', - 'pyyaml'] + 'pyyaml>=5.1'] ) diff --git a/sourcespec/spectrum.py b/sourcespec/spectrum.py index 9e7bda58..7df5f225 100644 --- a/sourcespec/spectrum.py +++ b/sourcespec/spectrum.py @@ -15,6 +15,7 @@ import warnings import math import h5py +import yaml import numpy as np @@ -414,14 +415,57 @@ def _write_hdf5(self, group): group.create_dataset( 'data_mag_logspaced', data=self.data_mag_logspaced) + def _write_ascii(self, filename): + """ + Write the spectrum to an ASCII file. + + :param filename: The name of the file to write to. + """ + with open(filename, 'w', encoding='utf-8') as fp: + fp.write('# %SOURCESPEC ASCII SPECTRUM FORMAT 1.0\n') + fp.write('# %BEGIN STATS YAML\n') + stats_str = yaml.dump(dict(self.stats), sort_keys=False).rstrip() + for line in stats_str.split('\n'): + fp.write(f'# {line}\n') + fp.write( + '# %END STATS YAML\n' + '# %BEGIN LINSPACED DATA\n' + '# frequency(Hz) data data_mag\n' + ) + if self.data_mag.size: + data_mag = self.data_mag + else: + data_mag = np.ones_like(self.data) * np.nan + for freq, data, data_mag in zip(self.freq, self.data, data_mag): + fp.write(f'{freq:.6f} {data:.6f} {data_mag:.6f}\n') + fp.write( + '# %END LINSPACED DATA\n' + '# %BEGIN LOGSPACED DATA\n' + '# frequency_logspaced(Hz) data_logspaced data_mag_logspaced\n' + ) + if self.data_mag_logspaced.size: + data_mag_logspaced = self.data_mag_logspaced + else: + data_mag_logspaced = np.ones_like(self.data_logspaced) * np.nan + for freq_logspaced, data_logspaced, data_mag_logspaced in zip( + self.freq_logspaced, self.data_logspaced, + data_mag_logspaced): + fp.write( + f'{freq_logspaced:.6f} {data_logspaced:.6f} ' + f'{data_mag_logspaced:.6f}\n' + ) + fp.write('# %END LOGSPACED DATA\n') + # pylint: disable=redefined-builtin def write(self, filename, format='HDF5', append=False): """ Write the spectrum to a file. :param filename: The name of the file to write to. - :param format: The format to use. Currently, only 'HDF5' is supported. + :param format: The format to use. One of 'HDF5' or 'ASCII'. + Default is 'HDF5'. :param append: If True, append the spectrum to an existing file. + Only valid for HDF5 format. """ if format == 'HDF5': if append: @@ -433,6 +477,10 @@ def write(self, filename, format='HDF5', append=False): newgroup = 'spectrum_0001' self._write_hdf5(fp.create_group(newgroup)) fp.close() + elif format == 'ASCII': + if append: + raise ValueError('Cannot append to an ASCII file') + self._write_ascii(filename) else: raise ValueError(f'Unsupported format: {format}')