rotation_matrix: Create specialized versions.

src/poliastro/core/elements.py

@@ -8,7 +8,7 @@
 from numpy.linalg import norm
 
 from poliastro.core.angles import E_to_nu, F_to_nu
-from poliastro.core.util import rotation_matrix
+from poliastro.core.util import rotation_matrix_x, rotation_matrix_z
 
 
 @jit
@@ -76,9 +76,9 @@ def rv_pqw(k, p, ecc, nu):
 @jit
 def coe_rotation_matrix(inc, raan, argp):
  """Create a rotation matrix for coe transformation"""
- r = rotation_matrix(raan, 2)
- r = r @ rotation_matrix(inc, 0)
- r = r @ rotation_matrix(argp, 2)
+ r = rotation_matrix_z(raan)
+ r = r @ rotation_matrix_x(inc)
+ r = r @ rotation_matrix_z(argp)
  return r
 
 

src/poliastro/core/iod.py

@@ -189,7 +189,7 @@ def izzo(k, r1, r2, tof, M, numiter, rtol):
  """
 
  assert tof > 0, "Assert tof must be positive"
- 
+
  # Check collinearity of r1 and r2
  if not cross(r1, r2).any():
  raise ValueError("Lambert solution cannot be computed for collinear vectors")
@@ -374,17 +374,15 @@ def _initial_guess(T, ll, M):
  x_0l = (((M * pi + pi) / (8 * T)) ** (2 / 3) - 1) / (
  ((M * pi + pi) / (8 * T)) ** (2 / 3) + 1
  )
- x_0r = (((8 * T) / (M * pi)) ** (2 / 3) - 1) / (
- ((8 * T) / (M * pi)) ** (2 / 3) + 1
- )
+ x_0r = (((8 * T) / (M * pi)) ** (2 / 3) - 1) / (((8 * T) / (M * pi)) ** (2 / 3) + 1)
 
  return [x_0l, x_0r]
 
 
 @jit
 def _initial_guess_M0(T, ll):
  """Initial guess for single revolution."""
- T_0 = np.arccos(ll) + ll * np.sqrt(1 - ll ** 2) # Equation 19
+ T_0 = np.arccos(ll) + ll * np.sqrt(1 - ll ** 2)  # Equation 19
  T_1 = 2 * (1 - ll ** 3) / 3 # Equation 21
  if T >= T_0:
  x_0 = (T_0 / T) ** (2 / 3) - 1

src/poliastro/core/thrust/change_ecc_inc.py

@@ -10,15 +10,24 @@
 from numpy import cross
 from numpy.linalg import norm
 
+from poliastro.core.angles import nu_to_E
 from poliastro.core.elements import rv2coe
 from poliastro.core.util import circular_velocity, eccentricity_vector
-from poliastro.core.angles import nu_to_E
+
 
 @jit
 def beta(ecc_0, ecc_f, inc_0, inc_f, argp):
- return np.arctan(np.sign(ecc_f - ecc_0) * np.sign(inc_f - inc_0) *
- (3 * np.pi / 4) * ((inc_f - inc_0) / np.cos(argp)) /
- (ecc_0 - ecc_f + np.log((ecc_f + 1) * (ecc_0 - 1) / ((ecc_f - 1) * (ecc_0 + 1)))))
+ return np.arctan(
+ np.sign(ecc_f - ecc_0)
+ * np.sign(inc_f - inc_0)
+ * (3 * np.pi / 4)
+ * ((inc_f - inc_0) / np.cos(argp))
+ / (
+ ecc_0
+ - ecc_f
+ + np.log((ecc_f + 1) * (ecc_0 - 1) / ((ecc_f - 1) * (ecc_0 + 1)))
+ )
+ )
  # # Note: "The argument of perigee will vary during the orbit transfer
  # # due to the natural drift and because e may approach zero.
  # # However, [the equation] still gives a good estimate of the desired
@@ -73,10 +82,10 @@ def a_d(t0, u_, k_):
  v_ = u_[3:]
  p, ecc, inc, raan, argp, nu = rv2coe(k_, r_, v_)
  beta_ = np.abs(beta_0) * np.sign(inc_f - inc_0)
- E = nu_to_E(nu, ecc) + .5 * np.pi
- if .5 * np.pi <= E < 1.5 * np.pi:
+ E = nu_to_E(nu, ecc) + 0.5 * np.pi
+ if 0.5 * np.pi <= E < 1.5 * np.pi:
  beta_ = -beta_
- w_ = cross(r_, v_) # Specific angular momentum vector
+ w_ = cross(r_, v_)  # Specific angular momentum vector
  w_ = w_ / norm(w_)
  accel_v = f * (np.cos(beta_) * thrust_unit + np.sin(beta_) * w_)
  return accel_v

src/poliastro/core/util.py

@@ -32,14 +32,14 @@ def circular_velocity(k, a):
 @jit
 def rotation_matrix(angle, axis):
  r"""Rotation matrix. Astropy uses the alias or passive transformation, whereas poliastro uses the alibi or active transformation
