def test_pyramid(self):
        """A simple right-pyramid triangular mesh"""
        # Face set:
        verts = np.vstack((np.zeros(3), np.eye(3))) # origin + unit along each axis
        faces = np.array([[0, 1, 2], [0, 1, 3], [0, 2, 3], [1, 2, 3]])
        assembly = Assembly(
            objects=[TriangulatedSurface(verts, faces, perfect_mirror)])
        
        # Ray bundle:
        pos = np.c_[[1.5, 0.5, 0.5], [-0.5, 0.5, 0.5],
            [0.5, 1.5, 0.5], [0.5, -0.5, 0.5],
            [0.5, 0.5, -0.5], [0.5, 0.5, 1.5]]
        direct = np.c_[[-1., 0., 0.], [1., 0., 0.],
            [0., -1., 0.], [0., 1., 0.],
            [0., 0., 1.], [0., 0., -1.]]
        rayb = RayBundle(pos, direct, energy=np.ones(6))

        engine = TracerEngine(assembly)
        verts = engine.ray_tracer(rayb, 1, .05)[0]
        
        p = engine.tree[-1].get_parents()
        zrays = (p >= 4)
        np.testing.assert_array_equal(verts[:,zrays],
            np.tile(np.c_[[0.5, 0.5, 0.]], (1,4)) )
        yrays = (p == 2) | (p ==3) # Only 2 rays here. Edge degeneracy? maybe.
        np.testing.assert_array_equal(verts[:,yrays],
            np.tile(np.c_[[0.5, 0., 0.5]], (1,4)) )
        xrays = (p < 2)
        np.testing.assert_array_equal(verts[:,xrays],
            np.tile(np.c_[[0., 0.5, 0.5]], (1,4)) )

class TestObjectBuilding2(unittest.TestCase):
    """Tests an object composed of two flat surfaces"""
    flat_only = True
    def setUp(self):
        self.assembly = Assembly()
        surface1 = Surface(flat_surface.FlatGeometryManager(),
            optics_callables.RefractiveHomogenous(1., 1.5),
            location=N.array([0,0,-1.]))
        surface2 = Surface(flat_surface.FlatGeometryManager(),
            optics_callables.RefractiveHomogenous(1.5, 1.),
            location=N.array([0,0,1.]))
        
        self.object = AssembledObject(surfs=[surface1, surface2])
        self.assembly.add_object(self.object)
        
        x = 1/(math.sqrt(2))
        dir = N.c_[[0,-x,x]]
        position = N.c_[[0,1,-2.]]
        self._bund = RayBundle(position, dir, energy=N.r_[1.], ref_index=N.r_[1.])

    def test_refraction1(self):
        """Tests the refractive functions after a single intersection"""
        self.engine = TracerEngine(self.assembly)
        ans =  self.engine.ray_tracer(self._bund,1,.05)
        params = N.arctan(ans[1][1]/ans[1][2])
        correct_params = N.r_[0.785398163, -.4908826]
        N.testing.assert_array_almost_equal(params, correct_params)

    def test_refraction2(self):
        """Tests the refractive functions after two intersections"""
        self.engine = TracerEngine(self.assembly)
        ans = self.engine.ray_tracer(self._bund,2,.05)
        params = N.arctan(ans[1][1]/ans[1][2])
        correct_params = N.r_[-0.7853981]
        N.testing.assert_array_almost_equal(params, correct_params)

	def trace(self, rph, iters = 10000, minE = 1e-9, render = False):
		"""Commences raytracing using (rph) number of rays per heliostat, for a maximum of 
		   (iters) iterations, discarding rays with energy less than (minE). If render is
		   True, a 3D scene will be displayed which would need to be closed to proceed."""
		# Get the solar vector using azimuth and elevation
		sun_vec = solar_vector(self.sun_az*degree, self.sun_elev*degree)
        	# Calculate number of rays used. Rays per heliostat * number of heliostats.
        	num_rays = rph*len(self.field.get_heliostats())
		self.no_of_rays += num_rays
		# Generates the ray bundle
        	rot_sun = rotation_to_z(-sun_vec)
        	direct = N.dot(rot_sun, pillbox_sunshape_directions(num_rays, 0.00465))
        
        	xy = N.random.uniform(low=-0.25, high=0.25, size=(2, num_rays))
        	base_pos = N.tile(self.pos, (rph, 1)).T #Check if its is rph or num_rays
       		base_pos += N.dot(rot_sun[:,:2], xy)
        	
        	base_pos -= direct
        	rays = RayBundle(base_pos, direct, energy=N.ones(num_rays))

		# Perform the raytracing
		e = TracerEngine(self.plant)
		e.ray_tracer(rays, iters, minE, tree=True)
		e.minener = minE
		rays_in = sum(e.tree._bunds[0].get_energy())
		self.helio_hits = sum(e.tree._bunds[1].get_energy())


		# Optional rendering
		if render == True:
			trace_scene = Renderer(e)
			trace_scene.show_rays()

 def test_absorbed_to_back(self):
     """Absorbed rays moved to back of recorded bundle"""
     engine = TracerEngine(self.assembly)
     engine.ray_tracer(self.bund, 300, .05)
     
     parents = engine.tree.ordered_parents()
     N.testing.assert_equal(parents, [N.r_[2,3, 0, 1]])

