from numpy import cos, cosh, sin, sinh
from numpy.testing import assert_allclose, assert_equal

from poliastro._math.special import stumpff_c2 as c2, stumpff_c3 as c3


def test_stumpff_functions_near_zero():
    psi = 0.5
    expected_c2 = (1 - cos(psi**0.5)) / psi
    expected_c3 = (psi**0.5 - sin(psi**0.5)) / psi**1.5

    assert_allclose(c2(psi), expected_c2)
    assert_allclose(c3(psi), expected_c3)


def test_stumpff_functions_above_zero():
    psi = 3.0
    expected_c2 = (1 - cos(psi**0.5)) / psi
    expected_c3 = (psi**0.5 - sin(psi**0.5)) / psi**1.5

    assert_equal(c2(psi), expected_c2)
    assert_equal(c3(psi), expected_c3)


def test_stumpff_functions_under_zero():
    psi = -3.0
    expected_c2 = (cosh((-psi) ** 0.5) - 1) / (-psi)
    expected_c3 = (sinh((-psi) ** 0.5) - (-psi) ** 0.5) / (-psi) ** 1.5

    assert_equal(c2(psi), expected_c2)
    assert_equal(c3(psi), expected_c3)

from numpy import cos, cosh, sin, sinh
from numpy.testing import assert_allclose

from poliastro._math.special import stumpff_c2 as c2, stumpff_c3 as c3


def test_stumpff_functions_near_zero():
    psi = 0.5
    expected_c2 = (1 - cos(psi**0.5)) / psi
    expected_c3 = (psi**0.5 - sin(psi**0.5)) / psi**1.5

    assert_allclose(c2(psi), expected_c2)
    assert_allclose(c3(psi), expected_c3)


def test_stumpff_functions_above_zero():
    psi = 3.0
    expected_c2 = (1 - cos(psi**0.5)) / psi
    expected_c3 = (psi**0.5 - sin(psi**0.5)) / psi**1.5

    assert_allclose(c2(psi), expected_c2, rtol=1e-10)
    assert_allclose(c3(psi), expected_c3, rtol=1e-10)


def test_stumpff_functions_under_zero():
    psi = -3.0
    expected_c2 = (cosh((-psi) ** 0.5) - 1) / (-psi)
    expected_c3 = (sinh((-psi) ** 0.5) - (-psi) ** 0.5) / (-psi) ** 1.5

    assert_allclose(c2(psi), expected_c2, rtol=1e-10)
    assert_allclose(c3(psi), expected_c3, rtol=1e-10)

Use numpy.testing.assert_allclose with small relative tolerance
test: Use numpy.testing.assert_allclose with small relative tolerance

To avoid having numpy.testing.assert_equal fail in some cases at the
level of machine precision, use numpy.testing.assert_allclose instead
with a relative tolerance set small, 1e-10, to still be restrictive in
the level of equality accepted.

  from numpy import cos, cosh, sin, sinh
- from numpy.testing import assert_allclose, assert_equal
+ from numpy.testing import assert_allclose
  
  from poliastro._math.special import stumpff_c2 as c2, stumpff_c3 as c3
  
  
  def test_stumpff_functions_near_zero():
      psi = 0.5
      expected_c2 = (1 - cos(psi**0.5)) / psi
      expected_c3 = (psi**0.5 - sin(psi**0.5)) / psi**1.5
  
      assert_allclose(c2(psi), expected_c2)
      assert_allclose(c3(psi), expected_c3)
  
  
  def test_stumpff_functions_above_zero():
      psi = 3.0
      expected_c2 = (1 - cos(psi**0.5)) / psi
      expected_c3 = (psi**0.5 - sin(psi**0.5)) / psi**1.5
  
-     assert_equal(c2(psi), expected_c2)
+     assert_allclose(c2(psi), expected_c2, rtol=1e-10)
-     assert_equal(c3(psi), expected_c3)
+     assert_allclose(c3(psi), expected_c3, rtol=1e-10)
  
  
  def test_stumpff_functions_under_zero():
      psi = -3.0
      expected_c2 = (cosh((-psi) ** 0.5) - 1) / (-psi)
      expected_c3 = (sinh((-psi) ** 0.5) - (-psi) ** 0.5) / (-psi) ** 1.5
  
-     assert_equal(c2(psi), expected_c2)
+     assert_allclose(c2(psi), expected_c2, rtol=1e-10)
-     assert_equal(c3(psi), expected_c3)
+     assert_allclose(c3(psi), expected_c3, rtol=1e-10)
  
Use numpy.testing.assert_allclose with small relative tolerance
## Code Before:
from numpy import cos, cosh, sin, sinh
from numpy.testing import assert_allclose, assert_equal

from poliastro._math.special import stumpff_c2 as c2, stumpff_c3 as c3


def test_stumpff_functions_near_zero():
    psi = 0.5
    expected_c2 = (1 - cos(psi**0.5)) / psi
    expected_c3 = (psi**0.5 - sin(psi**0.5)) / psi**1.5

    assert_allclose(c2(psi), expected_c2)
    assert_allclose(c3(psi), expected_c3)


def test_stumpff_functions_above_zero():
    psi = 3.0
    expected_c2 = (1 - cos(psi**0.5)) / psi
    expected_c3 = (psi**0.5 - sin(psi**0.5)) / psi**1.5

    assert_equal(c2(psi), expected_c2)
    assert_equal(c3(psi), expected_c3)


def test_stumpff_functions_under_zero():
    psi = -3.0
    expected_c2 = (cosh((-psi) ** 0.5) - 1) / (-psi)
    expected_c3 = (sinh((-psi) ** 0.5) - (-psi) ** 0.5) / (-psi) ** 1.5

    assert_equal(c2(psi), expected_c2)
    assert_equal(c3(psi), expected_c3)

## Instruction:
Use numpy.testing.assert_allclose with small relative tolerance
## Code After:
from numpy import cos, cosh, sin, sinh
from numpy.testing import assert_allclose

from poliastro._math.special import stumpff_c2 as c2, stumpff_c3 as c3


def test_stumpff_functions_near_zero():
    psi = 0.5
    expected_c2 = (1 - cos(psi**0.5)) / psi
    expected_c3 = (psi**0.5 - sin(psi**0.5)) / psi**1.5

    assert_allclose(c2(psi), expected_c2)
    assert_allclose(c3(psi), expected_c3)


def test_stumpff_functions_above_zero():
    psi = 3.0
    expected_c2 = (1 - cos(psi**0.5)) / psi
    expected_c3 = (psi**0.5 - sin(psi**0.5)) / psi**1.5

    assert_allclose(c2(psi), expected_c2, rtol=1e-10)
    assert_allclose(c3(psi), expected_c3, rtol=1e-10)


def test_stumpff_functions_under_zero():
    psi = -3.0
    expected_c2 = (cosh((-psi) ** 0.5) - 1) / (-psi)
    expected_c3 = (sinh((-psi) ** 0.5) - (-psi) ** 0.5) / (-psi) ** 1.5

    assert_allclose(c2(psi), expected_c2, rtol=1e-10)
    assert_allclose(c3(psi), expected_c3, rtol=1e-10)
