Skip to content

Commit 44519c9

Browse files
fancaiyufancaiyu
committed
update examples
1 parent 1cca1bd commit 44519c9

File tree

3 files changed

+322
-161
lines changed

3 files changed

+322
-161
lines changed

‎examples/example_1.py

Lines changed: 124 additions & 69 deletions
Original file line numberDiff line numberDiff line change
@@ -1,70 +1,125 @@
1-
# Example 1: Unit Circle and Half-Sphere Integrands Comparison
2-
31
import torch
4-
from MCintegration import MonteCarlo, MarkovChainMonteCarlo, Vegas, set_seed, get_device
5-
6-
set_seed(42)
7-
device = get_device()
8-
9-
10-
def unit_circle_integrand(x, f):
11-
f[:, 0] = (x[:, 0] ** 2 + x[:, 1] ** 2 < 1).double()
12-
return f[:, 0]
13-
14-
15-
def half_sphere_integrand(x, f):
16-
f[:, 0] = torch.clamp(1 - (x[:, 0] ** 2 + x[:, 1] ** 2), min=0) * 2
17-
return f[:, 0]
18-
19-
20-
dim = 2
21-
bounds = [(-1, 1), (-1, 1)]
22-
n_eval = 6400000
23-
batch_size = 10000
24-
n_therm = 20
25-
26-
vegas_map = Vegas(dim, device=device, ninc=10)
27-
28-
# Monte Carlo and MCMC for Unit Circle
29-
mc_integrator = MonteCarlo(
30-
f=unit_circle_integrand, bounds=bounds, batch_size=batch_size
31-
)
32-
mcmc_integrator = MarkovChainMonteCarlo(
33-
f=unit_circle_integrand, bounds=bounds, batch_size=batch_size, nburnin=n_therm
34-
)
35-
36-
print("Unit Circle Integration Results:")
37-
print("Plain MC:", mc_integrator(n_eval))
38-
print("MCMC:", mcmc_integrator(n_eval, mix_rate=0.5))
39-
40-
# Train VEGAS map for Unit Circle
41-
vegas_map.adaptive_training(batch_size, unit_circle_integrand, alpha=0.5)
42-
vegas_integrator = MonteCarlo(
43-
bounds, f=unit_circle_integrand, maps=vegas_map, batch_size=batch_size
44-
)
45-
vegasmcmc_integrator = MarkovChainMonteCarlo(
46-
bounds,
47-
f=unit_circle_integrand,
48-
maps=vegas_map,
49-
batch_size=batch_size,
50-
nburnin=n_therm,
51-
)
52-
53-
print("VEGAS:", vegas_integrator(n_eval))
54-
print("VEGAS-MCMC:", vegasmcmc_integrator(n_eval, mix_rate=0.5))
55-
56-
# Monte Carlo and MCMC for Half-Sphere
57-
mc_integrator.f = half_sphere_integrand
58-
mcmc_integrator.f = half_sphere_integrand
59-
60-
print("\nHalf-Sphere Integration Results:")
61-
print("Plain MC:", mc_integrator(n_eval))
62-
print("MCMC:", mcmc_integrator(n_eval, mix_rate=0.5))
63-
64-
vegas_map.make_uniform()
65-
vegas_map.adaptive_training(batch_size, half_sphere_integrand, epoch=10, alpha=0.5)
66-
vegas_integrator.f = half_sphere_integrand
67-
vegasmcmc_integrator.f = half_sphere_integrand
68-
69-
print("VEGAS:", vegas_integrator(n_eval))
70-
print("VEGAS-MCMC:", vegasmcmc_integrator(n_eval, mix_rate=0.5))
2+
import torch.distributed as dist
3+
import torch.multiprocessing as mp
4+
import os
5+
import sys
6+
import traceback
7+
from MCintegration import MonteCarlo, MarkovChainMonteCarlo, Vegas
8+
os.environ["NCCL_DEBUG"] = "OFF"
9+
os.environ["TORCH_DISTRIBUTED_DEBUG"] = "OFF"
10+
os.environ["GLOG_minloglevel"] = "2"
11+
os.environ["MASTER_ADDR"] = os.getenv("MASTER_ADDR", "localhost")
12+
os.environ["MASTER_PORT"] = os.getenv("MASTER_PORT", "12355")
13+
14+
backend = "nccl"
15+
16+
def init_process(rank, world_size, fn, backend=backend):
17+
try:
18+
dist.init_process_group(backend, rank=rank, world_size=world_size)
19+
fn(rank, world_size)
20+
except Exception as e:
21+
traceback.print_exc()
22+
if dist.is_initialized():
23+
dist.destroy_process_group()
24+
raise e
25+
26+
def run_mcmc(rank, world_size):
27+
try:
28+
if rank != 0:
29+
sys.stderr = open(os.devnull, "w")
30+
sys.stdout = open(os.devnull, "w")
31+
torch.manual_seed(42 + rank)
32+
33+
def unit_circle_integrand(x, f):
34+
f[:, 0] = (x[:, 0] ** 2 + x[:, 1] ** 2 < 1).double()
35+
return f[:, 0]
36+
37+
38+
def half_sphere_integrand(x, f):
39+
f[:, 0] = torch.clamp(1 - (x[:, 0] ** 2 + x[:, 1] ** 2), min=0) * 2
40+
return f[:, 0]
41+
42+
dim = 2
43+
bounds = [(-1, 1), (-1, 1)]
44+
n_eval = 6400000
45+
batch_size = 40000
46+
n_therm = 20
47+
48+
device = torch.device(f"cuda:{rank}")
49+
50+
vegas_map = Vegas(dim, device=device, ninc=10)
51+
52+
# Monte Carlo and MCMC for Unit Circle
53+
mc_integrator = MonteCarlo(
54+
f=unit_circle_integrand, bounds=bounds, batch_size=batch_size,
55+
device=device
56+
)
57+
mcmc_integrator = MarkovChainMonteCarlo(
58+
f=unit_circle_integrand, bounds=bounds, batch_size=batch_size, nburnin=n_therm,
59+
device=device
60+
)
61+
62+
print("Unit Circle Integration Results:")
63+
print("Plain MC:", mc_integrator(n_eval))
64+
print("MCMC:", mcmc_integrator(n_eval, mix_rate=0.5))
65+
66+
# Train VEGAS map for Unit Circle
67+
vegas_map.adaptive_training(batch_size, unit_circle_integrand, alpha=0.5)
68+
vegas_integrator = MonteCarlo(
69+
bounds, f=unit_circle_integrand, maps=vegas_map, batch_size=batch_size,
70+
device=device
71+
)
72+
vegasmcmc_integrator = MarkovChainMonteCarlo(
73+
bounds,
74+
f=unit_circle_integrand,
75+
maps=vegas_map,
76+
batch_size=batch_size,
77+
nburnin=n_therm,
78+
device=device
79+
)
80+
81+
print("VEGAS:", vegas_integrator(n_eval))
82+
print("VEGAS-MCMC:", vegasmcmc_integrator(n_eval, mix_rate=0.5))
83+
84+
# Monte Carlo and MCMC for Half-Sphere
85+
mc_integrator.f = half_sphere_integrand
86+
mcmc_integrator.f = half_sphere_integrand
87+
88+
print("\nHalf-Sphere Integration Results:")
89+
print("Plain MC:", mc_integrator(n_eval))
90+
print("MCMC:", mcmc_integrator(n_eval, mix_rate=0.5))
91+
92+
vegas_map.make_uniform()
93+
vegas_map.adaptive_training(batch_size, half_sphere_integrand, epoch=10, alpha=0.5)
94+
vegas_integrator.f = half_sphere_integrand
95+
vegasmcmc_integrator.f = half_sphere_integrand
96+
97+
print("VEGAS:", vegas_integrator(n_eval))
98+
print("VEGAS-MCMC:", vegasmcmc_integrator(n_eval, mix_rate=0.5))
99+
100+
101+
except Exception as e:
102+
print(f"Error in run_mcmc for rank {rank}: {e}")
103+
traceback.print_exc()
104+
raise e
105+
finally:
106+
if dist.is_initialized():
107+
dist.destroy_process_group()
108+
109+
110+
def test_mcmc(world_size):
111+
world_size = min(world_size, mp.cpu_count())
112+
try:
113+
mp.spawn(
114+
init_process,
115+
args=(world_size, run_mcmc),
116+
nprocs=world_size,
117+
join=True,
118+
daemon=False,
119+
)
120+
except Exception as e:
121+
print(f"Error in test_mcmc: {e}")
122+
123+
if __name__ == "__main__":
124+
mp.set_start_method("spawn", force=True)
125+
test_mcmc(4)

