temp2_contract.qlua

package.path = '/global/homes/s/srijitp/build/qlua-build-2017/build2/parts/qlua/tree/lhpc/?.qlua;/global/homes/s/srijitp/build/qlua-build-2017/build2/parts/qlua/tree/qlib/?.qlua;' .. package.path
require "stdlib"
require "gauge"
require "ape-smear"
require "wuppertal-smear"
require "stout-smear"
require "qdpc_io"
require "setup_util"
require "lhpc-std"
require "latvec"

-- ************************************************************************************
-- * we use the following numbering
-- *
-- * N.M.X
-- *
-- * N in
-- * 1 --- M   -     M   2-point functions
-- * 2 --- MxM -     M   2-point functions
-- * 3 --- MxM -     MxM 2-point functions
-- * 4 --- M   - J - M   3-point functions
-- * 5 --- Mxm - J - M   3-point functions
-- *
-- * M in
-- * 1 --- light-light
-- * 2 --- light-charm
-- * 3 --- light-bottom
-- *
-- * X in
-- * a --- prepare eta_xi
-- * b --- prepare eta_phi
-- * c --- combine eta_xi and eta_phi
-- ************************************************************************************

-- ************************************************************************************
-- * INITIALIZE
-- ************************************************************************************

total_timer = timer("total time")

-- print parameters
printf("# [invert_contract_v3] Dimensions                = %d %d %d %d\n", dims[1], dims[2], dims[3], dims[4])
printf("# [invert_contract_v3] prefix to conf name       = %s\n", conf_name_prefix)
printf("# [invert_contract_v3] suffix to conf name       = %s\n", conf_name_input_suffix)
printf("# [invert_contract_v3] suffix to conf name       = %s\n", conf_name_output_suffix)
printf("# [invert_contract_v3] Clover Inverter\n")
printf("#\t light quark mass               = %g\n", mq_l)
printf("#\t light quark kappa              = %g\n", kappa_l)
printf("#\t strange quark mass               = %g\n", mq_s)
printf("#\t strange quark kappa              = %g\n", kappa_s)
printf("#\t Boundary condition = %d, %d, %d, %d\n", quark_bc[1], quark_bc[2], quark_bc[3], quark_bc[4])
-- TODO: what are relevant parameters for the multigrid solver?
-- printf("#\t Epsilon inner / final    = %e / %e\n", eps_i, eps_f)
-- printf("#\t Solver max. iters inner / total = %d\n", max_iter_i / dirac_max_iter)


  -- initialize the lattice
L = qcd.lattice(dims, {network=network_dims})

  Nd = #L

  volume = 1
  for i = 0, #L - 1 do
  volume = volume * L[i]
  end
  printf("# [invert_contract] lattice volume = %d\n", volume)

  -- initialize param set for writing propagators
  write_param = {
    ["m_q"]         = mq_l,
    ["c_sw"]        = csw,
    ["ferm_bc"]     = quark_bc,
    ["ensemble"]    = "beta3.31_2hex_24c48_ml-0.09530_mh-0.04",
    ["cfg"]         = nconf,
    ["ape_alpha"]   = ape_alpha,
    ["ape_n"]       = ape_n,
    ["ape_accu"]    = ape_accu,
    ["ape_maxiter"] = ape_maxiter,
    ["wup_alpha"]   = wup_alpha,
    ["wup_n"]       = wup_n,
    ["action_type"] = "clover",
    ["stout_rho"]   = stout_rho,
    ["stout_n"]     = stout_n,
    ["solver_eps"]  = mg_max_relerr,
    ["src_pos"]     = {-1, -1, -1},
    ["src_t"]       = -1
  }


if stout_smear then
-- read configuration
conf_name = string.format("%s%d%s", conf_name_prefix, nconf, conf_name_input_suffix)
printf("# [invert_contract] reading original gauge configuration from file %s\n", conf_name)

  u_load_timer = timer("u load")
