Skip to content

Commit

Permalink
python refactor
Browse files Browse the repository at this point in the history
  • Loading branch information
@hannorein
hannorein committed Oct 27, 2023
1 parent 074f250 commit 57c92ff
Show file tree
Hide file tree
Showing 3 changed files with 262 additions and 250 deletions.
105 changes: 105 additions & 0 deletions rebound/orbit.py
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,105 @@
import ctypes
from .tools import M_to_E
from .vectors import Vec3dBasic

class Orbit(ctypes.Structure):
"""
A class containing orbital parameters for a particle.
This is an abstraction of the reb_orbit data structure in C.
When using the various REBOUND functions using Orbits, all angles are in radians.
The following image illustrated the most important angles used.
In REBOUND the reference direction is the positive x direction, the reference plane
is the xy plane.
.. image:: images/orbit.png
:width: 500px
:height: 450px
Image from wikipedia. CC-BY-SA-3.
Attributes
----------
d : float
radial distance from reference
v : float
velocity relative to central object's velocity
h : float
specific angular momentum
P : float
orbital period (negative if hyperbolic)
n : float
mean motion (negative if hyperbolic)
a : float
semimajor axis
e : float
eccentricity
inc : float
inclination
Omega : float
longitude of ascending node
omega : float
argument of pericenter
pomega : float
longitude of pericenter
f : float
true anomaly
M : float
mean anomaly
E : float
eccentric anomaly (requires solving Kepler's equation - only calculated when needed)
l : float
mean longitude = Omega + omega + M
theta : float
true longitude = Omega + omega + f
T : float
time of pericenter passage
rhill : float
Hill radius ( =a*pow(m/(3M),1./3.) )
pal_h : float
Cartesian component of the eccentricity ( = e*sin(pomega) )
pal_k : float
Cartesian component of the eccentricity ( = e*cos(pomega) )
pal_ix : float
Cartesian component of the inclination ( = 2*sin(i/2)*cos(Omega) )
pal_iy : float
Cartesian component of the inclination ( = 2*sin(i/2)*sin(Omega) )
hvec : Vec3d
Specific angular momentum vector
evec : Vec3d
Eccentricity vector (mag = eccentricity, points toward pericenter)
"""
_fields_ = [("d", ctypes.c_double),
("v", ctypes.c_double),
("h", ctypes.c_double),
("P", ctypes.c_double),
("n", ctypes.c_double),
("a", ctypes.c_double),
("e", ctypes.c_double),
("inc", ctypes.c_double),
("Omega", ctypes.c_double),
("omega", ctypes.c_double),
("pomega", ctypes.c_double),
("f", ctypes.c_double),
("M", ctypes.c_double),
("l", ctypes.c_double),
("theta", ctypes.c_double),
("T", ctypes.c_double),
("rhill", ctypes.c_double),
("pal_h", ctypes.c_double),
("pal_k", ctypes.c_double),
("pal_ix", ctypes.c_double),
("pal_iy", ctypes.c_double),
("hvec", Vec3dBasic),
("evec", Vec3dBasic)]

def __str__(self):
"""
Returns a string with the semi-major axis and eccentricity of the orbit.
"""
return "<rebound.Orbit instance, a={0} e={1} inc={2} Omega={3} omega={4} f={5}>".format(str(self.a),str(self.e), str(self.inc), str(self.Omega), str(self.omega), str(self.f))

@property
def E(self):
return M_to_E(self.e, self.M)

255 changes: 5 additions & 250 deletions rebound/simulation.py
View file
Original file line number Diff line number Diff line change
@@ -1,8 +1,9 @@
from ctypes import Structure, c_double, POINTER, c_uint32, c_float, c_int, c_uint, c_uint32, c_int64, c_long, c_ulong, c_ulonglong, c_void_p, c_char_p, CFUNCTYPE, byref, create_string_buffer, addressof, pointer, cast, c_char, c_size_t, string_at, sizeof
from . import clibrebound, Escape, NoParticles, Encounter, Collision, SimulationError, ParticleNotFound, M_to_E
from . import clibrebound, Escape, NoParticles, Encounter, Collision, SimulationError, ParticleNotFound
from .citations import cite
from .particle import Particle
from .particles import *
from .particles import Particles
from .orbit import Orbit
from .units import units_convert_particle, check_units, convert_G, hash_to_unit
from .hash import hash as rebhash, HashPointerPair
from .vectors import *
Expand Down Expand Up @@ -48,107 +49,6 @@ def __repr__(self):



