2nd Vassiliev invariant for a 100-vertex trefoil knot

knots.py

import numpy as np
import time
import matplotlib.pyplot as plt

# See: https://x.com/AlexanderRKlotz/status/1848152886443753574
# With h/t to chatgpt.

"""
Output should look sth like:

> Original: 1.613 seconds
> [warning that only affects the diagonal of wmat so shouldn't matter]
> New: 0.074 seconds
"""

def torgen(N,p,q):
    step=2*np.pi/(N)
    th=np.arange(0,2*np.pi,step)
    r=np.cos(th*q)+2
    x=4*r*np.cos(p*th)
    y=4*r*np.sin(p*th)
    z=-4*np.sin(q*th)
    knot=np.zeros((N,3))
    knot[:,0]=np.transpose(x)
    knot[:,1]=np.transpose(y)
    knot[:,2]=np.transpose(z)
    return knot
    
def og_vas(knot, closed):
    """
    Not sure if I copied this right, but it runs, at least.
    """
    N=len(knot)
    wmat=np.zeros((N,N))
    k2=np.roll(knot, 1, axis=0)
    k3=np.roll(knot, -1, axis=0)
    dk=(k2-k3)/2
    if closed==0:
        dk[-1,:]=knot[-1,:]-knot[-2,:]
        dk[0,:]=k2[0,:]-knot[0,:]
    for i in range(0,N):
        for j in range(0,N):
            if i>j:
                wmat[i,j]=np.dot(np.cross(dk[i,:],dk[j,:]),knot[i,:]-knot[j,:])/np.linalg.norm(knot[i,:]-knot[j,:])**3
    
    SLL=0
    for i in range(3,N):
        for j in range(2,i):
            for k in range(1,j):
                t1=wmat[i,k]
                for l in range(0,k):
                    t2=wmat[j,l]
                    SLL=SLL+t1*t2/(8*np.pi)
                    # q=q+1
    v2=SLL/np.sqrt(12*np.pi)
    return v2


def SLL_indices(N):
    indices = []
    for i in range(3, N):
        for j in range(2, i):
            for k in range(1, j):
                indices.append((i, j, k))
    indices = np.array(indices)
    return indices[:, 0], indices[:, 1], indices[:, 2]

def new_vas(knot, closed):
    N = len(knot)
    k2 = np.roll(knot, 1, axis=0)
    k3 = np.roll(knot, -1, axis=0)
    dk = (k2 - k3) / 2
    if closed == 0:
        dk[-1, :] = knot[-1, :] - knot[-2, :]
        dk[0, :] = k2[0, :] - knot[0, :]

    # Vectorized computation of the wmat matrix
    i_indices, j_indices = np.triu_indices(N, k=-1)
    cross_products = np.cross(dk[i_indices], dk[j_indices])
    vector_diffs = knot[i_indices] - knot[j_indices]
    norms = np.linalg.norm(vector_diffs, axis=1) ** 3
    dot_products = np.einsum('ij,ij->i', cross_products, vector_diffs)
    wmat = np.zeros((N, N))
    wmat[i_indices, j_indices] = dot_products / norms
    wmat = wmat.T
    wmat[np.isnan(wmat)] = 0

    # Precompute cumulative sums for each row of wmat, this speeds things up a lot.
    cumsum_wmat = np.cumsum(wmat, axis=1)

    # Initialize SLL
    SLL = 0.0
    i_indices, j_indices, k_indices = SLL_indices(N)
    t1_values = wmat[i_indices, k_indices]
    sum_t2_values = cumsum_wmat[j_indices, k_indices-1]  # Adjusted index for cumulative sum
    SLL = np.sum(t1_values * sum_t2_values)

    # Final normalization for the Vassiliev invariant
    final_value = SLL / np.sqrt(12 * np.pi) / (8 * np.pi)

    return final_value


knot = torgen(100,3,2)

start = time.perf_counter()
og_vas(knot, 1)  # ~1.4 sec
stop = time.perf_counter()
print(f"Original: {(stop - start):.3f} seconds")

start = time.perf_counter()
new_vas(knot, 1)  # ~0.15 sec
stop = time.perf_counter()
print(f"New: {(stop - start):.3f} seconds")