schubmult 2.0.2__py3-none-any.whl → 2.0.4__py3-none-any.whl

schubmult/schubmult_q/__main__.py

@@ -1,5 +1,6 @@
 import sys
+
 from ._script import main
 
 if __name__ == "__main__":
-    sys.exit(main())
+    sys.exit(main(sys.argv))

schubmult/schubmult_q/_funcs.py

@@ -1,54 +1,39 @@
-from ._vars import (
-    var_x,
-)
-from symengine import Add, Mul, Pow
+from functools import cached_property
+
+from symengine import Add, Mul, Pow, symarray
+
 from schubmult.perm_lib import (
-    elem_sym_perms_q,
     add_perm_dict,
+    code,
     compute_vpathdicts,
+    double_elem_sym_q,
+    elem_sym_perms_q,
+    inv,
     inverse,
-    strict_theta,
     medium_theta,
-    permtrim,
-    inv,
     mulperm,
-    code,
+    permtrim,
+    strict_theta,
     uncode,
-    double_elem_sym_q,
-    q_var,
 )
 
-# from symengine import sympify, Add, Mul, Pow, symarray, Symbol, expand
-# from schubmult._base_argparse import schub_argparse
-# from schubmult.perm_lib import (
-#     trimcode,
-#     elem_sym_perms_q,
-#     add_perm_dict,
-#     compute_vpathdicts,
-#     inverse,
-#     strict_theta,
-#     medium_theta,
-#     permtrim,
-#     inv,
-#     mulperm,
-#     code,
-#     uncode,
-#     double_elem_sym_q,
-#     longest_element,
-#     check_blocks,
-#     is_parabolic,
-#     q_vector,
-#     omega,
-#     count_less_than,
-#     q_var,
-#     sg,
-#     n,
-# )
-# import numpy as np
-# from schubmult.schubmult_q_double import factor_out_q_keep_factored
 
+class _gvars:
+    @cached_property
+    def n(self):
+        return 100
+
+    @cached_property
+    def var_x(self):
+        return tuple(symarray("x", self.n).tolist())
+
+    @cached_property
+    def q_var(self):
+        return tuple(symarray("q", self.n).tolist())
 
-def single_variable(coeff_dict, varnum, var_q=q_var):
+_vars = _gvars()
+
+def single_variable(coeff_dict, varnum, var_q=_vars.q_var):
     ret = {}
     for u in coeff_dict:
         new_perms_k = elem_sym_perms_q(u, 1, varnum, var_q)
@@ -64,33 +49,32 @@ def single_variable(coeff_dict, varnum, var_q=q_var):
     return ret
 
 
-def mult_poly(coeff_dict, poly, var_x=var_x, var_q=q_var):
+def mult_poly(coeff_dict, poly, var_x=_vars.var_x, var_q=_vars.q_var):
     if poly in var_x:
         return single_variable(coeff_dict, var_x.index(poly), var_q=var_q)
-    elif isinstance(poly, Mul):
+    if isinstance(poly, Mul):
         ret = coeff_dict
         for a in poly.args:
             ret = mult_poly(ret, a, var_x, var_q=var_q)
         return ret
-    elif isinstance(poly, Pow):
+    if isinstance(poly, Pow):
         base = poly.args[0]
         exponent = int(poly.args[1])
         ret = coeff_dict
         for i in range(int(exponent)):
             ret = mult_poly(ret, base, var_x, var_q=var_q)
         return ret
-    elif isinstance(poly, Add):
+    if isinstance(poly, Add):
         ret = {}
         for a in poly.args:
             ret = add_perm_dict(ret, mult_poly(coeff_dict, a, var_x, var_q=var_q))
         return ret
-    else:
-        ret = {}
-        for perm in coeff_dict:
-            ret[perm] = poly * coeff_dict[perm]
-        return ret
+    ret = {}
+    for perm in coeff_dict:
+        ret[perm] = poly * coeff_dict[perm]
+    return ret
 