‎examples/example_2.py

Lines changed: 108 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -1,54 +1,109 @@
1-
# Example 2: Sharp Peak Integrand in Higher Dimensions
2-
31
import torch
4-
from MCintegration import MonteCarlo, MarkovChainMonteCarlo, Vegas, set_seed, get_device
5-
6-
set_seed(42)
7-
device = get_device()
8-
9-
10-
def sharp_integrands(x, f):
11-
f[:, 0] = torch.sum((x - 0.5) ** 2, dim=-1)
12-
f[:, 0] *= -200
13-
f[:, 0].exp_()
14-
f[:, 1] = f[:, 0] * x[:, 0]
15-
f[:, 2] = f[:, 0] * x[:, 0] ** 2
16-
return f.mean(dim=-1)
17-
18-
19-
dim = 4
20-
bounds = [(0, 1)] * dim
21-
n_eval = 500000
22-
batch_size = 10000
23-
n_therm = 20
24-
25-
vegas_map = Vegas(dim, device=device, ninc=1000)
26-
27-
# Plain MC and MCMC
28-
mc_integrator = MonteCarlo(
29-
f=sharp_integrands, f_dim=3, bounds=bounds, batch_size=batch_size
30-
)
31-
mcmc_integrator = MarkovChainMonteCarlo(
32-
f=sharp_integrands, f_dim=3, bounds=bounds, batch_size=batch_size, nburnin=n_therm
33-
)
34-
35-
print("Sharp Peak Integration Results:")
36-
print("Plain MC:", mc_integrator(n_eval))
37-
print("MCMC:", mcmc_integrator(n_eval, mix_rate=0.5))
38-
39-
# Train VEGAS map
40-
vegas_map.adaptive_training(batch_size, sharp_integrands, f_dim=3, epoch=10, alpha=2.0)
41-
vegas_integrator = MonteCarlo(
42-
bounds, f=sharp_integrands, f_dim=3, maps=vegas_map, batch_size=batch_size
43-
)
44-
vegasmcmc_integrator = MarkovChainMonteCarlo(
45-
bounds,
46-
f=sharp_integrands,
47-
f_dim=3,
48-
maps=vegas_map,
49-
batch_size=batch_size,
50-
nburnin=n_therm,
51-
)
52-
53-
print("VEGAS:", vegas_integrator(n_eval))
54-
print("VEGAS-MCMC:", vegasmcmc_integrator(n_eval, mix_rate=0.5))
2+
import torch.distributed as dist
3+
import torch.multiprocessing as mp
4+
import os
5+
import sys
6+
import traceback
7+
from MCintegration import MonteCarlo, MarkovChainMonteCarlo, Vegas
8+
os.environ["NCCL_DEBUG"] = "OFF"
9+
os.environ["TORCH_DISTRIBUTED_DEBUG"] = "OFF"
10+
os.environ["GLOG_minloglevel"] = "2"
11+
os.environ["MASTER_ADDR"] = os.getenv("MASTER_ADDR", "localhost")
12+
os.environ["MASTER_PORT"] = os.getenv("MASTER_PORT", "12355")
13+
14+
backend = "nccl"
15+
16+
def init_process(rank, world_size, fn, backend=backend):
17+
try:
18+
dist.init_process_group(backend, rank=rank, world_size=world_size)
19+
fn(rank, world_size)
20+
except Exception as e:
21+
traceback.print_exc()
22+
if dist.is_initialized():
23+
dist.destroy_process_group()
24+
raise e
25+
26+
def run_mcmc(rank, world_size):
27+
try:
28+
if rank != 0:
29+
sys.stderr = open(os.devnull, "w")
30+
sys.stdout = open(os.devnull, "w")
31+
torch.manual_seed(42 + rank)
32+
33+
def sharp_integrands(x, f):
34+
f[:, 0] = torch.sum((x - 0.5) ** 2, dim=-1)
35+
f[:, 0] *= -200
36+
f[:, 0].exp_()
37+
f[:, 1] = f[:, 0] * x[:, 0]
38+
f[:, 2] = f[:, 0] * x[:, 0] ** 2
39+
return f.mean(dim=-1)
40+
41+
42+
dim = 4
43+
bounds = [(0, 1)] * dim
44+
n_eval = 6400000
45+
batch_size = 40000
46+
n_therm = 20
47+
48+
device = torch.device(f"cuda:{rank}")
49+
50+
print(f"Process {rank} using device: {device}")
51+
52+
vegas_map = Vegas(dim, device=device, ninc=1000)
53+
54+
# Plain MC and MCMC
55+
mc_integrator = MonteCarlo(
56+
f=sharp_integrands, f_dim=3, bounds=bounds, batch_size=batch_size,device=device
57+
)
58+
mcmc_integrator = MarkovChainMonteCarlo(
59+
f=sharp_integrands, f_dim=3, bounds=bounds, batch_size=batch_size, nburnin=n_therm,device=device
60+
)
61+
62+
print("Sharp Peak Integration Results:")
63+
print("Plain MC:", mc_integrator(n_eval))
64+
print("MCMC:", mcmc_integrator(n_eval, mix_rate=0.5))
65+
66+
# Train VEGAS map
67+
vegas_map.adaptive_training(batch_size, sharp_integrands, f_dim=3, epoch=10, alpha=2.0)
68+
vegas_integrator = MonteCarlo(
69+
bounds, f=sharp_integrands, f_dim=3, maps=vegas_map, batch_size=batch_size,device=device
70+
)
71+
vegasmcmc_integrator = MarkovChainMonteCarlo(
72+
bounds,
73+
f=sharp_integrands,
74+
f_dim=3,
75+
maps=vegas_map,
76+
batch_size=batch_size,
77+
nburnin=n_therm,
78+
device=device
79+
)
80+
81+
print("VEGAS:", vegas_integrator(n_eval))
82+
print("VEGAS-MCMC:", vegasmcmc_integrator(n_eval, mix_rate=0.5))
83+
84+
85+
except Exception as e:
86+
print(f"Error in run_mcmc for rank {rank}: {e}")
87+
traceback.print_exc()
88+
raise e
89+
finally:
90+
if dist.is_initialized():
91+
dist.destroy_process_group()
92+
93+
94+
def test_mcmc(world_size):
95+
world_size = min(world_size, mp.cpu_count())
96+
try:
97+
mp.spawn(
98+
init_process,
99+
args=(world_size, run_mcmc),
100+
nprocs=world_size,
101+
join=True,
102+
daemon=False,
103+
)
104+
except Exception as e:
105+
print(f"Error in test_mcmc: {e}")
106+
107+
if __name__ == "__main__":
108+
mp.set_start_method("spawn", force=True)
109+
test_mcmc(4)

0 commit comments

Comments
 (0)