In [ ]:

    

%%javascript
//Imports the javascript code - not required if the extensions are already installed
require(['/tree/py3dmol/py3dmol/molviz3d.js','/tree/py3dmol/py3dmol/molviz2d.js',
         '/tree/py3dmol/py3dmol/3Dmol-min.js','/tree/py3dmol/py3dmol/d3.v3.min.js'],
        function(){console.log('JS extensions loaded.')})


    

In [ ]:

    

%pylab inline


    

In [ ]:

    

import cStringIO as cio

import numpy as np
import pandas
import numba

import buckyball as bb
import buckyball.units as u


    

In [ ]:

    

"""Set up our unit system"""
u.ureg.define('earth_mass = 5.97219e24 * kg')
u.em = u.ureg.earth_mass
u.ureg.define('earth_radius = 7917.5/2.0 * miles')
u.au = u.ureg.au
u.yj = u.ureg.yottajoule
u.G = u.ureg.newtonian_constant_of_gravitation

u.default.length = u.au
u.default.mass = u.em
u.default.energy = u.yj
u.default.time = u.ureg.day


    

In [ ]:

    

"""Set up astronomical data"""
data = pandas.read_csv('data/solar.csv')
planets = {}
for name,mass,orbit,velocity in zip(data.Name,
                                    data['Mass (Earth=1)'],
                                    data['Mean Distance from Sun (AU)'],
                                    data['Mean Orbital Velocity (km/sec)']):
    planets[name] = {'mass':mass*u.em,
                    'orbit':orbit*u.au,
                    'speed':velocity*u.ureg.km/u.ureg.seconds}
def make_planet(name):
    planet = planets[name]
    atom = bb.Atom(name=name, mass=planet['mass'], atnum=0)
    atom.x = planet['orbit']
    atom.momentum[1] = planet['speed'] * atom.mass
    return atom


    

In [ ]:

    

"""Define the potential"""
class Gravity(bb.basemethods.EnergyModelBase):
    def prep(self):
        n = self.mol.num_atoms
        self.indices = np.triu_indices(n,k=1)
        self.forcebuffer = np.zeros((n,3)) * u.default.force
        self._prepped = True
        mg = u.to_units_array([u.G*self.mol.atoms[i].mass*self.mol.atoms[j].mass
                                    for i,j in zip(*self.indices)])
        self.mass_grav = mg


    def calculate(self,requests=None):
        if not self._prepped: self.prep()

        energy,forces = self.fast_gravity()
        
        return {'potential_energy':u.default.convert(energy),
               'forces':u.default.convert(forces)}
    
    def fast_gravity(self):
        dists = self.mol.calc_distance_array(self.mol)[self.indices]
        disps = self.mol.calc_displacements()[self.indices]
        base_array = self.mass_grav/dists
        potential = -np.sum(base_array)
        f1 = base_array/(dists**2)
        f2 = np.multiply(f1,disps.T).T
        self.forcebuffer *= 0.0
        for ind,(i,j) in enumerate(zip(*self.indices)):
            self.forcebuffer[3*i:3*i+3] += f2[ind]
            self.forcebuffer[3*j:3*j+3] -= f2[ind]
        return potential,self.forcebuffer


    

In [ ]:

    

"""Timing..."""
twobody = bb.Molecule([make_planet('Earth'),make_planet('Sun')],
                       name='EarthAndSun')
twobody.set_energy_model(Gravity())
#twobody.energy_model.calculate()
#%timeit twobody.energy_model.calculate()


    

In [ ]:

    

"""Run a simulation of the earth and sun"""
twobody = bb.Molecule([make_planet('Earth'),make_planet('Sun')],
                       name='EarthAndSun')
twobody.set_energy_model(Gravity())
print 'Starting potential energy: ',twobody.calc_potential_energy()
vv = bb.integrators.VelocityVerlet(frame_interval=20, timestep=1.0*u.ureg.day)
twobody.set_integrator(vv)
duration = 5.0*u.ureg.year
traj = twobody.run(duration)
print 'Created trajectory with %d frames over %s'%(len(traj),duration)


    

In [ ]:

    

twobody.energy_model.mass_grav


    

In [ ]:

    

"""Plot the results"""
print 'There are %d frames.'%len(traj)
plot(traj.time,traj.positions[:,0],label='x')
plot(traj.time,traj.positions[:,1],label='y')
xlabel('time / days'); ylabel('coordinate / Ast. Unit');grid(); legend(loc='lower left')
title("Earth's location relative to the sun")


    

In [ ]:

    

"""A bigger simulation"""
energy_model=Gravity()
integrator=bb.integrators.VelocityVerlet(timestep=1.0*u.ureg.day,
                                        frame_interval=10)
solar_system = bb.Molecule([make_planet(n) for n in data.Name],
                           energy_model=energy_model,
                          integrator=integrator)
traj = solar_system.run(0.5*u.ureg.year)


    

In [ ]: