<commit_before>"""
    Test the core Potential classes
"""

# Standard library
from collections import OrderedDict

# Third party
import pytest
import numpy as np
from astropy.constants import G
import astropy.units as u
from matplotlib import cm

# This package
from ..core import PotentialBase, CompositePotential, PotentialParameter
from ....units import UnitSystem

units = [u.kpc, u.Myr, u.Msun, u.radian]
usys = UnitSystem(u.au, u.yr, u.Msun, u.radian)
G = G.decompose(units)


def test_new_simple():

    class MyPotential(PotentialBase):
        def _energy(self, r, t=0.):
            return -1/r

        def _gradient(self, r, t=0.):
            return r**-2

    p = MyPotential()
    assert p(0.5) == -2.
    assert p.energy(0.5) == -2.
    assert p.acceleration(0.5) == -4.

    p(np.arange(0.5, 11.5, 0.5).reshape(1, -1))
    p.energy(np.arange(0.5, 11.5, 0.5).reshape(1, -1))
    p.acceleration(np.arange(0.5, 11.5, 0.5).reshape(1, -1))


class MyPotential(PotentialBase):
    m = PotentialParameter('m', 'mass')
    x0 = PotentialParameter('x0', 'length')

    def _energy(self, x, t):
        m = self.parameters['m']
        x0 = self.parameters['x0']
        r = np.sqrt(np.sum((x-x0[None])**2, axis=1))
        return -m/r

    def _gradient(self, x, t):
        m = self.parameters['m']
        x0 = self.parameters['x0']
        r = np.sqrt(np.sum((x-x0[None])**2, axis=1))
        return m*(x-x0[None])/r**3


def test_repr():
    p = MyPotential(m=1.E10*u.Msun, x0=0., units=usys)
    _repr = p.__repr__()<commit_msg>get value of params<commit_after>"""
    Test the core Potential classes
"""

# Standard library
from collections import OrderedDict

# Third party
import pytest
import numpy as np
from astropy.constants import G
import astropy.units as u
from matplotlib import cm

# This package
from ..core import PotentialBase, CompositePotential, PotentialParameter
from ....units import UnitSystem

units = [u.kpc, u.Myr, u.Msun, u.radian]
usys = UnitSystem(u.au, u.yr, u.Msun, u.radian)
G = G.decompose(units)


def test_new_simple():

    class MyPotential(PotentialBase):
        def _energy(self, r, t=0.):
            return -1/r

        def _gradient(self, r, t=0.):
            return r**-2

    p = MyPotential()
    assert p(0.5) == -2.
    assert p.energy(0.5) == -2.
    assert p.acceleration(0.5) == -4.

    p(np.arange(0.5, 11.5, 0.5).reshape(1, -1))
    p.energy(np.arange(0.5, 11.5, 0.5).reshape(1, -1))
    p.acceleration(np.arange(0.5, 11.5, 0.5).reshape(1, -1))


class MyPotential(PotentialBase):
    m = PotentialParameter('m', 'mass')
    x0 = PotentialParameter('x0', 'length')

    def _energy(self, x, t):
        m = self.parameters['m'].value
        x0 = self.parameters['x0'].value
        r = np.sqrt(np.sum((x-x0[None])**2, axis=1))
        return -m/r

    def _gradient(self, x, t):
        m = self.parameters['m'].value
        x0 = self.parameters['x0'].value
        r = np.sqrt(np.sum((x-x0[None])**2, axis=1))
        return m*(x-x0[None])/r**3


def test_repr():
    p = MyPotential(m=1.E10*u.Msun, x0=0., units=usys)
    _repr = p.__repr__()
{'fraction_non_alphanumeric': 0.1676300578034682, 'fraction_numerical': 0.03789338471419396, 'mean_word_length': 2.7803398058252426, 'pattern_counts': {'":': 0, '<': 3, '<?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': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2224241', 'n_tokens_mistral': 1270, 'n_tokens_neox': 1197, 'n_words': 276}