Fix casting bug in Tristan v2 spectra loading

Also, fixed the test data type and tests for the new gamma factor in loading spectra
PrincetonUniversity · Aug 19, 2024 · aec8bcc · aec8bcc
1 parent b9b973b
commit aec8bcc
Show file tree

Hide file tree

Showing 4 changed files with 3 additions and 3 deletions.
diff --git a/src/data_loading.py b/src/data_loading.py
@@ -205,7 +205,7 @@ def __handle_tristan_v2_spectra(spectra_file_path: pathlib.Path, spectra_file: h
     if not data_log_scale:
         spectral_data = np.log10(spectral_data)
 
-    spectral_data /= spectra_file['ebins']
+    spectral_data /= (spectra_file['ebins'][:])[:, np.newaxis]
 
     return spectral_data
 # =============================================================================

diff --git a/tests/data/tristan_v2/single_directory/spec.tot.00070 b/tests/data/tristan_v2/single_directory/spec.tot.00070
diff --git a/tests/data/tristan_v2/standard_structure/spec/spec.tot.00070 b/tests/data/tristan_v2/standard_structure/spec/spec.tot.00070
diff --git a/tests/test_data_loading.py b/tests/test_data_loading.py
@@ -204,7 +204,7 @@ def test_handle_tristan_v2_spectra():
         for i,dataset_name in enumerate(dataset_names):
             test_spectra = data_loading.__handle_tristan_v2_spectra(file_path, file, dataset_name, cli_args)
 
-            assert np.array_equiv(test_spectra, np.full(5,i+1))
+            assert np.array_equiv(test_spectra, np.full(5,i+2))
 
 def test_handle_tristan_v2_spectra_cli_args():
     parser = argparse.ArgumentParser()
@@ -219,7 +219,7 @@ def test_handle_tristan_v2_spectra_cli_args():
         for i,dataset_name in enumerate(dataset_names):
             test_spectra = data_loading.__handle_tristan_v2_spectra(file_path, file, dataset_name, cli_args)
 
-            assert np.array_equiv(test_spectra, np.full(5,i+1))
+            assert np.array_equiv(test_spectra, np.full(5,i+2))
 
 def test_handle_tristan_v2_spectra_raise_ValueError():
     with pytest.raises(ValueError):