Add remaining tests

Author: pjanowsk

tests/background_estimation/test_nf_unbinned_background.py

@@ -0,0 +1,168 @@
+import pytest
+
+import cosipy
+if not cosipy.with_ml:
+    pytest.skip(reason="Optional [ml] dependencies not installed", allow_module_level=True)
+
+import numpy as np
+import torch
+from unittest.mock import MagicMock
+
+from astropy import units as u
+from astropy.time import Time
+
+from cosipy.data_io.EmCDSUnbinnedData import TimeTagEmCDSEventInSCFrameInterface
+from cosipy.background_estimation.ml.NFBackground import NFBackground
+from cosipy.interfaces.data_interface import TimeTagEmCDSEventDataInSCFrameInterface
+from cosipy.interfaces.event import EventInterface
+
+from cosipy.background_estimation.ml.nf_unbinned_background import FreeNormNFUnbinnedBackground
+
+
+@pytest.fixture
+def mock_sc_history():
+    """Provides a realistic SpacecraftHistory mock using real Astropy units/times."""
+    sc_history = MagicMock()
+    
+    obstime = Time([1000.0, 1010.0, 1020.0], format='unix')
+    sc_history.obstime = obstime
+    sc_history.tstart = obstime[0]
+    sc_history.tstop = obstime[-1]
+    
+    sc_history.livetime = [9.0, 8.5] * u.s
+    sc_history.intervals_duration = [10.0, 10.0] * u.s
+
+    sc_history.cumulative_livetime.return_value = 17.5 * u.s
+    
+    return sc_history
+
+@pytest.fixture
+def mock_data():
+    """Provides a realistic Data mock using real Astropy times."""
+    data = MagicMock(spec=TimeTagEmCDSEventDataInSCFrameInterface)
+    
+    data.energy_keV = [500.0, 1000.0]
+    data.scattering_angle_rad = [0.5, 1.0]
+    data.scattered_lon_rad_sc = [0.1, 0.2]
+    data.scattered_lat_rad_sc = [0.3, 0.4]
+    
+    data.time = Time([1005.0, 1015.0], format='unix')
+    
+    return data
+
+@pytest.fixture
+def mock_model():
+    """Provides a mocked NFBackground model that returns predictable tensors."""
+    model = MagicMock(spec=NFBackground)
+    
+    model.evaluate_rate.return_value = torch.tensor([2.0, 3.0])
+    model.evaluate_density.return_value = torch.tensor([0.5, 0.6])
+    
+    model.active_pool = True
+    return model
+
+
+@pytest.fixture
+def background_instance(mock_model, mock_data, mock_sc_history):
+    """Instantiates the background class with all dependencies."""
+    return FreeNormNFUnbinnedBackground(
+        model=mock_model,
+        data=mock_data,
+        sc_history=mock_sc_history,
+        label="test_bkg_norm"
+    )
+
+
+class TestFreeNormNFUnbinnedBackground:
+
+    def test_init_and_properties(self, background_instance):
+        """Test initial state, event type, and parameters property."""
+        assert background_instance.event_type == TimeTagEmCDSEventInSCFrameInterface
+        assert background_instance.offset == 1e-12
+        assert background_instance._label == "test_bkg_norm"
+        assert background_instance._accum_livetime == 17.5 
+        
+        params = background_instance.parameters
+        assert "test_bkg_norm" in params
+        assert isinstance(params["test_bkg_norm"], u.Quantity)
+
+    def test_offset_validation(self, background_instance):
+        """Test the offset setter logic, including negative value guardrails."""
+        background_instance.offset = 0.0
+        assert background_instance.offset == 0.0
+        
+        background_instance.offset = None
+        assert background_instance.offset is None
+        
+        with pytest.raises(ValueError, match="The offset cannot be negative."):
+            background_instance.offset = -1.0
+
+    def test_integrate_rate(self, background_instance, mock_model):
+        """Test that the integration correctly multiplies rates by livetime."""
+        expected_counts = background_instance._integrate_rate()
+        
+        assert expected_counts == 43.5
+        mock_model.evaluate_rate.assert_called_once()
+        
+        passed_times = mock_model.evaluate_rate.call_args[0][0]
+        assert torch.allclose(passed_times, torch.tensor([[1005.0], [1015.0]], dtype=torch.float64))
+
+    def test_compute_density(self, background_instance, mock_model):
+        """Test the complex density computation and bin mapping logic."""
+        densities = background_instance._compute_density()
+        
+        assert isinstance(densities, np.ndarray)
+        assert densities.dtype == np.float64
+        np.testing.assert_allclose(densities, [0.9, 1.53])
+        
+        passed_source = mock_model.evaluate_density.call_args[0][1]
+        assert passed_source.shape == (2, 4)
+        assert torch.allclose(passed_source[0, 3], torch.tensor(np.pi/2 - 0.3, dtype=torch.float32))
+
+    def test_compute_density_time_bounds_error(self, background_instance, mock_data):
+        """Ensure evaluating events outside the spacecraft history raises an error."""
+        mock_data.time = Time([900.0, 1050.0], format='unix')
+        
+        with pytest.raises(ValueError, match="Input times are outside the spacecraft history range"):
+            background_instance._compute_density()
+
+    def test_norm_setter_and_getter(self, background_instance):
+        """Test scaling logic when norm is manipulated."""
+        
+        initial_norm_qty = background_instance.norm
+        assert initial_norm_qty.unit == u.Hz
+        np.testing.assert_allclose(initial_norm_qty.value, 43.5 / 17.5)
+        
+        background_instance.norm = 5.0 * u.Hz
+        
+        np.testing.assert_allclose(background_instance._norm, 5.0 * (17.5 / 43.5))
+        
+        background_instance.set_parameters(test_bkg_norm=10.0 * u.Hz)
+        assert background_instance.norm.value == 10.0
+
+    def test_expected_counts_method(self, background_instance):
+        """Test that expected counts scale by the internal norm."""
+        assert background_instance.expected_counts() == 43.5
+        
+        background_instance._norm = 2.0
+        assert background_instance.expected_counts() == 87.0
+
+    def test_expectation_density_method(self, background_instance):
+        """Test final density calculations, including norm scaling and offset addition."""
+        
+        res1 = background_instance.expectation_density()
+        np.testing.assert_allclose(res1, [0.9 + 1e-12, 1.53 + 1e-12])
+        
+        background_instance._norm = 2.0
+        background_instance.offset = None
+        res2 = background_instance.expectation_density()
+        np.testing.assert_allclose(res2, [1.8, 3.06])
+
+    def test_compute_pool_management(self, background_instance, mock_model):
+        """Ensure lazy evaluation accurately initializes and shutdowns compute pools."""
+        mock_model.active_pool = False
+        
+        background_instance.expectation_density()
+        
+        mock_model.init_compute_pool.assert_called_once()
+        mock_model.shutdown_compute_pool.assert_called_once()