class TestTraceProtocol1(unittest.TestCase):
    """ 
    Tests intersect_ray and the bundle driver with a single flat surface, not rotated, with 
    a single interation 
    """
    def setUp(self):
        dir = N.array([[1,1,-1],[-1,1,-1],[-1,-1,-1],[1,-1,-1]]).T/math.sqrt(3)
        position = N.c_[[0,0,1],[1,-1,1],[1,1,1],[-1,1,1]]
        self._bund = RayBundle(position, dir, energy=N.ones(4))
        
        self.assembly = Assembly()
        object = AssembledObject()
        object.add_surface(Surface(FlatGeometryManager(), opt.perfect_mirror))
        self.assembly.add_object(object)
        self.engine = TracerEngine(self.assembly)
        
    def test_intersect_ray1(self):
        surfaces = self.assembly.get_surfaces()
        objects = self.assembly.get_objects()
        surfs_per_obj = [len(obj.get_surfaces()) for obj in objects]
        surf_ownership = N.repeat(N.arange(len(objects)), surfs_per_obj)
        ray_ownership = -1*N.ones(self._bund.get_num_rays())
        surfs_relevancy = N.ones((len(surfaces), self._bund.get_num_rays()), dtype=N.bool)
        
        params = self.engine.intersect_ray(self._bund, surfaces, objects, \
            surf_ownership, ray_ownership, surfs_relevancy)[0]
        self.failUnless(params.all())

    def test_ray_tracer(self):
        """Ray tracer after one iteration returns what the surface would have"""
        params = self.engine.ray_tracer(self._bund,1,.05)[0]
        correct_pts = N.zeros((3,4))
        correct_pts[:2,0] = 1
        
        N.testing.assert_array_almost_equal(params, correct_pts)

    def test_tetrahedron(self):
        """Triangular mesh with oblique triangles"""
        # Face set:
        theta = np.arange(np.pi/2., np.pi*2, 2*np.pi/3)
        base_verts = np.vstack(( np.cos(theta), np.sin(theta), np.ones(3) )).T
        verts = np.vstack((np.zeros(3), base_verts))
        
        faces = np.array([[0, 1, 2], [0, 1, 3], [0, 2, 3], [1, 2, 3]])
        fset = TriangulatedSurface(verts, faces, perfect_mirror)
        
        # Flat floor:
        floor = rect_one_sided_mirror(5., 5., 1.)
        floor.set_location(np.r_[0., 0., 1.])
        assembly = Assembly(objects=[fset, floor])
        
        # Ray bundle of 3 rays starting at equal angles around the tetrahedron:
        theta -= np.pi/3.
        pos = np.vstack((np.cos(theta), np.sin(theta), np.ones(3)*0.2)) * 0.2
        direct = np.vstack(( np.zeros((2,3)), np.ones(3) ))
        rayb = RayBundle(pos, direct, energy=np.ones(6))
        
        # Check that the points on the floor describe an isosceles.
        engine = TracerEngine(assembly)
        engine.ray_tracer(rayb, 2, .05)[0]
        verts = engine.tree[-1].get_vertices()
        sizes = np.sqrt(
            np.sum((verts - np.roll(verts, 1, axis=1))**2, axis=0))
        
        self.assertAlmostEqual(sizes[0], sizes[1])
        self.assertAlmostEqual(sizes[2], sizes[1])

	'''

    def trace(self):
        """Generate a flux map using much more rays than drawn"""
        # Generate a large ray bundle using a radial stagger much denser
        # than the field.
        sun_vec = solar_vector(self.sun_az*degree, self.sun_elev*degree)
        
        hstat_rays = 20
        num_rays = hstat_rays*len(self.field.get_heliostats())
        rot_sun = rotation_to_z(-sun_vec)
        direct = N.dot(rot_sun, pillbox_sunshape_directions(num_rays, 0.00465))
        
        xy = N.random.uniform(low=-0.25, high=0.25, size=(2, num_rays))
        base_pos = N.tile(self.pos, (hstat_rays, 1)).T
        base_pos += N.dot(rot_sun[:,:2], xy)
        
        base_pos -= direct
        rays = RayBundle(base_pos, direct, energy=N.ones(num_rays))
        
        # Perform the trace:
        e = TracerEngine(self.plant)
        e.ray_tracer(rays, 100, 0.05, tree=True)
        e.minener = 1e-5

		# Render:
        trace_scene = Renderer(e)
        trace_scene.show_rays()

def test_case(focus, num_rays=100, h_depth=0.7, side=0.4):
	# Case parameters:
	D = 5.
	
	center = N.c_[[0, 7., 7.]]
	x = -1/(math.sqrt(2))
	direction = N.array([0,x,x])
	radius_sun = 3. 
	ang_range = 0.005
	
	iterate = 100
	min_energy = 1e-6
	
	# Model:
	assembly = MiniDish(D, focus, 0.9, focus + h_depth, side, h_depth, 0.9)
	assembly.set_transform(rotx(-N.pi/4))
	
	# Rays:
	sun = solar_disk_bundle(num_rays, center, direction, radius_sun, ang_range,
		flux=1000.)
	
	# Do the tracing:
	engine = TracerEngine(assembly)
	engine.ray_tracer(sun, iterate, min_energy)

	resolution = 20
 
	# Render a subset of the total rays:
	v = R(engine)
	v.show_rays(max_rays=100, resolution=resolution, fluxmap=True)	

	# Plot, scale in suns:
	f = plot_hits(assembly.histogram_hits(bins=resolution)[0]/(side/resolution)**2/1000., (-side/2., side/2., -side/2., side/2.))
	f.show()

class TestObjectBuilding1(unittest.TestCase):
    """Tests an object composed of sphere surfaces"""
    def setUp(self):
        self.assembly = Assembly()
        surface1 = Surface(HemisphereGM(3.), optics_callables.perfect_mirror, 
            location=N.array([0,0,-1.]),
            rotation=general_axis_rotation(N.r_[1,0,0], N.pi))
        surface2 = Surface(HemisphereGM(3.), optics_callables.perfect_mirror, 
            location=N.array([0,0,1.]))
        
        self.object = AssembledObject()
        self.object.add_surface(surface1)
        self.object.add_surface(surface2)
        self.assembly.add_object(self.object)

        dir = N.c_[[0,0,1.],[0,0,1.]]
        position = N.c_[[0,0,-3.],[0,0,-1.]]
    
        self._bund = RayBundle(position, dir, energy=N.ones(2))
    
    def test_object(self):
        """Tests that the assembly heirarchy works at a basic level"""
        self.engine = TracerEngine(self.assembly)

        inters = self.engine.ray_tracer(self._bund,1,.05)[0]
        correct_inters = N.c_[[0,0,2],[0,0,-2]]

        N.testing.assert_array_almost_equal(inters, correct_inters)
    
    def test_translation(self):
        """Tests an assembly that has been translated"""
        trans = N.array([[1,0,0,0],[0,1,0,0],[0,0,1,1],[0,0,0,1]])
        self.assembly.transform_children(trans)

        self.engine = TracerEngine(self.assembly)

        params =  self.engine.ray_tracer(self._bund,1,.05)[0]
        correct_params = N.c_[[0,0,3],[0,0,-1]]

        N.testing.assert_array_almost_equal(params, correct_params)

    def test_rotation_and_translation(self):
        """Tests an assembly that has been translated and rotated"""
        self._bund = RayBundle()
        self._bund.set_vertices(N.c_[[0,-5,1],[0,5,1]])
        self._bund.set_directions(N.c_[[0,1,0],[0,1,0]])
        self._bund.set_energy(N.r_[[1,1]])
        self._bund.set_ref_index(N.r_[[1,1]])

        trans = generate_transform(N.r_[[1,0,0]], N.pi/2, N.c_[[0,0,1]])
        self.assembly.transform_children(trans)

        self.engine = TracerEngine(self.assembly)

        params =  self.engine.ray_tracer(self._bund,1,.05)[0]
        correct_params = N.c_[[0,-2,1]]

        N.testing.assert_array_almost_equal(params, correct_params)

    def test_cylinder_height(self):
        """The bounding cylinder exists for planoconvex lens"""
        f = self.lens.focal_length()
        rb = RayBundle(N.c_[[0., 0., -0.01]], N.c_[[1., 0., 0.]],
            energy=N.r_[1.], ref_index=N.r_[1.5])

        e = TracerEngine(Assembly([self.lens]))
        verts, dirs = e.ray_tracer(rb, 1, 1e-6)

        N.testing.assert_array_equal(verts, N.array([]).reshape(3,0))

    def test_cylinder(self):
        """The bounding cylinder exists for biconcave lens"""
        f = self.lens.focal_length()
        rb = RayBundle(N.c_[[0., 0., 0.08]], N.c_[[1., 0., 0.]],
            energy=N.r_[1.], ref_index=N.r_[1.5])
        
        e = TracerEngine(Assembly([self.lens]))
        verts, dirs = e.ray_tracer(rb, 1, 1e-6)

        N.testing.assert_array_equal(verts, N.tile(N.c_[[0.5, 0., 0.08]], (1,2)))
        N.testing.assert_array_equal(dirs, N.c_[[-1., 0., 0.], [1., 0., 0.]])

 def test_paraxial_ray(self):
     """A paraxial ray reaches the focus"""
     rb = RayBundle(N.c_[[0., 0.001, 1.]], N.c_[[0., 0., -1.]], 
         energy=N.r_[1.], ref_index=N.r_[1.])
     screen = rect_one_sided_mirror(5, 5)
     f = self.lens.focal_length()
     screen.set_transform(translate(0, 0, -f))
     
     e = TracerEngine(Assembly([self.lens, screen]))
     vert, _ = e.ray_tracer(rb, 3, 1e-6)
     self.failUnlessAlmostEqual(vert[1,2], 0, 4)

 def test_regular(self):
     """One-sided plate without rotation"""
     e = TracerEngine(Assembly(objects=[self.mirror]))
     e.ray_tracer(self.bund, 1, 0.05)
     outg = e.tree[-1]
     
     correct_verts = N.zeros((3,2))
     correct_verts[0] = N.r_[0, 0.5]
     N.testing.assert_array_equal(
         outg.get_vertices()[:,outg.get_energy() > 0], correct_verts)
     N.testing.assert_array_almost_equal(
         outg.get_energy(), N.r_[100., 100., 0, 0])

 def test_aim(self):
     """Aiming heliostats works"""
     elev = N.pi/4
     az = N.pi/2
     self.field.aim_to_sun(az, elev)
     
     e = TracerEngine(self.field)
     v, d = e.ray_tracer(self.rays, 1, 0.05)
     
     N.testing.assert_array_almost_equal(d[1, :self.pos.shape[0]/2], 0)
     N.testing.assert_array_almost_equal(d[0, self.pos.shape[0]/2:], 0)
     N.testing.assert_array_almost_equal(abs(d[2]*(v[0] + v[1])/(d[0] + d[1])), 85.5)

class TestAssemblyBuilding3(unittest.TestCase):
    """Tests an assembly composed of objects that are transformed rel. the assembly"""
    def setUp(self):  
        self.assembly = Assembly()

        surface1 = Surface(flat_surface.FlatGeometryManager(), 
            optics_callables.RefractiveHomogenous(1., 1.5),
            location=N.array([0,0,-1.]))
        surface2 = Surface(flat_surface.FlatGeometryManager(), 
            optics_callables.RefractiveHomogenous(1., 1.5),
            location=N.array([0,0,1.]))
        
        self.object1 = AssembledObject() 
        self.object1.add_surface(surface1)
        self.object1.add_surface(surface2)
        
        boundary = BoundarySphere(location=N.r_[0,0.,3], radius=3.)
        surface3 = Surface(CutSphereGM(2., boundary), optics_callables.perfect_mirror)
        self.object2 = AssembledObject()
        self.object2.add_surface(surface3)
    
        self.transform = generate_transform(N.r_[1,0.,0],0.,N.c_[[0.,0,2]])
        self.assembly.add_object(self.object1)
        self.assembly.add_object(self.object2, self.transform)
        
        x = 1./(math.sqrt(2))
        dir = N.c_[[0,1.,0.],[0,x,x],[0,0,1.]]
        position = N.c_[[0,0,2.],[0,0,2.],[0,0.,2.]]
        self._bund = RayBundle(position, dir, energy=N.ones(3), ref_index=N.ones(3))
        
    def test_assembly1(self):
        """Tests the assembly after one iteration"""
        self.engine = TracerEngine(self.assembly)
        ans =  self.engine.ray_tracer(self._bund,1,.05)
        params = N.arctan(ans[1][1]/ans[1][2])
        correct_params = N.r_[0.7853981, 0]

        N.testing.assert_array_almost_equal(params, correct_params)

    def test_assembly2(self):
        """Tests the assembly after two iterations"""
        self.engine = TracerEngine(self.assembly)
        params = self.engine.ray_tracer(self._bund,2,.05)[0]
        correct_params = N.c_[[0,-1,1], [0,-1,1],[0,0,1]]
        N.testing.assert_array_almost_equal(params, correct_params)

    def test_assembly3(self):      
        """Tests the assembly after three iterations"""  
        self.engine = TracerEngine(self.assembly)
        params = self.engine.ray_tracer(self._bund, 3,.05)[0]
        correct_params = N.c_[[0,-2.069044,-1],[0,0,-1]]

        N.testing.assert_array_almost_equal(params, correct_params)

 def test_rotated(self):
     """One-sided plate with rotation"""
     rot = sp.roty(N.pi/4.)
     self.mirror.set_transform(rot)
     
     e = TracerEngine(Assembly(objects=[self.mirror]))
     e.ray_tracer(self.bund, 1, 0.05)
     outg = e.tree[-1]
     
     correct_verts = N.array([[0., 0.5], [0., 0.], [0., -0.5]])
     N.testing.assert_array_almost_equal(
         outg.get_vertices()[:,outg.get_energy() > 0], correct_verts)
     N.testing.assert_array_almost_equal(
         outg.get_energy(), N.r_[100., 100., 0, 0])

class TestNestedAssemblies(unittest.TestCase):
    def setUp(self):
        """
        Prepare an assembly with two subassemblies: one assembly representing
        a spherical lens behind a flat screen, and one asssembly representing a
        perfect mirror.
        The mirror will be placed at the two subassemblies' focus, so a paraxial
        ray will come back on the other side of the optical axis.
        
        Reference:
        In [1], the lensmaker equation
        """
        # focal length = 1, thickness = 1/6
        R = 1./6.
        back_surf = Surface(HemisphereGM(R), opt.RefractiveHomogenous(1., 1.5), location=N.r_[0., 0., -R/2.])
        front_surf = Surface(HemisphereGM(R), opt.RefractiveHomogenous(1., 1.5),location=N.r_[0., 0., R/2.], rotation=rotx(N.pi/2.)[:3,:3])
        
        front_lens = AssembledObject(surfs=[back_surf, front_surf])

        back_surf = Surface(RoundPlateGM(R), opt.RefractiveHomogenous(1., 1.5),
            location=N.r_[0., 0., -0.01])
        front_surf = Surface(RoundPlateGM(R), opt.RefractiveHomogenous(1., 1.5),
            location=N.r_[0., 0., 0.01])

        glass_screen = AssembledObject(surfs=[back_surf, front_surf],
            transform=translate(0., 0., 0.5))
        
        lens_assembly = Assembly(objects=[glass_screen, front_lens])
        lens_assembly.set_transform(translate(0., 0., 1.))
        full_assembly = Assembly(objects=[rect_one_sided_mirror(1., 1., 0.)],
            subassemblies = [lens_assembly])
        
        self.engine = TracerEngine(full_assembly)
    
    def test_paraxial_ray(self):
        """A paraxial ray in reflected correctly"""
        bund = RayBundle()
        bund.set_vertices(N.c_[[0.01, 0., 2.]])
        bund.set_directions(N.c_[[0., 0., -1.]])
        bund.set_energy(N.r_[100.])
        bund.set_ref_index(N.r_[1])
        
        self.engine.ray_tracer(bund, 15, 10.)

        non_degenerate = self.engine.tree[-1].get_energy() > 10.
        v = self.engine.tree[-1].get_vertices()[:,non_degenerate]
        d = self.engine.tree[-1].get_directions()[:,non_degenerate]
        # Not high equality demanded, because of spherical aberration.
        N.testing.assert_array_almost_equal(v, N.c_[[-0.01, 0., 1.5]], 2)
        N.testing.assert_array_almost_equal(d, N.c_[[0., 0., 1.]], 2)

    def test_assembly3(self):      
        """Tests the assembly after three iterations"""  
        self.engine = TracerEngine(self.assembly)
        params = self.engine.ray_tracer(self._bund, 3,.05)[0]
        correct_params = N.c_[[0,-2.069044,-1],[0,0,-1]]

        N.testing.assert_array_almost_equal(params, correct_params)

class TestHomogenizer(unittest.TestCase):
    def setUp(self):
        """A homogenizer transforms a bundle correctly"""
        hmg = rect_homogenizer(5., 3., 10., 0.9)
        self.engine = TracerEngine(hmg)
        self.bund = RayBundle()
        
        # 4 rays starting somewhat above (+z) the homogenizer
        pos = N.zeros((3,4))
        pos[2] = N.r_[11, 11, 11, 11]
        self.bund.set_vertices(pos)
        
        # One ray going to each wall:
        dir = N.c_[[1, 0, -1], [-1, 0, -1], [0, 1, -1], [0, -1, -1]]/N.sqrt(2)
        self.bund.set_directions(dir)
        
        # Laborious setup details:
        self.bund.set_energy(N.ones(4)*4.)
        self.bund.set_ref_index(N.ones(4))
    
    def test_first_hits(self):
        """Test bundle enters homogenizer correctly"""
        v, d = self.engine.ray_tracer(self.bund, 1, 0.05)
        
        out_dirs = N.c_[[-1, 0, -1], [1, 0, -1], [0, -1, -1], [0, 1, -1]]/N.sqrt(2)
        N.testing.assert_array_almost_equal(d, out_dirs)
        
        out_hits = N.c_[
            [2.5, 0, 8.5], 
            [-2.5, 0, 8.5], 
            [0, 1.5, 9.5], 
            [0, -1.5, 9.5]]
        N.testing.assert_array_almost_equal(v, out_hits)

 def setUp(self):
     """
     Prepare an assembly with two subassemblies: one assembly representing
     a spherical lens behind a flat screen, and one asssembly representing a
     perfect mirror.
     The mirror will be placed at the two subassemblies' focus, so a paraxial
     ray will come back on the other side of the optical axis.
     
     Reference:
     In [1], the lensmaker equation
     """
     # focal length = 1, thickness = 1/6
     R = 1./6.
     back_surf = Surface(HemisphereGM(R), opt.RefractiveHomogenous(1., 1.5),
         location=N.r_[0., 0., -R/2.])
     front_surf = Surface(HemisphereGM(R), opt.RefractiveHomogenous(1., 1.5),
         location=N.r_[0., 0., R/2.], rotation=rotx(N.pi/2.)[:3,:3])
     front_lens = AssembledObject(surfs=[back_surf, front_surf])
     
     back_surf = Surface(FlatGeometryManager(), opt.RefractiveHomogenous(1., 1.5),
         location=N.r_[0., 0., -0.01])
     front_surf = Surface(FlatGeometryManager(), opt.RefractiveHomogenous(1., 1.5),
         location=N.r_[0., 0., 0.01])
     glass_screen = AssembledObject(surfs=[back_surf, front_surf],
         transform=translate(0., 0., 0.5))
     
     lens_assembly = Assembly(objects=[glass_screen, front_lens])
     lens_assembly.set_transform(translate(0., 0., 1.))
     full_assembly = Assembly(objects=[rect_one_sided_mirror(1., 1., 0.)],
         subassemblies = [lens_assembly])
     
     self.engine = TracerEngine(full_assembly)