From 47a32d6e3225508272275b1239ae00affe51a256 Mon Sep 17 00:00:00 2001 From: Luchang Jin Date: Wed, 25 Dec 2024 04:27:24 -0500 Subject: [PATCH] add auto contract example --- examples-py/Makefile | 2 +- ...auto-contract.log => auto-contract-01.log} | 0 examples-py/auto-contract-01.log.json | 124 ++++++++++++++++++ examples-py/auto-contract-01.py | 95 ++++++++++++++ examples-py/auto-contract.log.json | 1 - examples-py/auto-contract.py | 33 ----- 6 files changed, 220 insertions(+), 35 deletions(-) rename examples-py/{auto-contract.log => auto-contract-01.log} (100%) create mode 100644 examples-py/auto-contract-01.log.json create mode 100755 examples-py/auto-contract-01.py delete mode 100644 examples-py/auto-contract.log.json delete mode 100755 examples-py/auto-contract.py diff --git a/examples-py/Makefile b/examples-py/Makefile index 6bfe485ba..05a8ae7ea 100644 --- a/examples-py/Makefile +++ b/examples-py/Makefile @@ -14,7 +14,7 @@ tests = \ auto-cexprs-kpipi-qed-cy.log \ auto-cexprs-uu-cy.log \ auto-cexprs-distillation.log \ - auto-contract.log \ + auto-contract-01.log \ mat-op-distillation.log \ mat-op-qcd.log \ qlat-id-node.log \ diff --git a/examples-py/auto-contract.log b/examples-py/auto-contract-01.log similarity index 100% rename from examples-py/auto-contract.log rename to examples-py/auto-contract-01.log diff --git a/examples-py/auto-contract-01.log.json b/examples-py/auto-contract-01.log.json new file mode 100644 index 000000000..47c4614d5 --- /dev/null +++ b/examples-py/auto-contract-01.log.json @@ -0,0 +1,124 @@ +[ + [ + "Expr([Term([],[],1)],'(1)')" + ], + [ + "Expr([Term([],[Qv('u','x1','s1','c1'), Qb('u','x2','s2','c2')],1)],\"Qv('u','x1','s1','c1') * Qb('u','x2','s2','c2')\")" + ], + [ + "Expr([Term([],[Qv('u','x1','s1','c1'), Qb('u','x2','s2','c2')],1), Term([],[Qb('u','x2','s2','c2'), Qv('u','x1','s1','c1')],1)],\"+Qv('u','x1','s1','c1') * Qb('u','x2','s2','c2') + Qb('u','x2','s2','c2') * Qv('u','x1','s1','c1')\")" + ], + [ + "Expr([Term([],[],1)],'< (1) >')" + ], + [ + "Expr([Term([S('l','x1','x2','s1','s2','c1','c2')],[],1)],\"< Qv('u','x1','s1','c1') * Qb('u','x2','s2','c2') >\")" + ], + [ + "Expr([],\"< Qv('u','x1','s1','c1') * Qb('u','x2','s2','c2') + Qb('u','x2','s2','c2') * Qv('u','x1','s1','c1') >\")" + ], + [ + "# Begin CExpr\ndiagram_type_dict = dict()\ndiagram_type_dict[()] = 'ADT0'\ndiagram_type_dict[((('x1', 'x2'), 1),)] = 'ADT1'\n# Positions:\nx1, x2 = ['x1', 'x2']\n# Diagram type coef:\ncoef_ADT0 = 1\ncoef_ADT1 = 1\n# Named terms:\nterm_ADT0_0 = coef_ADT0 * \nterm_ADT1_0 = coef_ADT1 * S_l(x1,x2)\nterms = [ 0 for i in range(2) ]\nterms[0] = term_ADT0_0\nterms[1] = term_ADT1_0\n# Named exprs:\nexprs = [ 0 for i in range(3) ]\n# < (1) > exprs[0]\nexprs[0] += 1*term_ADT0_0\n# < Qv('u','x1','s1','c1') * Qb('u','x2','s2','c2') > exprs[1]\nexprs[1] += 1*term_ADT1_0\n# < Qv('u','x1','s1','c1') * Qb('u','x2','s2','c2') + Qb('u','x2','s2','c2') * Qv('u','x1','s1','c1') > exprs[2]\n# End CExpr" + ], + [ + "[('ADT0', ()), ('ADT1', ((('x1', 'x2'), 1),))]" + ], + [ + "['x1', 'x2']" + ], + [ + "[]" + ], + [ + "[]" + ], + [ + "[]" + ], + [ + "[]" + ], + [ + "[]" + ], + [ + "[]" + ], + [ + "[('term_ADT0_0', Term([],[],1)), ('term_ADT1_0', Term([S('l','x1','x2','s1','s2','c1','c2')],[],1))]" + ], + [ + "[('< (1) > exprs[0]', [(1, 'term_ADT0_0')]), (\"< Qv('u','x1','s1','c1') * Qb('u','x2','s2','c2') > exprs[1]\", [(1, 'term_ADT1_0')]), (\"< Qv('u','x1','s1','c1') * Qb('u','x2','s2','c2') + Qb('u','x2','s2','c2') * Qv('u','x1','s1','c1') > exprs[2]\", [])]" + ], + [ + "# Begin CExpr\ndiagram_type_dict = dict()\ndiagram_type_dict[()] = 'ADT0'\ndiagram_type_dict[((('x1', 'x2'), 1),)] = 'ADT1'\n# Positions:\nx1, x2 = ['x1', 'x2']\n# Variables prop:\nV_S_0 = S_l(x1,x2)\n# Variables factor:\nV_factor_coef_0 = 1\n# Diagram type coef:\ncoef_ADT0 = 1\ncoef_ADT1 = 1\n# Named terms:\nterm_ADT0_0 = coef_ADT0 * \nterm_ADT1_0 = coef_ADT1 * V_S_0\nterms = [ 0 for i in range(2) ]\nterms[0] = term_ADT0_0\nterms[1] = term_ADT1_0\n# Named exprs:\nexprs = [ 0 for i in range(3) ]\n# < (1) > exprs[0]\nexprs[0] += (V_factor_coef_0)*term_ADT0_0\n# < Qv('u','x1','s1','c1') * Qb('u','x2','s2','c2') > exprs[1]\nexprs[1] += (V_factor_coef_0)*term_ADT1_0\n# < Qv('u','x1','s1','c1') * Qb('u','x2','s2','c2') + Qb('u','x2','s2','c2') * Qv('u','x1','s1','c1') > exprs[2]\n# End CExpr" + ], + [ + "[('ADT0', ()), ('ADT1', ((('x1', 'x2'), 1),))]" + ], + [ + "['x1', 'x2']" + ], + [ + "[]" + ], + [ + "[('V_factor_coef_0', ea.Expr([ea.Term([],1)]))]" + ], + [ + "[('V_S_0', S('l','x1','x2','s1','s2','c1','c2'))]" + ], + [ + "[]" + ], + [ + "[]" + ], + [ + "[]" + ], + [ + "[('term_ADT0_0', Term([],[],1)), ('term_ADT1_0', Term([Var('V_S_0')],[],1))]" + ], + [ + "[('< (1) > exprs[0]', [(ea.Expr([ea.Term([ea.Factor(V_factor_coef_0,[],Var)],1)]), 'term_ADT0_0')]), (\"< Qv('u','x1','s1','c1') * Qb('u','x2','s2','c2') > exprs[1]\", [(ea.Expr([ea.Term([ea.Factor(V_factor_coef_0,[],Var)],1)]), 'term_ADT1_0')]), (\"< Qv('u','x1','s1','c1') * Qb('u','x2','s2','c2') + Qb('u','x2','s2','c2') * Qv('u','x1','s1','c1') > exprs[2]\", [])]" + ], + [ + "from auto_contractor.runtime import *\n\n### ----\n\n@timer\ndef cexpr_function(*, positions_dict, get_prop, is_ama_and_sloppy=False):\n # get_props\n props, cms, factors = cexpr_function_get_prop(positions_dict, get_prop)\n # eval\n ama_val = cexpr_function_eval(positions_dict, props, cms, factors)\n # extract sloppy val\n val_sloppy = ama_extract(ama_val, is_sloppy = True)\n # extract AMA val\n val_ama = ama_extract(ama_val)\n # return\n if is_ama_and_sloppy:\n # return both AMA corrected results and sloppy results\n return val_ama, val_sloppy\n else:\n # return AMA corrected results by default\n return val_ama\n\n### ----\n\n@timer_flops\ndef cexpr_function_get_prop(positions_dict, get_prop):\n # set positions\n x1 = positions_dict['x1']\n x2 = positions_dict['x2']\n # get prop\n V_S_0 = get_prop('l', x1, x2)\n # get color matrix\n # set props for return\n props = [\n V_S_0,\n ]\n # set color matrix for return\n cms = [\n ]\n # set intermediate factors\n # declare factors\n factors = np.zeros(1, dtype=np.complex128)\n factors_view = factors # Python only\n # set factors\n # V_factor_coef_0\n V_factor_coef_0 = (1) # Python only\n factors_view[0] = V_factor_coef_0\n # set flops\n total_flops = len(props) * 144 * 2 * 8 + len(cms) * 9 * 2 * 8 + len(factors) * 2 * 8\n # return\n return total_flops, (props, cms, factors,)\n\n### ----\n\n@timer_flops\ndef cexpr_function_eval(positions_dict, props, cms, factors):\n # load AMA props with proper format\n props = [ load_prop(p) for p in props ]\n # join the AMA props\n ama_props = ama_list(*props)\n # apply eval to the factors and AMA props\n ama_val = ama_apply1(lambda x_props: cexpr_function_eval_with_props(positions_dict, x_props, cms, factors), ama_props)\n # set flops\n total_flops = ama_counts(ama_val) * total_sloppy_flops\n # return\n return total_flops, ama_val\n\n### ----\n\n@timer_flops\ndef cexpr_function_eval_with_props(positions_dict, props, cms, factors_view): # Python only\n # set positions\n x1 = positions_dict['x1']\n x2 = positions_dict['x2']\n # set props\n p_V_S_0 = props[0] # Python only\n # set cms\n # set factors\n V_factor_coef_0 = factors_view[0] # Python only\n # compute products\n # compute traces\n # set terms\n term_ADT0_0 = 1 # Python only\n term_ADT1_0 = p_V_S_0 # Python only\n # declare exprs\n exprs = np.empty(3, dtype=object)\n exprs_view = exprs\n # set exprs\n expr_V_G = SpinMatrix() # Python only\n expr_V_S = WilsonMatrix() # Python only\n expr_V_U = ColorMatrix() # Python only\n # 0 name='< (1) > exprs[0]' \n expr_V_a = 0\n expr_V_a += (((V_factor_coef_0))) * term_ADT0_0\n exprs_view[0] = expr_V_a # Python only\n # 1 name='< Qv('u','x1','s1','c1') * Qb('u','x2','s2','c2') > exprs[1]' \n expr_V_S.set_zero() # Python only\n expr_V_S += mat_mul_a_wm(((V_factor_coef_0)), term_ADT1_0)\n exprs_view[1] = expr_V_S.copy() # Python only\n # 2 name='< Qv('u','x1','s1','c1') * Qb('u','x2','s2','c2') + Qb('u','x2','s2','c2') * Qv('u','x1','s1','c1') > exprs[2]' \n exprs_view[2] = 0\n # set flops\n total_flops = total_sloppy_flops\n # return\n return total_flops, exprs\n\n### ----\n\n# Total flops per sloppy call is: 0\ntotal_sloppy_flops = 0" + ], + [ + "get_cexpr_test benchmark_eval_cexpr check get_data_sig is_cython=False", + { + "real": 15.325849886336291, + "imag": -15.659291332109987, + "__extended_json_type__": "complex128" + } + ], + [ + "get_cexpr_test benchmark_eval_cexpr check_ama get_data_sig is_cython=False", + { + "real": 11.492942956036883, + "imag": -13.926566927529471, + "__extended_json_type__": "complex128" + } + ], + [ + "from auto_contractor.runtime import *\nimport cython # Cython\ncimport qlat_utils.everything as cc # Cython\ncimport qlat_utils.all as qu # Cython\ncimport libcpp.complex # Cython\ncimport numpy # Cython\n\n### ----\n\n@timer\ndef cexpr_function(*, positions_dict, get_prop, is_ama_and_sloppy=False):\n # get_props\n props, cms, factors = cexpr_function_get_prop(positions_dict, get_prop)\n # eval\n ama_val = cexpr_function_eval(positions_dict, props, cms, factors)\n # extract sloppy val\n val_sloppy = ama_extract(ama_val, is_sloppy = True)\n # extract AMA val\n val_ama = ama_extract(ama_val)\n # return\n if is_ama_and_sloppy:\n # return both AMA corrected results and sloppy results\n return val_ama, val_sloppy\n else:\n # return AMA corrected results by default\n return val_ama\n\n### ----\n\n@timer_flops\n@cython.boundscheck(False) # Cython\n@cython.wraparound(False) # Cython\ndef cexpr_function_get_prop(positions_dict, get_prop):\n # set positions\n x1 = positions_dict['x1']\n x2 = positions_dict['x2']\n # get prop\n V_S_0 = get_prop('l', x1, x2)\n # get color matrix\n # set props for return\n props = [\n V_S_0,\n ]\n # set color matrix for return\n cms = [\n ]\n # set intermediate factors\n # declare factors\n cdef numpy.ndarray[numpy.complex128_t] factors # Cython\n factors = np.zeros(1, dtype=np.complex128)\n cdef cc.PyComplexD[:] factors_view = factors # Cython\n # set factors\n # V_factor_coef_0\n cdef cc.PyComplexD V_factor_coef_0 = (1) # Cython\n factors_view[0] = V_factor_coef_0\n # set flops\n total_flops = len(props) * 144 * 2 * 8 + len(cms) * 9 * 2 * 8 + len(factors) * 2 * 8\n # return\n return total_flops, (props, cms, factors,)\n\n### ----\n\n@timer_flops\ndef cexpr_function_eval(positions_dict, props, cms, factors):\n # load AMA props with proper format\n props = [ load_prop(p) for p in props ]\n # join the AMA props\n ama_props = ama_list(*props)\n # apply eval to the factors and AMA props\n ama_val = ama_apply1(lambda x_props: cexpr_function_eval_with_props(positions_dict, x_props, cms, factors), ama_props)\n # set flops\n total_flops = ama_counts(ama_val) * total_sloppy_flops\n # return\n return total_flops, ama_val\n\n### ----\n\n@timer_flops\n@cython.boundscheck(False) # Cython\n@cython.wraparound(False) # Cython\ndef cexpr_function_eval_with_props(dict positions_dict, list props, list cms, cc.PyComplexD[:] factors_view): # Cython\n # set positions\n x1 = positions_dict['x1']\n x2 = positions_dict['x2']\n # set props\n cdef cc.WilsonMatrix* p_V_S_0 = &(props[0]).xx # Cython\n # set cms\n # set factors\n cdef cc.PyComplexD V_factor_coef_0 = factors_view[0] # Cython\n # compute products\n # compute traces\n # set terms\n cdef cc.PyComplexD term_ADT0_0 = 1 # Cython\n cdef cc.WilsonMatrix term_ADT1_0 = p_V_S_0[0] # Cython\n # declare exprs\n exprs = np.empty(3, dtype=object)\n exprs_view = exprs\n # set exprs\n cdef cc.PyComplexD expr_V_a # Cython\n cdef cc.SpinMatrix expr_V_G # Cython\n cdef cc.WilsonMatrix expr_V_S # Cython\n cdef cc.ColorMatrix expr_V_U # Cython\n cdef qu.SpinMatrix expr_V_G_box # Cython\n cdef qu.WilsonMatrix expr_V_S_box # Cython\n cdef qu.ColorMatrix expr_V_U_box # Cython\n # 0 name='< (1) > exprs[0]' \n expr_V_a = 0\n expr_V_a += (((V_factor_coef_0))) * term_ADT0_0\n exprs_view[0] = expr_V_a # Cython\n # 1 name='< Qv('u','x1','s1','c1') * Qb('u','x2','s2','c2') > exprs[1]' \n cc.set_zero(expr_V_S) # Cython\n expr_V_S += cc.ccpy_d(((V_factor_coef_0))) * term_ADT1_0\n expr_V_S_box = WilsonMatrix() # Cython\n expr_V_S_box.xx = expr_V_S # Cython\n exprs_view[1] = expr_V_S_box # Cython\n # 2 name='< Qv('u','x1','s1','c1') * Qb('u','x2','s2','c2') + Qb('u','x2','s2','c2') * Qv('u','x1','s1','c1') > exprs[2]' \n exprs_view[2] = 0\n # set flops\n total_flops = total_sloppy_flops\n # return\n return total_flops, exprs\n\n### ----\n\n# Total flops per sloppy call is: 0\ntotal_sloppy_flops = 0" + ], + [ + "get_cexpr_test benchmark_eval_cexpr check get_data_sig is_cython=True", + { + "real": 15.325849886336291, + "imag": -15.659291332109987, + "__extended_json_type__": "complex128" + } + ], + [ + "get_cexpr_test benchmark_eval_cexpr check_ama get_data_sig is_cython=True", + { + "real": 11.492942956036883, + "imag": -13.926566927529471, + "__extended_json_type__": "complex128" + } + ] +] \ No newline at end of file diff --git a/examples-py/auto-contract-01.py b/examples-py/auto-contract-01.py new file mode 100755 index 000000000..45af3ad66 --- /dev/null +++ b/examples-py/auto-contract-01.py @@ -0,0 +1,95 @@ +#!/usr/bin/env python3 + +json_results = [] +check_eps = 1e-14 + +def json_results_append(*args): + json_results.append(args) + +import sys +import numpy as np +import qlat as q + +import auto_contractor as qac + +size_node_list = [ + [1, 1, 1, 1], + [1, 1, 1, 2], + [1, 1, 1, 4], + [1, 1, 1, 8], + [2, 2, 2, 2], + [2, 2, 2, 4], + ] + +q.begin_with_mpi(size_node_list) + +f = "u" # flavor, can be "u", "d", "s", "c" + +p1 = "x1" +s1 = "s1" +c1 = "c1" + +p2 = "x2" +s2 = "s2" +c2 = "c2" + +q1v = qac.Qv(f, p1, s1, c1) +q1b = qac.Qb(f, p1, s1, c1) +q2v = qac.Qv(f, p2, s2, c2) +q2b = qac.Qb(f, p2, s2, c2) + +exprs = [ + qac.mk_expr(1), + q1v * q2b, + q1v * q2b + q2b * q1v, +] + +for expr in exprs: + json_results_append(str(expr)) + +for expr in qac.contract_simplify(*exprs): + json_results_append(str(expr)) + +cexpr = qac.contract_simplify_compile(*exprs) + +json_results_append( + qac.display_cexpr(cexpr) +) + +for v in cexpr.list(): + json_results_append(str(v)) + +cexpr_opt = cexpr.copy() +cexpr_opt.optimize() + +json_results_append( + qac.display_cexpr(cexpr_opt) +) + +for v in cexpr_opt.list(): + json_results_append(str(v)) + +for is_cython in [ False, True, ]: + json_results_append( + qac.cexpr_code_gen_py(cexpr_opt, is_cython=is_cython) + ) + @q.timer + def get_cexpr_test(): + fn_base = f"cache/auto_contract_cexpr/get_cexpr_test" + def calc_cexpr(): + cexpr = qac.contract_simplify_compile(*exprs, is_isospin_symmetric_limit=True) + return cexpr + return qac.cache_compiled_cexpr(calc_cexpr, fn_base, is_cython=is_cython) + + cexpr = get_cexpr_test() + check, check_ama = qac.benchmark_eval_cexpr(cexpr) + json_results_append(f"get_cexpr_test benchmark_eval_cexpr check get_data_sig is_cython={is_cython}", q.get_data_sig(np.array(check, dtype=np.complex128), q.RngState())) + json_results_append(f"get_cexpr_test benchmark_eval_cexpr check_ama get_data_sig is_cython={is_cython}", q.get_data_sig(np.array(check_ama, dtype=np.complex128), q.RngState())) + +q.check_log_json(__file__, json_results, check_eps=check_eps) + +q.timer_display() + +q.end_with_mpi() + +q.displayln_info(f"CHECK: finished successfully.") diff --git a/examples-py/auto-contract.log.json b/examples-py/auto-contract.log.json deleted file mode 100644 index 0637a088a..000000000 --- a/examples-py/auto-contract.log.json +++ /dev/null @@ -1 +0,0 @@ -[] \ No newline at end of file diff --git a/examples-py/auto-contract.py b/examples-py/auto-contract.py deleted file mode 100755 index 8c42ef0ea..000000000 --- a/examples-py/auto-contract.py +++ /dev/null @@ -1,33 +0,0 @@ -#!/usr/bin/env python3 - -json_results = [] -check_eps = 1e-14 - -def json_results_append(*args): - json_results.append(args) - -import sys -import qlat as q - -import auto_contractor as qac - -size_node_list = [ - [1, 1, 1, 1], - [1, 1, 1, 2], - [1, 1, 1, 4], - [1, 1, 1, 8], - [2, 2, 2, 2], - [2, 2, 2, 4], - ] - -q.begin_with_mpi(size_node_list) - - - -q.check_log_json(__file__, json_results, check_eps=check_eps) - -q.timer_display() - -q.end_with_mpi() - -q.displayln_info(f"CHECK: finished successfully.")