<commit_before># coding: utf-8

from __future__ import division, print_function

# standard library
import warnings

# Third-party
from astropy.coordinates import SphericalRepresentation, Galactic
import astropy.units as u
import numpy as np
import pytest
import scipy.optimize as so

# Project
from ..core import PhaseSpacePosition
from ..orbit import Orbit
from ...integrate import DOPRI853Integrator
from ...potential import Hamiltonian, HernquistPotential, LogarithmicPotential, KeplerPotential
from ...potential.frame import StaticFrame, ConstantRotatingFrame
from ...units import galactic, solarsystem

# Tests below should be cleaned up a bit...
def test_initialize():

    with pytest.raises(ValueError):
        x = np.random.random(size=(3,10))
        v = np.random.random(size=(3,8))
        Orbit(pos=x, vel=v)

    with pytest.raises(ValueError):
        x = np.random.random(size=(3,10))
        v = np.random.random(size=(3,10))
        t = np.arange(8)
        Orbit(pos=x, vel=v, t=t)

    # TODO: always?
    # x = np.random.random(size=(3,10))
    # v = np.random.random(size=(3,10))
    # o = Orbit(pos=x, vel=v)
    # assert o.ndim == 3

    x = np.random.random(size=(3,10))*u.kpc
    v = np.random.random(size=(3,10))*u.km/u.s
    o = Orbit(pos=x, vel=v)
    assert o.xyz.unit == u.kpc
    assert o.v_x.unit == u.km/u.s

    # TODO: don't support < 3 dim?
    # x = np.random.random(size=(2,10))
    # v = np.random.random(size=(2,10))
    # o = Orbit(pos=x, vel=v)
    # assert o.ndim == 2
    # assert o.hamiltonian is None

    # Check that passing in frame and potential or Hamiltonian works
    x = np.random.random(size=(3,10))*u.kpc
    v = np.random.random(size=(3,10))*u.km/u.s
    frame = StaticFrame(galactic)
    potential = LogarithmicPotential(v_c=1., r_h=0.14, q1=1., q2=0.9, q3=1.,
                                     units=galactic)

    o = Orbit(pos=x, vel=v, frame=frame)
    assert o.hamiltonian is None
    assert o.potential is None

    o = Orbit(pos=x, vel=v, potential=potential)
    assert o.hamiltonian is None
    assert o.frame is None

    o = Orbit(pos=x, vel=v, potential=potential, frame=frame)
    o = Orbit(pos=x, vel=v,
              hamiltonian=Hamiltonian(potential, frame=frame))
    assert isinstance(o.hamiltonian, Hamiltonian)
    assert isinstance(o.potential, LogarithmicPotential)
    assert isinstance(o.frame, StaticFrame)

    # check that old class raises deprecation warning
    from ..orbit import CartesianOrbit
    warnings.simplefilter('always')
    with pytest.warns(DeprecationWarning):
        o = CartesianOrbit(pos=x, vel=v)

def test_from_w():

    w = np.random.random(size=(6,10))
    o = Orbit.from_w(w, galactic)
    assert o.xyz.unit == u.kpc
    assert o.v_x.unit == u.kpc/u.Myr