class Orbit(Structure):
"""
A class containing orbital parameters for a particle.
This is an abstraction of the reb_orbit data structure in C.
When using the various REBOUND functions using Orbits, all angles are in radians.
The following image illustrated the most important angles used.
In REBOUND the reference direction is the positive x direction, the reference plane
is the xy plane.
.. image:: images/orbit.png
:width: 500px
:height: 450px
Image from wikipedia. CC-BY-SA-3.
Attributes
----------
d : float
radial distance from reference
v : float
velocity relative to central object's velocity
h : float
specific angular momentum
P : float
orbital period (negative if hyperbolic)
n : float
mean motion (negative if hyperbolic)
a : float
semimajor axis
e : float
eccentricity
inc : float
inclination
Omega : float
longitude of ascending node
omega : float
argument of pericenter
pomega : float
longitude of pericenter
f : float
true anomaly
M : float
mean anomaly
E : float
eccentric anomaly (requires solving Kepler's equation - only calculated when needed)
l : float
mean longitude = Omega + omega + M
theta : float
true longitude = Omega + omega + f
T : float
time of pericenter passage
rhill : float
Hill radius ( =a*pow(m/(3M),1./3.) )
pal_h : float
Cartesian component of the eccentricity ( = e*sin(pomega) )
pal_k : float
Cartesian component of the eccentricity ( = e*cos(pomega) )
pal_ix : float
Cartesian component of the inclination ( = 2*sin(i/2)*cos(Omega) )
pal_iy : float
Cartesian component of the inclination ( = 2*sin(i/2)*sin(Omega) )
hvec : Vec3d
Specific angular momentum vector
evec : Vec3d
Eccentricity vector (mag = eccentricity, points toward pericenter)
"""
_fields_ = [("d", c_double),
("v", c_double),
("h", c_double),
("P", c_double),
("n", c_double),
("a", c_double),
("e", c_double),
("inc", c_double),
("Omega", c_double),
("omega", c_double),
("pomega", c_double),
("f", c_double),
("M", c_double),
("l", c_double),
("theta", c_double),
("T", c_double),
("rhill", c_double),
("pal_h", c_double),
("pal_k", c_double),
("pal_ix", c_double),
("pal_iy", c_double),
("hvec", Vec3dBasic),
("evec", Vec3dBasic)]

def __str__(self):
"""
Returns a string with the semi-major axis and eccentricity of the orbit.
"""
return "<rebound.Orbit instance, a={0} e={1} inc={2} Omega={3} omega={4} f={5}>".format(str(self.a),str(self.e), str(self.inc), str(self.Omega), str(self.omega), str(self.f))

@property
def E(self):
return M_to_E(self.e, self.M)

class Simulation(Structure):
"""
This is the REBOUND Simulation Class.
Expand Down Expand Up @@ -1585,153 +1485,6 @@ def tree_update(self):
"""
clibrebound.reb_simulation_update_tree(byref(self))

class Variation(Structure):
"""
REBOUND Variational Configuration Object.
This object encapsulated the configuration of one set of variational
equations in a REBOUND simulation. It is an abstraction of the
C struct reb_variational_configuration.
None of the fields in this struct should be changed after it has
been initialized.
One rebound simulation can include any number of first and second order
variational equations.
Note that variations are only encoded as particles for convenience.
A variational particle's position and velocity should be interpreted as a derivative, i.e. how much that position or velocity varies with respect to the first or second-order variation.
See ipython_examples/VariationalEquations.ipynb and Rein and Tamayo (2016) for details.
Examples
--------
>>> sim = rebound.Simulation() # Create a simulation
>>> sim.add(m=1.) # Add a star
>>> sim.add(m=1.e-3, a=1.) # a planet
>>> var_config = sim.add_variation() # Add a set of variational equations.
>>> var_config.particles[1].x = 1. # Initialize the variational particle corresponding to the planet
>>> sim.integrate(100.) # Integrate the simulation
>>> print(var_config.particles[0].vx) # Print the velocity of the variational particle corresponding to the star
"""
_fields_ = [
("_sim", POINTER(Simulation)),
("order", c_int),
("index", c_int),
("testparticle", c_int),
("index_1st_order_a", c_int),
("index_1st_order_b", c_int),
("_lrescale", c_double)]