-def schubmult_db(perm_dict, v, q_var=q_var):
+def schubmult_db(perm_dict, v, q_var=_vars.q_var):
     if v == (1, 2):
         return perm_dict
     th = medium_theta(inverse(v))
@@ -118,14 +102,12 @@ def schubmult_db(perm_dict, v, q_var=q_var):
             mx_th = 0
             for vp in vpathdicts[index]:
                 for v2, vdiff, s in vpathdicts[index][vp]:
-                    if th[index] - vdiff > mx_th:
-                        mx_th = th[index] - vdiff
+                    mx_th = max(mx_th, th[index] - vdiff)
             if index < len(th) - 1 and th[index] == th[index + 1]:
                 mx_th1 = 0
                 for vp in vpathdicts[index + 1]:
                     for v2, vdiff, s in vpathdicts[index + 1][vp]:
-                        if th[index + 1] - vdiff > mx_th1:
-                            mx_th1 = th[index + 1] - vdiff
+                        mx_th1 = max(mx_th1, th[index + 1] - vdiff)
                 newpathsums = {}
                 for up in vpathsums:
                     newpathsums0 = {}
@@ -142,7 +124,7 @@ def schubmult_db(perm_dict, v, q_var=q_var):
                                 if udiff1 + vdiff2 == th[index]:
                                     newpathsums0[(up1, udiff1, mul_val1)][v2] = (
                                         newpathsums0[(up1, udiff1, mul_val1)].get(
-                                            v2, 0
+                                            v2, 0,
                                         )
                                         + s2 * sumval * mul_val1
                                     )
@@ -189,7 +171,7 @@ def schubmult_db(perm_dict, v, q_var=q_var):
             vpathsums = newpathsums
         toget = tuple(vmu)
         ret_dict = add_perm_dict(
-            {ep: vpathsums[ep].get(toget, 0) for ep in vpathsums}, ret_dict
+            {ep: vpathsums[ep].get(toget, 0) for ep in vpathsums}, ret_dict,
         )
     return ret_dict
 
@@ -214,13 +196,12 @@ def schubmult(perm_dict, v):
             mx_th = 0
             for vp in vpathdicts[index]:
                 for v2, vdiff, s in vpathdicts[index][vp]:
-                    if th[index] - vdiff > mx_th:
-                        mx_th = th[index] - vdiff
+                    mx_th = max(mx_th, th[index] - vdiff)
             newpathsums = {}
             for up in vpathsums:
                 inv_up = inv(up)
                 newperms = elem_sym_perms_q(
-                    up, min(mx_th, (inv_mu - (inv_up - inv_u)) - inv_vmu), th[index]
+                    up, min(mx_th, (inv_mu - (inv_up - inv_u)) - inv_vmu), th[index],
                 )
                 for up2, udiff, mul_val in newperms:
                     if up2 not in newpathsums:
@@ -238,7 +219,7 @@ def schubmult(perm_dict, v):
             vpathsums = newpathsums
         toget = tuple(vmu)
         ret_dict = add_perm_dict(
-            {ep: vpathsums[ep].get(toget, 0) for ep in vpathsums}, ret_dict
+            {ep: vpathsums[ep].get(toget, 0) for ep in vpathsums}, ret_dict,
         )
     return ret_dict
 

schubmult/schubmult_q/_script.py

@@ -1,30 +1,33 @@
-from ._funcs import (
-    schubmult,
-    schubmult_db,
-    mult_poly,
-)
+import sys
+
+import numpy as np
 from symengine import sympify
+
 from schubmult._base_argparse import schub_argparse
 from schubmult.perm_lib import (
-    trimcode,
-    permtrim,
-    inv,
-    mulperm,
-    uncode,
-    longest_element,
     check_blocks,
+    count_less_than,
+    inv,
     is_parabolic,
-    q_vector,
+    longest_element,
+    mulperm,
     omega,
-    count_less_than,
+    permtrim,
     q_var,
+    q_vector,
     sg,
+    trimcode,
+    uncode,
+)
+from schubmult.schubmult_q._funcs import (
+    schubmult,
+    schubmult_db,
 )