def test_slice():

    # simple
    x = np.random.random(size=(3,10))
    v = np.random.random(size=(3,10))
    o = Orbit(pos=x, vel=v)
    new_o = o[:5]
    assert new_o.shape == (5,)

    x = np.random.random(size=(3,10))
    v = np.random.random(size=(3,10))
    t = np.linspace(0,10,10)
    o = Orbit(pos=x, vel=v, t=t)
    new_o = o[:5]
    assert new_o.shape == (5,)

    # 1d slice on 3d
    x = np.random.random(size=(3,10,8))
    v = np.random.random(size=(3,10,8))
    t = np.arange(x.shape[1])
    o = Orbit(pos=x, vel=v, t=t)
    new_o = o[:5]
    assert new_o.shape == (5,8)
    assert new_o.t.shape == (5,)

    # pick a single orbit
    new_o = o[:,0]
    assert isinstance(new_o, Orbit)
    assert new_o.shape == (10,)
    assert new_o.t.shape == (10,)

    # pick a single time
    new_o = o[3]
    assert isinstance(new_o, PhaseSpacePosition)
    assert new_o.shape == (8,)

    # 3d slice on 3d
    o = Orbit(pos=x, vel=v, t=t)
    new_o = o[:5,:4]
    assert new_o.shape == (5,4)
    assert new_o.t.shape == (5,)

    # boolean array
    x = np.random.random(size=(3,10))
    v = np.random.random(size=(3,10))
    t = np.arange(x.shape[1])
    o = Orbit(pos=x, vel=v, t=t)
    ix = np.array([0,0,0,0,0,1,1,1,1,1]).astype(bool)
    new_o = o[ix]
    assert new_o.shape == (sum(ix),)
    assert new_o.t.shape == (5,)

    # boolean array - 3D
    x = np.random.random(size=(3,10,4))
    v = np.random.random(size=(3,10,4))
    t = np.arange(x.shape[1])
    o = Orbit(pos=x, vel=v, t=t)
    ix = np.array([0,0,0,0,0,1,1,1,1,1]).astype(bool)
    new_o = o[ix]
    assert new_o.shape == (sum(ix),x.shape[-1])
    assert new_o.t.shape == (5,)

    # integer array
    x = np.random.random(size=(3,10))
    v = np.random.random(size=(3,10))
    t = np.arange(x.shape[1])
    o = Orbit(pos=x, vel=v, t=t)
    ix = np.array([0,3,5])
    new_o = o[ix]
    assert new_o.shape == (len(ix),)
    assert new_o.t.shape == (len(ix),)

def test_represent_as():

    # simple / unitless
    x = np.random.random(size=(3,10))
    v = np.random.random(size=(3,10))
    o = Orbit(pos=x, vel=v)
    sph = o.represent_as(SphericalRepresentation)

    assert sph.pos.distance.unit == u.one
    assert sph.vel.d_distance.unit == u.one

    # simple / with units
    x = np.random.random(size=(3,10))*u.kpc
    v = np.random.normal(0.,100.,size=(3,10))*u.km/u.s
    o = Orbit(pos=x, vel=v)
    sph = o.represent_as(SphericalRepresentation)
    assert sph.pos.distance.unit == u.kpc
    assert sph.vel.d_distance.unit == u.km/u.s

def test_to_coord_frame():
    # simple / unitless
    x = np.random.random(size=(3,10))
    v = np.random.random(size=(3,10))
    o = Orbit(pos=x, vel=v)

    with pytest.raises(u.UnitConversionError):
        o.to_coord_frame(Galactic)

    # simple / with units
    x = np.random.random(size=(3,10))*u.kpc
    v = np.random.normal(0.,100.,size=(3,10))*u.km/u.s
    o = Orbit(pos=x, vel=v)
    coo,vel = o.to_coord_frame(Galactic)
    assert coo.name == 'galactic'

    # simple / with units and time
    x = np.random.random(size=(3,128,10))*u.kpc
    v = np.random.normal(0.,100.,size=(3,128,10))*u.km/u.s
    o = Orbit(pos=x, vel=v)
    coo,vel = o.to_coord_frame(Galactic)
    assert coo.name == 'galactic'

def test_w():
    # simple / unitless
    x = np.random.random(size=(3,10))
    v = np.random.random(size=(3,10))
    o = Orbit(pos=x, vel=v)
    w = o.w()
    assert w.shape == (6,10)

    # simple / with units
    x = np.random.random(size=(3,10))*u.kpc
    v = np.random.normal(0.,100.,size=(3,10))*u.km/u.s
    o = Orbit(pos=x, vel=v)
    with pytest.raises(ValueError):
        o.w()
    w = o.w(units=galactic)
    assert np.allclose(x.value, w[:3,:])
    assert np.allclose(v.value, (w[3:,:]*u.kpc/u.Myr).to(u.km/u.s).value)<commit_msg>fix to_coord_frame on orbit class<commit_after># coding: utf-8

from __future__ import division, print_function

# standard library
import warnings

