update integrators_test

Author: houpc

MCintegration/integrators.py

@@ -233,7 +233,7 @@ def __call__(self, neval, nblock=32, verbose=-1, **kwargs):
 
         if verbose > 0:
             print(
-                f"nblock = {nblock}, n_steps_perblock = {epoch_perblock}, batch_size = {self.batch_size}, actual neval = {self.batch_size*epoch_perblock*nblock}"
+                f"nblock = {nblock}, n_steps_perblock = {epoch_perblock}, batch_size = {self.batch_size}, actual neval = {self.batch_size * epoch_perblock * nblock}"
             )
 
         config = Configuration(
@@ -356,13 +356,13 @@ def __call__(
         else:
             nblock = epoch // nsteps_perblock
         n_meas_perblock = nsteps_perblock // meas_freq
-        assert (
-            n_meas_perblock > 0
-        ), f"neval ({neval}) should be larger than batch_size * nblock * meas_freq ({self.batch_size} * {nblock} * {meas_freq})"
+        assert n_meas_perblock > 0, (
+            f"neval ({neval}) should be larger than batch_size * nblock * meas_freq ({self.batch_size} * {nblock} * {meas_freq})"
+        )
 
         if verbose > 0:
             print(
-                f"nblock = {nblock}, n_meas_perblock = {n_meas_perblock}, meas_freq = {meas_freq}, batch_size = {self.batch_size}, actual neval = {self.batch_size*nsteps_perblock*nblock}"
+                f"nblock = {nblock}, n_meas_perblock = {n_meas_perblock}, meas_freq = {meas_freq}, batch_size = {self.batch_size}, actual neval = {self.batch_size * nsteps_perblock * nblock}"
             )
 
         config = Configuration(

MCintegration/integrators_test.py

@@ -1,12 +1,11 @@
 import unittest
-from unittest.mock import patch
+from unittest.mock import patch, MagicMock
 import os
 import torch
 import numpy as np
 from integrators import Integrator, MonteCarlo, MarkovChainMonteCarlo
 from integrators import get_ip, get_open_port, setup
 
-# from base import LinearMap, Uniform
 from maps import Configuration
 
 
@@ -54,6 +53,34 @@ def test_initialization(self):
         self.assertTrue(hasattr(integrator.maps, "device"))
         self.assertTrue(hasattr(integrator.maps, "dtype"))
 
+    @patch("MCintegration.maps.CompositeMap")
+    @patch("MCintegration.base.LinearMap")
+    def test_initialization_with_maps(self, mock_linear_map, mock_composite_map):
+        # Mock the LinearMap and CompositeMap
+        mock_linear_map_instance = MagicMock()
+        mock_linear_map.return_value = mock_linear_map_instance
+        mock_composite_map_instance = MagicMock()
+        mock_composite_map.return_value = mock_composite_map_instance
+
+        # Create a mock map
+        mock_map = MagicMock()
+        mock_map.device = "cpu"
+        mock_map.dtype = torch.float32
+        mock_map.forward_with_detJ.return_value = (torch.rand(10, 2), torch.rand(10))
+
+        # Initialize Integrator with maps
+        integrator = Integrator(
+            bounds=self.bounds, f=self.f, maps=mock_map, batch_size=self.batch_size
+        )
+
+        # Assertions
+        self.assertEqual(integrator.dim, 2)
+        self.assertEqual(integrator.batch_size, 1000)
+        self.assertEqual(integrator.f_dim, 1)
+        self.assertTrue(hasattr(integrator.maps, "forward_with_detJ"))
+        self.assertTrue(hasattr(integrator.maps, "device"))
+        self.assertTrue(hasattr(integrator.maps, "dtype"))
+
     def test_bounds_conversion(self):
         # Test various input types
         test_cases = [
@@ -98,11 +125,11 @@ def test_invalid_bounds(self):
                 with self.assertRaises(error_type):
                     Integrator(bounds=bounds, f=self.f)
 
-    def test_device_handling(self):
-        if torch.cuda.is_available():
-            integrator = Integrator(bounds=self.bounds, f=self.f, device="cuda")
-            self.assertTrue(integrator.bounds.is_cuda)
-            self.assertTrue(integrator.maps.device == "cuda")
+    # def test_device_handling(self):
+    #     if torch.cuda.is_available():
+    #         integrator = Integrator(bounds=self.bounds, f=self.f, device="cuda")
+    #         self.assertTrue(integrator.bounds.is_cuda)
+    #         self.assertTrue(integrator.maps.device == "cuda")
 
     def test_dtype_handling(self):
         dtypes = [torch.float32, torch.float64]
@@ -266,30 +293,6 @@ def test_burnin_effect(self):
                     value = result
                 self.assertAlmostEqual(float(value), 1.0, delta=tolerance)
 
-    # def test_mix_rate_sensitivity(self):
-    #     # Modified mix rate test to be more robust
-    #     mix_rates = [0.0, 0.5, 1.0]
-    #     results = []
-
-    #     for mix_rate in mix_rates:
-    #         accumulated_error = 0
-    #         n_trials = 3  # Run multiple trials for each mix_rate
-
-    #         for _ in range(n_trials):
-    #             result = self.mcmc(neval=50000, mix_rate=mix_rate, nblock=10)
-    #             if hasattr(result, "mean"):
-    #                 value = result.mean
-    #                 error = result.sdev
-    #             else:
-    #                 value = result
-    #                 error = abs(float(value) - 1.0)
-    #             accumulated_error += error
-
-    #         results.append(accumulated_error / n_trials)
-
-    #     # We expect moderate mix rates to have lower average error
-    #     self.assertLess(results[1], max(results[0], results[2]))
-
 
 class TestDistributedFunctionality(unittest.TestCase):
     @unittest.skipIf(not torch.distributed.is_available(), "Distributed not available")
@@ -300,26 +303,26 @@ def test_distributed_initialization(self):
         self.assertEqual(integrator.rank, 0)
         self.assertEqual(integrator.world_size, 1)
 
-    @unittest.skipIf(not torch.distributed.is_available(), "Distributed not available")
-    def test_multi_gpu_consistency(self):
-        if torch.cuda.device_count() >= 2:
-            bounds = torch.tensor([[0.0, 1.0]], dtype=torch.float64)
-            f = lambda x, fx: torch.ones_like(x)
+    # @unittest.skipIf(not torch.distributed.is_available(), "Distributed not available")
+    # def test_multi_gpu_consistency(self):
+    #     if torch.cuda.device_count() >= 2:
+    #         bounds = torch.tensor([[0.0, 1.0]], dtype=torch.float64)
+    #         f = lambda x, fx: torch.ones_like(x)
 
-            # Create two integrators on different devices
-            integrator1 = Integrator(bounds=bounds, f=f, device="cuda:0")
-            integrator2 = Integrator(bounds=bounds, f=f, device="cuda:1")
+    #         # Create two integrators on different devices
+    #         integrator1 = Integrator(bounds=bounds, f=f, device="cuda:0")
+    #         integrator2 = Integrator(bounds=bounds, f=f, device="cuda:1")
 
-            # Results should be consistent across devices
-            result1 = integrator1(neval=10000)
-            result2 = integrator2(neval=10000)
+    #         # Results should be consistent across devices
+    #         result1 = integrator1(neval=10000)
+    #         result2 = integrator2(neval=10000)
 
-            if hasattr(result1, "mean"):
-                value1, value2 = result1.mean, result2.mean
-            else:
-                value1, value2 = result1, result2
+    #         if hasattr(result1, "mean"):
+    #             value1, value2 = result1.mean, result2.mean
+    #         else:
+    #             value1, value2 = result1, result2
 
-            self.assertAlmostEqual(float(value1), float(value2), places=1)
+    #         self.assertAlmostEqual(float(value1), float(value2), places=1)
 
 
 if __name__ == "__main__":