uorig = load_nersc(L, conf_name)
  u_load_timer("done")
  printf("# [invert_contract] Plaquette (original) = %12.10f\n", plaquette_full(uorig))

  -- smear the gauge configuration
  stout_timer = timer("stout smear")
  printf("# [invert_contract] Stout-smearing gauge field\n")
ustout = stout_smear_4d(uorig, stout_rho, stout_n)
  printf("# [invert_contract] Plaquette (smeared) = %12.10f\n", plaquette_full(ustout))
  stout_timer("done")
  else
  -- read configuration
  --  conf_name = string.format("%s%d%s", conf_name_prefix, nconf, conf_name_output_suffix)
  --  printf("# [invert_contract_v3] reading 4d-stout-smeared gauge configuration from file %s\n", conf_name)

  --  u_load_timer = timer("u load")
--  ustout = load_ildg(L, conf_name)
  --  u_load_timer("done")
  --  printf("# [invert_contract_v3] Plaquette (4d-stout-smeared) = %12.10f\n", plaquette_full(ustout))
  printf("[invert_contract] Error, need original gauge configuration for heavy quark sector\n")
os.exit(2)
  end

  --[[
  -- TEST
  ustout = {}
  for i = 0, #L - 1 do
ustout[i+1] = L:ColorMatrix(complex(1,0))
  end
  uorig = ustout
  --]]


  -- APE-smear the configuration
  --if ape_n > 0 then
--uape = APE_smear(ustout, ape_alpha, ape_n, ape_accu, ape_maxiter, ape_orthdir)
  --  printf("# [invert_contract] Plaquette after (light quark) APE smearing = %12.10f\n", plaquette_full(uape))
  --end

--ustout_heavy = stout_smear_full(uorig, heavy_stout_rho, heavy_stout_n, 3)
  -- TEST
  -- ustout_heavy = uorig
  --printf("# [invert_contract] Plaquette after (heavy quark) stout smearing = %12.10f\n", plaquette_full(ustout_heavy))


  -- TEST
  --solver_light = make_mg_solver (L, ustout, mg_block, mg_max_abserr, mg_max_relerr, quark_bc, kappa_l, csw )
--solver_heavy = make_aicl_solver(L, uorig, charm, bottom)

  --[[
  -- TEST

function solver_light (v)
  L = v.lattice
s = tostring(v)
  if string.find(s, "^QDP:DiracFermion") then
dv = L:DiracFermion() 
  elseif string.find(s, "^QDP:DiracPropagator") then
dv = L:DiracPropagator() 
  else
  printf("# [solver_light] Bad argument to solver: %s\n", s)
os.exit(1)
  end
dv:set(v)
  return dv
  end

  -- TEST
function solver_heavy(v, flavor)
  L = v.lattice
s = tostring(v)
  if string.find(s, "^QDP:DiracFermion") then
dv = L:DiracFermion() 
  elseif string.find(s, "^QDP:DiracPropagator") then
dv = L:DiracPropagator() 
  else
  printf("# [solver_heavy] Bad argument to solver: %s\n", s)
os.exit(1)
  end
dv:set(v)
  return dv
  end
  --]]


-- prepare a list with all source locations (base + coherent)
  source_locations_all = {}
  for i_src, v_src in ipairs(source_locations) do

  local t   = v_src["t"]
  local pos = v_src["pos"]

  --  for i_coherent = 0, n_coherent_source - 1 do
  --              
  local i_prop = i_src
  --              
  --    local t = ( t + i_coherent * (L[3]/n_coherent_source) ) % L[3]
  --              
  --    local pos = {}
  --    pos[1] = ( pos[1]  + i_coherent * (L[0] / n_coherent_source) ) % L[0]
  --    pos[2] = ( pos[2]  + i_coherent * (L[1] / n_coherent_source) ) % L[1]
  --    pos[3] = ( pos[3]  + i_coherent * (L[2] / n_coherent_source) ) % L[2]
  --
  source_locations_all[i_prop] = {t=t, pos=pos}
  --     
  --  end
  end

  -- TEST
  for i = 1, #source_locations_all do
  printf("# [invert_contract_v3] src[%d] = (%d, %d, %d, t = %d\n", i, 
      source_locations_all[i].pos[1], source_locations_all[i].pos[2],
      source_locations_all[i].pos[3], source_locations_all[i].t)
  end


  -- random number generator