def vary(self, particle_index, variation, variation2=None, primary=None):
"""
This function can be used to initialize the variational particles that are
part of a Variation.
Note that rather than using this convenience function, one can
also directly manipulate the particles' coordinate using the following
syntax:
>>> var = sim.add_variation()
>>> var.particles[0].x = 1.
The ``vary()`` function is useful for initializing variations corresponding to
changes in one of the orbital parameters for a particle on a bound
Keplerian orbit.
The function supports both first and second order variations in the following
classical orbital parameters:
a, e, inc, omega, Omega, f
as well as the Pal (2009) coordinates:
a, h, k, ix, iy, lambda
and in both cases the mass m of the particle. The advantage of the Pal coordinate
system is that all derivatives are well behaved (infinitely differentiable).
Classical orbital parameters on the other hand exhibit coordinate singularities,
for example when e=0.
The following example initializes the variational particles corresponding to a
change in the semi-major axis of the particle with index 1:
>>> var = sim.add_variation()
>>> var.vary(1,"a")
Parameters
----------
particle_index : int
The index of the particle that should be varied. The index starts at 0 and runs through N-1. The first particle added to a simulation receives the index 0, the second 1, and the on.
variation : string
This parameter determines which orbital parameter is varied.
variation2: string, optional
This is only used for second order variations which can depend on two varying parameters. If omitted, then it is assumed that the parameter variation is variation2.
primary: Particle, optional
By default, variational particles are created in the Heliocentric frame.
Set this parameter to use any other particles as a primary (e.g. the center of mass).
"""
if self.order==2 and variation2 is None:
variation2 = variation
if self._sim is not None:
sim = self._sim.contents
particles = sim.particles
else:
raise RuntimeError("Something went wrong. Cannot seem to find simulation corresponding to variation.")
if self.testparticle >= 0:
particles[self.index] = Particle(simulation=sim,particle=particles[particle_index], variation=variation, variation2=variation2, primary=primary)
else:
particles[self.index + particle_index] = Particle(simulation=sim,particle=particles[particle_index], variation=variation, variation2=variation2, primary=primary)

@property
def particles(self):
"""
Access the variational particles corresponding to this set of variational equations.
The function returns a list of particles which are sorted in the same way as those in
sim.particles
The particles are pointers and thus can be modified.
If there are N real particles, this function will also return a list of N particles (all of which are
variational particles).
"""
sim = self._sim.contents
ps = []
if self.testparticle>=0:
N = 1
else:
N = sim.N-sim.N_var

ParticleList = Particle*N
ps = ParticleList.from_address(addressof(sim._particles.contents)+self.index*sizeof(Particle))
return ps

@property
def lrescale(self):
"""
Access the lrescale parameter.
This is a property because sim.add_variation() returns a copy of the struct, so need to find up-to-date reb_variational_configuration struct in simulation.
"""
sim = self._sim.contents
for i in range(sim.N_var_config):
if sim.var_config[i].index == self.index:
return sim.var_config[i]._lrescale
raise SimulationError("An error occured while trying to find variational struct in simulation.")

@lrescale.setter
def lrescale(self, value):
"""
Set the lrescale parameter.
This is a property because sim.add_variation() returns a copy of the struct, so need to find up-to-date reb_variational_configuration struct in simulation.
"""
sim = self._sim.contents
for i in range(sim.N_var_config):
if sim.var_config[i].index == self.index:
self._lrescale = c_double(value)
sim.var_config[i]._lrescale = c_double(value)
return

raise SimulationError("An error occured while trying to find variational struct in simulation.")


class timeval(Structure):
_fields_ = [("tv_sec",c_long),("tv_usec",c_long)]
Expand Down Expand Up @@ -1780,6 +1533,8 @@ class DisplayData(Structure):
from .integrators.saba import IntegratorSABA
from .integrators.mercurius import IntegratorMercurius

from .variation import Variation

# Setting up fields after class definition (because of self-reference)
Simulation._fields_ = [
("t", c_double),
Expand Down
Loading

0 comments on commit 57c92ff

Please sign in to comment.