# Third-party
from astropy.coordinates import SphericalRepresentation, Galactic
import astropy.units as u
import numpy as np
import pytest
import scipy.optimize as so

# Project
from ..core import PhaseSpacePosition
from ..orbit import Orbit
from ...integrate import DOPRI853Integrator
from ...potential import Hamiltonian, HernquistPotential, LogarithmicPotential, KeplerPotential
from ...potential.frame import StaticFrame, ConstantRotatingFrame
from ...units import galactic, solarsystem

# Tests below should be cleaned up a bit...
def test_initialize():

    with pytest.raises(ValueError):
        x = np.random.random(size=(3,10))
        v = np.random.random(size=(3,8))
        Orbit(pos=x, vel=v)

    with pytest.raises(ValueError):
        x = np.random.random(size=(3,10))
        v = np.random.random(size=(3,10))
        t = np.arange(8)
        Orbit(pos=x, vel=v, t=t)

    # TODO: always?
    # x = np.random.random(size=(3,10))
    # v = np.random.random(size=(3,10))
    # o = Orbit(pos=x, vel=v)
    # assert o.ndim == 3

    x = np.random.random(size=(3,10))*u.kpc
    v = np.random.random(size=(3,10))*u.km/u.s
    o = Orbit(pos=x, vel=v)
    assert o.xyz.unit == u.kpc
    assert o.v_x.unit == u.km/u.s

    # TODO: don't support < 3 dim?
    # x = np.random.random(size=(2,10))
    # v = np.random.random(size=(2,10))
    # o = Orbit(pos=x, vel=v)
    # assert o.ndim == 2
    # assert o.hamiltonian is None

    # Check that passing in frame and potential or Hamiltonian works
    x = np.random.random(size=(3,10))*u.kpc
    v = np.random.random(size=(3,10))*u.km/u.s
    frame = StaticFrame(galactic)
    potential = LogarithmicPotential(v_c=1., r_h=0.14, q1=1., q2=0.9, q3=1.,
                                     units=galactic)

    o = Orbit(pos=x, vel=v, frame=frame)
    assert o.hamiltonian is None
    assert o.potential is None

    o = Orbit(pos=x, vel=v, potential=potential)
    assert o.hamiltonian is None
    assert o.frame is None

    o = Orbit(pos=x, vel=v, potential=potential, frame=frame)
    o = Orbit(pos=x, vel=v,
              hamiltonian=Hamiltonian(potential, frame=frame))
    assert isinstance(o.hamiltonian, Hamiltonian)
    assert isinstance(o.potential, LogarithmicPotential)
    assert isinstance(o.frame, StaticFrame)

    # check that old class raises deprecation warning
    from ..orbit import CartesianOrbit
    warnings.simplefilter('always')
    with pytest.warns(DeprecationWarning):
        o = CartesianOrbit(pos=x, vel=v)

def test_from_w():

    w = np.random.random(size=(6,10))
    o = Orbit.from_w(w, galactic)
    assert o.xyz.unit == u.kpc
    assert o.v_x.unit == u.kpc/u.Myr

