Generate forward kinematics data for neural network training

generate_training_data.py

def transformMatrix(theta, d, a, alpha):
    return np.array([[np.cos(theta), -np.sin(theta)*np.cos(alpha), np.sin(theta)*np.sin(alpha), a*np.cos(theta)], 
                     [np.sin(theta), np.cos(theta)*np.cos(alpha), -np.cos(theta)*np.sin(alpha), a*np.sin(theta)], 
                     [0, np.sin(alpha), np.cos(alpha), d], 
                     [0, 0, 0, 1]])

def forwardKinematics_2(theta1, theta2):
    T00 = transformMatrix(theta1,0,1,0)
    T01 = transformMatrix(theta2,0,1,0)
    pos = [0, 0, 0, 1]
    Etip = np.matmul(np.matmul(T00, T01), pos)
    return T00, T01, Etip

def get_positions_2(theta):
    # assuming theta is already in radian
    theta1 = theta[0]
    theta2 = theta[1]
    
    T00, T01, Etip = forwardKinematics_2(theta1, theta2)
    t = np.transpose(np.array([[0, 0, 0, 1]]))
    pos_1 = np.matmul(T00, t)
    
    # only return first 2 elements
    return np.array([pos_1[:2], np.reshape(Etip[:2], (2, 1))])

def get_xy_and_theta_2(num):
    xy = np.zeros((num, 2))
    theta = np.zeros((num, 3))

    theta[:,0] = (np.random.random((num)) * 2 * np.pi) - np.pi
    theta[:,1] = (np.random.random((num)) * np.pi) - (0.5 * np.pi)
    for i in range(num):
        _, _, temp = forwardKinematics_2(theta[i,0], theta[i,1])
        xy[i, :] = temp[:2]
    return xy, theta