AMBER/tests/test_background.py
2025-04-09 16:46:45 +02:00

177 lines
6.2 KiB
Python

import unittest
import torch
import numpy as np
import os
import sys
maindir = os.getcwd()
main_path = maindir[:maindir.find('ds4ms/code')]
sys.path.append(main_path+"/ds4ms/code/src")
from background import background
class TestBackground(unittest.TestCase):
def setUp(self):
self.bg = background(data_path='/mydata/ds4ms/victor/Data/', str_dataset='VanadiumOzide', str_option='6p9T')
self.bg.load_data(verbose=True)
self.bg.set_grid_volume(dqx=0.03, dqy=0.03, dE=0.08)
self.bg.set_binned_data()
self.bg.set_radial_bins(max_radius=6.0, n_bins=10
self.bg.Ygrid = torch.tensor(np.random.rand(10, 10), dtype=torch.float64)
def test_load_data(self):
self.bg.load_data(verbose=True)
self.assertIsNotNone(self.bg.ds)
def test_set_dataset(self):
ds = "dummy_dataset"
self.bg.set_dataset(ds)
self.assertEqual(self.bg.ds, ds)
def test_mask_preprocessing(self):
self.bg.mask_preprocessing()
self.assertIsNotNone(self.bg.ds.mask)
def test_set_grid_volume(self):
self.bg.set_grid_volume(dqx=0.03, dqy=0.03, dE=0.08)
self.assertEqual(self.bg.dqx, 0.03)
self.assertEqual(self.bg.dqy, 0.03)
self.assertEqual(self.bg.dE, 0.08)
def test_set_binned_data(self):
self.bg.set_binned_data()
self.assertIsNotNone(self.bg.data)
self.assertIsNotNone(self.bg.bins)
def test_set_variables(self):
self.bg.set_variables()
self.assertIsNotNone(self.bg.X)
self.assertIsNotNone(self.bg.b)
self.assertIsNotNone(self.bg.b_grid)
def test_R_operator(self):
Y_r = torch.tensor(np.random.rand(10, 10), dtype=torch.float64)
b = torch.tensor(np.random.rand(10, 10), dtype=torch.float64)
result = self.bg.R_operator(Y_r, b)
self.assertIsNotNone(result)
def test_Rstar_operator(self):
Y_r = torch.tensor(np.random.rand(10, 10), dtype=torch.float64)
X = torch.tensor(np.random.rand(10, 10), dtype=torch.float64)
result = self.bg.Rstar_operator(Y_r, X)
self.assertIsNotNone(result)
def test_mask_nans(self):
x = torch.tensor(np.random.rand(10, 10), dtype=torch.float64)
result = self.bg.mask_nans(x)
self.assertIsNotNone(result)
def test_S_lambda(self):
x = torch.tensor(np.random.rand(10, 10), dtype=torch.float64)
lambda_ = torch.tensor(np.random.rand(10), dtype=torch.float64)
result = self.bg.S_lambda(x, lambda_)
self.assertIsNotNone(result)
def test_L_b(self):
self.bg.set_radial_nans()
result = self.bg.L_b(e_cut=0)
self.assertIsNotNone(result)
def test_gamma_matrix(self):
self.bg.set_radial_nans()
result = self.bg.gamma_matrix(e_cut=0)
self.assertIsNotNone(result)
def test_set_radial_nans(self):
self.bg.set_radial_nans()
self.assertIsNotNone(self.bg.u)
def test_set_b_design_matrix(self):
beta_ = torch.tensor(1.0, dtype=torch.float64)
alpha_ = torch.tensor(np.random.rand(10), dtype=torch.float64)
result = self.bg.set_b_design_matrix(beta_, alpha_, e_cut=0)
self.assertIsNotNone(result)
def test_compute_laplacian(self):
Y_r = torch.tensor(np.random.rand(10, 10), dtype=torch.float64)
result = self.bg.compute_laplacian(Y_r)
self.assertIsNotNone(result)
def test_compute_all_laplacians(self):
Y_r = torch.tensor(np.random.rand(10, 10), dtype=torch.float64)
self.bg.compute_all_laplacians(Y_r)
self.assertIsNotNone(self.bg.L_list)
def test_TV_denoising(self):
gamma_ = torch.tensor(np.random.rand(10), dtype=torch.float64)
result = self.bg.TV_denoising(gamma_, n_epochs=1, verbose=False)
self.assertIsNotNone(result)
def test_L_e(self):
result = self.bg.L_e()
self.assertIsNotNone(result)
def test_set_e_design_matrix(self):
result = self.bg.set_e_design_matrix(mu_=1.0)
self.assertIsNotNone(result)
def test_MAD_lambda(self):
result = self.bg.MAD_lambda()
self.assertIsNotNone(result)
def test_mu_estimator(self):
result = self.bg.mu_estimator()
self.assertIsNotNone(result)
def test_denoising(self):
Y_r = torch.tensor(np.random.rand(10, 10), dtype=torch.float64)
lambda_ = torch.tensor(np.random.rand(10), dtype=torch.float64)
beta_ = torch.tensor(np.random.rand(10), dtype=torch.float64)
alpha_ = torch.tensor(np.random.rand(10, 10), dtype=torch.float64)
self.bg.denoising(Y_r, lambda_, beta_, alpha_, mu_=1.0, n_epochs=1, verbose=False)
self.assertIsNotNone(self.bg.X)
self.assertIsNotNone(self.bg.b)
def test_applyBackground(self):
self.bg.applyBackground(median=False)
self.assertIsNotNone(self.bg.ds.backgroundModel)
def test_median_bg(self):
result = self.bg.median_bg(self.bg.Ygrid)
self.assertIsNotNone(result)
def test_compute_signal_to_noise(self):
b = torch.tensor(np.random.rand(10, 10), dtype=torch.float64)
result = self.bg.compute_signal_to_noise(b)
self.assertIsNotNone(result)
def test_plot_snr(self):
b = torch.tensor(np.random.rand(10, 10), dtype=torch.float64)
self.bg.plot_snr(b, e_cut=range(2, 4), fmin=0.0, fmax=0.1)
def test_save_arrays(self):
self.bg.save_arrays(median=True)
self.assertTrue(os.path.exists('arrays'))
def test_compute_signal_to_obs(self):
b = torch.tensor(np.random.rand(10, 10), dtype=torch.float64)
result = self.bg.compute_signal_to_obs(b)
self.assertIsNotNone(result)
def test_cross_validation(self):
lambda_range = torch.tensor([1.0])
alpha_range = torch.tensor([1.0])
beta_range = torch.tensor([1.0])
mu_range = torch.tensor([1.0])
result = self.bg.cross_validation(lambda_range=lambda_range, alpha_range=alpha_range, beta_range=beta_range, mu_range=mu_range, n_epochs=1, verbose=False)
self.assertIsNotNone(result)
def test_compute_mask(self):
result = self.bg.compute_mask(q=0.75, e_cut=None)
self.assertIsNotNone(result)
if __name__ == '__main__':
unittest.main()