def test_slice():

    # simple
    x = np.random.random(size=(3,10))
    v = np.random.random(size=(3,10))
    o = Orbit(pos=x, vel=v)
    new_o = o[:5]
    assert new_o.shape == (5,)

    x = np.random.random(size=(3,10))
    v = np.random.random(size=(3,10))
    t = np.linspace(0,10,10)
    o = Orbit(pos=x, vel=v, t=t)
    new_o = o[:5]
    assert new_o.shape == (5,)

    # 1d slice on 3d
    x = np.random.random(size=(3,10,8))
    v = np.random.random(size=(3,10,8))
    t = np.arange(x.shape[1])
    o = Orbit(pos=x, vel=v, t=t)
    new_o = o[:5]
    assert new_o.shape == (5,8)
    assert new_o.t.shape == (5,)

    # pick a single orbit
    new_o = o[:,0]
    assert isinstance(new_o, Orbit)
    assert new_o.shape == (10,)
    assert new_o.t.shape == (10,)

    # pick a single time
    new_o = o[3]
    assert isinstance(new_o, PhaseSpacePosition)
    assert new_o.shape == (8,)

    # 3d slice on 3d
    o = Orbit(pos=x, vel=v, t=t)
    new_o = o[:5,:4]
    assert new_o.shape == (5,4)
    assert new_o.t.shape == (5,)

    # boolean array
    x = np.random.random(size=(3,10))
    v = np.random.random(size=(3,10))
    t = np.arange(x.shape[1])
    o = Orbit(pos=x, vel=v, t=t)
    ix = np.array([0,0,0,0,0,1,1,1,1,1]).astype(bool)
    new_o = o[ix]
    assert new_o.shape == (sum(ix),)
    assert new_o.t.shape == (5,)

    # boolean array - 3D
    x = np.random.random(size=(3,10,4))
    v = np.random.random(size=(3,10,4))
    t = np.arange(x.shape[1])
    o = Orbit(pos=x, vel=v, t=t)
    ix = np.array([0,0,0,0,0,1,1,1,1,1]).astype(bool)
    new_o = o[ix]
    assert new_o.shape == (sum(ix),x.shape[-1])
    assert new_o.t.shape == (5,)

    # integer array
    x = np.random.random(size=(3,10))
    v = np.random.random(size=(3,10))
    t = np.arange(x.shape[1])
    o = Orbit(pos=x, vel=v, t=t)
    ix = np.array([0,3,5])
    new_o = o[ix]
    assert new_o.shape == (len(ix),)
    assert new_o.t.shape == (len(ix),)

def test_represent_as():

    # simple / unitless
    x = np.random.random(size=(3,10))
    v = np.random.random(size=(3,10))
    o = Orbit(pos=x, vel=v)
    sph = o.represent_as(SphericalRepresentation)

    assert sph.pos.distance.unit == u.one
    assert sph.vel.d_distance.unit == u.one

    # simple / with units
    x = np.random.random(size=(3,10))*u.kpc
    v = np.random.normal(0.,100.,size=(3,10))*u.km/u.s
    o = Orbit(pos=x, vel=v)
    sph = o.represent_as(SphericalRepresentation)
    assert sph.pos.distance.unit == u.kpc
    assert sph.vel.d_distance.unit == u.km/u.s

def test_to_coord_frame():
    # simple / unitless
    x = np.random.random(size=(3,10))
    v = np.random.random(size=(3,10))
    o = Orbit(pos=x, vel=v)

    with pytest.raises(u.UnitConversionError):
        o.to_coord_frame(Galactic)

    # simple / with units
    x = np.random.random(size=(3,10))*u.kpc
    v = np.random.normal(0.,100.,size=(3,10))*u.km/u.s
    o = Orbit(pos=x, vel=v)
    coo = o.to_coord_frame(Galactic)
    assert coo.name == 'galactic'

    # simple / with units and time
    x = np.random.random(size=(3,128,10))*u.kpc
    v = np.random.normal(0.,100.,size=(3,128,10))*u.km/u.s
    o = Orbit(pos=x, vel=v)
    coo = o.to_coord_frame(Galactic)
    assert coo.name == 'galactic'

def test_w():
    # simple / unitless
    x = np.random.random(size=(3,10))
    v = np.random.random(size=(3,10))
    o = Orbit(pos=x, vel=v)
    w = o.w()
    assert w.shape == (6,10)

    # simple / with units
    x = np.random.random(size=(3,10))*u.kpc
    v = np.random.normal(0.,100.,size=(3,10))*u.km/u.s
    o = Orbit(pos=x, vel=v)
    with pytest.raises(ValueError):
        o.w()
    w = o.w(units=galactic)
    assert np.allclose(x.value, w[:3,:])
    assert np.allclose(v.value, (w[3:,:]*u.kpc/u.Myr).to(u.km/u.s).value)
{'fraction_non_alphanumeric': 0.17075773745997866, 'fraction_numerical': 0.03445647202317426, 'mean_word_length': 2.8710534080849808, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '951862', 'n_tokens_mistral': 5531, 'n_tokens_neox': 5227, 'n_words': 1029}