tests/background_estimation/test_nfbackgroundmodels.py

@@ -0,0 +1,187 @@
+import pytest
+
+import cosipy
+if not cosipy.with_ml:
+    pytest.skip(reason="Optional [ml] dependencies not installed", allow_module_level=True)
+
+import numpy as np
+import torch
+from unittest.mock import MagicMock, patch
+
+from cosipy.background_estimation.ml.NFBackgroundModels import (
+    TotalDC4BackgroundRate,
+    TotalBackgroundDensityCMLPDGaussianCARQSFlow
+)
+
+@pytest.fixture
+def dummy_rate_input():
+    return {
+        "slew_duration": 10.0,
+        "obs_duration": 50.0,
+        "start_time": 1000.0,
+        "offset": 5.0,
+        "slope": 0.1,
+        "buildup": ((1.0, 2.0), (10.0, 20.0)),
+        "scale": 0.5,
+        "cutoff": (90.0, (1.0, 1.0, 1.0), (0.1, 0.1, 0.1), (0.0, 30.0, 60.0)),
+        "outlocs": torch.tensor([500.0, 900.0, 1500.0]),
+        "saa_decay": ((2.0, 3.0), (15.0, 30.0))
+    }
+
+@pytest.fixture
+def dummy_density_input():
+    return {
+        "model_state_dict": {},
+        "bins": 8,
+        "hidden_units": 64,
+        "residual_blocks": 2,
+        "total_layers": 3,
+        "context_size": 5,
+        "mlp_hidden_units": 32,
+        "mlp_hidden_layers": 2,
+        "menergy_cuts": (100.0, 10000.0),
+        "phi_cuts": (0.0, np.pi),
+        "start_time": 1000.0,
+        "total_time": 10000.0,
+        "period": 5400.0, 
+        "slew_duration": 600.0,
+        "obs_duration": 3000.0,
+        "outlocs": torch.tensor([500.0, 900.0, 1500.0])
+    }
+
+class TestTotalDC4BackgroundRate:
+    
+    def test_context_dim_property(self, dummy_rate_input):
+        model = TotalDC4BackgroundRate(dummy_rate_input)
+        assert model.context_dim == 1
+
+    def test_unpack_rate_input(self, dummy_rate_input):
+        """Ensure all elements from the dictionary are mapped to the correct instance variables."""
+        model = TotalDC4BackgroundRate(dummy_rate_input)
+        
+        assert model._slew_duration == 10.0
+        assert model._offset == 5.0
+        assert model._buildup_A == (1.0, 2.0)
+        assert model._cutoff_mu == (0.0, 30.0, 60.0)
+        assert torch.allclose(model._outlocs, torch.tensor([500.0, 900.0, 1500.0]))
+
+    def test_static_math_methods(self):
+        """Test the pure mathematical equations for buildup and decay."""
+        t = torch.tensor([10.0])
+        
+        buildup_res = TotalDC4BackgroundRate._buildup(t, A=4.0, T=10.0)
+        assert torch.allclose(buildup_res, torch.tensor([2.0]))
+        
+        decay_res = TotalDC4BackgroundRate._decay(t, A=4.0, T=10.0)
+        assert torch.allclose(decay_res, torch.tensor([2.0]))
+        
+        vm_res = TotalDC4BackgroundRate._von_mises(torch.tensor([0.0]), T=10.0, A=2.0, kappa=1.0, mu=0.0)
+        assert torch.allclose(vm_res, torch.tensor([2.0 * np.exp(1.0)], dtype=torch.float32))
+
+    def test_pointing_scale(self, dummy_rate_input):
+        """Test the sigmoid boundary logic inside pointing scale."""
+        model = TotalDC4BackgroundRate(dummy_rate_input)
+        
+        res = model._pointing_scale(torch.tensor([0.0, 60.0, 120.0]), scale=0.5, k0=10.0)
+        assert res.shape == (3,)
+        assert np.isclose(res[0], res[2])
+        
+    def test_saa_decay(self, dummy_rate_input):
+        """Test SAA decay correctly identifies the proper last exit time using searchsorted."""
+        model = TotalDC4BackgroundRate(dummy_rate_input)
+        
+        time_mins = torch.tensor([0.0, 10.0]) 
+        
+        decay = model._saa_decay(time_mins, A=(2.0, 3.0), T=(15.0, 30.0))
+        assert decay.shape == (2,)
+        assert np.all(np.isclose(decay,
+                                 model._decay(time_mins - torch.tensor([-100/60, 500/60]), 2.0, 15.0) +
+                                 model._decay(time_mins - torch.tensor([-100/60, 500/60]), 3.0, 30.0)))
+        assert torch.all(decay > 0) 
+
+    def test_evaluate_rate(self, dummy_rate_input):
+        """Test the full aggregation method."""
+        model = TotalDC4BackgroundRate(dummy_rate_input)
+        
+        abs_times = torch.tensor([1000.0, 1060.0, 1120.0])
+        rates = model.evaluate_rate(abs_times)
+        
+        assert rates.shape == (3,)
+        assert rates.dtype == torch.float32
+
+class TestTotalBackgroundDensity:
+
+    @patch('cosipy.background_estimation.ml.NFBackgroundModels.NNDensityInferenceWrapper')
+    @patch('cosipy.background_estimation.ml.NFBackgroundModels.build_c_arqs_flow')
+    @patch('cosipy.background_estimation.ml.NFBackgroundModels.build_cmlp_diaggaussian_base')
+    def test_init_and_properties(self, mock_base_builder, mock_flow_builder, mock_wrapper, dummy_density_input):
+        """Test that the flow builds correctly from dict parameters and properties read out correctly."""
+        model = TotalBackgroundDensityCMLPDGaussianCARQSFlow(
+            density_input=dummy_density_input,
+            worker_device="cpu",
+            batch_size=128,
+            compile_mode=None
+        )
+        
+        assert model.context_dim == 1
+        assert model.source_dim == 4
+        
+        assert model._menergy_cuts == (100.0, 10000.0)
+        assert model._total_time == 10000.0
+        
+        mock_base_builder.assert_called_once()
+        mock_flow_builder.assert_called_once()
+        mock_wrapper.assert_called_once()
+
+    @patch('cosipy.background_estimation.ml.NFBackgroundModels.TotalBackgroundDensityCMLPDGaussianCARQSFlow._load_model', return_value=None)
+    def test_inverse_transform_coordinates(self, mock_load, dummy_density_input):
+        """Test the physics to normalized-coordinate inverse mappings."""
+        model = TotalBackgroundDensityCMLPDGaussianCARQSFlow(dummy_density_input, "cpu", 128)
+        
+        nem = torch.tensor([0.0])  
+        nphi = torch.tensor([0.5]) 
+        npsi = torch.tensor([0.25])
+        nchi = torch.tensor([0.5]) 
+        dummy = torch.tensor([0.0])
+        
+        res = model._inverse_transform_coordinates(nem, nphi, npsi, nchi, dummy)
+        
+        assert res.shape == (1, 4)
+        np.testing.assert_allclose(res[0, 0].item(), 100.0)
+        np.testing.assert_allclose(res[0, 1].item(), np.pi / 2)
+        np.testing.assert_allclose(res[0, 2].item(), np.pi / 2)
+        np.testing.assert_allclose(res[0, 3].item(), np.pi / 2)
+
+    @patch('cosipy.background_estimation.ml.NFBackgroundModels.TotalBackgroundDensityCMLPDGaussianCARQSFlow._load_model', return_value=None)
+    def test_transform_coordinates(self, mock_load, dummy_density_input):
+        """Test calculation of transformed context, source, and jacobian."""
+        model = TotalBackgroundDensityCMLPDGaussianCARQSFlow(dummy_density_input, "cpu", 128)
+        
+        time = torch.tensor([1000.0])
+        em = torch.tensor([1000.0])
+        phi = torch.tensor([np.pi])
+        scatt_az = torch.tensor([np.pi])
+        scatt_pol = torch.tensor([np.pi/2])
+        
+        ctx, src, jac = model._transform_coordinates(time, em, phi, scatt_az, scatt_pol)
+        
+        assert ctx.shape == (1, 5)
+        assert src.shape == (1, 4) 
+        assert jac.shape == (1,)
+        
+        np.testing.assert_allclose(src[0, 0].item(), 0.5)
+
+    @patch('cosipy.background_estimation.ml.NFBackgroundModels.TotalBackgroundDensityCMLPDGaussianCARQSFlow._load_model', return_value=None)
+    def test_valid_samples(self, mock_load, dummy_density_input):
+        """Test the logical masking bounds for validation checks."""
+        model = TotalBackgroundDensityCMLPDGaussianCARQSFlow(dummy_density_input, "cpu", 128)
+        
+        nem = torch.tensor([0.0, -1.0, 0.5])   
+        nphi = torch.tensor([0.5, 1.5, 0.5])   
+        npsi = torch.tensor([0.5, 0.5, -0.1])  
+        nchi = torch.tensor([0.5, 0.5, 0.5])
+        dummy = torch.tensor([0.0, 0.0, 0.0])
+        
+        mask = model._valid_samples(nem, nphi, npsi, nchi, dummy)
+        
+        assert mask.tolist() == [True, False, False]