Source code for irfpy.util.keplernumeric

r""" Kepler's trajectory using numeric solvers

Kepler's trajectory is determined by parameters together with
the initial conditions.

A generic form of the equation of motion is

.. math::

    \frac{d\mathbf{v}}{dt} = -GM\frac{\hat{\mathbf{r}}}{r^2}

where G is the constant of gravity and M is the mass of the center body
(or center of mass).

:class:`Summary` reports the details about two-body problems

:class:`Kepler2RungeKutta` is the most reliable one for
non-reltivistic Kepler motion with centrifugal force supported.

.. math::

        \frac{dx}{dt} = v_x \\
        \frac{dy}{dt} = v_y \\
        \frac{dz}{dt} = v_z \\
        \frac{dv_x}{dt} = -Gm_0 x / r^3 + \omega^2 x +
                                2v_y\omega_z - 2 v_z\omega_y\\
        \frac{dv_y}{dt} = -Gm_0 y / r^3 + \omega^2 y +
                                2v_z\omega_x - 2 v_x\omega_z\\
        \frac{dv_z}{dt} = -Gm_0 z / r^3 + \omega^2 z +
                                2v_x\omega_y - 2 v_y\omega_x

here *r* is :math:`\sqrt{x^2+y^2+z^2}`.


.. autosummary::

    Kepler2RungeKutta
    Kepler
    Kepler_SunEarth
    Summary
"""

import math
import numpy as np
from scipy import interpolate

from irfpy.util.leapflog import LeapFlog