S = random_state(L)

  -- propagator lists
  light_propagator_list = {}

  --charm_propagator_list = {}
  --bottom_propagator_list = {}


  light_sequential_propagator_list = {}

  -- ************************************************************************************
  -- * LOOP ON BASE SOURCE TIMESLICES
  -- ************************************************************************************
  for i_src, v_src in ipairs(source_locations) do

  local t   = v_src["t"]
  local pos = v_src["pos"]
  printf("# [invert_contract_v3] base source location[%d] = (%d, %d, %d) t=%d\n", i_src, pos[1], pos[2], pos[3], t)

  --  for i_coherent = 0, n_coherent_source - 1 do
  --   for i_coherent = 0, 0 do

  local i_prop = i_src
  printf("# [invert_contract_v3] setting forward propagator no. %d\n", i_prop)

-- * set original (i_coherent = 0) or coherent source location (i_coherent > 0)

  --  local t = ( t + i_coherent * (L[3]/n_coherent_source) ) % L[3]

  -- set the shifted, coherent source location
  --    local pos = {}
  --    pos[1] = ( pos[1]  + i_coherent * (L[0] / n_coherent_source) ) % L[0]
  --    pos[2] = ( pos[2]  + i_coherent * (L[1] / n_coherent_source) ) % L[1]
  --    pos[3] = ( pos[3]  + i_coherent * (L[2] / n_coherent_source) ) % L[2]

  local source_timeslice = t
  --   printf("# [invert_contract_v3] coherent source timeslice = %d\n", source_timeslice)


local source_position = vector.int(4)
  source_position[0] = pos[1]
  source_position[1] = pos[2]
  source_position[2] = pos[3]
  source_position[3] = t


  -- ************************************************************************************
  -- * INVERT POINT SOURCE
  -- ************************************************************************************

  -- set a point source and smear the source

local source = L:DiracPropagator()

  if not read_forward_propagator then

  for icol  = 0, 2 do
  for ispin = 0, 3 do
  printf("# [invert_contract_v3] setting / smearing point source  %d / %d - (c=%2d, s=%2d)\n", i_src, i_prop, icol, ispin)
  local dv  = L:DiracFermion()
dv[{pos[1], pos[2], pos[3], t, c=icol, d=ispin}] = complex(1, 0)
  --if wup_n > 0 then
  -- source smearing
--source[{c=icol,d=ispin}] = wuppertal_smear(uape, wup_alpha, wup_n, dv, 3)
  -- else
  source[{c=icol,d=ispin}] = dv
  -- end
  end
  end

  -- propagator
  local i_timer = timer("light forward inversion")
  light_propagator_list[i_prop] = solver_light(source)
  i_timer("done")

  -- write the propagator UNSMEARED
  if write_forward_propagator then
  write_param["name"]    = light_propfile_prefix
  write_param["src_pos"] = pos
  write_param["src_t"]   = t

  write_propagator (light_propagator_list[i_prop], write_param)

  finfo = {["$tag"] = "test", "QDP/C propagator test file"}

  w = qcd.qdpc.Writer(L, string.format("%s.%.4d.x%.2dy%.2dz%.2dt%.2d.lime",light_propfile_prefix,nconf,pos[1],pos[2],pos[3],t), xml.unparse(finfo))

  --w:DiracFermion(src, xml.unparse({["$tag"] = "source"}))
  for icol  = 0, 2 do
  for ispin = 0, 3 do

  w:DiracFermion(light_propagator_list[i_prop][{c=icol, d=ispin}], xml.unparse({["$tag"] = "Fermion"}))

  end
  end
  w:close()
  end

  else
  printf("# [invert_contract_v3] reading light forward propagator\n")
  write_param["name"]    = light_propfile_prefix
  write_param["path"]    = light_propfile_path
  write_param["src_pos"] = pos
  write_param["src_t"]   = t
  light_propagator_list[i_prop] = read_propagator_lime (L, write_param)
  end

