update feynmandiag example

Author: houpc

examples/example_effective_mass.py

@@ -12,7 +12,8 @@
 
 maxOrder = 3
 beta = 10.0
-rs = 1.0
+rs = 2.0
+mass2 = 0.5
 
 device = get_device()
 batch_size = 5000
@@ -26,7 +27,7 @@ def zfactor_inv(maxOrder):
     sig_imag = [1.0]
 
     for order in range(1, maxOrder + 1):
-        feynfunc = init_feynfunc(order, rs, beta, batch_size, is_real=False)
+        feynfunc = init_feynfunc(order, rs, beta, mass2, batch_size, is_real=False)
         feynfunc.to(device)
         f_dim = num_roots[order - 1]
 
@@ -58,8 +59,11 @@ def meff_inv(maxOrder):
     sig_real = [1.0]
 
     for order in range(1, maxOrder + 1):
-        feynfunc = init_feynfunc(order, rs, beta, batch_size, is_real=True, has_dk=True)
-        feynfunc.to(device)
+        feynfunc = init_feynfunc(
+            order, rs, beta, mass2, batch_size, is_real=True, has_dk=True
+        )
+        # feynfunc.to(device)
+        device = torch.device("cpu")
         f_dim = num_roots[order - 1]
 
         vegas_map = Vegas(feynfunc.ndims, ninc=1000, device=device)
@@ -94,5 +98,5 @@ def meff_inv(maxOrder):
 
 meff = sum(z_factor_inv) / sum(m_eff_inv)
 
-print(f"\nEffective mass (up to order {maxOrder}):")
-print(meff)
+print(f"\n z-factor (up to order {maxOrder}): ", 1 / sum(z_factor_inv))
+print(f"\n Effective mass (up to order {maxOrder}): ", meff, "\n")

examples/feynman_test.py

@@ -16,8 +16,11 @@
 
 num_roots = [1, 2, 3, 4, 5, 6]
 order = 2
+rs = 2.0
 beta = 10.0
-feynfunc = init_feynfunc(order, beta, batch_size)
+mass2 = 0.5
+
+feynfunc = init_feynfunc(order, rs, beta, mass2, batch_size)
 feynfunc.to(device)
 f_dim = num_roots[order]
 

examples/feynmandiag.py

@@ -46,6 +46,7 @@ def __init__(
         order,
         rs,
         beta,
+        mass2,
         loopBasis,
         leafstates,
         leafvalues,
@@ -80,7 +81,7 @@ def __init__(
         self.register_buffer("mass2", torch.tensor(0.5))
         self.register_buffer("me", torch.tensor(0.5))
         self.register_buffer("spin", torch.tensor(2.0))
-        self.register_buffer("rs", torch.tensor(2.0))
+        self.register_buffer("rs", torch.tensor(rs))
         self.dim = 3
 
         # Derived constants
@@ -319,7 +320,7 @@ def __call__(self, var, root):
         return root.sum(dim=1)
 
 
-def init_feynfunc(order, rs, beta, batch_size, is_real=True, has_dk=False):
+def init_feynfunc(order, rs, beta, mass2, batch_size, is_real=True, has_dk=False):
     if has_dk:
         name = "sigmadk"
         root_dir = os.path.join(os.path.dirname(__file__), "funcs_sigmadk/")
@@ -349,6 +350,7 @@ def init_feynfunc(order, rs, beta, batch_size, is_real=True, has_dk=False):
         order,
         rs,
         beta,
+        mass2,
         loopBasis,
         leafstates[0],
         leafvalues[0],