YAML read/write support for Spectrum and SpectrumStream classes

setup.py

@@ -92,5 +92,5 @@
         'obspy>=1.2.0',
         'pyproj',
         'tzlocal',
-        'pyyaml']
+        'pyyaml>=5.1']
 )

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}')