From 57c92ff890d106ccb5831f915b4daaff220b7bc0 Mon Sep 17 00:00:00 2001 From: Hanno Rein Date: Fri, 27 Oct 2023 09:47:08 -0400 Subject: [PATCH] python refactor --- rebound/orbit.py | 105 +++++++++++++++++ rebound/simulation.py | 255 +----------------------------------------- rebound/variation.py | 152 +++++++++++++++++++++++++ 3 files changed, 262 insertions(+), 250 deletions(-) create mode 100644 rebound/orbit.py create mode 100644 rebound/variation.py diff --git a/rebound/orbit.py b/rebound/orbit.py new file mode 100644 index 000000000..5d9c896b9 --- /dev/null +++ b/rebound/orbit.py @@ -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 "".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) + diff --git a/rebound/simulation.py b/rebound/simulation.py index a116f06ba..6bc5d8e1e 100644 --- a/rebound/simulation.py +++ b/rebound/simulation.py @@ -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 * @@ -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 "".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. @@ -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)] @@ -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), diff --git a/rebound/variation.py b/rebound/variation.py new file mode 100644 index 000000000..fb351caba --- /dev/null +++ b/rebound/variation.py @@ -0,0 +1,152 @@ +import ctypes +from . import SimulationError +from .particle import Particle +from .simulation import Simulation + +class Variation(ctypes.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", ctypes.POINTER(Simulation)), + ("order", ctypes.c_int), + ("index", ctypes.c_int), + ("testparticle", ctypes.c_int), + ("index_1st_order_a", ctypes.c_int), + ("index_1st_order_b", ctypes.c_int), + ("_lrescale", ctypes.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(ctypes.addressof(sim._particles.contents)+self.index*ctypes.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 = ctypes.c_double(value) + sim.var_config[i]._lrescale = ctypes.c_double(value) + return + + raise SimulationError("An error occured while trying to find variational struct in simulation.") +