updata parameters in example

Author: houpc

examples/example_1.py

@@ -28,6 +28,7 @@
 os.environ["MASTER_PORT"] = os.getenv("MASTER_PORT", "12355")
 
 backend = "nccl"
+# backend = "gloo"
 
 
 def init_process(rank, world_size, fn, backend=backend):
@@ -60,11 +61,18 @@ def half_sphere_integrand(x, f):
         bounds = [(-1, 1), (-1, 1)]
         n_eval = 6400000
         batch_size = 40000
+        alpha = 2.0
+        ninc = 1000
         n_therm = 20
 
-        device = torch.device(f"cuda:{rank}")
+        if backend == "gloo":
+            device = torch.device("cpu")
+        elif backend == "nccl":
+            device = torch.device(f"cuda:{rank}")
+        else:
+            raise ValueError(f"Invalid backend: {backend}")
 
-        vegas_map = Vegas(dim, device=device, ninc=10)
+        vegas_map = Vegas(dim, device=device, ninc=ninc)
 
         # Monte Carlo and MCMC for Unit Circle
         mc_integrator = MonteCarlo(
@@ -83,7 +91,7 @@ def half_sphere_integrand(x, f):
         print("MCMC:", mcmc_integrator(n_eval, mix_rate=0.5))
 
         # Train VEGAS map for Unit Circle
-        vegas_map.adaptive_training(batch_size, unit_circle_integrand, alpha=0.5)
+        vegas_map.adaptive_training(batch_size, unit_circle_integrand, alpha=alpha)
         vegas_integrator = MonteCarlo(
             bounds,
             f=unit_circle_integrand,
@@ -113,7 +121,7 @@ def half_sphere_integrand(x, f):
 
         vegas_map.make_uniform()
         vegas_map.adaptive_training(
-            batch_size, half_sphere_integrand, epoch=10, alpha=0.5
+            batch_size, half_sphere_integrand, epoch=10, alpha=alpha
         )
         vegas_integrator.f = half_sphere_integrand
         vegasmcmc_integrator.f = half_sphere_integrand

examples/example_2.py

@@ -28,6 +28,7 @@
 os.environ["MASTER_PORT"] = os.getenv("MASTER_PORT", "12355")
 
 backend = "nccl"
+# backend = "gloo"
 
 
 def init_process(rank, world_size, fn, backend=backend):
@@ -60,13 +61,20 @@ def sharp_integrands(x, f):
         bounds = [(0, 1)] * dim
         n_eval = 6400000
         batch_size = 40000
+        alpha = 2.0
+        ninc = 1000
         n_therm = 20
 
-        device = torch.device(f"cuda:{rank}")
+        if backend == "gloo":
+            device = torch.device("cpu")
+        elif backend == "nccl":
+            device = torch.device(f"cuda:{rank}")
+        else:
+            raise ValueError(f"Invalid backend: {backend}")
 
         print(f"Process {rank} using device: {device}")
 
-        vegas_map = Vegas(dim, device=device, ninc=1000)
+        vegas_map = Vegas(dim, device=device, ninc=ninc)
 
         # Plain MC and MCMC
         mc_integrator = MonteCarlo(
@@ -91,7 +99,7 @@ def sharp_integrands(x, f):
 
         # Train VEGAS map
         vegas_map.adaptive_training(
-            batch_size, sharp_integrands, f_dim=3, epoch=10, alpha=2.0
+            batch_size, sharp_integrands, f_dim=3, epoch=10, alpha=alpha
         )
         vegas_integrator = MonteCarlo(
             bounds,

examples/example_3.py

@@ -5,13 +5,16 @@
 import sys
 import traceback
 from MCintegration import MonteCarlo, MarkovChainMonteCarlo, Vegas
+
 os.environ["NCCL_DEBUG"] = "OFF"
 os.environ["TORCH_DISTRIBUTED_DEBUG"] = "OFF"
 os.environ["GLOG_minloglevel"] = "2"
 os.environ["MASTER_ADDR"] = os.getenv("MASTER_ADDR", "localhost")
 os.environ["MASTER_PORT"] = os.getenv("MASTER_PORT", "12355")
 
 backend = "nccl"
+# backend = "gloo"
+
 
 def init_process(rank, world_size, fn, backend=backend):
     try:
@@ -23,6 +26,7 @@ def init_process(rank, world_size, fn, backend=backend):
             dist.destroy_process_group()
         raise e
 
+
 def run_mcmc(rank, world_size):
     try:
         if rank != 0:
@@ -42,7 +46,12 @@ def func(x, f):
         ninc = 1000
         n_therm = 20
 
-        device = torch.device(f"cuda:{rank}")
+        if backend == "gloo":
+            device = torch.device("cpu")
+        elif backend == "nccl":
+            device = torch.device(f"cuda:{rank}")
+        else:
+            raise ValueError(f"Invalid backend: {backend}")
 
         print(f"Process {rank} using device: {device}")
 
@@ -54,26 +63,32 @@ def func(x, f):
 
         print("Integration Results for log(x)/sqrt(x):")
 
-
         # Plain MC Integration
-        mc_integrator = MonteCarlo(bounds, func, batch_size=batch_size,device=device)
+        mc_integrator = MonteCarlo(bounds, func, batch_size=batch_size, device=device)
         print("Plain MC Integral Result:", mc_integrator(n_eval))
 
         # MCMC Integration
         mcmc_integrator = MarkovChainMonteCarlo(
-            bounds, func, batch_size=batch_size, nburnin=n_therm,device=device
+            bounds, func, batch_size=batch_size, nburnin=n_therm, device=device
         )
         print("MCMC Integral Result:", mcmc_integrator(n_eval, mix_rate=0.5))
 
         # Perform VEGAS integration
-        vegas_integrator = MonteCarlo(bounds, func, maps=vegas_map, batch_size=batch_size,device=device)
+        vegas_integrator = MonteCarlo(
+            bounds, func, maps=vegas_map, batch_size=batch_size, device=device
+        )
         res = vegas_integrator(n_eval)
 
         print("VEGAS Integral Result:", res)
 
         # VEGAS-MCMC Integration
         vegasmcmc_integrator = MarkovChainMonteCarlo(
-            bounds, func, maps=vegas_map, batch_size=batch_size, nburnin=n_therm,device=device
+            bounds,
+            func,
+            maps=vegas_map,
+            batch_size=batch_size,
+            nburnin=n_therm,
+            device=device,
         )
         res_vegasmcmc = vegasmcmc_integrator(n_eval, mix_rate=0.5)
         print("VEGAS-MCMC Integral Result:", res_vegasmcmc)
@@ -100,6 +115,7 @@ def test_mcmc(world_size):
     except Exception as e:
         print(f"Error in test_mcmc: {e}")
 
+
 if __name__ == "__main__":
     mp.set_start_method("spawn", force=True)
-    test_mcmc(4)
+    test_mcmc(4)