[docs]class Kepler: r''' Simulate Kepler motion using Leap Flog method. ''' def __init__(self, mass0, mass1, xinit, vinit, dt, gravity_constant=6.674e-11): ''' Initalizer :param mass0: Mass of the central body (or equivalent mass in the center of mass frame) MKSA system is assumed, unless you specify gravity_constant. :param mass1: Mass of the body of interest. It does not affect to the motion, but kept it by historical reason :( :param xinit: Initial position of mass1 at t=0 :type xinit: ``np.array`` with (3,) :param vinit: Initial velocity of mass1 at t=0 :type vinit: ``np.array`` with (3,) :param dt: Time resolution. Negative *dt* will be backtracing. :keyword gravity_constant: Gravity constant. The MKSA value for default. USAGE: Let's see the Io motion around Jupiter. The mass of Jupiter is 1.89686e27 kg, and Io is 8.932e22 (which does not affect the calculation). Put the Io at 4.217e5 km with velocity of 17.334 km/s at t=0. Time resolution is set 3600 s (1 hour). >>> kepler = Kepler(1.89686e27, 8.932e22, ... np.array((4.217e8, 0, 0)), ... np.array((0, 1.7334e4, 0)), ... 3600) Then, you can proceed the calculation by 100 steps. This means that the calculation was done until 3.6e5 s (100 hours). >>> kepler.proceed(100) # 100 step forward. You may find the position at t=7.6e4 s, which is about the half of the orbital period of Io (full orbital period is 1.52853e5 s). Thus, the position should be close to [-4.217e8, 0, 0] with velocity of (0, -1.7334e4, 0). >>> pos = kepler.get_pos(7.6e4) >>> print(pos) # doctest: +NORMALIZE_WHITESPACE [-4.26590115e+08 1.92925628e+07 0.00000000e+00] >>> vel = kepler.get_vel(7.6e4) >>> print(vel) # doctest: +NORMALIZE_WHITESPACE [ -778.52664693 -17145.37459877 0. ] ''' self._m0 = mass0 self._m1 = mass1 self._xinit = xinit self._vinit = vinit self._dt = dt self._G = gravity_constant self._xlist = [xinit] self._xtlist = [0] self._vlist = [vinit] self._vtlist = [0] self._force = lambda r: ( -self._G * self._m0 * r / ((r * r).sum() ** 1.5)) vnow = LeapFlog.progress_velocity( self._vinit, self._force(self._xinit), 0.5 * self._dt) self._vlist.append(vnow) self._vtlist.append(0.5 * self._dt) self._interpos = None self._intervel = None
[docs] def proceed(self, niter): ''' Proceed n steps. ''' for i in range(niter): self.step()
[docs] def step(self): ''' Proceed 1 step. ''' x0 = self._xlist[-1] # Position at t=t0 v0 = self._vlist[-1] # Velocity at t=t0+dt/2 dt = self._dt x1 = LeapFlog.progress_position(x0, v0, dt) # Position at t=t0+dt v1 = LeapFlog.progress_velocity(v0, self._force(x1), dt) # V at t0+dt*2 self._xlist.append(x1) self._xtlist.append(self._xtlist[-1] + dt) self._vlist.append(v1) self._vtlist.append(self._vtlist[-1] + dt) self._interpos = None self._intervel = None
[docs] def get_xlist(self): ''' Return the list of positions. :returns: Tuple. (xtlist, xlist), where xtlist is (N,) shape and xlist is (N, 3) shape. ''' return np.array(self._xtlist), np.array(self._xlist)
[docs] def get_vlist(self): ''' Return the list of velocity. :returns: Tuple. (vtlist, vlist), where vtlist is (N,) shape and vlist is (N, 3) shape. ''' return np.array(self._vtlist), np.array(self._vlist)
[docs] def get_pos(self, t): ''' Return the position, interpolated. :param t: Scalar or tuple of the time. Should be interpolated. :returns: np.array with (3,) or (3, N) shape. >>> kep = Kepler_SunEarth(dt=86400) >>> kep.proceed(100) # 100 day move. >>> tx, xyz = kep.get_xlist() # Return results in list format. >>> print(tx[38]) 3283200 >>> x, y, z = xyz[38, :] # Position of 38th data. >>> print('%.5e %.5e' % (x, y)) 1.18759e+11 9.09623e+10 >>> x, y, z = kep.get_pos(3300000) >>> print('%.5e %.5e' % (x, y)) 1.18454e+11 9.13590e+10 >>> kep.proceed(100) # 100 day more. >>> x, y, z = kep.get_pos(86400 * 180.38) # At 180.38 days >>> print('%.5e %.5e' % (x, y)) -1.49434e+11 5.65757e+09 ''' if self._interpos is None: tx, xyz = self.get_xlist() if self._dt > 0: self._interpos = [ # kind=3 for 3-d spline interpolate.interp1d(tx, xyz[:, 0], kind=3), interpolate.interp1d(tx, xyz[:, 1], kind=3), interpolate.interp1d(tx, xyz[:, 2], kind=3)] else: self._interpos = [ # kind=3 for 3-d spline interpolate.interp1d(tx[::-1], xyz[::-1, 0], kind=3), interpolate.interp1d(tx[::-1], xyz[::-1, 1], kind=3), interpolate.interp1d(tx[::-1], xyz[::-1, 2], kind=3)] x = self._interpos[0](t) y = self._interpos[1](t) z = self._interpos[2](t) return np.array([x, y, z])
[docs] def get_vel(self, t): ''' Return the velocity. ''' if self._intervel is None: tv, vxyz = self.get_vlist() if self._dt > 0: self._interpos = [ # kind=3 for 3-d spline interpolate.interp1d(tv, vxyz[:, 0], kind=3), interpolate.interp1d(tv, vxyz[:, 1], kind=3), interpolate.interp1d(tv, vxyz[:, 2], kind=3)] else: self._interpos = [ # kind=3 for 3-d spline interpolate.interp1d(tv[::-1], vxyz[::-1, 0], kind=3), interpolate.interp1d(tv[::-1], vxyz[::-1, 1], kind=3), interpolate.interp1d(tv[::-1], vxyz[::-1, 2], kind=3)] vx = self._interpos[0](t) vy = self._interpos[1](t) vz = self._interpos[2](t) return np.array([vx, vy, vz])
[docs]def Kepler_SunEarth(dt=3600): r''' Return :class:`Kepler` class for Sun-Earth system >>> kse = Kepler_SunEarth() >>> kse.proceed(24 * 365) # About 1 year. >>> tx, x = kse.get_xlist() >>> print('%.2e %.2e %.2e' % (x[-1, 0], x[-1, 1], x[-1, 2])) # Last data 1.50e+11 -3.11e+08 0.00e+00 ''' msun = 1.9884e30 # kg mearth = 5.972e24 # kg rearth = 149597870700. # m vearth = 29780. # m/s dphiearth = vearth / rearth # rad/s kep = Kepler(msun, mearth, np.array([rearth, 0, 0]), np.array([0, vearth, 0]), dt) return kep
[docs]def Kepler_EarthMoon(): r''' Return :class:`Kepler` class for Earth-moon system ''' raise NotImplementedError('Not yet implemented:')
[docs]class KeplerRungeKutta: r''' Kepler motion solver using Runge-Kutta. Similar functionality to :class:`Kepler`, but using classical RK solver. No leap flog used. Equation of motion is .. math:: \frac{dx}{dt} = v_x \\ \frac{dy}{dt} = v_y \\ \frac{dz}{dt} = v_z \\ \frac{dv_x}{dt} = -Gm_0 x / r^3 \\ \frac{dv_y}{dt} = -Gm_0 y / r^3 \\ \frac{dv_z}{dt} = -Gm_0 z / r^3 here r is :math:`\sqrt{x^2+y^2+z^2}`. >>> rkkep = KeplerRungeKutta(1e27, 1.5, (5e8, 0, 0), (0, 20e3, 0)) >>> x, y, z, vx, vy, vz = rkkep.get_posvel(np.arange(0, 1000, 30)).T >>> print('%.4e %.4e %.4e' % (x[-1], y[-1], z[-1])) 4.9987e+08 1.9798e+07 0.0000e+00 Compare with :class:`Kepler` class. >>> kep = Kepler(1e27, 1.5, np.array((5e8, 0, 0)), ... np.array((0, 20e3, 0)), 10) >>> kep.proceed(150) >>> x, y, z = kep.get_pos(990) >>> print('%.4e %.4e %.4e' % (x, y, z)) 4.9987e+08 1.9798e+07 0.0000e+00 ''' def __init__(self, mass0, mass1, xinit, vinit, gravity_constant=6.674e-11): ''' :param mass0: The center object. :param mass1: The flying object. It does not affect the results. :param xinit: (3,) shaped numpy.array to specify the initial position of mass1. :param vinit: (3,) shpaed numpy.array to specify the initial velocity of mass1. :param gravity_constant: The gravity constant. You may change it if the equation and variables are normalized. ''' self.mass0 = mass0 self.mass1 = mass1 self.xinit = xinit self.vinit = vinit self.GM = mass0 * gravity_constant def _rhs(self, variables): x, y, z, vx, vy, vz = variables r2 = x * x + y * y + z * z sq3r2 = np.sqrt(r2) ** 3 rhs1 = vx rhs2 = vy rhs3 = vz rhs4 = -self.GM * x / sq3r2 rhs5 = -self.GM * y / sq3r2 rhs6 = -self.GM * z / sq3r2 return np.array((rhs1, rhs2, rhs3, rhs4, rhs5, rhs6))
[docs] def get_posvel(self, tlist): ''' Return the pos, vel ''' x0, y0, z0 = self.xinit vx0, vy0, vz0 = self.vinit initcond = np.array([x0, y0, z0, vx0, vy0, vz0]) f = lambda y, t: self._rhs(y) import scipy.integrate y = scipy.integrate.odeint(f, initcond, tlist) return y
[docs] def get_pos(self, tlist): return self.get_posvel(tlist)[:, :3]
[docs] def get_vel(self, tlist): return self.get_posvel(tlist)[:, 3:]
[docs]class Kepler2RungeKutta: r''' Kepler motion seen in a rotation frame. Kepler motion in a rotation frame. The center of the rotation MUST be at the origin, where massive object (mass0) is located. Note that the :math:`\omega=0` case, the results should return exactly the same as :class:`KeplerRungeKutta`. Probably it can handle the "center of gravity" frame, but not yet checked... Equation of motion is: .. math:: \frac{dx}{dt} = v_x \\ \frac{dy}{dt} = v_y \\ \frac{dz}{dt} = v_z \\ \frac{dv_x}{dt} = -Gm_0 x / r^3 + \omega^2 x + 2v_y\omega_z - 2 v_z\omega_y\\ \frac{dv_y}{dt} = -Gm_0 y / r^3 + \omega^2 y + 2v_z\omega_x - 2 v_x\omega_z\\ \frac{dv_z}{dt} = -Gm_0 z / r^3 + \omega^2 z + 2v_x\omega_y - 2 v_y\omega_x here r is :math:`\sqrt{x^2+y^2+z^2}`. >>> rkkep = Kepler2RungeKutta(1e27, 1.5, (5e8, 0, 0), (0, 20e3, 0), ... omega=np.array([0.0, 0.0, 0.0])) The simplest way of obtaining the position and velocity is to give a list of the time to integrate. Usually, it starts from 0 and increases monotonically with a specific time step. Note that the integral starts with the first element of the time step. >>> x, y, z, vx, vy, vz = rkkep.get_posvel(np.arange(0, 30001, 30)).T >>> print('%.4e %.4e %.4e' % (x[0], y[0], z[0])) 5.0000e+08 0.0000e+00 0.0000e+00 >>> print('%.4e %.4e %.4e' % (x[-1], y[-1], z[-1])) # Data on t=30000. 4.0260e+08 5.6589e+08 0.0000e+00 If you switch on the rotation, you can instance as >>> rkkep_rott = Kepler2RungeKutta(1e27, 1.5, (5e8, 0, 0), (0, 15e3, 0), ... omega=np.array([0, 0, 5e3 / 5e8])) Here the rotation frame is rotating with angular frequency of 5e3 / 5e8 = 1e-5 rad/s. This means that the initial velocity of the particle, measured in rotation frame, is 5e3 less than that seen in inertia. This corresponds to 15e3. >>> x, y, z, vx, vy, vz = rkkep_rott.get_posvel(np.arange(0, 30001, 30)).T >>> print('%.4e %.4e %.4e' % (x[0], y[0], z[0])) 5.0000e+08 0.0000e+00 0.0000e+00 >>> print('%.4e %.4e %.4e' % (x[-1], y[-1], z[-1])) # Data on t=30000. 5.5185e+08 4.2164e+08 0.0000e+00 The final coordinates at t=30000 in rotation frame can be (have been) validated from the rotational conversion using :math:`\omega t=0.3` rad. ''' def __init__(self, mass0, mass1, xinit, vinit, omega=np.array([0.0, 0.0, 0.0]), gravity_constant=6.674e-11): ''' :param mass0: The center object. :param mass1: The flying object. It does not affect the results. :param xinit: (3,) shaped numpy.array to specify the initial position of mass1. Given in the rotation frame. :param vinit: (3,) shpaed numpy.array to specify the initial velocity of mass1. Given in the rotation frame. :param omega: (3,) shpaed the angular velocity vector of the rotation frame seen from inertia frame. :param gravity_constant: The gravity constant. You may change it if the equation and variables are normalized. ''' self.mass0 = mass0 self.mass1 = mass1 self.xinit = xinit self.vinit = vinit self.GM = mass0 * gravity_constant self.w = omega def _rhs(self, variables): x, y, z, vx, vy, vz = variables r2 = x * x + y * y + z * z wx, wy, wz = self.w w2 = wx * wx + wy * wy + wz * wz sq3r2 = np.sqrt(r2) ** 3 rhs1 = vx rhs2 = vy rhs3 = vz rhs4 = -self.GM * x / sq3r2 + w2 * x + 2 * vy * wz - 2 * vz * wy rhs5 = -self.GM * y / sq3r2 + w2 * y + 2 * vz * wx - 2 * vx * wz rhs6 = -self.GM * z / sq3r2 + w2 * z + 2 * vx * wy - 2 * vy * wx return np.array((rhs1, rhs2, rhs3, rhs4, rhs5, rhs6))
[docs] def get_posvel(self, tlist): ''' Return the pos, vel Get the posistion and velocity as (N, 3) shape where N is the shape of tlist. Note that the tlist should start from t=0. If tlist[0] is not 0, the initial conditions are considered at the time of tlist[0]. Thus, the results will be "shifted". ''' x0, y0, z0 = self.xinit vx0, vy0, vz0 = self.vinit initcond = np.array([x0, y0, z0, vx0, vy0, vz0]) f = lambda y, t: self._rhs(y) import scipy.integrate y = scipy.integrate.odeint(f, initcond, tlist) return y
[docs] def get_pos(self, tlist): return self.get_posvel(tlist)[:, :3]
[docs] def get_vel(self, tlist): return self.get_posvel(tlist)[:, 3:]
[docs]class Summary: ''' Report summary of the two body problem. See http://www.mns.kyutech.ac.jp/~kishine/notes/sentan_note101130e.pdf for summary of conservables. >>> msun = 1.9884e30 # kg >>> mearth = 5.972e24 # kg >>> rearth = 149597870700. #m >>> vearth = 29780. # m/s >>> summary = Summary(msun, mearth, [rearth, 0, 0], [0, vearth, 0]) >>> # summary.print_summary() # will show the summary ''' def __init__(self, mass0, mass1, xinit, vinit, gravity_constant=6.674e-11): self._m0 = mass0 self._m1 = mass1 self._xinit = np.array(xinit) self._vinit = np.array(vinit) self._G = gravity_constant
[docs] def print_summary(self): ''' Print summary of the two-body system. ''' l = self.get_angular_momentum() print('angular momentum=', l, 'kg m^2 / s') arev = self.get_areal_velocity() print('areal velocity=', arev, 'm^2 / s (or equivalent)') ene = self.get_total_energy() print('ene=', ene, 'J (or equivalent)') lv = self.get_laplace_vector() print('laplace vector=', lv) ecc = self.get_eccentricity() print('ecc=', ecc)
[docs] def get_eccentricity_vector(self): ''' Return the eccentricity vector. Use :meth:`get_eccentricity` to get scalar data. ''' return self.get_laplace_vector() / self._G / self._m0
[docs] def get_eccentricity(self): ''' Return the eccentricity. If eccentricity, e, is 0, the motion is circular. For e between 0 and 1, the motion is eliptic. For e=1, the motion is parabolic. For e>1, the motion is hyperbolic. ''' return np.linalg.norm(self.get_eccentricity_vector())
[docs] def get_total_energy(self): ''' Return the total energy Usually the unit is in Joule, if you instance this object in MKSA system. ''' m0 = self._m0 m1 = self._m1 r_sq = (self._xinit ** 2).sum() v_sq = (self._vinit ** 2).sum() em = m1 * v_sq / 2.0 ep = self._G * m0 * m1 / np.sqrt(r_sq) return em - ep
[docs] def get_angular_momentum(self): ''' Return the angular momentum. If you specify the initial values in MKSA system, the angular momentum vector is in the unit of :math:`kg m^2/s`. ''' return self._m1 * np.cross(self._xinit, self._vinit)
[docs] def get_areal_velocity(self): ''' Return the areal velcoity vector ''' return self.get_angular_momentum() / self._m1
[docs] def get_laplace_vector(self): ''' Return the Laplace vector ''' h = self.get_areal_velocity() a0 = np.cross(h, self._vinit) a1 = self._G * self._m0 * self._xinit / np.linalg.norm(self._xinit) return a0 + a1
import unittest import doctest
[docs]def doctests(): return unittest.TestSuite(( doctest.DocTestSuite(), ))
if __name__ == '__main__': unittest.main(defaultTest='doctests')