Function get_random_direction_within_cone

plib/rigid_motion.py:365–407

Get uniformly sampled random directions within a cone centered at (0,0,1). Args: n: number of samples theta: the half angle of the cone (in degree) rng: numpy random state method: 'random' Returns:

(
        n: int,
        theta: float,
        rng: np.random.RandomState = None,
        method: str = 'random',
)

def get_random_direction_within_cone(
        n: int,
        theta: float,
        rng: np.random.RandomState = None,
        method: str = 'random',
):
    """
    Get uniformly sampled random directions within a cone centered at (0,0,1).

    Args:
        n:
            number of samples
        theta:
            the half angle of the cone (in degree)
        rng:
            numpy random state
        method:
            'random'

    Returns:
        (n, 3), float64

    We are to use Archimedes' Hat-Box Theorem to sample the directions.
    """
    assert 0 < theta <= 180.
    t_max = 1
    t_min = np.cos(theta / 180. * np.pi)

    if rng is None:
        rng = np.random

    # sample uniformly within [0,1]^2
    samples = sample_utils.get_samples(total_samples=n, d=2, method=method, rng=rng)  # (n, 2)
    # adjust range
    wzs = samples[..., 0] * (t_max - t_min) + t_min  # (n,)
    phis = samples[..., 1] * (2 * np.pi)  # (n,)

    rs = np.sqrt(1 - np.power(wzs, 2))
    wxs = rs * np.cos(phis)
    wys = rs * np.sin(phis)

    ds = np.stack((wxs, wys, wzs), axis=-1)  # (n,3)
    return ds


def construct_coord_frame(