# Path: tracer/spatial_geometry.py
# def rotation_to_z(vecs):
#     """
#     Generate a rotation into a frame whose Z zxis points along the direction
#     indicated by ``vecs``. The rest of the directions are determined by
#     requiring the new X to lie on the XY plane of the original frame, and by
#     the right-hand rule.
#     
#     In the singular case that the required direction is the Z axis, the
#     original frame is retained.
#     
#     Arguments:
#     vec - a 3-component 1D array representing a unit direction in 3D space.
#     
#     Returns:
#     A 3x3 array representing the global-to-local rotation to get the desired
#         frame.
#     """
#     vecs = N.atleast_2d(vecs)
#     perp = N.hstack((vecs[:,1][:,None],  -vecs[:,0][:,None],
#         N.zeros((vecs.shape[0], 1)) ))
#     perp[N.all(perp == 0, axis=1)] = N.r_[1.,  0.,  0.]
#     perp /= N.sqrt(N.sum(perp**2, axis=1))[:,None]
#     perp_rot = N.concatenate((perp[...,None], N.cross(vecs, perp)[...,None],
#         vecs[...,None]), axis=2)
#     return N.squeeze(perp_rot)


from numpy import random, linalg as LA
from .ray_bundle import RayBundle
from .spatial_geometry import rotation_to_z
import numpy as N
    xi1 = random.uniform(size=num_rays)
    thetas = random.uniform(high=2*N.pi, size=num_rays)
    rs = radius*N.sqrt(xi1)
    xs = rs * N.cos(thetas)
    ys = rs * N.sin(thetas)

    # Rotate locations to the plane defined by <direction>:
    vertices_local = N.vstack((xs,  ys,  N.zeros(num_rays)))
    vertices_global = N.dot(perp_rot,  vertices_local)

    rayb = RayBundle(vertices=vertices_global + center, directions=directions)
    if flux is not None:
        rayb.set_energy(N.pi*radius**2/num_rays*flux*N.ones(num_rays))
    
    return rayb

def square_bundle(num_rays, center, direction, width):
    """
    Generate a ray bundles whose rays are equally spaced along a square grid,
    and all pointing in the same direction.
    
    Arguments:
    num_rays - number of rays to generate.
    center - a column 3-array with the 3D coordinate of the disk's center
    direction - a 1D 3-array with the unit direction vector for the bundle.
    width - of the square of starting points.
    
    Returns: 
    A RayBundle object with the above charachteristics set.
    """

# Path: tracer/spatial_geometry.py
# def rotation_to_z(vecs):
#     """
#     Generate a rotation into a frame whose Z zxis points along the direction
#     indicated by ``vecs``. The rest of the directions are determined by
#     requiring the new X to lie on the XY plane of the original frame, and by
#     the right-hand rule.
#     
#     In the singular case that the required direction is the Z axis, the
#     original frame is retained.
#     
#     Arguments:
#     vec - a 3-component 1D array representing a unit direction in 3D space.
#     
#     Returns:
#     A 3x3 array representing the global-to-local rotation to get the desired
#         frame.
#     """
#     vecs = N.atleast_2d(vecs)
#     perp = N.hstack((vecs[:,1][:,None],  -vecs[:,0][:,None],
#         N.zeros((vecs.shape[0], 1)) ))
#     perp[N.all(perp == 0, axis=1)] = N.r_[1.,  0.,  0.]
#     perp /= N.sqrt(N.sum(perp**2, axis=1))[:,None]
#     perp_rot = N.concatenate((perp[...,None], N.cross(vecs, perp)[...,None],
#         vecs[...,None]), axis=2)
#     return N.squeeze(perp_rot)


# Path: tracer/sources.py
from numpy import random, linalg as LA
from .ray_bundle import RayBundle
from .spatial_geometry import rotation_to_z
import numpy as N

    xi1 = random.uniform(size=num_rays)
    thetas = random.uniform(high=2*N.pi, size=num_rays)
    rs = radius*N.sqrt(xi1)
    xs = rs * N.cos(thetas)
    ys = rs * N.sin(thetas)

    # Rotate locations to the plane defined by <direction>:
    vertices_local = N.vstack((xs,  ys,  N.zeros(num_rays)))
    vertices_global = N.dot(perp_rot,  vertices_local)

    rayb = RayBundle(vertices=vertices_global + center, directions=directions)
    if flux is not None:
        rayb.set_energy(N.pi*radius**2/num_rays*flux*N.ones(num_rays))
    
    return rayb

def square_bundle(num_rays, center, direction, width):
    """
    Generate a ray bundles whose rays are equally spaced along a square grid,
    and all pointing in the same direction.
    
    Arguments:
    num_rays - number of rays to generate.
    center - a column 3-array with the 3D coordinate of the disk's center
    direction - a 1D 3-array with the unit direction vector for the bundle.
    width - of the square of starting points.
    
    Returns: 
    A RayBundle object with the above charachteristics set.
    """

    rot = rotation_to_z(direction)