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 [ ]: