udpate integrators and utils tests

Author: houpc

MCintegration/integrators_test.py

@@ -5,8 +5,10 @@
 import numpy as np
 from integrators import Integrator, MonteCarlo, MarkovChainMonteCarlo
 from integrators import get_ip, get_open_port, setup
+from integrators import gaussian, random_walk
 
 from maps import Configuration
+from base import EPSILON
 
 
 class TestIntegrators(unittest.TestCase):
@@ -34,6 +36,50 @@ def test_setup(self, mock_set_device, mock_init_process_group):
         mock_init_process_group.assert_called_once_with(backend="gloo")
         mock_set_device.assert_called_once_with(0)
 
+    def test_random_walk(self):
+        # Test random_walk function with default parameters
+        dim = 3
+        device = "cpu"
+        dtype = torch.float32
+        u = torch.rand(dim, device=device, dtype=dtype)
+        step_size = 0.2
+
+        new_u = random_walk(dim, device, dtype, u, step_size=step_size)
+
+        # Validate output shape and range
+        self.assertEqual(new_u.shape, u.shape)
+        self.assertTrue(torch.all(new_u >= 0) and torch.all(new_u <= 1))
+
+        # Test with custom step size
+        custom_step_size = 0.5
+        new_u_custom = random_walk(dim, device, dtype, u, step_size=custom_step_size)
+        self.assertEqual(new_u_custom.shape, u.shape)
+        self.assertTrue(torch.all(new_u_custom >= 0) and torch.all(new_u_custom <= 1))
+
+    def test_gaussian(self):
+        # Test gaussian function with default parameters
+        dim = 3
+        device = "cpu"
+        dtype = torch.float32
+        u = torch.rand(dim, device=device, dtype=dtype)
+
+        mean = torch.zeros_like(u)
+        std = torch.ones_like(u)
+
+        new_u = gaussian(dim, device, dtype, u, mean=mean, std=std)
+
+        # Validate output shape
+        self.assertEqual(new_u.shape, u.shape)
+
+        # Test with custom mean and std
+        custom_mean = torch.full_like(u, 0.5)
+        custom_std = torch.full_like(u, 0.1)
+        new_u_custom = gaussian(dim, device, dtype, u, mean=custom_mean, std=custom_std)
+
+        self.assertEqual(new_u_custom.shape, u.shape)
+        self.assertTrue(torch.all(new_u_custom > custom_mean - 3 * custom_std))
+        self.assertTrue(torch.all(new_u_custom < custom_mean + 3 * custom_std))
+
 
 class TestIntegrator(unittest.TestCase):
     def setUp(self):
@@ -81,6 +127,18 @@ def test_initialization_with_maps(self, mock_linear_map, mock_composite_map):
         self.assertTrue(hasattr(integrator.maps, "device"))
         self.assertTrue(hasattr(integrator.maps, "dtype"))
 
+        integrator = Integrator(
+            bounds=self.bounds,
+            f=self.f,
+            maps=mock_map,
+            batch_size=self.batch_size,
+            device="cpu",
+        )
+
+        # Assertions
+        self.assertEqual(integrator.device, "cpu")
+        self.assertTrue(hasattr(integrator.maps, "device"))
+
     def test_bounds_conversion(self):
         # Test various input types
         test_cases = [
@@ -223,6 +281,27 @@ def test_batch_size_handling(self):
                     # Should not raise warning
                     self.mc(neval=neval, nblock=nblock)
 
+    def test_block_size_warning(self):
+        mc = MonteCarlo(bounds=self.bounds, f=self.simple_integral, batch_size=1000)
+        with self.assertWarns(UserWarning):
+            mc(neval=500, nblock=10)  # neval too small for nblock
+
+    def test_varying_nblock(self):
+        test_cases = [
+            (10000, 10),  # Standard case
+            (10000, 1),  # Single block
+            (10000, 100),  # Many blocks
+        ]
+
+        for neval, nblock in test_cases:
+            with self.subTest(neval=neval, nblock=nblock):
+                result = self.mc(neval=neval, nblock=nblock)
+                if hasattr(result, "mean"):
+                    value = result.mean
+                else:
+                    value = result
+                self.assertAlmostEqual(float(value), 1.0, delta=0.1)
+
 
 class TestMarkovChainMonteCarlo(unittest.TestCase):
     def setUp(self):
@@ -293,6 +372,41 @@ def test_burnin_effect(self):
                     value = result
                 self.assertAlmostEqual(float(value), 1.0, delta=tolerance)
 
+    def test_sample_acceptance(self):
+        config = Configuration(
+            self.mcmc.batch_size,
+            self.mcmc.dim,
+            self.mcmc.f_dim,
+            self.mcmc.device,
+            self.mcmc.dtype,
+        )
+        config.u, config.detJ = self.mcmc.q0.sample_with_detJ(self.mcmc.batch_size)
+        config.x, detj = self.mcmc.maps.forward_with_detJ(config.u)
+        config.detJ *= detj
+        config.weight = torch.rand(self.mcmc.batch_size, device=self.mcmc.device)
+
+        self.mcmc.sample(config, nsteps=1, mix_rate=0.5)
+
+        # Validate acceptance logic
+        self.assertTrue(torch.all(config.weight >= EPSILON))
+        self.assertEqual(config.u.shape, config.x.shape)
+
+    def test_varying_mix_rate(self):
+        test_cases = [
+            (0.1, 0.2),  # Low mix rate
+            (0.5, 0.1),  # Medium mix rate
+            (0.9, 0.05),  # High mix rate
+        ]
+
+        for mix_rate, tolerance in test_cases:
+            with self.subTest(mix_rate=mix_rate):
+                result = self.mcmc(neval=50000, mix_rate=mix_rate, nblock=10)
+                if hasattr(result, "mean"):
+                    value = result.mean
+                else:
+                    value = result
+                self.assertAlmostEqual(float(value), 1.0, delta=tolerance)
+
 
 class TestDistributedFunctionality(unittest.TestCase):
     @unittest.skipIf(not torch.distributed.is_available(), "Distributed not available")

MCintegration/utils_test.py

@@ -142,21 +142,39 @@ def test_converged(self):
         self.weighted_ravg.add(gvar.gvar(1.0, 0.01))
         self.assertTrue(self.weighted_ravg.converged(0.1, 0.1))
 
-    # def test_multiplication(self):
-    #     ravg1 = RAvg(weighted=True)
-    #     ravg1.update(2.0, 0.1)
-    #     ravg2 = RAvg(weighted=True)
-    #     ravg2.update(3.0, 0.1)
-    #     result = ravg1 * ravg2
-    #     self.assertAlmostEqual(result.mean, 6.0)
-
-    # def test_division(self):
-    #     ravg1 = RAvg(weighted=True)
-    #     ravg1.update(6.0, 0.1)
-    #     ravg2 = RAvg(weighted=True)
-    #     ravg2.update(3.0, 0.1)
-    #     result = ravg1 / ravg2
-    #     self.assertAlmostEqual(result.mean, 2.0)
+    def test_reduce_ex_serialization(self):
+        ravg = RAvg(weighted=True)
+        ravg.add(gvar.gvar(1.0, 0.1))
+        reduced = ravg.__reduce_ex__(2)
+        restored = reduced[0](*reduced[1])
+        self.assertEqual(restored.mean, ravg.mean)
+        self.assertEqual(restored.sdev, ravg.sdev)
+
+    def test_summary_output(self):
+        ravg = RAvg(weighted=True)
+        ravg.add(gvar.gvar(1.0, 0.1))
+        summary = ravg.summary()
+        self.assertIn("itn", summary)
+        self.assertIn("integral", summary)
+        self.assertIn("wgt average", summary)
+
+    def test_converged_criteria(self):
+        ravg = RAvg(weighted=True)
+        ravg.add(gvar.gvar(1.0, 0.1))
+        self.assertTrue(ravg.converged(0.1, 0.1))
+        self.assertFalse(ravg.converged(0.001, 0.001))
+
+    def test_multiplication_with_another_ravg(self):
+        ravg1 = RAvg(weighted=True)
+        ravg1.update(2.0, 0.1)
+        ravg2 = RAvg(weighted=True)
+        ravg2.update(3.0, 0.1)
+
+        result = ravg1 * ravg2
+        self.assertAlmostEqual(result.mean, 6.0)
+        sdev = (0.1 / 2**2 + 0.1 / 3**2) ** 0.5 * 6.0
+        self.assertAlmostEqual(result.sdev, sdev)
+
     def test_multiplication(self):
         ravg1 = RAvg(weighted=True)
         # Test multiplication by another RAvg object
@@ -198,6 +216,17 @@ def test_multiplication(self):
             np.allclose([r.sdev for r in result], [2.0 * ravg1.sdev, 3.0 * ravg1.sdev])
         )
 
+    def test_division_with_another_ravg(self):
+        ravg1 = RAvg(weighted=True)
+        ravg1.update(6.0, 0.1)
+        ravg2 = RAvg(weighted=True)
+        ravg2.update(3.0, 0.1)
+
+        result = ravg1 / ravg2
+        self.assertAlmostEqual(result.mean, 2.0)
+        sdev = (0.1 / 6.0**2 + 0.1 / 3.0**2) ** 0.5 * 2.0
+        self.assertAlmostEqual(result.sdev, sdev)
+
     def test_division(self):
         ravg1 = RAvg(weighted=True)
         ravg1.update(6.0, 0.1)
@@ -254,25 +283,35 @@ def test_set_seed_cpu(self):
         # Test set_seed on a CPU-only environment
         set_seed(42)
         self.assertEqual(torch.initial_seed(), 42)
-
-    @unittest.skipIf(not torch.cuda.is_available(), "CUDA is not available")
-    def test_set_seed_cuda(self):
-        # Test set_seed on a CUDA-enabled environment
+        u1 = torch.rand(10)
         set_seed(42)
-        self.assertEqual(torch.initial_seed(), 42)
-        self.assertEqual(torch.cuda.initial_seed(), 42)
+        u2 = torch.rand(10)
+        self.assertTrue(torch.all(u1 == u2))
+
+    # @unittest.skipIf(not torch.cuda.is_available(), "CUDA is not available")
+    # def test_set_seed_cuda(self):
+    #     # Test set_seed on a CUDA-enabled environment
+    #     set_seed(42)
+    #     self.assertEqual(torch.initial_seed(), 42)
+    #     self.assertEqual(torch.cuda.initial_seed(), 42)
 
     @unittest.skipIf(torch.cuda.is_available(), "CUDA is available")
     def test_get_device_cpu(self):
         # Test get_device when CUDA is not available
         device = get_device()
         self.assertEqual(device, torch.device("cpu"))
 
-    @unittest.skipIf(not torch.cuda.is_available(), "CUDA is not available")
-    def test_get_device_cuda(self):
-        # Test get_device when CUDA is available
-        device = get_device()
-        self.assertEqual(device, torch.cuda.current_device())
+    # @unittest.skipIf(not torch.cuda.is_available(), "CUDA is not available")
+    # def test_get_device_cuda(self):
+    #     # Test get_device when CUDA is available
+    #     device = get_device()
+    #     self.assertEqual(device, torch.cuda.current_device())
+
+    def test_get_device_cuda_inactive(self):
+        if not torch.cuda.is_available():
+            torch.cuda.set_device(-1)  # Simulate inactive CUDA
+            device = get_device()
+            self.assertEqual(device, torch.device("cpu"))
 
 
 if __name__ == "__main__":