coord_transform.py

def acos2(num, denom, disamb):
    cosine = num/denom
    
    return torch.where((cosine > -1) & (cosine < 1.),
        torch.acos(cosine) * torch.where(disamb < 0.0, -1, 1),
        torch.where(cosine <= -1, np.pi, 0.0)
    )

def coord_transform(x):
    # Assumes in CoM frame
    G = 1.0
    primary_m = 1.0
    
    m = x[:, -1]
    mu = G*(m + primary_m);
    
    dx = x[:, 0]
    dy = x[:, 1]
    dz = x[:, 2]
    
    dvx = x[:, 3]
    dvy = x[:, 4]
    dvz = x[:, 5]
    
    d = torch.sqrt( dx*dx + dy*dy + dz*dz )
    vsquared = dvx*dvx + dvy*dvy + dvz*dvz
    v = torch.sqrt(vsquared)
    vcircsquared = mu/d
    a = -mu/(vsquared - 2.0*vcircsquared)
    
    hx = (dy*dvz - dz*dvy);                     #angular momentum vector
    hy = (dz*dvx - dx*dvz);
    hz = (dx*dvy - dy*dvx);
    h = torch.sqrt( hx*hx + hy*hy + hz*hz );        # abs value of angular momentum
    
    vdiffsquared = vsquared - vcircsquared
    vr = (dx*dvx + dy*dvy + dz*dvz)/d
    rvr = d*vr
    muinv = 1./mu;
    
    ex = muinv*( vdiffsquared*dx - rvr*dvx );
    ey = muinv*( vdiffsquared*dy - rvr*dvy );
    ez = muinv*( vdiffsquared*dz - rvr*dvz );
    
    e = torch.sqrt( ex*ex + ey*ey + ez*ez );        # eccentricity
    
    inc = acos2(hz, h, torch.ones_like(h))
    return (a, e, inc)