- 
-  Parameters
-  ----------
- 
-  angle : numpy.array
- 
-  axis : int
-  0, 1, or 2
+
+ Parameters
+ ----------
+
+ angle : numpy.array
+
+ axis : int
+ 0, 1, or 2
  """
  assert axis in (0, 1, 2)
  angle = np.asarray(angle)
@@ -57,6 +57,48 @@ def rotation_matrix(angle, axis):
  return R
 
 
+@jit
+def rotation_matrix_x(angle):
+ angle = np.asarray(angle)
+ c = cos(angle)
+ s = sin(angle)
+ R = np.zeros(angle.shape + (3, 3))
+ R[..., 0, 0] = 1.0
+ R[..., 1, 1] = c
+ R[..., 1, 2] = -s
+ R[..., 2, 1] = s
+ R[..., 2, 2] = c
+ return R
+
+
+@jit
+def rotation_matrix_y(angle):
+ angle = np.asarray(angle)
+ c = cos(angle)
+ s = sin(angle)
+ R = np.zeros(angle.shape + (3, 3))
+ R[..., 1, 1] = 1.0
+ R[..., 2, 2] = c
+ R[..., 2, 0] = -s
+ R[..., 0, 2] = s
+ R[..., 0, 0] = c
+ return R
+
+
+@jit
+def rotation_matrix_z(angle):
+ angle = np.asarray(angle)
+ c = cos(angle)
+ s = sin(angle)
+ R = np.zeros(angle.shape + (3, 3))
+ R[..., 2, 2] = 1.0
+ R[..., 0, 0] = c
+ R[..., 0, 1] = -s
+ R[..., 1, 0] = s
+ R[..., 1, 1] = c
+ return R
+
+
 @jit
 def alinspace(start, stop=None, num=50, endpoint=True):
  """Return increasing, evenly spaced angular values over a specified interval."""

src/poliastro/maneuver.py

@@ -1,6 +1,7 @@
 """Orbital maneuvers.
 
 """
+import numpy as np
 from astropy import units as u
 
 from poliastro.core.maneuver import (
@@ -10,7 +11,7 @@
 )
 from poliastro.iod.izzo import lambert as lambert_izzo
 from poliastro.util import norm
-import numpy as np
+
 
 class Maneuver:
  r"""Class to represent a Maneuver.
@@ -185,8 +186,8 @@ def lambert(cls, orbit_i, orbit_f, method=lambert_izzo, short=True, **kwargs):
 
  # Time of flight is solved by subtracting both orbit epochs
  tof = orbit_f.epoch - orbit_i.epoch
- assert tof > 0, f'Time of flight={tof} must be positive'
- 
+ assert tof > 0, f"Time of flight={tof} must be positive"
+
  # Compute all possible solutions to the Lambert transfer
  sols = list(method(k, r_i, r_f, tof, **kwargs))
 
@@ -210,7 +211,7 @@ def get_total_cost(self):
  tc2 = np.linalg.norm(self._dvs.to_value(u.km / u.s), axis=1).sum() * u.km / u.s
  assert np.allclose(tc.value, tc2.value)
  return tc2
- 
+
  @classmethod
  @u.quantity_input(max_delta_r=u.km)
  def correct_pericenter(cls, orbit, max_delta_r):

src/poliastro/twobody/orbit.py

@@ -291,8 +291,7 @@ def apply_impulse(self, v):
  raise ValueError(f"Vectors must have dimension 1, got {v.ndim}")
 
  self._state = RVState(self.attractor, self.r, self.v + v, self.plane)
-
-
+
  @classmethod
  def from_coords(cls, attractor, coord, plane=Planes.EARTH_EQUATOR):
  """Creates an `Orbit` from an attractor and astropy `SkyCoord`

tests/tests_core/test_core_util.py

@@ -10,8 +10,8 @@
 
 from poliastro.core.util import (
  alinspace,
- rotation_matrix as rotation_matrix_poliastro,
  cartesian_to_spherical,
+ rotation_matrix as rotation_matrix_poliastro,
  spherical_to_cartesian,
 )
 
@@ -26,12 +26,20 @@ def _test_rotation_matrix(angle, axis):
  expected = rotation_matrix_astropy(-np.rad2deg(angle), "xyz"[axis])
  result = rotation_matrix_poliastro(angle, axis)
  assert_allclose(expected, result)
+ if axis == 0:
+ expected = rotation_matrix_x(angle)
+ elif axis == 1:
+ expected = rotation_matrix_y(angle)
+ elif axis == 2:
+ expected = rotation_matrix_z(angle)
+ assert_allclose(expected, result)
 
 
 def test_rotation_matrix():
  v = np.array([-0.30387748, -1.4202498, 0.24305655])
  for angle in (0.5, np.array([-np.pi, np.pi])):
  for axis in (0, 1, 2):
+ _test_rotation_matrix(angle, axis)
  _test_rotation_matrix_with_v(v, angle, axis)