diff --git a/pySNOM/spectra.py b/pySNOM/spectra.py index e980e8c..a921b54 100644 --- a/pySNOM/spectra.py +++ b/pySNOM/spectra.py @@ -91,6 +91,19 @@ def parameters(self): """Property - scantype (ScanType Enum name)""" return self._parameters + def add_channel(self, values, channelname, zerofilling=1): + """Adds a new channel to data dictionary""" + if channelname not in list(self._data.keys()): + self._data[channelname] = np.reshape( + values, + ( + int(self.parameters["PixelArea"][0]), + int(self.parameters["PixelArea"][1]), + int(self.parameters["PixelArea"][2]*zerofilling), + ), + ) + else: + raise ValueError class SingleChannelSpectrum(NeaSpectrum): def __init__( @@ -228,6 +241,10 @@ def reshape_spectrum_data(data, params): spectral_depth = len(np.unique(data["Index"])) elif "Omega" in allchannels: spectral_depth = len(np.unique(data["Omega"])) + elif "Wavenumber" in allchannels: + spectral_depth = len(np.unique(data["Wavenumber"])) + elif "Wavelength" in allchannels: + spectral_depth = len(np.unique(data["Wavelength"])) else: spectral_depth = params["PixelArea"][2] * n diff --git a/pySNOM/tests/test_spectra.py b/pySNOM/tests/test_spectra.py index 873eac8..a7d4f49 100644 --- a/pySNOM/tests/test_spectra.py +++ b/pySNOM/tests/test_spectra.py @@ -20,6 +20,18 @@ def test_pointspectrum_object(self): np.testing.assert_string_equal(s.scantype, "Point") np.testing.assert_equal(np.shape(s.data["O2A"])[0], 2048) + def test_add_channel(self): + f = "datasets/testspectrum_singlepoint.txt" + file_reader = readers.NeaSpectralReader(os.path.join(pySNOM.__path__[0], f)) + data, params = file_reader.read() + + newchannel = np.zeros(np.shape(data["O3A"])) + s = spectra.NeaSpectrum(data, params) + s.add_channel(newchannel, "O6A",zerofilling=2) + + np.testing.assert_almost_equal(s.data["O6A"][0], 0) + + def test_multipointspectrum_object(self): f = "datasets/testspectrum_multipoint.txt" file_reader = readers.NeaSpectralReader(os.path.join(pySNOM.__path__[0], f))