-import numpy as np
 from schubmult.schubmult_q_double import factor_out_q_keep_factored
 
 
 def _display_full(coeff_dict, args, formatter):
+    raw_result_dict = {}
     ascode = args.ascode
     parabolic_index = [int(s) for s in args.parabolic]
     parabolic = len(parabolic_index) != 0
@@ -63,7 +66,7 @@ def _display_full(coeff_dict, args, formatter):
                         q_var[index + 1 - count_less_than(parabolic_index, index + 1)] ** qv[index]
                         for index in range(len(qv))
                         if index + 1 not in parabolic_index
-                    ]
+                    ],
                 )
 
                 try:
@@ -81,25 +84,28 @@ def _display_full(coeff_dict, args, formatter):
         val = sympify(coeff_dict[perm]).expand()
         if val != 0:
             if ascode:
-                print(f"{str(trimcode(perm))}  {formatter(val)}")
+                raw_result_dict[tuple(trimcode(perm))] = val
+                if formatter:
+                    print(f"{trimcode(perm)!s}  {formatter(val)}")
             else:
-                print(f"{str(perm)}  {formatter(val)}")
+                raw_result_dict[perm] = val
+                if formatter:
+                    print(f"{perm!s}  {formatter(val)}")
+    return raw_result_dict
 
 
-def main():
+def main(argv=None):
+    if argv is None:
+        argv = sys.argv
     try:
         args, formatter = schub_argparse(
             "schubmult_q",
             "Compute products of quantum Schubert polynomials",
+            argv=argv[1:],
             quantum=True,
         )
 
-        mulstring = ""
 
-        mult = False
-        if args.mult is not None:
-            mult = True
-            mulstring = " ".join(args.mult)
 
         perms = args.perms
 
@@ -130,7 +136,7 @@ def main():
                 index = parabolic_index[i] - 1
                 if sg(index, perms[0]) == 1 or sg(index, perms[1]) == 1:
                     print(
-                        "Parabolic given but elements are not minimal length coset representatives."
+                        "Parabolic given but elements are not minimal length coset representatives.",
                     )
                     exit(1)
 
@@ -143,15 +149,18 @@ def main():
             for perm in perms[1:]:
                 coeff_dict = schubmult(coeff_dict, tuple(permtrim([*perm])))
 
-        if mult:
-            mul_exp = sympify(mulstring)
-            coeff_dict = mult_poly(coeff_dict, mul_exp)
+        # if mult:
+        #     mul_exp = sympify(mulstring)
+        #     coeff_dict = mult_poly(coeff_dict, mul_exp)
 
-        if pr:
-            _display_full(coeff_dict, args, formatter)
+        if pr or formatter is None:
+            raw_result_dict = _display_full(coeff_dict, args, formatter)
+        if formatter is None:
+            return raw_result_dict
     except BrokenPipeError:
         pass
 
 
 if __name__ == "__main__":
-    main()
+    import sys
+    sys.exit(main(sys.argv))

schubmult/schubmult_q_double/__init__.py

@@ -1,16 +1,10 @@
-from ._funcs import (
-    schubmult,
-    schubmult_db,
-    mult_poly,
-    factor_out_q_keep_factored,
-    schubpoly_quantum
-)
-
+from ._funcs import factor_out_q_keep_factored, mult_poly, schubmult, schubmult_db, schubpoly_quantum
 
 __all__ = [
+    "factor_out_q_keep_factored",
+    "mult_poly",
     "schubmult",
     "schubmult_db",
-    "mult_poly",
-    "factor_out_q_keep_factored",
-    "schubpoly_quantum"
+    "schubpoly_quantum",
 ]
+

schubmult/schubmult_q_double/__main__.py

@@ -1,5 +1,6 @@
 import sys
+
 from ._script import main
 
 if __name__ == "__main__":
-    sys.exit(main())
+    sys.exit(main(sys.argv))