<commit_before># coding: utf-8

""" Read and write potentials to text (YAML) files. """

from __future__ import division, print_function

__author__ = "adrn <<EMAIL>>"

# Standard library
import os

# Third-party
import astropy.units as u
from astropy.utils import isiterable
import numpy as np
import six
import yaml

from .. import potential as gp

__all__ = ['load', 'save']

def _unpack_params(p):
    params = p.copy()
    for key,item in six.iteritems(p):
        if '_unit' in key:
            continue

        if isiterable(item) and not isinstance(item, str):
            params[key] = np.array(item).astype(float)
        else:
            params[key] = float(item)

        if key+'_unit' in params:
            params[key] = params[key] * u.Unit(params[key+'_unit'])
            del params[key+'_unit']

    return params

def _parse_component(component, module):
    try:
        class_name = component['class']
    except KeyError:
        raise KeyError("Potential dictionary must contain a key 'class' for "
                       "specifying the name of the Potential class.")

    if 'units' not in component:
        unitsys = None
    else:
        try:
            unitsys = [u.Unit(unit) for ptype,unit in component['units'].items()]
        except KeyError:
            raise KeyError("Potential dictionary must contain a key 'units' with "
                           "a list of strings specifying the unit system.")

    params = component.get('parameters', {})

    # need to crawl the dictionary structure and unpack quantities
    params = _unpack_params(params)

    if module is None:
        from .. import potential
    else:
        potential = module

    Potential = getattr(potential, class_name)
    return Potential(units=unitsys, **params)

def from_dict(d, module=None):
    """
    Convert a dictionary potential specification into a
    :class:`~gary.potential.PotentialBase` subclass object.

    Parameters
    ----------
    d : dict
        Dictionary specification of a potential.
    module : namespace (optional)

    """

    if 'type' in d and d['type'] == 'composite':
        p = getattr(gp, d['class'])()
        for i,component in enumerate(d['components']):
            c = _parse_component(component, module)
            name = component.get('name', str(i))
            p[name] = c

    else:
        p = _parse_component(d, module)

    return p

# ----------------------------------------------------------------------------

def _pack_params(p):
    params = p.copy()
    for key,item in six.iteritems(p):
        if hasattr(item,'unit'):
            params[key] = item.value
            params[key+'_unit'] = str(item.unit)

        if hasattr(params[key], 'tolist'): # convert array to list
            params[key] = params[key].tolist()

    return params

def _to_dict_help(potential):
    d = dict()

    d['class'] = potential.__class__.__name__

    if potential.units is not None:
        d['units'] = dict([(str(ptype),str(unit)) for ptype,unit in potential.units.to_dict().items()])

    if len(potential.parameters) > 0:
        params = _pack_params(potential.parameters)
        d['parameters'] = params

    return d

def to_dict(potential):
    """
    Turn a potential object into a dictionary that fully specifies the
    state of the object.

    Parameters
    ----------
    potential : :class:`~gary.potential.PotentialBase`
        The instantiated :class:`~gary.potential.PotentialBase` object.

    """

    if isinstance(potential, gp.CompositePotential):
        d = dict()
        d['type'] = 'composite'
        d['class'] = potential.__class__.__name__
        d['components'] = []
        for k,p in potential.items():
            comp_dict = _to_dict_help(p)
            comp_dict['name'] = k
            d['components'].append(comp_dict)

    else:
        d = _to_dict_help(potential)

    return d

# ----------------------------------------------------------------------------

def load(f, module=None):<commit_msg>move imports<commit_after># coding: utf-8

""" Read and write potentials to text (YAML) files. """

from __future__ import division, print_function

__author__ = "adrn <<EMAIL>>"

# Standard library
import os

# Third-party
import astropy.units as u
from astropy.utils import isiterable
import numpy as np
import six
import yaml

__all__ = ['load', 'save']

def _unpack_params(p):
    params = p.copy()
    for key,item in six.iteritems(p):
        if '_unit' in key:
            continue

        if isiterable(item) and not isinstance(item, str):
            params[key] = np.array(item).astype(float)
        else:
            params[key] = float(item)

        if key+'_unit' in params:
            params[key] = params[key] * u.Unit(params[key+'_unit'])
            del params[key+'_unit']

    return params

def _parse_component(component, module):
    try:
        class_name = component['class']
    except KeyError:
        raise KeyError("Potential dictionary must contain a key 'class' for "
                       "specifying the name of the Potential class.")

    if 'units' not in component:
        unitsys = None
    else:
        try:
            unitsys = [u.Unit(unit) for ptype,unit in component['units'].items()]
        except KeyError:
            raise KeyError("Potential dictionary must contain a key 'units' with "
                           "a list of strings specifying the unit system.")

    params = component.get('parameters', {})

    # need to crawl the dictionary structure and unpack quantities
    params = _unpack_params(params)

    if module is None:
        from .. import potential
    else:
        potential = module

    Potential = getattr(potential, class_name)
    return Potential(units=unitsys, **params)

def from_dict(d, module=None):
    """
    Convert a dictionary potential specification into a
    :class:`~gary.potential.PotentialBase` subclass object.

    Parameters
    ----------
    d : dict
        Dictionary specification of a potential.
    module : namespace (optional)

    """

    if 'type' in d and d['type'] == 'composite':
        from .. import potential as gp
        p = getattr(gp, d['class'])()
        for i,component in enumerate(d['components']):
            c = _parse_component(component, module)
            name = component.get('name', str(i))
            p[name] = c

    else:
        p = _parse_component(d, module)

    return p

# ----------------------------------------------------------------------------

def _pack_params(p):
    params = p.copy()
    for key,item in six.iteritems(p):
        if hasattr(item,'unit'):
            params[key] = item.value
            params[key+'_unit'] = str(item.unit)

        if hasattr(params[key], 'tolist'): # convert array to list
            params[key] = params[key].tolist()

    return params

def _to_dict_help(potential):
    d = dict()

    d['class'] = potential.__class__.__name__

    if potential.units is not None:
        d['units'] = dict([(str(ptype),str(unit)) for ptype,unit in potential.units.to_dict().items()])

    if len(potential.parameters) > 0:
        params = _pack_params(potential.parameters)
        d['parameters'] = params

    return d

def to_dict(potential):
    """
    Turn a potential object into a dictionary that fully specifies the
    state of the object.

    Parameters
    ----------
    potential : :class:`~gary.potential.PotentialBase`
        The instantiated :class:`~gary.potential.PotentialBase` object.

    """
    from .. import potential as gp
    if isinstance(potential, gp.CompositePotential):
        d = dict()
        d['type'] = 'composite'
        d['class'] = potential.__class__.__name__
        d['components'] = []
        for k,p in potential.items():
            comp_dict = _to_dict_help(p)
            comp_dict['name'] = k
            d['components'].append(comp_dict)

    else:
        d = _to_dict_help(potential)

    return d

# ----------------------------------------------------------------------------

def load(f, module=None):
{'fraction_non_alphanumeric': 0.1511166253101737, 'fraction_numerical': 0.0004962779156327543, 'mean_word_length': 2.1365758754863813, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 9, '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': '6773674', 'n_tokens_mistral': 2340, 'n_tokens_neox': 2221, 'n_words': 653}