witten.py

r"""
Code for computations involving Witten's r-spin class.
Authors: Felix Janda (main author), Aaron Pixton (code improvements), Johannes Schmitt (integration into admcycles)
"""

from __future__ import absolute_import, print_function

import itertools

from six.moves import range

from admcycles.admcycles import Tautv_to_tautclass, tautclass
from admcycles.DR import *

# from sage.combinat.subset import Subsets
from sage.arith.all import factorial, lcm
from sage.functions.other import floor, ceil
from sage.misc.misc_c import prod
from sage.rings.all import PolynomialRing, QQ, ZZ
from sage.modules.free_module_element import vector
# from sage.rings.polynomial.polynomial_ring import polygen
from sage.rings.polynomial.multi_polynomial_element import MPolynomial
from collections import Iterable
from sage.rings.integer import Integer
from sage.calculus.var import var
from sage.misc.functional import symbolic_sum


# Computation of Witten's rspin class for large r

def rspin_leg_factor(d,a):
  if a < 0: a = X + a
  return Pmpolynomial(d).subs(a=a)

def rspin_edge_factor(w1,m,d1,d2):
  R=PolynomialRing(QQ,1,'X')
  X = R.gens()[0]
  d = d1+d2+1
  S = 0
  for i in range(d+1):
    S += R(rspin_leg_factor(i, (w1 + i) % (m-1))(X=m)*rspin_leg_factor(d-i, (m-2-i-w1) % (m-1))(X=m)*X**i)
  S /= X+1
  assert(S.denominator() == 1)
  S = S.numerator()
  #print(-S[d1])
  return -S[d1]

def rspin_coeff_setup(num,g,r,n=0,dvector=(),moduli_type=MODULI_ST):
  markings = tuple(range(1,n+1))
  G = single_stratum(num,g,r,markings,moduli_type)
  nr = G.M.nrows()
  nc = G.M.ncols()
  edge_list = []
  exp_list = []
  scalar_factor = 1/autom_count(num,g,r,markings,moduli_type)
  given_weights = [1 - G.M[i+1,0][0] for i in range(nr-1)]
  for i in range(1,nr):
    for j in range(1,G.M[i,0].degree()+1):
      scalar_factor /= factorial(G.M[i,0][j])
      scalar_factor *= (-1)**G.M[i,0][j]
      scalar_factor *= (rspin_leg_factor(j, 0))**(G.M[i,0][j])
      given_weights[i-1] -= j*G.M[i,0][j]
  for j in range(1,nc):
    ilist = [i for i in range(1,nr) if G.M[i,j] != 0]
    if G.M[0,j] == 0:
      if len(ilist) == 1:
        i1 = ilist[0]
        i2 = ilist[0]
        exp1 = G.M[i1,j][1]
        exp2 = G.M[i1,j][2]
      else:
        i1 = ilist[0]
        i2 = ilist[1]
        exp1 = G.M[i1,j][1]
        exp2 = G.M[i2,j][1]
      edge_list.append([i1-1,i2-1])
      exp_list.append(exp1)
      exp_list.append(exp2)
    else:
      exp1 = G.M[ilist[0],j][1]
      scalar_factor *= rspin_leg_factor(exp1, dvector[G.M[0,j][0]-1])
      given_weights[ilist[0] - 1] += dvector[G.M[0,j][0] - 1] - exp1
  return edge_list,exp_list,given_weights,scalar_factor

def rspin_coeff(num,g,r,n=0,dvector=(),r_coeff=None,step=1,m0given=-1,deggiven=-1,moduli_type=MODULI_ST):
  markings = tuple(range(1,n+1))
  G = single_stratum(num,g,r,markings,moduli_type)
  nr = G.M.nrows()
  nc = G.M.ncols()
  edge_list,exp_list,given_weights,scalar_factor = rspin_coeff_setup(num,g,r,n,dvector,moduli_type)
  if m0given == -1:
    m0 = (ceil(sum([abs(i.subs(X=0)) for i in dvector])/2) + g*1+1)*step
  else:
    m0 = m0given
  h0 = nc - nr - n + 1
  if deggiven == -1:
    deg = 2*sum(exp_list) + 2*len(edge_list)
  else:
    deg = deggiven
  if r_coeff is None:
    mrange = list(range(m0 + step, m0 + step*deg + step + 1, step))
  else:
    mrange = [r_coeff+1]  # just evaluate at a single value m = r_coeff  
  mvalues = []
  for m in mrange:
    given_weights_m = [ZZ(i.subs(X=m)) for i in given_weights]
    total = 0
    for weight_data in itertools.product(*[list(range(m-1)) for i in range(len(edge_list))]):
      vertex_weights = copy(given_weights_m)
      for i in range(len(edge_list)):
        vertex_weights[edge_list[i][0]] += weight_data[i]
        vertex_weights[edge_list[i][1]] -= weight_data[i]+2+exp_list[2*i]+exp_list[2*i+1]
      if len([i for i in vertex_weights if i % (m-1) != 0]) > 0:
        continue
      term = 1
      for i in range(len(edge_list)):
        term *= rspin_edge_factor(weight_data[i], m, exp_list[2*i], exp_list[2*i+1])
      total += term
    if r_coeff is None:
      mvalues.append(total*ZZ(m-1)**(-h0))
    else:
      #print(m-1)
      mvalues.append(total*ZZ(m-1)**(-r-h0)) # undo the rescaling by r**degree
  #print(mrange,mvalues, scalar_factor, r, h0)
  mpoly = ZZ(-1)**(r-g)*(interpolate(mrange, mvalues)*scalar_factor).simplify_rational()
  if r_coeff is not None:
    mpoly = QQ(mpoly.subs(X=r_coeff))
  return mpoly

