<commit_before># coding: utf-8

""" Utilities for Potential classes """

from __future__ import division, print_function

__author__ = "adrn <<EMAIL>>"

# Standard library
import os
import sys

# Third-party
import astropy.units as u
import numpy as np
from sympy.utilities.lambdify import lambdify

# Project
from .core import PotentialBase

def _classnamify(s):
    s = [x.lower() for x in str(s).split()]
    words = []
    for word in s:
        words.append(word[0].upper() + word[1:])
    return "".join(words)

def from_equation(expr, vars, pars, name=None):
    """
    Create a potential class from an expression for the potential.

    Parameters
    ----------
    expr : sympy., str
    vars : iterable
        Variables, e.g., x, y, z. Should be an iterable of strings containing
        the names of the variables.
    pars : iterable
        Parameters of the potential. For example, ... TODO
    name : str (optional)
        The name of the potential class returned.
    """
    try:
        import sympy
    except ImportError:
        raise ImportError("sympy is required to use 'from_equation()' "
                          "potential class creation.")

    # convert all input to Sympy objects
    expr = sympy.sympify(expr)
    vars = [sympy.sympify(v) for v in vars]
    var_names = [v.name for v in vars]
    pars = [sympy.sympify(p) for p in pars]
    par_names = [p.name for p in pars]

    class MyPotential(PotentialBase):

        def __init__(self, units=None, **kwargs):
            self.parameters = kwargs
            for par in par_names:
                if par not in self.parameters:
                    raise ValueError("You must specify a value for "
                                     "parameter '{}'.".format(par))

            super(MyPotential,self).__init__(units)

    if name is not None:
        name = _classnamify(name)
        if "potential" not in name.lower():
            name = name + "Potential"
        MyPotential.__name__ = name

    # Energy / value
    valuefunc = lambdify(vars + pars, expr, dummify=False)
    def _value(self, w, t):
        kw = self.parameters.copy()
        for i,name in enumerate(var_names):
            kw[name] = w[i]
        return valuefunc(**kw)
    MyPotential._value = _value

    # Gradient
    gradfuncs = []
    for var in vars:
        gradfuncs.append(lambdify(vars + pars, sympy.diff(expr,var), dummify=False))

    def _gradient(self, w, t):
        kw = self.parameters.copy()
        for i,name in enumerate(var_names):
            kw[name] = w[i]
        return np.vstack([f(**kw) for f in gradfuncs])
    MyPotential._gradient = _gradient

    return MyPotential<commit_msg>docstring example<commit_after># coding: utf-8

""" Utilities for Potential classes """

from __future__ import division, print_function

__author__ = "adrn <<EMAIL>>"

# Third-party
import numpy as np
from sympy.utilities.lambdify import lambdify

# Project
from .core import PotentialBase

def _classnamify(s):
    s = [x.lower() for x in str(s).split()]
    words = []
    for word in s:
        words.append(word[0].upper() + word[1:])
    return "".join(words)

def from_equation(expr, vars, pars, name=None, hessian=False):
    """
    Create a potential class from an expression for the potential.

    .. note::

        This utility requires having :ref:`sympy` installed.

    Parameters
    ----------
    expr : :ref:`sympy.core.expr.Expr`, str
        Either a ``sympy`` expression, or a string that can be converted to
        a ``sympy`` expression.
    vars : iterable
        An iterable of variable names in the expression.
    pars : iterable
        An iterable of parameter names in the expression.
    name : str (optional)
        The name of the potential class returned.
    hessian : bool (optional)
        Generate a function to compute the Hessian.

    Returns
    -------
    CustomPotential : `~gary.potential.PotentialBase`
        A potential class that represents the input equation. To instantiate the
        potential, use just like a normal class with parameters.

    Examples
    --------
    Here we'll create a potential class for the harmonic oscillator
    potential, :math:`\Phi(x) = \frac{1}{2}\,k\,x^2`::

        >>> Potential = from_equation("1/2*k*x**2", vars="x", pars="k",
        ...                           name='HarmonicOscillator')
        >>> p1 = Potential(k=1.)
        >>> p1
        <HarmonicoscillatorPotential: k=1.00 (dimensionless)>

    The potential class (and object) is a fully-fledged subclass of
    `~gary.potential.PotentialBase` and therefore has many useful methods.
    For example, to integrate an orbit::

        >>> t,w = p1.integrate_orbit([1.,0], dt=0.01, nsteps=1000)

    """
    try:
        import sympy
    except ImportError:
        raise ImportError("sympy is required to use 'from_equation()' "
                          "potential class creation.")

    # convert all input to Sympy objects
    expr = sympy.sympify(expr)
    vars = [sympy.sympify(v) for v in vars]
    var_names = [v.name for v in vars]
    pars = [sympy.sympify(p) for p in pars]
    par_names = [p.name for p in pars]

    class MyPotential(PotentialBase):

        def __init__(self, units=None, **kwargs):
            self.parameters = kwargs
            for par in par_names:
                if par not in self.parameters:
                    raise ValueError("You must specify a value for "
                                     "parameter '{}'.".format(par))
            super(MyPotential,self).__init__(units)

    if name is not None:
        name = _classnamify(name)
        if "potential" not in name.lower():
            name = name + "Potential"
        MyPotential.__name__ = name

    # Energy / value
    valuefunc = lambdify(vars + pars, expr, dummify=False)
    def _value(self, w, t):
        kw = self.parameters.copy()
        for i,name in enumerate(var_names):
            kw[name] = w[i]
        return valuefunc(**kw)
    MyPotential._value = _value

    # Gradient
    gradfuncs = []
    for var in vars:
        gradfuncs.append(lambdify(vars + pars, sympy.diff(expr,var), dummify=False))

    def _gradient(self, w, t):
        kw = self.parameters.copy()
        for i,name in enumerate(var_names):
            kw[name] = w[i]
        return np.vstack([f(**kw) for f in gradfuncs])
    MyPotential._gradient = _gradient

    # Hessian
    if hessian:
        raise NotImplementedError("Hessian functionality doesn't exist yet...sorry.")
        hessfuncs = []

    return MyPotential
{'fraction_non_alphanumeric': 0.10912698412698413, 'fraction_numerical': 0.004273504273504274, 'mean_word_length': 2.283066132264529, 'pattern_counts': {'":': 0, '<': 8, '<?xml version=': 0, '>': 20, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 3, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7282677', 'n_tokens_mistral': 2049, 'n_tokens_neox': 1931, 'n_words': 618}