In [ ]:
#For new lenses defined in terms of Radius of Curvature and Thickness

from vpython import *
from scipy.optimize import fsolve  #You will need to install the scipy package for python
from math import sqrt  #some of the equations require the use of the sqrt function
import numpy  #Package for scientific computing in Python

#####
# Start of World
#####

class world():
    '''Creates a virtual environment for rays and objects to interact in; a world'''
    def __init__(self):
        self.rays=[] #An array to store rays that will interact in the 'world'
        self.objects=[] #An array to store oobjects that will become obstacles to rays
        self.dL = 0.1   #Ray step size
        self.n = 1   #Default index of refraction for the world
        self.MAX_LENGTH = 10  ### Set to 10 ###
        self.current_length = 0 ### Length of current ray being drawn ###
        
        
    def add_ray(self,new_ray): #Works for Ray
        '''Add a ray to the world'''
        self.rays.append(new_ray)
        new_ray.length=self.dL
        new_ray.world_n = self.n
        
    def add_object(self,new_object): #Works for Ball, Block, Lens
        '''Add a new object to the world'''
        self.objects.append(new_object)
    
    def add_light(self,light): #Works for Point Source
        '''Add a point source to the world'''
        for i in light.Rays:
            self.add_ray(i)
    
    def add_beam(self, beam): #Works for Beam
        '''Add a beam to the world'''
        for i in beam.Rays:
            self.add_ray(i)
        
    def draw_rays(self):
        '''Draws rays in short steps'''
        #Loop for checking Boundaries
        for i in self.rays:
            i.current_length = 0
            while i.current_length < self.MAX_LENGTH:  ### Runs until we reach a max length
                for j in self.objects:  ### Loop through all objects ###
                    j.check_boundaries(i)   ### Call check_boundaries method for each object ###
                i.draw_ray()
                i.current_length=i.current_length+self.dL
                
class ray():
    '''A light ray class'''
    def __init__(self,startPoint,directionVector,color):
        self.startPoint = startPoint
        self.oldPoint = startPoint   ### Position of light ray prior to current time step ###
        self.currentPoint = startPoint
        self.directionVector = directionVector.norm()
        self.color = color
        self.world_n=1   ### Value of n for the medium ray is currently in ###
        self.ray = curve(pos=self.currentPoint, color=self.color)
        
        #self.draw_ray()  ### Commented out to let world() draw ray
    
    def draw_ray(self):
        '''Draw a single light ray'''
        self.oldPoint = self.currentPoint ### Set old point to the current position BEFORE the update ###
        newPoint = self.currentPoint + self.length*self.directionVector #add on one time step of distance
        self.currentPoint = newPoint #Update the current position of the ray to be what we just calculated
        self.ray.append(newPoint, color=self.color) #Add this point to an array holding the previous steps of this ray
        
class beam():
    '''A Beam Class'''
    def __init__(self, position, direction, width):  #User-input parameters, applied to their own variables in the lines following
        self.direction = direction
        self.position = position
        self.centerRay = ray(position, direction, color.green) #Creates a ray with the same parameters that the beam will have
        self.currentPoint = self.centerRay.currentPoint
        self.beamDirection = direction
        self.width = width
        self.color = self.centerRay.color
        self.Rays = [] #Creates an array that will store parallel rays (built off the center ray) to form the beam
        
        ### Begin changes to fix perpendicular rays ###
        self.find_perp()   #Find two rays perpendicular to centerRay; establishes a set of local coordinate axes for the beam to be drawn in
        self.draw_beam()
          
    def find_perp(self):
        '''Find two vectors perpendicular to the center ray'''
        if self.centerRay.directionVector.y==0 and self.centerRay.directionVector.x==0:
            self.perp1=vec(1,0,0)
        else:
            self.perp1 = norm(vec(self.centerRay.directionVector.y,-self.centerRay.directionVector.x,0))
        self.perp2 = norm(vec(self.centerRay.directionVector.cross(self.perp1)))
        ### End changes to fix perpendicular rays ###

    def draw_beam(self):
        '''Draw a beam; lots of parallel rays'''#Loops through a range adding parallel rays in rings around the center ray
        for i in numpy.arange(0,self.width,.5): #Increasing beam width by stepping out radially
            for k in numpy.arange(0,2*pi+0.1,.5): #Adding ring by stepping around center ray in a circle
                self.newStart = self.centerRay.startPoint + i*(self.perp1*cos(k)+self.perp2*sin(k))
                beamRay = ray(self.newStart, self.beamDirection, self.color)                  
                self.Rays.append(beamRay) #add new ray to the array in line 55 above
        
class pointSource():
    '''A point source class'''
    def __init__(self, position):
        self.position = position
        self.centerRay = ray(self.position, vec(1,0,0), color.green) #Starts with a set ray in the x-direction
        self.color = self.centerRay.color
        self.Rays = [] #An array to store other rays that will form the point source
        self.source = sphere(pos = vec(self.centerRay.startPoint), radius = 0.01, color = self.color) #Visual object to represent the source
        
        self.draw_pointSource()
    
    def draw_pointSource(self):
        '''Draw a point source''' #Steps through a set of parametric equations for the surface of a sphere in cartesian coordinates
        self.source
        c = pi/7 
        
        for i in numpy.arange(0,2*pi+c,c): #loops a half circle around in the xy-plane
            for k in numpy.arange(0,2*pi+c,c): #Loops a full circle around in the yz-plane, carrying the last loop into the z-dimension
                
                x = self.centerRay.length*(cos(k)*cos(i)) #Parametric equation for a sphere
                y = self.centerRay.length*(cos(k)*sin(i))
                z = self.centerRay.length*(sin(k))
                
                pointSourceRay = ray(self.position, vec(x,y,z), self.centerRay.color) #Creates a new ray based on the values obtained from the loop (A new point on the sphere)
                
                self.Rays.append(pointSourceRay) #Adds the newly created ray to the array in line 113
                
class lens():
     '''a lens to focus light rays in terms of Radius of Curvature'''
    def __init__ (self, position, axis, n, world_n, curveRad, perpRad): #User-input parameters