check_det_without_rotat.py

""" trying to understand the determinant of sphere-to-cyl """

import numpy as np
import jax.numpy as jnp
from jax import jacfwd


# map from cylinder to sphere
def phi(zr):
    d = zr.shape[0]
    w = jnp.full(d-1, jnp.sqrt(1 - zr[-1]**2))
    w = jnp.append(w, 1)
    return w * zr


def phi_inv(x):
    r = x[-1]  # generally, dot(x, mu)
    z = x[:-1] / jnp.sqrt(1 - r**2)
    return jnp.append(z, r)


def main():

    d = 4

    # get some point on d-sphere
    x = np.random.randn(d+1)
    x /= np.linalg.norm(x)

    # map to cylinder

    zr = phi_inv(x)

    # jacobian
    J = jacfwd(phi)(zr)

    z, r = zr[:-1], zr[-1]

    # expected end result is
    res = (1 - r ** 2) ** (d-2)
    print(res)

    # find a basis on T_z S^d-1

    P = np.eye(d) - np.outer(z, z)
    eigvals, eigvecs = np.linalg.eigh(P)
    E = eigvecs[:, np.abs(eigvals - 1) < 1e-5]
    # print(E @ E.T - P)  # == 0

    # append a 1 to the bottom-right of E
    E = np.block([[E, np.zeros((4, 1))], [np.zeros((1, 3)), np.eye(1)]])


    JE = J @ E

    print(jnp.linalg.det(JE.T @ JE))


if __name__ == '__main__':
    main()