<jupyter_start><jupyter_text>Symbolic computation of the modified quantities from  "Structure preserving discretization of time-reparametrized Hamiltonian systems with application to nonholonomic mechanics"  by <NAME> and <NAME>.Written in SageMath 8.1 https://www.sagemath.org.<jupyter_code>### Initialization and auxiliary methods ###

function('x')
function('y')
var('t,h,Ed')

# Maximum order of derivative to consider.
# Values up to 7 are implemented.
maxorder = 7

# Display elapsed time
showtiming = True

# Define x,y and their derivatives as algebraic variables (jet bundle coordinates)
var('x0', latex_name='x')
var('y0', latex_name='y')
for k in [0..maxorder]:
    var('x' + str(k))
    var('y' + str(k))
    def xder(k):
        return var('x' + str(k))
    def yder(k):
        return var('y' + str(k))

# Replace derivatives of the functions x,y by the jet bundle coordinates defined above
def dertovar(func,order):
    for k in [0..order+1]:
        func = expand(func).subs( {x(t).diff(t,k): xder(k), y(t).diff(t,k): yder(k)})
    return func

# Replace jet bundle coordinates by derivatives of the functions x,y
def vartoder(func,order):
    func = func.subs( {x0: x(t), y0: y(t)})
    for k in [1..order+1]:
        func = func.subs( {xder(k): x(t).diff(t,k), yder(k): y(t).diff(t,k)})
    return func

# Derivative acting on the jet bundle
def vdiff(func,order,*var):
    func = vartoder(func,order+1)
    func = diff(func,*var)
    func = dertovar(func,order+1)
    return func

# Sustitute expression in a truncated power series (polynomial) in h.
# Terms of higher order are systematically deleted, unlike with the subs() method.
def series_subs(expr,replace,order):
    expr = expr.series(h,order).truncate()
    out = 0
    for i in [0..order-1]:
        replace = [ r.lhs() == r.rhs().series(h,order-i).truncate() for r in replace]
        new = expr.coefficient(h,i).subs(replace).series(h,order-i).truncate()
        out += (h^i*new).series(h,order).truncate() #######
    return out.series(h,order).truncate() ########

# Solve a truncated power series equation up to a desired order.
# Proceeds iteratively through the orders.
def series_solve(eqn,order, *x):
    if type(eqn) == type([]): # If list of equations
        d = len(eqn)
        newx = var('newx',n=d)
        leading = [e.series(h,1).truncate() for e in eqn]
        higher = [(e.series(h,order).truncate() - e.series(h,1).truncate()).collect(h) for e in eqn]
        sol = [y.subs(solve([e == 0 for e in leading], *x)) for y in x]
        for i in [1..order-1]:
            coeffeqn = [leading[k].subs([x[j] == sol[j] + h^i*newx[j] for j in [0..d-1]])
                        + series_subs( higher[k] , [x[l] == sol[l] for l in [0..d-1]] , i+1) for k in [0..d-1]]
            coeffeqn = [e.series(h,i+1).coefficient(h,i) for e in coeffeqn]
            sol = [sol[k] + h^i*( newx[k].subs(solve([e == 0 for e in coeffeqn], newx)).simplify_full() ) for k in [0..d-1]]
    else: # If single equation
        if type(eqn) == type(1):
            print("\n HELP! \n")
            print(eqn)
        newx = var('newx')
        leading =  eqn.series(h,1).truncate()
        higher = (eqn.series(h,order).truncate() - leading).collect(h)
        sol = x[0].subs(solve(leading == 0, *x))
        for i in [1..order-1]:
            coeffeqn = leading.subs({x[0]: sol + h^i*newx}) + series_subs( higher , [x[0] == sol], i+1)
            coeffeqn = coeffeqn.series(h,i+1).truncate()
            sol = sol + h^i*( newx.subs(solve(coeffeqn == 0, newx)).expand() )
    return sol

# Convert output to python syntax
def convert_py(string):
    string = string.replace("^","**")
    string = string.replace("sqrt","np.sqrt")
    string = string.replace("x0","x")
    string = string.replace("x1","dx")
    string = string.replace("y0","y")
    string = string.replace("y1","dy")
    return string
### Calculate the modified Lagrangian corresponding to a discrete Lagrangian Ldisc ###
#
# Using the algorithm described in
# <NAME>. Modified equations for variational integrators. Numerische Mathematik 137:1001, 2017.