-- ************************************************************************************
-- * CONTRACTIONS 
-- ************************************************************************************
  printf("# [invert_contract_baryon] B - B contraction for source location[%d / %d] = (%d, %d, %d) t=%d\n", i_src, i_prop, 
      source_position[0], source_position[1], source_position[2], source_position[3]) 

    local name = string.format("%s.%.4d.t%.2d", "[u(Cgout*u*Cgin,u)]sp.contract_baryon_2pt", nconf, t)
    printf("# [invert_contract_baryon] B - B 2pt output file %s\n", name)
    aw_b_b_2pt, am = qcd.aff.Writer(name)
    name = nil
  local ctimer = timer("contract B-B")
  contract_baryon_2pt(L, light_propagator_list[i_prop], source_position, sink_momentum_list, aw_b_b_2pt)
  ctimer("done")

  --    vector_gamma_index = {1,2,4,15}
  --
  --    ctimer = timer("contract m-m")
  --    contract_vector_meson_2pt(L, light_propagator_smeared,       light_propagator_smeared,       source_position, sink_momentum_list, aw_b_b_2pt, "m-m")
  --    ctimer("done")
  --
aw_b_b_2pt:close()


 end  -- of loop on coherent source locations

  -- ************************************************************************************
  -- * INVERT SEQUENTIAL SOURCE
  -- ************************************************************************************