def rspin_compute(g,r,n=0,dvector=(),r_coeff=None,step=1,m0=-1,deg=-1,moduli_type=MODULI_ST):
  answer = []
  markings = tuple(range(1,n+1))
  for i in range(num_strata(g,r,markings,moduli_type)):
    answer.append(rspin_coeff(i,g,r,n,dvector,r_coeff,step,m0,deg,moduli_type))
  return vector(answer)

def rspin_degree_test(g,r,n=0,dvector=(),step=1,m0=-1,deg=-1,moduli_type=MODULI_ST):
  answer = []
  markings = tuple(range(1,n+1))
  for i in range(num_strata(g,r,markings,moduli_type)):
    answer.append(rspin_coeff(i,g,r,n,dvector,step,m0,deg,moduli_type).degree(X))
  return answer

def rspin_constant(g,r,n=0,dvector=(),step=1,m0=-1,deg=-1,moduli_type=MODULI_ST):
  answer = []
  markings = tuple(range(1,n+1))
  for i in range(num_strata(g,r,markings,moduli_type)):
    answer.append(rspin_coeff(i,g,r,n,dvector,step,m0,deg,moduli_type).subs(X=0))
  return vector(answer)

def Wittenrspin(g, Avector, r_coeff = None, d = None, rpoly = False):
  r"""
  Returns the polynomial limit of Witten's r-spin class in genus g with input Avector,
  as discussed in the appendix of [Pandharipande-Pixton-Zvonkine '16].
  
  More precisely, Avector is expected to be a vector of linear polynomials in QQ[r],
  with leading coefficients being rational numbers in the interval [0,1] and constant
  coefficients being integers. Elements of Avector should sum to C * r + 2g-2 for some 
  nonnegative integer C. Then Wittenrspin returns the tautological class obtained as
  the limit of
  
  r^(g-1+C) W_{g,n}^r(Avector)
  
  for r >> 0 sufficiently large and divisible, evaluated at r=0.
  
  INPUT::

  - ``g`` -- integer; underlying genus

  - ``Avector``   -- tuple ; a tuple of either integers, elements of a polynomial
    ring QQ[r] in some variable r or tuples (C_i, D_i) which are interpreted as
    polynomials C_i * r + D_i. For the entries with C_i = 1 we assume D_i <=-2.

  - ``r_coeff``  -- integer or None (default: `None`); if a particular integer 
    r_coeff = r is specified, the function will return the (unscaled) class W_{g,n}^r(Avector),
    not taking a limit for large r. 

  - ``d`` -- integer; desired degree in tautological ring; will be set to g-1+C by 
    default

  - ``rpoly``  -- bool (default: `False`); if True, return the limit of 
    r^(g-1+C) W_{g,n}^r(Avector) without evaluating at r=0, as a tautclass with 
    coefficients being polynomials in r
  
  EXAMPLES:
  
  We start by verifying the conjecture from the appendix of [Pandharipande-Pixton-Zvonkine '16]
  for g = 2 and mu = (2)::
  
    sage: from admcycles import Wittenrspin, Strataclass
    sage: H1 = Wittenrspin(2, (2,))
    sage: H2 = Strataclass(2, 1, (2,))
    sage: (H1-H2).is_zero()
    True

  We can also verify a new conjecture for classes of strata of meromorphic differentials for 
  g = 1 and mu = (3,-1,-2). The argument (1/2,-1) stands for an insertion 1/2 * r -1::
  
    sage: H1 = Wittenrspin(1, (3, (1/2,-1), (1/2,-2)))
    sage: H2 = Strataclass(1, 1, (3,-1,-2))
    sage: (H1+H2).is_zero()
    True

  As a variant of this, we also verify that insertions r-b stand for poles of order b with
  vanishing residues::
  
    sage: R.<r> = PolynomialRing(QQ,1)
    sage: H1 = Wittenrspin(1, (5,r-2, r-3))
    sage: H2 = Strataclass(1, 1, (5,-2,-3), res_cond=(2,))
    sage: (H1+H2).is_zero()
    True   

  We can also compute the (scaled) Witten's class without substituting r=0::
  
    sage: Wittenrspin(1,(2,r-2),rpoly=True)
    Graph :      [1] [[1, 2]] []
    Polynomial : (1/12*r^2 - 5/24*r + 1/12)*(kappa_1^1 )_0 
    <BLANKLINE>
    Graph :      [1] [[1, 2]] []
    Polynomial : (-1/12*r^2 + 29/24*r - 37/12)*psi_1^1 
    <BLANKLINE>
    Graph :      [1] [[1, 2]] []
    Polynomial : (-1/12*r^2 + 17/24*r - 13/12)*psi_2^1 
    <BLANKLINE>
    Graph :      [0, 1] [[1, 2, 4], [5]] [(4, 5)]
    Polynomial : (1/12*r^2 - 17/24*r + 13/12)*
    <BLANKLINE>
    Graph :      [0] [[4, 5, 1, 2]] [(4, 5)]
    Polynomial : (-1/48*r + 1/24)*
  
  Instead of calculating the asymptotic, polynomial behaviour of Witten's class, 
  we can also input a concrete value for r, using the option r_coeff. Below we
  verify the CohFT property of Witten's class::
  
    sage: R.<r> = PolynomialRing(QQ,1)
    sage: H1 = Wittenrspin(1, (5,r-2, r-3))
    sage: H2 = Strataclass(1, 1, (5,-2,-3), res_cond=(2,))
    sage: (H1+H2).is_zero()
    True       
  """
  n = len(Avector)
  
  if r_coeff is None:
    polyentries = [a for a in Avector if isinstance(a,MPolynomial)]
    if len(polyentries) > 0:
      R = polyentries[0].parent()
      r = R.gens()[0]
    else:
      R = PolynomialRing(QQ,1,'r')
      r = R.gens()[0]
    
    X = SR.var('X')
    AvectorX = []
    Cvector = []
    Dvector = []
    
    # Extract entries of Avector into a unified format (AvectorX)
    # Collect linear and constant coefficients of entries of AvectorX in Cvector, Dvector
    for a in Avector:
      if isinstance(a,Integer):
        AvectorX.append(a)
        Cvector.append(QQ(0))
        Dvector.append(a)
      elif isinstance(a,MPolynomial):
        AvectorX.append(a[1]*X+a[0])
        Cvector.append(QQ(a[1]))
        Dvector.append(a[0])
      elif isinstance(a,Iterable):
        AvectorX.append(a[0]*X+a[1])
        Cvector.append(QQ(a[0]))
        Dvector.append(a[1])
      else:
        raise ValueError('Entries of Avector must be Integers, Polynomials or tuples (C_i, D_i)')
    
    # For C_i = 1, D_i = -1 we require a simple pole with vanishing residue => return zero class
    if any((Cvector[i]==1) and (Dvector[i]==-1) for i in range(n)):
      return tautclass([])
    
    step = lcm(c.denom() for c in Cvector)
    C = sum(Cvector)
    assert C in ZZ
    assert sum(Dvector) == 2*g-2
  else:
    AvectorX=Avector
    step=0    
  
  if d is None:
    if r_coeff is None:
      d = ZZ(g-1+C)
    else:
      denom = ZZ((r_coeff-2)*(g-1) + sum(Avector))
      if denom % r_coeff != 0:
        return tautclass([]) # Witten's class vanishes unless this congruence is satisfied
      else:  
        d = ZZ(denom/ZZ(r_coeff))
  rvect = rspin_compute(g,d,n,tuple(AvectorX),r_coeff,step,m0=-1,deg=-1,moduli_type=MODULI_ST)
  if not rpoly:
    rvect = rvect.subs(X=0)
  rvect = convert_vector_to_monomial_basis(rvect, g, d, tuple(range(1, n+1)), MODULI_ST)
  
  if rpoly:
    rvect = vector((b.subs(X=r) for b in rvect))
  
  return Tautv_to_tautclass(rvect, g, n, d)  

