Extended rms beamsize calculation

src/madl_gphys.mad

@@ -1960,27 +1960,46 @@ local function is_ord_even (m)
   return true
 end
 
-function gphys.beamsize(x, sig)
+function gphys.beamsize(x, sig, mo)
   assert(is_tpsa(x), "invalid argument #1 (real TPSA expected)")
   assert(is_indexable(sig) and #sig >= 5,
                      "invalid argument #2 (iterable with 5+ values expected")
-  local nv, xs = x:nv(), 0
+  mo = mo or x.mo
+  assert(is_number(mo) and mo <= x.mo, "invalid argument #3 (order must not exceed TPSA order")
+
+  local nv, xs_4, xs_5 = x:nv(), 0, 0
   local mi, mj, mm = monomial(nv), monomial(nv), monomial(nv)
 
   for i,xi in x:iter(mi)   do
   for j,xj in x:iter(mj,i) do
-    if mj >= mi and is_ord_even(mi:add(mj,mm)) then
+    if mj >= mi and is_ord_even(mi:add(mj,mm)) and mj:ord() <= mo and mi:ord() <= mo then
+
       local sk, xx = 0, (i == j and 1 or 2)*xi*xj
       for k=1,4 do
         xx = xx * sig[k]^mm[k] * tgamma((1+mm[k])/2)
         sk = sk + mm[k]
       end
+
       xx = xx * sig[5]^mm[6] * 2^(sk/2-mm[6]) / ((1+mm[6])*pi^2)
-      xs = xs + xx
+
+      xs_4 = xs_4 + (mm[6] == 0 and xx or 0)
+      xs_5 = xs_5 + xx
     end
   end end
+
+  return sqrt(xs_4), sqrt(xs_5)
+end
+
+function gphys.beamsize_oo(x, sig)
+  assert(is_tpsa(x), "invalid argument #1 (real TPSA expected)")
+  assert(is_indexable(sig) and #sig >= 5,
+                     "invalid argument #2 (iterable with 5+ values expected")
+  local sg, sd = {}, {} 
+  for i=1,x.mo do
+    sg[i], sd[i] = gphys.beamsize(x, sig, i)
+  end
 
-  return sqrt(xs)
+  return sg, sd
 end
 
 -- env ------------------------------------------------------------------------o