--  light_sequential_propagator_list[i_src] = {}
--
--  for i_seq_mom, seq_mom in ipairs(sequential_source_momentum_list) do
--
--    local seq_source_mom_timer = timer("seq_source_mom")
--    printf("#[invert_contract_v2]  sequential source momentum(%d) = (%d, %d, %d)\n", i_seq_mom, seq_mom[1], seq_mom[2], seq_mom[3])
--
--    local seq_source_mom = L:DiracPropagator()
--
--
--    if not read_sequential_propagator then
--
--      -- local field exp(i p x)
--      --   p = seq_mom sequential source momentum
--      --   x = sink of light propagator
--      local px = (2 * math.pi * seq_mom[1]) / L[0]
--      local py = (2 * math.pi * seq_mom[2]) / L[1]
--      local pz = (2 * math.pi * seq_mom[3]) / L[2]
--      local phase = ( (L:Real(L:pcoord(0)) ) * px +
--                      (L:Real(L:pcoord(1)) ) * py +
--                      (L:Real(L:pcoord(2)) ) * pz):expi()
--            
--      -- Dirac gamma matrix at sequential vertex, is always g5 (binary gamma id 15)
--      local seq_gamma = sequential_source_gamma["val"]
--          
--      -- sequential source at t = propagator at t
----      for i_coherent = 0, n_coherent_source - 1 do
----        local i_prop = ( i_src - 1) * n_coherent_source + i_coherent + 1
----        local t_coherent = ( t + i_coherent * (L[3]/n_coherent_source) ) % L[3]
--        printf("# [invert_contract_v2] adding source timeslice %d to sequential source list from propagator number %d\n", t, i_prop)
--        L:Subset{axis=3, position = t}:where(function()
--          seq_source_mom:set( light_propagator_list[i_prop] )
--        end)
----     end
--      solver_light = make_mg_solver (L, ustout, mg_block, mg_max_abserr, mg_max_relerr, quark_bc, kappa_l, csw )
--      -- seq_source_mom <- J exp(i p x) g5 J seq_source_mom
--      for icol  = 0, 2 do
--      for ispin = 0, 3 do
--          local dv  = L:DiracFermion()
--          dv[{pos[1], pos[2], pos[3], t, c=icol, d=ispin}] = complex(1, 0)
--
----        if wup_n > 0 then
----          printf("# [invert_contract_v2] smearing source + setting gamma x moentum + smearing source %d / %d - (c=%2d, s=%2d)\n", i_src, i_seq_mom, icol, ispin)
----          local dv = seq_source_mom[{c=icol,d=ispin}]
----          local dv2 = wuppertal_smear(uape, wup_alpha, wup_n, dv, 3)
----          dv = phase * ( seq_gamma * dv2 )
----          seq_source_mom[{c=icol,d=ispin}] = wuppertal_smear(uape, wup_alpha, wup_n, dv, 3)
----        else
--          printf("# [invert_contract_v2] setting gamma x momentum %d / %d - (c=%2d, s=%2d)\n", i_src, i_seq_mom, icol, ispin)
--          --seq_source_mom[{c=icol,d=ispin}] = phase * ( seq_gamma * seq_source_mom[{c=icol,d=ispin}] )
--          local sol = L:DiracFermion() 
--          sol = solver_light(dv)
--          
--          light_sequential_propagator_list[i_src][i_seq_mom][{c=icol, d=ispin}] = sol
--          seq_source_mom[{c=icol, d=ispin}] = dv
--          printf("seq_source = %g\n", seq_source_mom[{c=icol,d=ispin}]:norm2())
--    
----        end
--
--      end
--      end
--      seq_source_mom_timer("done")
--     -- solver_light = make_mg_solver (L, ustout, mg_block, mg_max_abserr, mg_max_relerr, quark_bc, kappa_l, csw )
--       
--      -- invert on the sequential momentum source
--      local i_timer = timer("light sequential inversion")
--      --light_sequential_propagator_list[i_src][i_seq_mom] = solver_light(seq_source_mom)
--      i_timer("done")
--
--      -- write sequential propagator to file
--      if write_sequential_propagator then
--        write_param["name"]    = seq_propfile_prefix
--        write_param["src_pos"] = pos
--        write_param["src_t"]   = t
--        write_param["mom"]     = seq_mom
--        write_propagator (light_sequential_propagator_list[i_src][i_seq_mom], write_param)
--      end  -- of if write_sequential_propagator
--
--    else
--      printf("# [invert_contract_v2] reading light sequential propagator\n")
--      write_param["name"]    = seq_propfile_prefix
--      write_param["path"]    = seq_propfile_path
--      write_param["src_pos"] = pos
--      write_param["src_t"]   = t
--      write_param["mom"]     = seq_mom
--      light_sequential_propagator_list[i_src][i_seq_mom] = read_propagator (L, write_param)
--    end
--
--     
--
--
--    
--    -- sink smearing of sequential propagator
--    --   can be smeared at this point, since for 2-point and 3-point functions we
--    --   always use the sequential propagator smeared at source and sink
----    if wup_n > 0 then
----      for icol  = 0, 2 do
----      for ispin = 0, 3 do
----        printf("# [invert_contract_v2] smearing sequential propagator %d / %d - (c=%2d, s=%2d)\n", i_src, i_seq_mom, icol, ispin)
----        local dv = light_sequential_propagator_list[i_src][i_seq_mom][{c=icol, d=ispin}]
----        light_sequential_propagator_list[i_src][i_seq_mom][{c=icol, d=ispin}] = wuppertal_smear(uape, wup_alpha, wup_n, dv, 3)
----      end
----      end
----
----    end
----
--    seq_source_mom = nil
--
--  end  -- end of loop on sequential source momenta
--
--end  -- end of loop on base source locations
  total_timer("done")