@cached_function
def Pmpolynomial(m):
  r"""
  Returns the expression P_m(X,a) as defined in [Pandharipande-Pixton-Zvonkine '16, Section 4.5].
  
  TESTS::

    sage: from admcycles.witten import Pmpolynomial
    sage: Pmpolynomial(0)
    1  
    sage: Pmpolynomial(1)
    -1/12*X^2 + 1/2*(X - 1)*a - 1/2*a^2 + 5/24*X - 1/12
    sage: Pmpolynomial(2)
    1/288*X^4 - 1/12*(5*X - 1)*a^3 + 1/8*a^4 + 7/288*X^3 + 1/48*(20*X^2 - 5*X - 4)*a^2 - 29/384*X^2 - 1/48*(6*X^3 + X^2 - 9*X + 2)*a + 7/288*X + 1/288
  """
  (b,X,a) = var('b,X,a')
  if m==0:
    return 1+0*X
  return ( QQ(1)/2 * symbolic_sum((2*m*X-X-2*b)*Pmpolynomial(m-1).subs(a=b-1),b,1,a) - QQ(1)/(4*m*X*(X-1)) * symbolic_sum((X-1-b)*(2*m*X-b)*(2*m*X-X-2*b)*Pmpolynomial(m-1).subs(a=b-1),b,1,X-2) ).simplify_rational()