<commit_before># coding: utf-8
"""
    Test the Potential classes
"""

from __future__ import absolute_import, unicode_literals, division, print_function

__author__ = "adrn <<EMAIL>>"

import os
import pytest
import numpy as np
from astropy.constants.si import G
import astropy.units as u
import matplotlib.pyplot as plt

from ..core import *
from ..common import *

plot_path = "plots/tests/potential"
if not os.path.exists(plot_path):
    os.makedirs(plot_path)
else:
    for plot in os.listdir(plot_path):
        os.remove(os.path.join(plot_path,plot))

class TestPointMass(object):
    usys = UnitSystem(u.au, u.M_sun, u.yr)
    
    def test_pointmass_creation(self):
        potential = PointMassPotential(unit_system=self.usys,
                                       m=1.*u.M_sun, 
                                       r_0=[0.,0.,0.]*u.au)
        
        # no mass provided
        with pytest.raises(AssertionError):
            potential = PointMassPotential(unit_system=self.usys, 
                                           r_0=[0.,0.,0.]*u.au)
        
        # no r_0 provided
        with pytest.raises(AssertionError):
            potential = PointMassPotential(unit_system=self.usys, 
                                           m=1.*u.M_sun)
    
    
    def test_pointmass_eval(self):
        potential = PointMassPotential(unit_system=self.usys,
                                       m=1.*u.M_sun, 
                                       r_0=[0.,0.,0.]*u.au)
                                       
        # Test with a single position
        r = [1.,0.,0.]*u.au
        pot_val = potential.value_at(r)
        assert pot_val.unit.is_equivalent(u.J/u.kg)
        _pot_val = potential._value_at(r.value)
        
        np.testing.assert_array_almost_equal(_pot_val, -39.487906, decimal=5) 
        
        acc_val = potential.acceleration_at(r)
        assert acc_val.unit.is_equivalent(u.m/u.s**2)
        _acc_val = potential._acceleration_at(r.value)
        
        np.testing.assert_array_almost_equal(_acc_val, 
                                             [-39.487906,0.,0.], decimal=5)
    
    def test_pointmass_plot(self):
        
        # 2-d case
        potential = PointMassPotential(unit_system=self.usys,
                                       m=1.*u.M_sun, 
                                       r_0=[0.,0.]*u.au)
        grid = np.linspace(-5.,5)*u.au
        fig,axes = potential.plot(ndim=2, grid=grid)
        fig.savefig(os.path.join(plot_path, "point_mass_2d.png"))
    
    
    
def test_composite(self):
    
    potential = CompositePotential(units=[u.au,u.yr,u.M_sun], 
                               origin=[0.,0.,0.]*u.au)
    potential["one"] = PointMassPotential(units=potential.units,
                                          origin=[-1.,0.,0.]*u.au,
                                          m=1.*u.M_sun)
    print(potential.value_at([0.17157,0.,0.]*u.au), 
          potential.acceleration_at([0.17157,0.,0.]*u.au))
    
    potential["two"] = PointMassPotential(units=potential.units,
                                            origin=[1.,0.,0.]*u.au,
                                            m=0.5*u.M_sun)
    
    # Where forces cancel
    np.testing.assert_array_almost_equal(
                    sum(potential.acceleration_at([0.17157,0.,0.]*u.au)),
                    0.0, decimal=3)
    
    grid = np.linspace(-5.,5)*u.au<commit_msg>add 3d plotting case<commit_after># coding: utf-8
"""
    Test the Potential classes
"""

from __future__ import absolute_import, unicode_literals, division, print_function

__author__ = "adrn <<EMAIL>>"

import os
import pytest
import numpy as np
from astropy.constants.si import G
import astropy.units as u
import matplotlib.pyplot as plt

from ..core import *
from ..common import *

plot_path = "plots/tests/potential"
if not os.path.exists(plot_path):
    os.makedirs(plot_path)
else:
    for plot in os.listdir(plot_path):
        os.remove(os.path.join(plot_path,plot))

class TestPointMass(object):
    usys = UnitSystem(u.au, u.M_sun, u.yr)
    
    def test_pointmass_creation(self):
        potential = PointMassPotential(unit_system=self.usys,
                                       m=1.*u.M_sun, 
                                       r_0=[0.,0.,0.]*u.au)
        
        # no mass provided
        with pytest.raises(AssertionError):
            potential = PointMassPotential(unit_system=self.usys, 
                                           r_0=[0.,0.,0.]*u.au)
        
        # no r_0 provided
        with pytest.raises(AssertionError):
            potential = PointMassPotential(unit_system=self.usys, 
                                           m=1.*u.M_sun)
    
    
    def test_pointmass_eval(self):
        potential = PointMassPotential(unit_system=self.usys,
                                       m=1.*u.M_sun, 
                                       r_0=[0.,0.,0.]*u.au)
                                       
        # Test with a single position
        r = [1.,0.,0.]*u.au
        pot_val = potential.value_at(r)
        assert pot_val.unit.is_equivalent(u.J/u.kg)
        _pot_val = potential._value_at(r.value)
        
        np.testing.assert_array_almost_equal(_pot_val, -39.487906, decimal=5) 
        
        acc_val = potential.acceleration_at(r)
        assert acc_val.unit.is_equivalent(u.m/u.s**2)
        _acc_val = potential._acceleration_at(r.value)
        
        np.testing.assert_array_almost_equal(_acc_val, 
                                             [-39.487906,0.,0.], decimal=5)
    
    def test_pointmass_plot(self):
        
        # 2-d case
        potential = PointMassPotential(unit_system=self.usys,
                                       m=1.*u.M_sun, 
                                       r_0=[0.,0.]*u.au)
        grid = np.linspace(-5.,5)*u.au
        fig,axes = potential.plot(ndim=2, grid=grid)
        fig.savefig(os.path.join(plot_path, "point_mass_2d.png"))
        
        # 3-d case
        potential = PointMassPotential(unit_system=self.usys,
                                       m=1.*u.M_sun, 
                                       r_0=[0.,0.,0.]*u.au)
        grid = np.linspace(-5.,5)*u.au
        fig,axes = potential.plot(ndim=3, grid=grid)
        fig.savefig(os.path.join(plot_path, "point_mass_3d.png"))
    
    
def test_composite(self):
    
    potential = CompositePotential(units=[u.au,u.yr,u.M_sun], 
                               origin=[0.,0.,0.]*u.au)
    potential["one"] = PointMassPotential(units=potential.units,
                                          origin=[-1.,0.,0.]*u.au,
                                          m=1.*u.M_sun)
    print(potential.value_at([0.17157,0.,0.]*u.au), 
          potential.acceleration_at([0.17157,0.,0.]*u.au))
    
    potential["two"] = PointMassPotential(units=potential.units,
                                            origin=[1.,0.,0.]*u.au,
                                            m=0.5*u.M_sun)
    
    # Where forces cancel
    np.testing.assert_array_almost_equal(
                    sum(potential.acceleration_at([0.17157,0.,0.]*u.au)),
                    0.0, decimal=3)
    
    grid = np.linspace(-5.,5)*u.au
{'fraction_non_alphanumeric': 0.12062742920599667, 'fraction_numerical': 0.026790671848972793, 'mean_word_length': 1.5130798744332055, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 44, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3618865', 'n_tokens_mistral': 2383, 'n_tokens_neox': 2251, 'n_words': 321}