# Taylor expansion of the discrete Lagrangian
def Lseries(Ldisc,order): #,mid=False):
    return Ldisc(x(t-h/2),x(t+h/2),y(t-h/2),y(t+h/2)).series(h,order).truncate()

# Use Euler-Maclaurin formula to turn action sum into integral
# WARNING: only implemented for orders less than 8.
def EMcorrection(func,order):
    series = func - h^2/24 * diff(func.series(h,order-2).truncate(),t,2) 
    if order > 4:
        series += 7*h^4/5760 * diff(func.series(h,order-4).truncate(),t,4)
    if order > 6:
        series += -31*h^6/967680 * diff(func.series(h,order-6).truncate(),t,6)
    return dertovar(series,order)

# Calculate higher order EL equations
def EL(L,order):
    ELlst_ders = []
    for locorder in [1..order]:
        pxseries = series_subs( diff(L,x0) - vdiff(diff(L,x1),locorder,t) , ELlst_ders , locorder )
        pyseries = series_subs( diff(L,y0) - vdiff(diff(L,y1),locorder,t) , ELlst_ders , locorder )
        leadingsol = series_solve( [pxseries,pyseries] , locorder, x2,y2 )
        ELlst = [x2 == leadingsol[0], y2 == leadingsol[1]]
        for k in [3..locorder]:
            ELlst += [ xder(k) == series_subs( vdiff(ELlst[-2].rhs(),order,t) , ELlst , order+1-k ).subs(ELlst) ]
            ELlst += [ yder(k) == series_subs( vdiff(ELlst[-2].rhs(),order,t) , ELlst , order+1-k ).subs(ELlst) ]
        ELlst_ders = ELlst[2:]
    return ELlst

# Calculate modified Lagrangian
def modLag(Ldisc,order,mid=False):
    series1 = Lseries(Ldisc,order)
    series2 = EMcorrection(series1,order)
    series3 = series_subs(series2,EL(series2,order-2),order) #Substitue higher derivatives using EL equations
    return series3

# Calculate energy in tangent bundle coordinates
def energy(L):
    px = diff(L,x1)
    py = diff(L,y1)
    return expand(px*x1 + py*y1 - L).collect(h)

# Calculate Hamiltonian (energy in cotangent bundle coordinates)
var("px",latex_name="p_x")
var("py",latex_name="p_y")
def hamiltonian(L,order):
    w = walltime()
    legendre = series_solve([px - diff(L,x1),py - diff(L,y1)],order,x1,y1)
    if showtiming: print(walltime(w))
    legendre = [x1 == legendre[0].combine(), y1 == legendre[1].combine()]
    if showtiming: print(walltime(w))
    out = series_subs( x1*px + y1*py - L, legendre, order)
    if showtiming: print(walltime(w))
    return out
### Specify conformal factor and Lagrangian ###
# The altered Lagrangian is then N*(L+E)

def N(y):
    return 1/sqrt(1+y^2)
def L(x,y,vx,vy):
    return 1/2*(1+y^2)*vx^2 + 1/2*vy^2 - 1/2*(x^2+y^2)
### 5 discretizations of the Lagrangian ###
# The parameter E (Ed) is left unspecified
def LdiscMM(x0,x1,y0,y1):
    return N( (y0+y1)/2 ) * ( L( (x0+x1)/2, (y0+y1)/2, (x1-x0)/h, (y1-y0)/h ) + Ed )

def LdiscTM(x0,x1,y0,y1):
    return (N(y0)+N(y1))/2 * ( L( (x0+x1)/2, (y0+y1)/2, (x1-x0)/h, (y1-y0)/h ) + Ed )

def LdiscMT(x0,x1,y0,y1):
    return N( (y0+y1)/2 ) * ( L(x0, y0, (x1-x0)/h, (y1-y0)/h )/2 + L(x1, y1, (x1-x0)/h, (y1-y0)/h )/2 + Ed )

def LdiscTT(x0,x1,y0,y1):
    return (N(y0)+N(y1))/2 * ( L(x0, y0, (x1-x0)/h, (y1-y0)/h )/2 + L(x1, y1, (x1-x0)/h, (y1-y0)/h )/2 + Ed )

def LdiscT(x0,x1,y0,y1):
    return N(y0)*(L(x0, y0, (x1-x0)/h, (y1-y0)/h ) + Ed)/2 + N(y1)*(L(x1, y1, (x1-x0)/h, (y1-y0)/h ) + Ed)/2

# To track elapsed time
w = walltime()

# Order to which to calculate modified quantities.
# Must not be larger than maxorder
order = 7

# Parallel computation setup
import os
os.environ["SAGE_NUM_THREADS"] = "8" #Indicate max number of cpu cores to be used

# Compute the 5 modified Lagrangians
@parallel
def paramodLag(lag):
    return modLag(lag[1],order)
gen = paramodLag([[1,LdiscMM]])#,[2,LdiscTM],[3,LdiscMT],[4,LdiscTT],[5,LdiscT]])
methodnames = ["MM"]#,"TM","MT","TT","T"]
Lmod = [0 for i in methodnames]
for i in gen:
    index = i[0][0][0][0]
    lag = i[1]
    Lmod[index-1] = [index,lag]

if showtiming: print(walltime(w))
### Calculate the Hamiltonians corresponding to the modified Lagrangians from the previous step

w = walltime()

os.environ["SAGE_NUM_THREADS"] = "4" #Limit number of cores to save RAM
@parallel
def paraham(lag):
    return hamiltonian(lag[1],order)
gen = paraham(Lmod)
K = [0 for i in methodnames]
for i in gen:
    index = i[0][0][0][0]
    ham = i[1]
    K[index-1] = ham

if showtiming: print(walltime(w))
### Get expressions for modified quantities
#
# Modified altered Hamiltonian K,
# modified conformal Hamiltonian E, and
# modified conformal factor N.
#
# Python-formatted output.

w = walltime()

@parallel 
def get_KEN(H):
    E = series_solve(expand(H[1]), order, Ed)
    Nmod = expand(-diff(H[1],Ed).subs(Ed==E)).series(h,order).truncate().simplify_full()
    print("")
    print("K " + methodnames[H[0]-1] + " =")
    print(convert_py(str(H[1].simplify_full())))
    print("E " + methodnames[H[0]-1] + " =")
    print(convert_py(str(E)))
    print("N " + methodnames[H[0]-1] + " =")
    print(convert_py(str(Nmod)))
    
gen = get_KEN( [[i,K[i-1]] for i in [1..len(K)]] )
for i in gen:
    1

if showtiming: print(walltime(w))
### Print Latex-formatted output

w = walltime()
var("E")

@parallel 
def get_E(H):
    E = series_solve(expand(H[1]), order, Ed)
    print(methodnames[H[0]-1])
    print("E")
    print(latex(E.coefficient(h,0).simplify_full()))
    print("E2")
    print(latex(E.coefficient(h,2).simplify_full()))

@parallel 
def get_N(H):
    Nmod = expand(-diff(H[1],Ed).subs(Ed==series_solve(expand(H[1]), order, Ed))).series(h,order).truncate().simplify_full()
    print(methodnames[H[0]-1])
    print("N")
    print(latex(Nmod.coefficient(h,0).simplify_full()))
    print("N2")
    print(latex(Nmod.coefficient(h,2).simplify_full()))

@parallel     
def get_K(H):
    K = H[1].subs(Ed==E)
    print(methodnames[H[0]-1])
    print("K")
    print(latex(K.coefficient(h,0).simplify_full()))
    print("K2")
    print(latex(K.coefficient(h,2).simplify_full()))
    
gen = get_K( [[i,K[i-1]] for i in [1..len(K)]] )
for i in gen: 1
gen = get_E( [[i,K[i-1]] for i in [1..len(K)]] )
for i in gen: 1
gen = get_N( [[i,K[i-1]] for i in [1..len(K)]] )
for i in gen: 1
    
if showtiming: print(walltime(w))<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1832270720941638, 'fraction_numerical': 0.03011280039234919, 'mean_word_length': 3.6052393857271907, 'pattern_counts': {'":': 0, '<': 8, '<?xml version=': 0, '>': 10, 'https://': 1, 'lorem ipsum': 0, 'www.': 1, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '381712', 'license': 'MIT', 'n_tokens_mistral': 4141, 'n_tokens_neox': 3994, 'n_words': 926, 'path': 'Modified Lagrangian.ipynb', 'repo_name': 'mvermeeren/conf-ham-sys-2020', 'size': '10224'}