text:
"""Classes for storing ionization state data."""

from typing import Dict, List, Optional, Tuple, Union

import astropy.units as u
import collections
import numpy as np

from plasmapy.atomic import atomic_number
from plasmapy.atomic.particle_class import Particle
from plasmapy.atomic.particle_input import particle_input
from plasmapy.atomic.symbols import particle_symbol
from plasmapy.utils import (AtomicError, ChargeError, InvalidParticleError, check_quantity)

State = collections.namedtuple(
    'State', [
        'integer_charge',
        'ionic_fraction',
        'ionic_symbol',
    ])

_number_density_errmsg = (
    "Number densities must be Quantity objects with units of inverse "
    "volume."
)

class IonizationState:
    """
    Representation of the ionization state distribution of a single
    element or isotope.

    Parameters
    ----------
    particle: str, int, np.integer, or ~plasmapy.atomic.Particle
        A `str` or `~plasmapy.atomic.Particle` instance representing
        an element or isotope, or an `int` representing the atomic
        number of an element.

    ionic_fractions: ~numpy.ndarray, list, tuple, or ~astropy.units.Quantity
        The ionization fractions of an element, where the indices
        correspond to integer charge.  This argument should contain the
        atomic number plus one items, and must sum to one within an
        absolute tolerance of `tol` if dimensionless.  Alternatively,
        this argument may be a `~astropy.units.Quantity` that represents
        the number densities of each neutral/ion.

    T_e: ~astropy.units.Quantity, keyword-only, optional
        The electron temperature or thermal energy per particle.

    n_e: ~astropy.units.Quantity, keyword-only, optional
        The electron number density.

    n_elem: ~astropy.units.Quantity, keyword-only, optional
        The number density of the element, including neutrals and all
        ions.

    tol: float or int, keyword-only, optional
        The absolute tolerance used by `~numpy.isclose` when testing
        normalizations and making comparisons.  Defaults to `1e-15`.

    Raises
    ------
    ~plasmapy.utils.AtomicError
        If the ionic fractions are not normalized or contain invalid
        values

    ~plasmapy.utils.InvalidParticleError
        If the particle is invalid.

    Examples
    --------
    >>> states = IonizationState('H', [0.6, 0.4], n_elem=1*u.cm**-3, T_e=11000*u.K)
    >>> states.ionic_fractions
    array([0.6, 0.4])
    >>> states.n_e  # electron number density
    <Quantity 400000. 1 / m3>
    >>> states.n_elem  # element number density
    <Quantity 1000000. 1 / m3>

    Notes
    -----
    Only one of `n_e` and `n_elem` may be set.

    """

    # TODO: Allow this class to (optionally?) handle negatively charged
    # ions.  There are instances where singly negatively charged ions
    # are important in astrophysical plasmas, such as H- in the
    # atmospheres of relatively cool stars.  There may be some rare
    # situations where doubly negatively charged ions show up too,
    # though triply negatively charged ions are very unlikely.

    # TODO: Add in functionality to find equilibrium ionization states.
    # How much data will this require?

    @check_quantity({
        "T_e": {"units": u.K, "none_shall_pass": True},
        "n_e": {"units": u.m ** -3, "none_shall_pass": True},
        "n_elem": {"units": u.m ** -3, "none_shall_pass": True},
    })
    @particle_input(require='element', exclude='ion')
    def __init__(self,
                 particle: Particle,
                 ionic_fractions=None,
                 *,
                 T_e=None,
                 n_e=None,
                 n_elem=None,
                 tol: Union[float, int] = 1e-15):
        """Initialize a `~plasmapy.atomic.IonizationState` instance."""

        self._particle = particle

        try:
            self.tol = tol
            self.T_e = T_e
            self.n_elem = n_elem
            self.n_e = n_e
            self.ionic_fractions = ionic_fractions

            # This functionality has not yet been implemented:

            # if self._ionic_fractions is None and self.T_e is not None:
            #     self.equilibrate()

        except Exception as exc:
            raise AtomicError(
                f"Unable to create IonizationState instance for "
                f"{particle.particle}.") from exc

    def __getitem__(self, value) -> State:
        """Return the ionic fraction(s)."""
        if isinstance(value, slice):
            raise TypeError("IonizationState instances cannot be sliced.")

        if isinstance(value, (int, np.integer)) and 0 <= value <= self.atomic_number:
            result = State(value, self.ionic_fractions[value], self.ionic_symbols[value])
        else:
            if not isinstance(value, Particle):
                try:
                    value = Particle(value)
                except InvalidParticleError as exc:
                    raise InvalidParticleError(
                        f"{value} is not a valid integer charge or "
                        f"particle.") from exc

            same_element = value.element == self.element
            same_isotope = value.isotope == self.isotope
            has_charge_info = value.is_category(any_of=["charged", "uncharged"])

            if same_element and same_isotope and has_charge_info:
                Z = value.integer_charge
                result = State(Z, self.ionic_fractions[Z], self.ionic_symbols[Z])
            else:
                if not same_element or not same_isotope:
                    raise AtomicError("Inconsistent element or isotope.")
                elif not has_charge_info:
                    raise ChargeError("No integer charge provided.")
        return result

    def __iter__(self):
        """Initialize the instance prior to an iteration."""
        self._charge_index = 0
        return self

    def __next__(self):
        """
        Return a `~State` `~collections.namedtuple` which contains
        ``integer_charge``, ``ionic_fraction``, and ``ionic_symbol``.
        """
        if self._charge_index <= self.atomic_number:
            result = State(
                self._charge_index,
                self._ionic_fractions[self._charge_index],
                self.ionic_symbols[self._charge_index],
            )
            self._charge_index += 1
            return result
        else:
            del self._charge_index
            raise StopIteration

    def __eq__(self, other):
        """
        Return `True` if the ionic fractions for two `IonizationState`
        instances are approximately equal to within the minimum `tol`
        specified by either, and `False` otherwise.

        Raises
        ------
        AtomicError
            If `other` is not an `~plasmapy.atomic.IonizationState`
            instance, or if `other` corresponds to a different element.

        Examples
        --------
        >>> IonizationState('H', [1, 0], tol=1e-6) == IonizationState('H', [1, 1e-6], tol=1e-6)
        True
        >>> IonizationState('H', [1, 0], tol=1e-8) == IonizationState('H', [1, 1e-6], tol=1e-5)
        False

        """
        if not isinstance(other, IonizationState):
            raise AtomicError(
                "Instances of the IonizationState class may only be "
                "compared with other IonizationState instances.")

        if self.element != other.element:
            raise AtomicError("Only ionization states of the same element may be compared.")

        # Use the tightest of the two absolute tolerances
        min_tol = np.min([self.tol, other.tol])

        return np.allclose(self.ionic_fractions, other.ionic_fractions, atol=min_tol)

    @property
    def ionic_fractions(self) -> np.ndarray:
        """
        Return the ionic fractions, where the index corresponds to
        the integer charge.

        Examples
        --------
        >>> hydrogen_states = IonizationState('H', [0.9, 0.1])
        >>> hydrogen_states.ionic_fractions
        array([0.9, 0.1])

        """
        return self._ionic_fractions

    @ionic_fractions.setter
    def ionic_fractions(self, fractions):
        """
        Set the ionic fractions, while checking that the new values are
        valid and normalized to one.
        """

        if fractions is None:
            self._ionic_fractions = np.full(self.atomic_number + 1, np.nan, dtype=np.float64)
        else:
            try:
                if not isinstance(fractions, np.ndarray) or 'float' not in str(fractions.dtype):
                    fractions = np.array(fractions, dtype=np.float)
            except Exception as exc:
                raise AtomicError(
                    f"Unable to set ionic fractions of {self.element} "
                    f"to {fractions}.") from exc

            if np.min(fractions) < 0:
                raise AtomicError("Ionic fractions cannot be negative.")

            if isinstance(fractions, u.Quantity):
                if self._n_e is not None or self._n_elem is not None:
                    raise AtomicError(
                        "The ionization state may be set using number "
                        "densities for each ion only if neither of the "
                        "electron density and element density has "
                        "already been set.")
                self.number_densities = fractions
            else:

                if not np.any(np.isnan(fractions)):

                    total = np.sum(fractions)
                    if not np.isclose(total, 1, atol=self.tol, rtol=0):
                        raise AtomicError(
                            f"The sum of the ionic fractions of {self.element} "
                            f"equals {total}, which is not approximately one.")
                    if not len(fractions) == self.atomic_number + 1:
                        raise AtomicError(
                            f"len(fractions) equals {len(fractions)}, but "
                            f"should equal {self.atomic_number + 1} which "
                            f"is the atomic number of {self.element} + 1."
                        )

                self._ionic_fractions = fractions

    @property
    def n_e(self):
        """
        Return the electron number density assuming a single species
        plasma.
        """
        if self._n_e is not None:
            return self._n_e
        elif self._n_elem is not None:
            return np.sum(self._n_elem * self.ionic_fractions * self.integer_charges)
        else:
            raise AtomicError(
                "Insufficient information to calculate the electron "
                "density.")

    @n_e.setter
    def n_e(self, value):
        """
        Return the electron density assuming a single-species
        plasma.
        """
        if value is None:
            self._n_e = None
            return
        elif self._n_elem is not None:
            raise AtomicError(
                "Only one of n_e and n_elem may be set for a "
                "single element, quasineutral plasma.")
        try:
            self._n_e = value.to(u.m ** -3)
            self._n_elem = None
        except (AttributeError, u.UnitConversionError):
            raise AtomicError(_number_density_errmsg)

    @property
    def n_elem(self):
        """
        Return the number density of atoms plus ions for this
        species.
        """
        if self._n_elem is not None:
            return self._n_elem
        elif self._n_e is not None:
            return self._n_e / (self.ionic_fractions * self.integer_charges)

    @n_elem.setter
    def n_elem(self, value):
        """The number density of atoms plus ions of this species."""
        if value is None:
            self._n_elem = None
        else:
            if '_n_e' in dir(self) and self._n_e is not None:
                raise AtomicError(
                    "Only one of n_e and n_elem may be set for a "
                    "single element, quasineutral plasma.")
            try:
                self._n_elem = value.to(u.m ** -3)
            except (AttributeError, u.UnitConversionError):
                raise AtomicError(_number_density_errmsg) from None

    @property
    def number_densities(self):
        """Return the number densities for each state."""
        if self._n_e is not None or self._n_elem is not None:
            return (self.n_elem * self.ionic_fractions).to(u.m ** -3)
        else:
            raise AtomicError(
                "Insufficient information to return number densities.")

    @number_densities.setter
    def number_densities(self, value):
        if self._n_elem is not None or self._n_e is not None:
            raise AtomicError
        if not isinstance(value, u.Quantity):
            raise AtomicError
        if np.any(value < 0):
            raise AtomicError("Number densities cannot be negative.")
        try:
            value = value.to(u.m ** -3)
        except (AttributeError, u.UnitsError):
            raise AtomicError

        self._n_elem = value.sum()
        self._ionic_fractions = value / self._n_elem

    @property
    def T_e(self):
        """Return the electron temperature."""
        if self._T_e is None:
            raise AtomicError("No electron temperature has been specified.")
        return self._T_e.to(u.K, equivalencies=u.temperature_energy())

    @T_e.setter
    def T_e(self, value):

        if value is None:
            self._T_e = None
        else:
            try:
                value = value.to(u.K, equivalencies=u.temperature_energy())
            except (AttributeError, u.UnitsError):
                raise AtomicError("Invalid temperature.") from None
            self._T_e = value

    @property
    def equil_ionic_fractions(self, T_e=None):
        """
        Return the equilibrium ionic fractions for temperature `T_e` or
        the temperature set in the IonizationState instance.  Not
        implemented.
        """
        raise NotImplementedError

    def equilibrate(self, T_e=None):
        """
        Set the ionic fractions to collisional ionization equilibrium
        for temperature `T_e`.  Not implemented.
        """
        self.ionic_fractions = self.equil_ionic_fractions

    @property
    def atomic_number(self) -> int:
        """Return the atomic number of the element."""
        return self._particle.atomic_number

    @property
    def element(self) -> str:
        """Return the atomic symbol of the element."""
        return self._particle.element

    @property
    def isotope(self) -> str:
        """
        Return the isotope symbol for an isotope, or `None` if the
        particle is not an isotope.
        """
        return self._particle.isotope

    @property
    def base_particle(self) -> str:
        """
        Return the element or isotope corresponding to this
        `~plasmapy.atomic.IonizationState` instance.
        """
        return self._particle.particle

    @property
    def particles(self) -> List[Particle]:
        """
        Return a list of the `~plasmapy.atomic.Particle` class
        instances.
        """
        return [Particle(self._particle.particle, Z=i) for i in range(self.atomic_number + 1)]

    @property
    def ionic_symbols(self) -> List[str]:
        """Return the ionic symbols for all charge states."""
        return [p.ionic_symbol for p in self.particles]

    @property
    def integer_charges(self) -> np.ndarray:
        """Return an array with the integer charges."""
        return np.arange(0, self.atomic_number + 1, dtype=np.int)

    @property
    def Z_mean(self) -> np.float64:
        """Return the mean integer charge"""
        if np.nan in self.ionic_fractions:
            raise ChargeError(
                "Z_mean cannot be found because no ionic fraction "
                f"information is available for {self.particle}.")
        return np.sum(self.ionic_fractions * np.arange(self.atomic_number + 1))

    @property
    def Z_rms(self) -> np.float64:
        return np.sqrt(np.sum(self.ionic_fractions * np.arange(self.atomic_number + 1) ** 2))

    def is_normalized(self, tol=None) -> bool:
        """
        Return `True` if the sum of the ionization fractions is equal to
        one within the allowed tolerance, and `False` otherwise.
        """

        tol = tol if tol is not None else self.tol

        if not isinstance(tol, (float, int, np.integer)):
            raise TypeError("tol must be an int or float.")

        if tol < 0 or tol > 1:
            raise ValueError("tol must be between 0 and 1, inclusive.")

        total = np.sum(self._ionic_fractions)

        return np.isclose(total, 1, atol=tol, rtol=0)

    def normalize(self):
        """
        Normalize the ionization state distribution so that the sum
        becomes equal to one.
        """
        self._ionic_fractions = self._ionic_fractions / np.sum(self._ionic_fractions)

    @property
    def tol(self) -> float:
        """Return the absolute tolerance for comparisons."""
        return self._tol

    @tol.setter
    def tol(self, atol):
        """Set the absolute tolerance for comparisons."""
        if not isinstance(atol, (float, int, np.integer)):
            raise TypeError("The attribute tol must be a float or an int.")
        if 0 <= atol <= 1.0:
            self._tol = atol
        else:
            raise ValueError("Need 0 <= tol <= 1.")

class IonizationStates:
    """
    Describe the ionization state distributions of multiple elements.

    Parameters
    ----------
    inputs: `list`, `tuple`, or `dict`
        A `list` or `tuple` of elements or isotopes (if `T_e` is
        provided); a `list` of `~plasmapy.atomic.IonizationState`
        instances; or a `dict` with elements or isotopes as keys and
        a `~numpy.ndarray` of ionic fractions as the values.

    T_e: `~astropy.units.Quantity`
        The electron temperature in units of temperature or thermal
        energy per particle.

    abundances: `dict` or `str`, optional
        The relative abundances of each element in the plasma.

    log_abundances: `dict`, optional
        The base 10 logarithm of the relative abundances of each element
        in the plasma.

    number_densities: `dict`, optional
        The number densities of elements (including both neutral atoms
        and ions) in units of inverse volume.

    n_H: ~astropy.units.Quantity, optional
        The number density of neutral and ionized hydrogen atoms.  May
        only be used if the ionization state of hydrogen is included.

    Raises
    ------
    AtomicError

    Notes
    -----
    Exactly one of `abundances`, `log_abundances`, and
    `number_densities` must be specified.

    """

    @check_quantity({
        "T_e": {"units": u.K, "none_shall_pass": True},
        "n_H": {"units": u.m ** -3, "none_shall_pass": True},
    })
    def __init__(
            self,
            inputs,
            *,
            T_e=None,
            abundances=None,
            log_abundances=None,
            n_H=None,
            tol=1e-15,
        ):

        self._pars = collections.defaultdict(lambda: None)
        self.T_e = T_e
        self.n_H = n_H

        self.tol = tol

        if isinstance(inputs, dict):
            self.ionic_fractions = inputs
        elif isinstance(inputs, (list, tuple)):
            self.ionic_fractions = inputs
        else:
            raise TypeError(f"{inputs} are invalid inputs.")

        self.abundances = abundances
        self.log_abundances = log_abundances

    def __str__(self) -> str:
        join_str = ", " if len(self.elements) <= 5 else ","
        return f"<IonizationStates for: {join_str.join(self.elements)}>"

    def __repr__(self) -> str:
        return self.__str__()

    @property
    def ionic_fractions(self):
        try:
            return self._ionic_fractions
        except AttributeError as exc:
            raise AttributeError("Ionic fractions were not set.") from exc

    @ionic_fractions.setter
    def ionic_fractions(self, inputs: Union[Dict, List, Tuple]):
        """Set the ionic fractions"""
        if isinstance(inputs, dict):
            original_keys = inputs.keys()

            ionfrac_types = {type(inputs[key]) for key in original_keys}
            if u.Quantity in ionfrac_types and len(ionfrac_types) != 1:
                raise TypeError(
                    "Ionic fraction information may only be inputted "
                    "as a Quantity object if all ionic fractions are "
                    "Quantity arrays with units of inverse volume.")

            # Create a dictionary of Particle instances
            particles = dict()
            for key in original_keys:
                try:
                    particles[key] = key if isinstance(key, Particle) else Particle(key)
                except (InvalidParticleError, TypeError) as exc:
                    raise AtomicError(
                        f"Unable to create IonizationStates instance "
                        f"because {key} is not a valid particle") from exc

            # The particles whose ionization states are to be recorded
            # should be elements or isotopes but not ions or neutrals.
            is_element = particles[key].is_category('element')
            has_charge_info = particles[key].is_category(any_of=["charged", "uncharged"])
            if not is_element or has_charge_info:
                raise AtomicError(
                    f"{key} is not an element or isotope without "
                    f"charge information.")

            # We are sorting the elements/isotopes by atomic number and
            # mass number since we will often want to plot and analyze
            # things and this is the most sensible order.

            sorted_keys = sorted(original_keys, key=lambda k: (
                particles[k].atomic_number,
                particles[k].mass_number if particles[k].isotope else 0,
            ))

            _elements = []
            _particles = []
            new_ionic_fractions = {}
            for key in sorted_keys:
                new_key = particles[key].particle
                _particles.append(particles[key])
                if new_key in _elements:
                    raise AtomicError("Repeated particles in IonizationStates.")

                _elements.append(new_key)
                if isinstance(inputs[key], u.Quantity):
                    try:
                        number_densities = inputs[key].to(u.m ** -3)
                        n_elem = np.sum(number_densities)
                        new_ionic_fractions[new_key] = np.array(number_densities / n_elem)
                    except u.UnitConversionError as exc:
                        raise AtomicError("Units are not inverse volume.") from exc
                elif isinstance(inputs[key], np.ndarray) and inputs[key].dtype.kind == 'f':
                    new_ionic_fractions[particles[key].particle] = inputs[key]
                else:
                    try:
                        new_ionic_fractions[particles[key].particle] = \
                            np.array(inputs[key], dtype=np.float)
                    except ValueError:
                        raise AtomicError(f"Inappropriate ionic fractions for {key}.")

            for key in _elements:
                if np.min(new_ionic_fractions[key]) < 0 or np.max(new_ionic_fractions[key]) > 1:
                    raise AtomicError(f"Ionic fractions for {key} are not between 0 and 1.")
                if not np.isclose(np.sum(new_ionic_fractions[key]), 1, atol=self.tol, rtol=0):
                    raise AtomicError(f"Ionic fractions for {key} are not normalized to 1.")

        elif isinstance(inputs, (list, tuple)):

            try:
                _particles = [Particle(particle) for particle in inputs]
            except (InvalidParticleError, TypeError):
                raise AtomicError("Invalid inputs to IonizationStates")

            _particles.sort(key=lambda p: (p.atomic_number, p.mass_number if p.isotope else 0))
            _elements = [particle.particle for particle in _particles]
            new_ionic_fractions = {
                particle.particle: np.full(
                    particle.atomic_number + 1,
                    fill_value=np.nan,
                    dtype=np.float64
                ) for particle in _particles
            }
        else:
            raise TypeError

        # Because this depends on _particles being sorted, we add in an
        # easy check that atomic numbers do not decrease.
        for i in range(1, len(_particles)):
            if _particles[i - 1].element == _particles[i].element:
                if not _particles[i - 1].isotope and _particles[i].isotope:
                    raise AtomicError("Cannot have an element and isotopes of that element.")
            elif _particles[i - 1].atomic_number > _particles[i].atomic_number:
                raise AtomicError("_particles has not been sorted.")

        self._particles = _particles
        self._elements = _elements
        self._ionic_fractions = new_ionic_fractions

    def __getitem__(self, *values):

        errmsg = f"Invalid indexing for IonizationStates instance: {values[0]}"

        one_input = not isinstance(values[0], tuple)
        two_inputs = len(values[0]) == 2

        if not one_input and not two_inputs:
            raise TypeError(errmsg)

        try:
            arg1 = values[0] if one_input else values[0][0]
            int_charge = None if one_input else values[0][1]
            particle = arg1 if arg1 in self.elements else particle_symbol(arg1)

            if int_charge is None:
                return IonizationState(
                    particle=particle,
                    ionic_fractions=self.ionic_fractions[particle],
                    T_e=self._pars["T_e"],
                    tol=self.tol,
                )
            else:
                if not isinstance(int_charge, (int, np.integer)):
                    raise TypeError(f"{int_charge} is not a valid charge for {particle}.")
                elif not 0 <= int_charge <= atomic_number(particle):
                    raise ChargeError(f"{int_charge} is not a valid charge for {particle}.")
                return State(
                    integer_charge=int_charge,
                    ionic_fraction=self.ionic_fractions[particle][int_charge],
                    ionic_symbol=particle,
                )
        except Exception as exc:
            raise AtomicError(errmsg) from exc

    def __setitem__(self, key, value):
        if isinstance(value, dict):
            raise NotImplementedError("Dictionary assignment not implemented.")
        else:
            try:
                particle = particle_symbol(key)
                if particle not in self.elements:
                    raise AtomicError(
                        f"{key} is not one of the particles kept track of "
                        f"by this IonizationStates instance.")
                new_fractions = np.array(value, dtype=np.float64)
                if new_fractions.min() < 0 or new_fractions.max() > 1:
                    raise ValueError("Ionic fractions must be between 0 and 1.")
                if not np.isclose(np.sum(new_fractions), 1):
                    raise ValueError("Ionic fractions are not normalized.")
                if len(new_fractions) != atomic_number(particle) + 1:
                    raise ValueError(f"Incorrect size of ionic fraction array for {key}.")
                self._ionic_fractions[particle][:] = new_fractions[:]
            except Exception as exc:
                raise AtomicError(
                    f"Cannot set item for this IonizationStates "
                    f"instance for key = {repr(key)} and value = "
                    f"{repr(value)}")

    def __iter__(self):
        self._element_index = 0
        return self

    def __next__(self):
        if self._element_index < len(self.elements):
            particle = self.elements[self._element_index]
            result = IonizationState(
                particle,
                self.ionic_fractions[particle],
                T_e=self.T_e,
                tol=self.tol,
            )
            self._element_index += 1
            return result
        else:
            del self._element_index
            raise StopIteration

    def __eq__(self, other):
        """
        Test that the ionic fractions are approximately equal to
        another `~plasmapy.atomic.IonizationStates` instance.
        """
        # TODO: Should we check temperatures, densities, and other
        # parameters too?

        if not isinstance(other, IonizationStates):
            raise TypeError(
                "IonizationStates instance can only be compared with "
                "other IonizationStates instances.")

        if self.elements != other.elements:
            raise AtomicError

        tol = np.min([self.tol, other.tol])

        for element in self.elements:

            are_all_close = np.allclose(
                self.ionic_fractions[element],
                other.ionic_fractions[element],
                atol=tol,
                rtol=0,
            )

            if not are_all_close:
                return False

        return True

    @property
    def elements(self) -> List[str]:
        """
        Return a list of the elements whose ionization states are being
        kept track of.
        """
        return self._elements

    @elements.setter
    def elements(self, elems):
        if hasattr(self, "_elements"):
            raise AtomicError("Cannot change elements once they have been set.")
        else:
            self._elements = elems

    @property
    def abundances(self) -> Dict:
        if self._pars['abundances'] is None:
            raise AtomicError("No abundances are available.")
        return self._pars['abundances']

    @abundances.setter
    def abundances(self, abundances_dict: Optional[Dict]):
        """
        Set the elemental (or isotopic) abundances.  The elements and
        isotopes must be the same as or a superset of the elements whose
        ionization states are being tracked.
        """
        if abundances_dict is None:
            self._pars['abundances'] = None
        elif not isinstance(abundances_dict, dict):
            raise TypeError(
                f"The abundances argument {abundances_dict} must be a dict with elements "
                "or isotopes as keys and ")
        else:
            old_keys = abundances_dict.keys()
            try:
                new_keys_dict = {particle_symbol(old_key): old_key for old_key in old_keys}
            except Exception:
                raise AtomicError(
                    "The key {repr(old_key)} in the abundances "
                    "dictionary is not a valid element or isotope.")

            new_elements = new_keys_dict.keys()

            old_elements_set = set(self.elements)
            new_elements_set = set(new_elements)

            if old_elements_set > new_elements_set:
                raise AtomicError(
                    f"The abundances of the following particles are "
                    f"missing: {old_elements_set - new_elements_set}")

            new_abundances_dict = {}

            for element in new_elements:
                inputted_abundance = abundances_dict[new_keys_dict[element]]
                try:
                    inputted_abundance = float(inputted_abundance)
                except Exception:
                    raise TypeError

                if inputted_abundance < 0:
                    raise AtomicError(f"The abundance of {element} is negative.")
                new_abundances_dict[element] = inputted_abundance

            self._pars['abundances'] = new_abundances_dict

    @property
    def log_abundances(self):
        if self._pars['abundances'] is not None:
            log_abundances_dict = {}
            for key in self.abundances.keys():
                log_abundances_dict[key] = np.log10(self.abundances[key])
            return log_abundances_dict
        else:
            raise AtomicError("No abundances are available.")

    @log_abundances.setter
    def log_abundances(self, value):

        if value is not None:
            try:
                new_abundances_input = {}
                for key in value.keys():
                    new_abundances_input[key] = 10 ** value[key]
                self.abundances = new_abundances_input
            except Exception as exc:
                raise AtomicError("Invalid log_abundances.")

    @property
    def T_e(self):
        """Return the electron temperature."""
        return self._pars['T_e']

    @T_e.setter
    def T_e(self, electron_temperature):
        if electron_temperature is None:
            self._pars['T_e'] = None
        else:
            try:
                temp = electron_temperature.to(u.K, equivalencies=u.temperature_energy())
            except (AttributeError, u.UnitsError):
                raise AtomicError("Invalid electron temperature.")
            else:
                if temp < 0 * u.K:
                    raise AtomicError("The electron temperature cannot be negative.")
                self._pars['T_e'] = temp

    def equilibrate(self, T_e=None):
        if T_e is not None:
            raise NotImplementedError
        elif self._T_e is not None:
            raise NotImplementedError
        else:
            raise AtomicError

    @property
    def tol(self) -> float:
        """Return the absolute tolerance for comparisons."""
        return self._tol

    @tol.setter
    def tol(self, atol):
        """Set the absolute tolerance for comparisons."""
        if not isinstance(atol, (float, int, np.integer)):
            raise TypeError("The attribute tol must be a float or an int.")
        if 0 <= atol <= 1.0:
            self._tol = atol
        else:
            raise ValueError("Need 0 <= tol <= 1.")

    def normalize(self):
        for particle in self.elements:
            tot = np.sum(self.ionic_fractions[particle])
            self.ionic_fractions[particle] = self.ionic_fractions[particle] / tot

    @property
    def number_densities(self):
        ...

    @property
    def n_e(self):
        raise NotImplementedError

    @property
    def n_H(self):
        """
        The number density of hydrogen neutrals and atoms of all
        isotopes, if defined.
        """
        if 'H' not in self.elements or self._pars['n_H'] is None:
            raise AtomicError("The number density of hydrogen is not ")
        return self._pars['n_H']

    @n_H.setter
    def n_H(self, n):
        if n is None:
            self._pars['n_H'] = n
        else:
            try:
                self._pars['n_H'] = n.to(u.m ** -3)
            except u.UnitConversionError:
                raise AtomicError("Units cannot be converted to u.m**-3.")
            except Exception:
                raise AtomicError(f"{n} is not a valid number density.") from None

max_stars_repo_name: EigenDev/NEI-Research

max_issues_repo_name: EigenDev/NEI-Research

max_forks_repo_name: EigenDev/NEI-Research

reason: import numpy,import astropy


text:
"""
Collapser Optimizer
===================

#. :class:`.Collapser`
#. :class:`.CollapserMC`

Optimizer for collapsing enlarged topologies.

"""

import logging
from itertools import combinations
import numpy as np
from collections import defaultdict
from scipy.spatial.distance import pdist
import random
import matplotlib.pyplot as plt
import uuid
import os
import shutil
from stk import PdbWriter

from .optimizers import Optimizer
from ..utilities import get_atom_distance

logger = logging.getLogger(__name__)

class Collapser(Optimizer):
    """
    Collapse stk.ConstructedMolecule to decrease enlarged bonds.

    It is recommended to use the MCHammer version of this code with
    :mod:`MCHammer` [1]_, where a much cleaner version is written.
    The utilities `get_long_bond_ids` will help generate sub units.

    This optimizer aims to bring extended bonds closer together for
    further optimisation.

    .. code-block:: python

        import stk
        import stko

        bb1 = stk.BuildingBlock('NCCN', [stk.PrimaryAminoFactory()])
        bb2 = stk.BuildingBlock(
            smiles='O=CC(C=O)C=O',
            functional_groups=[stk.AldehydeFactory()],
        )
        cage1 = stk.ConstructedMolecule(
            topology_graph=stk.cage.FourPlusSix((bb1, bb2)),
        )

        # Perform collapser optimisation.
        optimizer = stko.Collapser(
            output_dir='test_coll',
            step_size=0.05,
            distance_cut=2.0,
            scale_steps=True,
        )
        cage1 = optimizer.optimize(mol=cage1)

    References
    ----------
    .. [1] https://github.com/andrewtarzia/MCHammer

    """

    def __init__(
        self,
        output_dir,
        step_size,
        distance_cut,
        scale_steps=True,
    ):
        """
        Initialize a :class:`Collapser` instance.

        Parameters
        ----------
        output_dir : :class:`str`
            The name of the directory into which files generated during
            the calculation are written, if ``None`` then
            :func:`uuid.uuid4` is used.

        step_size : :class:`float`
            The relative size of the step to take during collapse.

        distance_cut : :class:`float`
            Distance between distinct building blocks to use as
            threshold for halting collapse in Angstrom.

        scale_steps : :class:`bool`, optional
            Whether to scale the step of each distict building block
            by their relative distance from the molecules centroid.
            Defaults to ``True``

        """

        self._output_dir = output_dir
        self._step_size = step_size
        self._distance_cut = distance_cut
        self._scale_steps = scale_steps

    def _get_inter_bb_distance(self, mol):
        """
        Yield The distances between building blocks in mol.

        Ignores H atoms.

        """

        position_matrix = mol.get_position_matrix()

        for atom1, atom2 in combinations(mol.get_atom_infos(), 2):
            chk1 = (
                atom1.get_atom().get_id() != atom2.get_atom().get_id()
            )
            chk2 = (
                atom1.get_atom().get_atomic_number() != 1
                and atom2.get_atom().get_atomic_number() != 1
            )
            chk3 = (
                atom1.get_building_block_id() !=
                atom2.get_building_block_id()
            )
            if chk1 and chk2 and chk3:
                dist = get_atom_distance(
                    position_matrix=position_matrix,
                    atom1_id=atom1.get_atom().get_id(),
                    atom2_id=atom2.get_atom().get_id()
                )
                yield dist

    def _has_short_contacts(self, mol):
        """
        Calculate if there are short contants in mol.

        """

        return any(
            dist < self._distance_cut
            for dist in self._get_inter_bb_distance(mol)
        )

    def _get_new_position_matrix(self, mol, step, vectors, scales):
        """
        Get the position matrix of the mol after translation.

        """

        new_position_matrix = mol.get_position_matrix()
        for atom in mol.get_atom_infos():
            bb_id = atom.get_building_block_id()
            _id = atom.get_atom().get_id()
            pos = mol.get_position_matrix()[_id]
            new_position_matrix[_id] = (
                pos - step*vectors[bb_id]*scales[bb_id]
            )

        return new_position_matrix

    def _get_new_cell_vectors(self, unit_cell, step, scales):
        """
        Get the unit cell vectors after collapse step.

        """

        vector_1 = unit_cell.get_vector_1()
        vector_2 = unit_cell.get_vector_2()
        vector_3 = unit_cell.get_vector_3()

        max_scale = max(list(scales.values()))

        vector_1 = vector_1 - vector_1*step*max_scale
        vector_2 = vector_2 - vector_2*step*max_scale
        vector_3 = vector_3 - vector_3*step*max_scale

        return vector_1, vector_2, vector_3

    def _get_bb_vectors(self, mol, bb_atom_ids):
        """
        Get the building block to COM vectors.

        Parameters
        ----------
        mol : :class:`.Molecule`
            The molecule to be optimized.

        bb_atom_ids : :class:`dict` mapping :class:`int`: :class:`list`
            Dictionary mapping building block ids (keys) to a list of
            atom ids (values) in each distinct building block in the
            molecule.

        Returns
        -------
        bb_cent_vectors :
            :class:`dict` mapping :class:`int`: :class:`numpy.ndarray`
            Dictionary mapping building block ids (keys) to centroid
            vectors (values) of each distinct building block in the
            molecule.

        bb_cent_scales :
            :class:`dict` mapping :class:`int`: :class:`float`
            Dictionary mapping building block ids (keys) to relative
            magnitude of centroid vectors (values) of each distinct
            building block in the molecule.

        """

        cent = mol.get_centroid()

        # Get bb COM vector to molecule COM.
        bb_cent_vectors = {
            i: mol.get_centroid(atom_ids=bb_atom_ids[i])-cent
            for i in bb_atom_ids
        }

        # Scale the step size based on the different distances of
        # bbs from the COM. Impacts anisotropic topologies.
        if self._scale_steps:
            norms = {
                i: np.linalg.norm(bb_cent_vectors[i])
                for i in bb_cent_vectors
            }
            max_distance = max(list(norms.values()))
            bb_cent_scales = {
                i: norms[i]/max_distance
                for i in norms
            }
        else:
            bb_cent_scales = {
                i: 1
                for i in bb_cent_vectors
            }

        return bb_cent_vectors, bb_cent_scales

    def optimize(self, mol):
        """
        Optimize `mol`.

        Parameters
        ----------
        mol : :class:`stk.ConstructedMolecule`
            The molecule to be optimized.

        Returns
        -------
        mol : :class:`stk.ConstructedMolecule`
            The optimized molecule.

        """

        # Handle output dir.
        if self._output_dir is None:
            output_dir = str(uuid.uuid4().int)
        else:
            output_dir = self._output_dir
        output_dir = os.path.abspath(output_dir)

        if os.path.exists(output_dir):
            shutil.rmtree(output_dir)
        os.mkdir(output_dir)

        bb_atom_ids = defaultdict(list)
        for i in mol.get_atom_infos():
            bb_atom_ids[i.get_building_block_id()].append(
                i.get_atom().get_id()
            )

        # Translate each building block along bb_COM_vectors by a
        # distance `step`. I.e. `step` is the proportion of the
        # bb_COM_vectors that the building block is moved.
        step_no = 0
        step = self._step_size
        while not self._has_short_contacts(mol):
            # Update each step the building block vectors and distance.
            bb_cent_vectors, bb_cent_scales = self._get_bb_vectors(
                mol=mol,
                bb_atom_ids=bb_atom_ids
            )

            new_pos = self._get_new_position_matrix(
                mol=mol,
                step=step,
                vectors=bb_cent_vectors,
                scales=bb_cent_scales
            )
            step_no += 1
            mol = mol.with_position_matrix(new_pos)
            mol.write(
                os.path.join(output_dir, f'collapsed_{step_no}.mol')
            )

        bb_cent_vectors, bb_cent_scales = self._get_bb_vectors(
            mol=mol,
            bb_atom_ids=bb_atom_ids
        )

        # Check that we have not gone too far.
        min_dist = min(
            dist for dist in self._get_inter_bb_distance(mol)
        )
        if min_dist < self._distance_cut / 2:
            # Revert to half the previous step if we have.
            step = -(self._step_size/2)
            new_pos = self._get_new_position_matrix(
                mol=mol,
                step=step,
                vectors=bb_cent_vectors,
                scales=bb_cent_scales
            )
            step_no += 1
            mol = mol.with_position_matrix(new_pos)
            mol.write(
                os.path.join(output_dir, 'collapsed_rev.mol')
            )

        out_file = os.path.join(output_dir, 'collapser.out')
        with open(out_file, 'w') as f:
            f.write(
                "Collapser algorithm.\n"
                "====================\n"
                f"Step size: {self._step_size}\n"
                f"Scale steps?: {self._scale_steps}\n"
                f"Distance cut: {self._distance_cut}\n"
                f"====================\n"
                f"Steps run: {step_no}\n"
                f"Minimum inter-bb distance: {min_dist}\n"
                f"====================\n"
            )
        return mol

    def p_optimize(self, mol, unit_cell):
        """
        Optimize `mol`.

        Parameters
        ----------
        mol : :class:`stk.ConstructedMolecule`
            The molecule to be optimized.

        unit_cell : :class:`.UnitCell`
            The cell to be optimized.

        Returns
        -------
        mol : :class:`stk.ConstructedMolecule`
            The optimized molecule.

        unit_cell : :class:`.UnitCell`
            The optimized cell.

        """

        # Handle output dir.
        if self._output_dir is None:
            output_dir = str(uuid.uuid4().int)
        else:
            output_dir = self._output_dir
        output_dir = os.path.abspath(output_dir)

        if os.path.exists(output_dir):
            shutil.rmtree(output_dir)
        os.mkdir(output_dir)

        bb_atom_ids = defaultdict(list)
        for i in mol.get_atom_infos():
            bb_atom_ids[i.get_building_block_id()].append(
                i.get_atom().get_id()
            )

        # Translate each building block along bb_COM_vectors by a
        # distance `step`. I.e. `step` is the proportion of the
        # bb_COM_vectors that the building block is moved.
        step_no = 0
        step = self._step_size
        while not self._has_short_contacts(mol):
            # Update each step the building block vectors and distance.
            bb_cent_vectors, bb_cent_scales = self._get_bb_vectors(
                mol=mol,
                bb_atom_ids=bb_atom_ids
            )

            new_pos = self._get_new_position_matrix(
                mol=mol,
                step=step,
                vectors=bb_cent_vectors,
                scales=bb_cent_scales,
            )
            step_no += 1
            mol = mol.with_position_matrix(new_pos)
            new_vector_1, new_vector_2, new_vector_3 = (
                self._get_new_cell_vectors(
                    unit_cell=unit_cell,
                    step=step,
                    scales=bb_cent_scales,
                )
            )
            unit_cell = unit_cell.with_cell_from_vectors(
                vector_1=new_vector_1,
                vector_2=new_vector_2,
                vector_3=new_vector_3,
            )
            PdbWriter().write(
                molecule=mol,
                path=os.path.join(
                    output_dir, f'collapsed_{step_no}.pdb'
                ),
                periodic_info=unit_cell,
            )

        bb_cent_vectors, bb_cent_scales = self._get_bb_vectors(
            mol=mol,
            bb_atom_ids=bb_atom_ids
        )

        # Check that we have not gone too far.
        min_dist = min(
            dist for dist in self._get_inter_bb_distance(mol)
        )
        if min_dist < self._distance_cut / 2:
            # Revert to half the previous step if we have.
            step = -(self._step_size/2)
            new_pos = self._get_new_position_matrix(
                mol=mol,
                step=step,
                vectors=bb_cent_vectors,
                scales=bb_cent_scales
            )
            step_no += 1
            mol = mol.with_position_matrix(new_pos)
            new_vector_1, new_vector_2, new_vector_3 = (
                self._get_new_cell_vectors(
                    unit_cell=unit_cell,
                    step=step,
                    scales=bb_cent_scales,
                )
            )
            unit_cell = unit_cell.with_cell_from_vectors(
                vector_1=new_vector_1,
                vector_2=new_vector_2,
                vector_3=new_vector_3,
            )
            PdbWriter().write(
                molecule=mol,
                path=os.path.join(
                    output_dir, 'collapsed_rev.pdb'
                ),
                periodic_info=unit_cell,
            )

        out_file = os.path.join(output_dir, 'collapser.out')
        with open(out_file, 'w') as f:
            f.write(
                "Collapser algorithm.\n"
                "====================\n"
                f"Step size: {self._step_size}\n"
                f"Scale steps?: {self._scale_steps}\n"
                f"Distance cut: {self._distance_cut}\n"
                f"====================\n"
                f"Steps run: {step_no}\n"
                f"Minimum inter-bb distance: {min_dist}\n"
                f"====================\n"
            )
        return mol, unit_cell

class CollapserMC(Collapser):
    """
    Collapse molecule to decrease enlarged bonds using MC algorithm.

    It is recommended to use the MCHammer version of this code with
    :mod:`MCHammer` [1]_, where a much cleaner version is written.
    The utilities `get_long_bond_ids` will help generate sub units.

    Smarter optimisation than Collapser using simple Monte Carlo
    algorithm to perform rigid translations of building blocks.

    References
    ----------
    .. [1] https://github.com/andrewtarzia/MCHammer

    """

    def __init__(
        self,
        output_dir,
        step_size,
        target_bond_length,
        num_steps,
        bond_epsilon=50,
        nonbond_epsilon=20,
        nonbond_sigma=1.2,
        nonbond_mu=3,
        beta=2,
        random_seed=None,
    ):
        """
        Initialize a :class:`Collapser` instance.

        Parameters
        ----------
        output_dir : :class:`str`
            The name of the directory into which files generated during
            the calculation are written, if ``None`` then
            :func:`uuid.uuid4` is used.

        step_size : :class:`float`
            The relative size of the step to take during step.

        target_bond_length : :class:`float`
            Target equilibrium bond length for long bonds to minimize
            to.

        num_steps : :class:`int`
            Number of MC moves to perform.

        bond_epsilon : :class:`float`, optional
            Value of epsilon used in the bond potential in MC moves.
            Determines strength of the bond potential.
            Defaults to 50.

        nonbond_epsilon : :class:`float`, optional
            Value of epsilon used in the nonbond potential in MC moves.
            Determines strength of the nonbond potential.
            Defaults to 20.

        nonbond_sigma : :class:`float`, optional
            Value of sigma used in the nonbond potential in MC moves.
            Defaults to 1.2.

        nonbond_mu : :class:`float`, optional
            Value of mu used in the nonbond potential in MC moves.
            Determines the steepness of the nonbond potential.
            Defaults to 3.

        beta : :class:`float`, optional
            Value of beta used in the in MC moves. Beta takes the
            place of the inverse boltzmann temperature.
            Defaults to 2.

        random_seed : :class:`int`, optional
            Random seed to use for MC algorithm. Should only be set
            if exactly reproducible results are required, otherwise
            a system-based random seed should be used for proper
            sampling.

        """

        self._output_dir = output_dir
        self._step_size = step_size
        self._target_bond_length = target_bond_length
        self._num_steps = num_steps
        self._bond_epsilon = bond_epsilon
        self._nonbond_epsilon = nonbond_epsilon
        self._nonbond_sigma = nonbond_sigma
        self._nonbond_mu = nonbond_mu
        self._beta = beta
        if random_seed is None:
            random.seed()
        else:
            random.seed(random_seed)

    def _get_bb_atom_ids(self, mol):

        bb_atom_ids = defaultdict(list)
        for i in mol.get_atom_infos():
            bb_atom_ids[i.get_building_block_id()].append(
                i.get_atom().get_id()
            )

        return bb_atom_ids

    def _get_bond_length(self, mol, bond):

        position_matrix = mol.get_position_matrix()
        return get_atom_distance(
            position_matrix=position_matrix,
            atom1_id=bond.get_atom1().get_id(),
            atom2_id=bond.get_atom2().get_id()
        )

    def _get_bond_vector(self, mol, bond):

        position_matrix = mol.get_position_matrix()
        atom1_pos = position_matrix[bond.get_atom1().get_id()]
        atom2_pos = position_matrix[bond.get_atom2().get_id()]
        return atom2_pos - atom1_pos

    def _get_long_bond_infos(self, mol):
        """
        Returns dict of long bond infos.

        """

        long_bond_infos = {}
        for bond_infos in mol.get_bond_infos():
            if bond_infos.get_building_block() is None:
                ids = (
                    bond_infos.get_bond().get_atom1().get_id(),
                    bond_infos.get_bond().get_atom2().get_id(),
                )
                long_bond_infos[ids] = bond_infos

        return long_bond_infos

    def _get_bb_centroids(self, mol, bb_atom_ids):
        """
        Returns dict of building block centroids.

        """

        bb_centroids = {
            i: mol.get_centroid(atom_ids=bb_atom_ids[i])
            for i in bb_atom_ids
        }

        return bb_centroids

    def _get_cent_to_lb_vector(
        self,
        mol,
        bb_centroids,
        long_bond_infos
    ):
        """
        Returns dict of long bond atom to bb centroid vectors.

        """

        position_matrix = mol.get_position_matrix()
        centroid_to_lb_vectors = {}
        for bb in bb_centroids:
            cent = bb_centroids[bb]
            for b_atom_ids, bond_info in long_bond_infos.items():
                for atom_id in b_atom_ids:
                    atom_info, = mol.get_atom_infos(atom_ids=atom_id)
                    atom_pos = position_matrix[atom_id]
                    if atom_info.get_building_block_id() == bb:
                        centroid_to_lb_vectors[(bb, atom_id)] = (
                            atom_pos - cent,
                        )
                        break

        return centroid_to_lb_vectors

    def _bond_potential(self, distance):
        """
        Define an arbitrary parabolic bond potential.

        This potential has no relation to an empircal forcefield.
        """

        potential = (distance - self._target_bond_length) ** 2
        potential = self._bond_epsilon * potential

        return potential

    def _non_bond_potential(self, distance):
        """
        Define an arbitrary repulsive nonbonded potential.

        This potential has no relation to an empircal forcefield.
        """

        return (
            self._nonbond_epsilon * (
                (self._nonbond_sigma/distance) ** self._nonbond_mu
            )
        )

    def _compute_non_bonded_potential(self, mol):

        # Get all pairwise distances.
        pair_dists = pdist(mol.get_position_matrix())
        nonbonded_potential = np.sum(
            self._non_bond_potential(pair_dists)
        )

        return nonbonded_potential

    def _compute_potential(self, mol, long_bond_infos):

        system_potential = self._compute_non_bonded_potential(mol)
        for long_bond_ids in long_bond_infos:
            long_bond = long_bond_infos[long_bond_ids]

            system_potential += self._bond_potential(
                distance=self._get_bond_length(
                    mol=mol,
                    bond=long_bond.get_bond()
                )
            )

        return system_potential

    def _translate_atoms_along_vector(self, mol, atom_ids, vector):

        new_position_matrix = mol.get_position_matrix()
        for atom in mol.get_atom_infos(atom_ids=atom_ids):
            _id = atom.get_atom().get_id()
            pos = mol.get_position_matrix()[_id]
            new_position_matrix[_id] = pos - vector

        return mol.with_position_matrix(new_position_matrix)

    def _test_move(self, curr_pot, new_pot):

        if new_pot < curr_pot:
            return True
        else:
            exp_term = np.exp(-self._beta*(new_pot-curr_pot))
            rand_number = random.random()

            if exp_term > rand_number:
                return True
            else:
                return False

    def _output_top_lines(self):

        string = (
            '====================================================\n'
            '                Collapser optimisation              \n'
            '                ----------------------              \n'
            '                                                    \n'
            f' step size = {self._step_size} \n'
            f' target bond length = {self._target_bond_length} \n'
            f' num. steps = {self._num_steps} \n'
            f' bond epsilon = {self._bond_epsilon} \n'
            f' nonbond epsilon = {self._nonbond_epsilon} \n'
            f' nonbond sigma = {self._nonbond_sigma} \n'
            f' nonbond mu = {self._nonbond_mu} \n'
            f' beta = {self._beta} \n'
            '====================================================\n\n'
        )

        return string

    def _plot_progess(self, steps, maxds, spots, npots, output_dir):

        fig, ax = plt.subplots(figsize=(8, 5))
        ax.plot(steps, maxds, c='k', lw=2)
        # Set number of ticks for x-axis
        ax.tick_params(axis='both', which='major', labelsize=16)
        ax.set_xlabel('step', fontsize=16)
        ax.set_ylabel('max long bond length [angstrom]', fontsize=16)
        ax.axhline(y=self._target_bond_length, c='r')
        fig.tight_layout()
        fig.savefig(
            os.path.join(output_dir, 'maxd_vs_step.pdf'),
            dpi=360,
            bbox_inches='tight'
        )
        plt.close()
        # Plot energy vs timestep.
        fig, ax = plt.subplots(figsize=(8, 5))
        ax.plot(steps, spots, c='k', lw=2, label='total')
        ax.plot(steps, npots, c='r', lw=2, label='non bonded')
        # Set number of ticks for x-axis
        ax.tick_params(axis='both', which='major', labelsize=16)
        ax.set_xlabel('step', fontsize=16)
        ax.set_ylabel('potential', fontsize=16)
        ax.legend(fontsize=16)
        fig.tight_layout()
        fig.savefig(
            os.path.join(output_dir, 'pot_vs_step.pdf'),
            dpi=360,
            bbox_inches='tight'
        )
        plt.close()

    def optimize(self, mol):
        """
        Optimize `mol`.

        Parameters
        ----------
        mol : :class:`stk.ConstructedMolecule`
            The molecule to be optimized.

        Returns
        -------
        mol : :class:`stk.ConstructedMolecule`
            The optimized molecule.

        """

        # Handle output dir.
        if self._output_dir is None:
            output_dir = str(uuid.uuid4().int)
        else:
            output_dir = self._output_dir
        output_dir = os.path.abspath(output_dir)

        if os.path.exists(output_dir):
            shutil.rmtree(output_dir)
        os.mkdir(output_dir)

        # Define long bonds to optimise.
        long_bond_infos = self._get_long_bond_infos(mol)

        # If no long bonds, then optimisation is done.
        if len(long_bond_infos) == 0:
            return mol

        # Define bb centroid - long bond atom vectors.
        # These are to be maintained during optimisation.
        # centroid_to_lb_vectors = self._get_cent_to_lb_vector(
        #     mol, bb_centroids, long_bond_infos
        # )
        print('###################################################')
        print('WARNING: centroid_to_lb_vectors not maintained yet.')
        print('###################################################')

        system_potential = self._compute_potential(
            mol=mol,
            long_bond_infos=long_bond_infos
        )

        with open(os.path.join(output_dir, 'coll.out'), 'w') as f:
            f.write(self._output_top_lines())
            mol.write(os.path.join(output_dir, 'coll_0.mol'))
            steps = [0]
            passed = []
            spots = [system_potential]
            npots = [self._compute_non_bonded_potential(mol=mol)]
            maxds = [max([
                self._get_bond_length(
                    mol, long_bond_infos[i].get_bond()
                )
                for i in long_bond_infos
            ])]
            f.write(
                'Step system_potential nonbond_potential max_dist '
                'opt_bbs updated?\n'
            )
            f.write(
                f'{steps[-1]} {spots[-1]} {npots[-1]} {maxds[-1]} '
                '-- --\n'
            )
            for step in range(1, self._num_steps):

                # Randomly select a long bond.
                lb_ids = random.choice(list(long_bond_infos.keys()))
                lb_info = long_bond_infos[lb_ids]

                lb_vector = self._get_bond_vector(
                    mol=mol,
                    bond=lb_info.get_bond()
                )

                bb_id_1, bb_id_2 = (
                    i.get_building_block_id()
                    for i in mol.get_atom_infos(atom_ids=lb_ids)
                )

                # Choose bb to move out of the two randomly.
                moving_bb = random.choice([bb_id_1, bb_id_2])
                moving_bb_atom_ids = tuple(
                    i.get_atom().get_id()
                    for i in mol.get_atom_infos()
                    if i.get_building_block_id() == moving_bb
                )

                # Randomly choose between translation along long bond
                # vector or along BB-COM vector.
                # Random number from -1 to 1
                rand = (random.random() - 0.5) * 2

                # Define translation along long bond vector where
                # direction is from force, magnitude is randomly
                # scaled.
                long_bond_translation = (
                    -lb_vector * self._step_size * rand
                )

                # Get bb COM vector to molecule COM.
                cent = mol.get_centroid()
                bb_cent_vector = (
                    mol.get_centroid(atom_ids=moving_bb_atom_ids)-cent
                )
                com_translation = (
                    bb_cent_vector * self._step_size * rand
                )

                translation_vector = random.choice([
                    long_bond_translation,
                    com_translation,
                ])

                # Translate building block.
                # Update atom position of building block.
                mol = self._translate_atoms_along_vector(
                    mol=mol,
                    atom_ids=moving_bb_atom_ids,
                    vector=translation_vector,
                )

                ###################################################
                # Here I want to add rotations
                # To maintain cent_vectors relative orientations.
                # cent_vector_1 = centroid_to_lb_vectors[
                #     (bb_id_1, lb_ids[0])
                # ]
                # cent_vector_2 = centroid_to_lb_vectors[
                #     (bb_id_2, lb_ids[1])
                # ]
                ###################################################

                new_system_potential = self._compute_potential(
                    mol=mol,
                    long_bond_infos=long_bond_infos
                )

                if self._test_move(
                    system_potential,
                    new_system_potential
                ):
                    updated = 'T'
                    system_potential = new_system_potential
                    passed.append(step)
                else:
                    updated = 'F'
                    # Reverse move.
                    mol = self._translate_atoms_along_vector(
                        mol=mol,
                        atom_ids=moving_bb_atom_ids,
                        vector=-translation_vector,
                    )

                mol.write(os.path.join(output_dir, f'coll_{step}.xyz'))
                steps.append(step)
                spots.append(system_potential)
                npots.append(
                    self._compute_non_bonded_potential(mol=mol)
                )
                maxds.append(max([
                    self._get_bond_length(
                        mol, long_bond_infos[i].get_bond()
                    )
                    for i in long_bond_infos
                ]))
                f.write(
                    f'{steps[-1]} {spots[-1]} '
                    f'{npots[-1]} {maxds[-1]} {lb_ids} {updated}\n'
                )
                step += 1

            f.write('\n============================================\n')
            f.write(
                'Optimisation done:\n'
                f'{len(passed)} steps passed: '
                f'{len(passed)/self._num_steps}'
            )

        self._plot_progess(steps, maxds, spots, npots, output_dir)

        return mol

ext: py

reason: import numpy,from scipy


text:
import pickle

import nltk
import numpy as np
from sentence_transformers import SentenceTransformer

with open("data/raw_cleaned/war_and_peace.txt", "r") as infile:
    corpus = infile.read().strip().replace("\n", " ")

corpus = corpus.replace("”", "").replace("‘", "").replace("“", "")
corpus = corpus.replace("....", ".")
corpus = corpus.replace("—", " ").replace("...", ".")
corpus = corpus.replace("..", ".")
corpus = corpus.replace("“", "")
corpus = corpus.replace("*", "")

def filter_sentence(sent):
    tokens = nltk.tokenize.word_tokenize(sent)
    wds = [e.lower() for e in tokens if e.isalpha()]
    return " ".join(wds)

sentences = nltk.sent_tokenize(corpus)
sentences2vectorize = [filter_sentence(s) for s in sentences]

model = SentenceTransformer("bert-base-nli-mean-tokens")
sentence_embeddings = model.encode(sentences2vectorize)

s2 = np.asarray(sentence_embeddings)
with open("data/embeddings/sentence_embeddings.npy", "wb") as outfile:
    np.save(outfile, s2)

with open("data/embeddings/sentences.pkl", "wb") as outfile:
    pickle.dump(sentences, outfile)

max_stars_repo_name: crow-intelligence/text-visualization

max_issues_repo_name: crow-intelligence/text-visualization

max_forks_repo_name: crow-intelligence/text-visualization

reason: import numpy


text:
import Image
import numpy as np
from Urutu import *

@Urutu("gpu")
def colortoblack(r,g,b,black):
	black[tx+bx*blockDim.x] = (r[tx+bx*blockDim.x] + b[tx+bx*blockDim.x] + g[tx+bx*blockDim.x])/3
	return black

def main():
	im = Image.open("image.jpg")
	u,v = im.size[0], im.size[1]
	r = np.zeros(u*v,dtype=np.int)
	b = np.zeros_like(r)
	g = np.zeros_like(r)
	black = np.zeros_like(r)
	for j in range(v):
		for i in range(u):
			r[i+j*u] = im.getpixel((i,j))[0]
			b[i+j*u] = im.getpixel((i,j))[1]
			g[i+j*u] = im.getpixel((i,j))[2]
	out = Image.new('RGB',(u,v),"black")
	colortoblack([u,1,1],[v,1,1],r,g,b,black)
	pixels = out.load()
	for j in range(v):
		for i in range(u):
			pixels[i,j] = (black[i+j*u],black[i+j*u],black[i+j*u])
	out.save("black_gpu.jpg")

if __name__ == '__main__':
	main()

max_stars_repo_name: adityaatluri/Urutu

max_issues_repo_name: adityaatluri/Urutu

max_forks_repo_name: adityaatluri/Urutu


text:
import unittest
import os
import numpy as np
from common_wrangler.common import capture_stdout, capture_stderr, DIHES, silent_remove
from gaussian_wrangler.gausslog_unique import main, compare_gausslog_info, print_results, DEF_OUT_NAME
from gaussian_wrangler.gw_common import process_gausslog_file, CONVERG_ERR, TS, CONVERG
import logging

# logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)
DISABLE_REMOVE = logger.isEnabledFor(logging.DEBUG)

__author__ = 'hmayes'

TEST_DIR = os.path.dirname(__file__)
MAIN_DIR = os.path.dirname(TEST_DIR)
DATA_DIR = os.path.join(os.path.dirname(__file__), 'test_data')
SUB_DATA_DIR = os.path.join(DATA_DIR, 'gausslog_unique')

LOG_LIST = os.path.join(SUB_DATA_DIR, 'list.txt')
LIGNIN_LIST = os.path.join(SUB_DATA_DIR, 'lignin_list.txt')

HEADER = '"File","Convergence","Energy","Enthalpy","Gibbs_Free_E"\n'
LME2_TS3 = '"lme2acetoxprpnt_ts3_ircf_opt.log",7.5369,-535.578434,-535.403311,-535.453798\n'
LME2_TS3_NO_FREQ = '"lme2acetoxprpnt_ts3_ircf_opt_no_freq.log",0.9367,-535.578434,nan,nan\n'
LME2_25_T = '"lme2acetoxypropionate_25_t.log",0.1303,-535.578443,-535.403293,-535.453677\n'
DIFF_HEADER = '"File","Convergence","Energy","Enthalpy","Gibbs_Free_E","Diff(kcal/mol)"\n'

EMPTY_LIST = os.path.join(SUB_DATA_DIR, 'empty_list.txt')
LIST_NO_FREQ = os.path.join(SUB_DATA_DIR, 'list_no_freq.txt')
MISSING_FILE_LIST = os.path.join(SUB_DATA_DIR, 'list_with_missing_files.txt')
TWO_MOL_LIST = os.path.join(SUB_DATA_DIR, 'list_two_molecules.txt')
TWO_MORE_MOL_LIST = os.path.join(SUB_DATA_DIR, 'list_two_more_molecules.txt')
SIMILAR_LIST = os.path.join(SUB_DATA_DIR, 'list_similar_molecules.txt')
CALCALL_LIST = os.path.join(SUB_DATA_DIR, 'list_calcall.txt')

class TestGausslogUniqueNoOut(unittest.TestCase):
    def testHelp(self):
        test_input = ['-h']
        if logger.isEnabledFor(logging.DEBUG):
            main(test_input)
        with capture_stderr(main, test_input) as output:
            self.assertFalse(output)
        with capture_stdout(main, test_input) as output:
            self.assertTrue("optional arguments" in output)

    # These all test failure cases
    def testNoArgs(self):
        test_input = []
        # main(test_input)
        with capture_stderr(main, test_input) as output:
            self.assertTrue("Problems reading file" in output)

    def testNoFiles(self):
        test_input = ["-l", EMPTY_LIST]
        # main(test_input)
        with capture_stderr(main, test_input) as output:
            self.assertTrue("This program expects at least two files" in output)

    def testMissingFiles(self):
        test_input = ["-l", MISSING_FILE_LIST]
        # main(test_input)
        with capture_stderr(main, test_input) as output:
            self.assertTrue("Could not find the following file" in output)

    def testNoSuchOption(self):
        test_input = ["-@", LOG_LIST]
        # main(test_input)
        with capture_stderr(main, test_input) as output:
            self.assertTrue("unrecognized arguments" in output)

    def testSortByEnthalpyWithCutoff(self):
        test_input = ["-l", LOG_LIST, '-n', "-m", "!.0"]
        with capture_stderr(main, test_input) as output:
            self.assertTrue("WARNING:  could not convert string to float" in output)

class TestGausslogUnique(unittest.TestCase):
    # These test/demonstrate different options
    def testStandard(self):
        test_input = ["-l", LOG_LIST]
        good_output = ''.join([HEADER, LME2_TS3, LME2_25_T]) + '\n'
        with capture_stdout(main, test_input) as output:
            self.assertTrue(good_output in output)
        with capture_stderr(main, test_input) as output:
            self.assertTrue('Check convergence' in output)

    # These test/demonstrate different options
    def testSortByEnergy(self):
        test_input = ["-l", LOG_LIST, "-e"]
        good_output = ''.join([HEADER, LME2_25_T, LME2_TS3]) + '\n'
        with capture_stdout(main, test_input) as output:
            self.assertTrue(good_output in output)
        with capture_stderr(main, test_input) as output:
            self.assertTrue('Check convergence' in output)

    def testSortByEnthalpy(self):
        test_input = ["-l", LOG_LIST, '-n']
        good_output = ''.join([HEADER, LME2_TS3, LME2_25_T]) + '\n'
        with capture_stdout(main, test_input) as output:
            self.assertTrue(good_output in output)
        with capture_stderr(main, test_input) as output:
            self.assertTrue('Check convergence' in output)

    def testSortByEnergyWithCutoff(self):
        test_input = ["-l", LOG_LIST, "-e", "-m", "5.0"]
        good_output = '"File","Convergence","Energy","Enthalpy","Gibbs_Free_E","Diff(kcal/mol)"\n' \
                      '"Files within SCF energy cutoff of 5.00 kcal/mol"\n' + LME2_25_T[:-1] + ',0.00\n' + \
                      LME2_TS3[:-1] + ',0.01'
        try:
            # main(test_input)
            with capture_stdout(main, test_input) as output:
                self.assertTrue(good_output in output)
            with open(DEF_OUT_NAME) as f:
                cutoff_files = [row.strip() for row in f.readlines()]
            self.assertEqual(cutoff_files, ['lme2acetoxypropionate_25_t.log', 'lme2acetoxprpnt_ts3_ircf_opt.log'])
        finally:
            silent_remove(DEF_OUT_NAME)
            pass

    def testSortByEnthalpyWithCutoffAlso(self):
        test_input = ["-l", LIGNIN_LIST, "-m", "1.0"]
        good_output = DIFF_HEADER + '"Files within SCF energy cutoff of 1.00 kcal/mol"\n' \
                                    '"g_dimer_8.log",1.4217,-1225.835244,nan,nan,0.00\n' \
                                    '"g_dimer_102.log",0.2522,-1225.834813,nan,nan,0.27\n' \
                                    '"Files outside of cutoff:"\n' \
                                    '"g_dimer_100.log",0.5761,-1225.831663,nan,nan,2.25\n' \
                                    '"g_dimer_901.log",0.3047,-1225.823048,nan,nan,7.65'
        try:
            with capture_stdout(main, test_input) as output:
                self.assertTrue(good_output in output)
            with open(DEF_OUT_NAME) as f:
                cutoff_files = [row.strip() for row in f.readlines()]
            self.assertEqual(cutoff_files, ['g_dimer_8.log', 'g_dimer_102.log'])
        finally:
            silent_remove(DEF_OUT_NAME)

    def testSortByEnthalpyWithCutoff(self):
        test_input = ["-l", LOG_LIST, '-n', "-m", "1.0"]
        good_output = (DIFF_HEADER + '"Files within enthalpy cutoff of 1.00 kcal/mol"\n' + LME2_TS3[:-1] + ',0.00\n'
                       + LME2_25_T[:-1] + ',0.01')
        try:
            with capture_stdout(main, test_input) as output:
                self.assertTrue(good_output in output)
        finally:
            silent_remove(DEF_OUT_NAME)

    def testNoFreq(self):
        # also tests that it can skip a blank line
        test_input = ["-l", LIST_NO_FREQ, "-n"]
        good_output = ''.join([HEADER, LME2_25_T, LME2_TS3_NO_FREQ]) + '\n'
        main(test_input)
        with capture_stdout(main, test_input) as output:
            self.assertTrue(good_output in output)
        with capture_stderr(main, test_input) as output:
            self.assertFalse('Check convergence' in output)

    def testTwoMolecules(self):
        pet_843 = '"pet_mono_843_tzvp.log",1.4694,-917.071861,-916.796704,-916.862565\n'
        pet_1 = '"pet_mono_1_tzvp.log",0.8478,-917.069491,-916.794649,-916.861929\n'
        test_input = ["-l", TWO_MOL_LIST, "-n"]
        good_output = ''.join([HEADER, pet_843, pet_1, LME2_TS3, LME2_25_T]) + '\n'
        # main(test_input)
        with capture_stdout(main, test_input) as output:
            self.assertTrue(good_output in output)
        with capture_stderr(main, test_input) as output:
            self.assertTrue('lme2acetoxprpnt_ts3_ircf_opt.log' in output)

    def testTwoMoreMolecules(self):
        # This test checks that the program can handle the case when Gaussian prints '********' for convergence,
        #    when the convergence is so bad that it can't fit in the space allotted
        good_result = '"lme2acetoxprpnt_ts4_ircf_opt.log",2.0767,-535.576027,-535.401005,-535.450649\n' \
                      '"lme2acetoxprpnt_ts4_b_ts_ircf_opt.log",8.2747,-535.575022,-535.399906,-535.449855\n'
        test_input = ["-l", TWO_MORE_MOL_LIST, "-n"]
        good_output = ''.join([HEADER, good_result]) + '\n'
        # main(test_input)
        with capture_stdout(main, test_input) as output:
            self.assertTrue(good_output in output)
        with capture_stderr(main, test_input) as output:
            self.assertTrue('lme2acetoxprpnt_ts4_ircf_opt.log' in output)
            self.assertTrue('lme2acetoxprpnt_ts4_b_ts_ircf_opt.log' in output)

    def testSimilarMolecules(self):
        # I was surprised that these weren't listed as the same; turns out, the difference in dihedral angle
        #  was almost, but not quite, 360; now, if within tolerance of 360 degrees, it will subtract 360, catching
        #  these similar conformations
        good_result = '"me2pheoxprpnt_30.log",0.0139,-613.945900,-613.726343,-613.779248\n\n'
        good_output = ''.join([HEADER, good_result])
        test_input = ["-l", SIMILAR_LIST, "-n"]
        with capture_stdout(main, test_input) as output:
            self.assertTrue(good_output in output)
        with capture_stderr(main, test_input) as output:
            self.assertTrue('' == output)

    def testCalcAllOutput(self):
        # make sure program sees the enthalpy when "CalcAll" is run instead of frequency
        good_result = '"hexyl_acrylate_239.log",1.1100,-503.005111,-502.751977,-502.807357\n' \
                      '"hexyl_acrylate_419.log",0.0706,-503.004423,-502.751021,-502.805958\n\n'
        good_output = ''.join([HEADER, good_result])
        test_input = ["-l", CALCALL_LIST, "-n"]
        # main(test_input)
        with capture_stdout(main, test_input) as output:
            self.assertTrue(good_output in output)
        with capture_stderr(main, test_input) as output:
            self.assertTrue('' == output)

class TestGausslogUniqueFunctions(unittest.TestCase):
    def testOneFileMissingDihedralInfo(self):
        # testing for some specific problem encountered: when there is a freq only job (no opt) there is no
        #     dihedral info (which caused trouble when tried to compare it) and also prevented finding the
        #     convergence. Thus, checking that this specific error is resolved with the expected warning.
        current_file_list = ["tieg5ipatse_ts_ircf_opt_noeg.log", "tieg5ipatse_ts_noeg_reacta.log",
                             "tieg5ipatse_ts_noeg_reactc.log"]
        full_info_dict = {}
        check_list = []
        gausslog_loc = os.path.join(SUB_DATA_DIR, current_file_list[0])
        with capture_stderr(process_gausslog_file, gausslog_loc, find_dih=True, find_converg=True) as output:
            self.assertTrue("Requested dihedral data not found for file: tieg5ipatse_ts_ircf_opt_noeg.log" in
                            output)
        for gausslog_fname in current_file_list:
            if gausslog_fname == "":
                continue
            gausslog_loc = os.path.join(SUB_DATA_DIR, gausslog_fname)
            gausslog_content = process_gausslog_file(gausslog_loc, find_dih=True, find_converg=True)
            full_info_dict[gausslog_fname] = gausslog_content
            if gausslog_content[CONVERG_ERR]:
                check_list.append(gausslog_fname)
        self.assertTrue(full_info_dict[current_file_list[0]][DIHES] is None)
        self.assertTrue(full_info_dict[current_file_list[1]][DIHES] is not None)
        self.assertFalse(full_info_dict[current_file_list[0]][TS])
        self.assertTrue(full_info_dict[current_file_list[0]][CONVERG_ERR])
        self.assertTrue(full_info_dict[current_file_list[1]][CONVERG_ERR])
        self.assertEqual(len(check_list), 3)
        list_of_conf_lists = compare_gausslog_info(full_info_dict, 5)
        self.assertEqual(len(list_of_conf_lists), 3)
        winner_str, warn_files_str = print_results(full_info_dict, list_of_conf_lists, False, True, print_winners=False)
        warn_files_list = warn_files_str.split('\n')
        self.assertEqual(len(warn_files_list), 4)

    def testRemoveConvErr(self):
        # sometimes Gaussian ignore convergence error stating:
        #     "Optimization completed on the basis of negligible forces."
        # in those cases, let's also ignore the error and change the error flag to false
        gausslog_loc = os.path.join(SUB_DATA_DIR, "tieg5pdc2tsc_ts_ircr_opt.log")
        gausslog_content = process_gausslog_file(gausslog_loc, find_dih=True, find_converg=True)
        self.assertTrue("False" in gausslog_content[CONVERG_ERR])
        self.assertAlmostEqual(gausslog_content[CONVERG], 1.5711111111111113)

    def testTSInfo(self):
        goodvibes_helper_folder = os.path.join(DATA_DIR, 'goodvibes_helper')
        file_list = ["co_gas.log", "hcoch3_gas.log", "tieg5ipatse_ts.log", "water.log", "h_gas_stable_t.log"]
        full_info_dict = {}
        check_list = []
        for fname in file_list:
            gausslog_loc = os.path.join(goodvibes_helper_folder, fname)
            gausslog_content = process_gausslog_file(gausslog_loc, find_dih=True, find_converg=True)
            full_info_dict[fname] = gausslog_content
            if gausslog_content[CONVERG_ERR]:
                check_list.append(fname)
        for fname in file_list:
            if fname == "tieg5ipatse_ts.log":
                self.assertTrue(full_info_dict[fname][TS])
            else:
                self.assertFalse(full_info_dict[fname][TS])

    def testCatchSortError(self):
        log_info = {'ti_eg5_dime_pdc1_tsb_ts.log': {'atoms_section': {}, 'base_name': 'ti_eg5_dime_pdc1_tsb_ts.log',
                                                    'Stoichiometry': 'C19H34O16Ti', 'Transition_State': True,
                                                    'Energy': -2797.66176465, 'Enthalpy': -2797.058267,
                                                    'Gibbs_Free_E': -2797.178818,
                                                    'converg_dict': {}, 'Charge': 0, 'Mult': 1, 'Dihedrals': {},
                                                    'Convergence': 2.055833333333333, 'Convergence_Error': True},
                    'ti_eg5_dime_pdc1_tsc_ts.log':
                        {'atoms_section': {},
                         'base_name': 'ti_eg5_dime_pdc1_tsc_ts.log',
                         'Stoichiometry': 'C19H34O16Ti',
                         'Transition_State': None,
                         'Energy': -2797.72551346,
                         'Enthalpy': np.nan,
                         'Gibbs_Free_E': np.nan,
                         'converg_dict': {},
                         'Charge': 0,
                         'Mult': 1,
                         'Dihedrals': {},
                         'Convergence': 171.12472222222223,
                         'Convergence_Error': True},
                    'ti_eg5_dime_pdc1_tsd_ts.log':
                        {'atoms_section': {},
                         'base_name': 'ti_eg5_dime_pdc1_tsd_ts.log',
                         'Stoichiometry': 'C19H34O16Ti',
                         'Transition_State': True,
                         'Energy': -2797.6659881,
                         'Enthalpy': -2797.062691,
                         'Gibbs_Free_E': -2797.182718,
                         'converg_dict': {},
                         'Charge': 0,
                         'Mult': 1,
                         'Dihedrals': {},
                         'Convergence': 0.8447222222222222,
                         'Convergence_Error': False},
                    'ti_eg5_dime_pdc1_tse_ts.log':
                        {'atoms_section': {},
                         'base_name': 'ti_eg5_dime_pdc1_tse_ts.log',
                         'Stoichiometry': 'C19H34O16Ti',
                         'Transition_State': None,
                         'Energy': np.nan,
                         'Enthalpy': np.nan,
                         'Gibbs_Free_E': np.nan,
                         'converg_dict': {},
                         'Charge': 0,
                         'Mult': 1,
                         'Dihedrals': {},
                         'Convergence': np.nan,
                         'Convergence_Error': None},
                    }
        expected_warn_str = '\n    ti_eg5_dime_pdc1_tsc_ts.log:  171.12' \
                            '\n    ti_eg5_dime_pdc1_tsb_ts.log:  2.06' \
                            '\n    ti_eg5_dime_pdc1_tse_ts.log:  Not found'

        list_of_conf_lists = [['ti_eg5_dime_pdc1_tsb_ts.log'], ['ti_eg5_dime_pdc1_tsc_ts.log'],
                              ['ti_eg5_dime_pdc1_tsd_ts.log'], ['ti_eg5_dime_pdc1_tse_ts.log']]
        sort_by_enthalpy = False
        sort_by_energy = True
        winner_str, warn_files_str = print_results(log_info, list_of_conf_lists, sort_by_enthalpy,
                                                   sort_by_energy, print_winners=True)
        self.assertTrue('"ti_eg5_dime_pdc1_tse_ts.log",nan,nan,nan' in winner_str)
        self.assertEqual(warn_files_str, expected_warn_str)
        pass


text:
# coding=utf-8
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tests of polymer evolutionary optimization tools."""

import jax.numpy as jnp
from jax import vmap, jit, random
from flatland import evolution as evo

def test_types():

  key = random.PRNGKey(0)

  pop = evo.generate_population(key=key)

  assert isinstance(pop, evo.Population)

  pop2 = evo.mutate_population(
      key=random.split(key)[0],
      population=pop,
  )

  assert isinstance(pop2, evo.Population)

  fitnesses = random.uniform(key, (pop2.shape[0],))

  pop3 = evo.resample_population(key=random.split(key)[0],
                                 population=pop2,
                                 fitnesses=fitnesses)

  assert isinstance(pop3, evo.Population)

def test_evolve_with_mutation_runs():
  """Evolve a population of polymers whose mean value is 1.0.

  E.g. [0,1,2,1,0,1,2,1,0,1,2,1,0]
  """

  alphabet_size = 3
  population_size = 10
  genome_length = 5
  mutation_rate = 0.15
  num_generations = 10

  def fitness_mean_value_target(v, target_value=1.0):
    return 1 - jnp.abs(jnp.mean(v) - target_value) / target_value

  @jit
  def batched_fitness_mean_value_target(population):
    return vmap(fitness_mean_value_target)(population)

  res = evo.evolve_with_mutation(fitness_fn=batched_fitness_mean_value_target,
                                 num_generations=num_generations,
                                 pop_size=population_size,
                                 genome_length=genome_length,
                                 mutation_rate=mutation_rate,
                                 alphabet_size=alphabet_size,
                                 keep_full_population_history=True,
                                 key=random.PRNGKey(1))

  fitness_history, population_history, population = res

max_stars_repo_name: cayley-group/flatland

max_issues_repo_name: cayley-group/flatland

max_forks_repo_name: cayley-group/flatland

reason: import jax,from jax


text:
import matplotlib as mpl
import numpy as np
import dill

mpl.use('Agg')

from pySDC.helpers.stats_helper import get_sorted
from pySDC.implementations.collocations import Collocation
from pySDC.implementations.controller_classes.controller_nonMPI import controller_nonMPI
from pySDC.implementations.problem_classes.Piline import piline
from pySDC.implementations.sweeper_classes.imex_1st_order import imex_1st_order
import pySDC.helpers.plot_helper as plt_helper

from pySDC.core.Hooks import hooks

class log_data(hooks):

    def post_step(self, step, level_number):

        super(log_data, self).post_step(step, level_number)

        # some abbreviations
        L = step.levels[level_number]

        L.sweep.compute_end_point()

        self.add_to_stats(process=step.status.slot, time=L.time, level=L.level_index, iter=0,
                          sweep=L.status.sweep, type='v1', value=L.uend[0])
        self.add_to_stats(process=step.status.slot, time=L.time, level=L.level_index, iter=0,
                          sweep=L.status.sweep, type='v2', value=L.uend[1])
        self.add_to_stats(process=step.status.slot, time=L.time, level=L.level_index, iter=0,
                          sweep=L.status.sweep, type='p3', value=L.uend[2])

def main():
    """
    A simple test program to do SDC/PFASST runs for the Piline model
    """

    # initialize level parameters
    level_params = dict()
    level_params['restol'] = 1E-10
    level_params['dt'] = 0.25

    # initialize sweeper parameters
    sweeper_params = dict()
    sweeper_params['collocation_class'] = Collocation
    sweeper_params['node_type'] = 'LEGENDRE'
    sweeper_params['quad_type'] = 'LOBATTO'
    sweeper_params['num_nodes'] = 3
    # sweeper_params['QI'] = 'LU'  # For the IMEX sweeper, the LU-trick can be activated for the implicit part

    # initialize problem parameters
    problem_params = dict()
    problem_params['Vs'] = 100.0
    problem_params['Rs'] = 1.0
    problem_params['C1'] = 1.0
    problem_params['Rpi'] = 0.2
    problem_params['C2'] = 1.0
    problem_params['Lpi'] = 1.0
    problem_params['Rl'] = 5.0

    # initialize step parameters
    step_params = dict()
    step_params['maxiter'] = 20

    # initialize controller parameters
    controller_params = dict()
    controller_params['logger_level'] = 20
    controller_params['hook_class'] = log_data

    # fill description dictionary for easy step instantiation
    description = dict()
    description['problem_class'] = piline                         # pass problem class
    description['problem_params'] = problem_params                # pass problem parameters
    description['sweeper_class'] = imex_1st_order                 # pass sweeper
    description['sweeper_params'] = sweeper_params                # pass sweeper parameters
    description['level_params'] = level_params                    # pass level parameters
    description['step_params'] = step_params                      # pass step parameters

    assert 'errtol' not in description['step_params'].keys(), "No exact or reference solution known to compute error"

    # set time parameters
    t0 = 0.0
    Tend = 20

    # instantiate controller
    controller = controller_nonMPI(num_procs=1, controller_params=controller_params, description=description)

    # get initial values on finest level
    P = controller.MS[0].levels[0].prob
    uinit = P.u_exact(t0)

    # call main function to get things done...
    uend, stats = controller.run(u0=uinit, t0=t0, Tend=Tend)

    fname = 'piline.dat'
    f = open(fname, 'wb')
    dill.dump(stats, f)
    f.close()

    # filter statistics by number of iterations
    iter_counts = get_sorted(stats, type='niter', sortby='time')

    # compute and print statistics
    min_iter = 20
    max_iter = 0

    f = open('piline_out.txt', 'w')
    niters = np.array([item[1] for item in iter_counts])
    out = '   Mean number of iterations: %4.2f' % np.mean(niters)
    f.write(out + '\n')
    print(out)
    for item in iter_counts:
        out = 'Number of iterations for time %4.2f: %1i' % item
        f.write(out + '\n')
        print(out)
        min_iter = min(min_iter, item[1])
        max_iter = max(max_iter, item[1])

    assert np.mean(niters) <= 10, "Mean number of iterations is too high, got %s" % np.mean(niters)
    f.close()

    plot_voltages()

def plot_voltages(cwd='./'):
    f = open(cwd + 'piline.dat', 'rb')
    stats = dill.load(f)
    f.close()

    # convert filtered statistics to list of iterations count, sorted by process
    v1 = get_sorted(stats, type='v1', sortby='time')
    v2 = get_sorted(stats, type='v2', sortby='time')
    p3 = get_sorted(stats, type='p3', sortby='time')

    times = [v[0] for v in v1]

    setup_mpl()
    fig, ax = plt_helper.plt.subplots(1, 1, figsize=(4.5, 3))
    ax.plot(times, [v[1] for v in v1], linewidth=1, label='$v_{C_1}$')
    ax.plot(times, [v[1] for v in v2], linewidth=1, label='$v_{C_2}$')
    ax.plot(times, [v[1] for v in p3], linewidth=1, label='$i_{L_\pi}$')
    ax.legend(frameon=False, fontsize=12, loc='center right')

    ax.set_xlabel('Time')
    ax.set_ylabel('Energy')

    fig.savefig('data/piline_model_solution.png', dpi=300, bbox_inches='tight')

def setup_mpl(fontsize=8):
    plt_helper.setup_mpl(reset=True)

    style_options = {
        "font.family": "sans-serif",
        "font.serif": "Computer Modern Sans Serif",
        "font.sans-serif": "Computer Modern Sans Serif",
        "font.monospace": "Computer Modern Sans Serif",
        "axes.labelsize": 12,  # LaTeX default is 10pt font.
        "legend.fontsize": 13,  # Make the legend/label fonts a little smaller
        "axes.xmargin": 0.03,
        "axes.ymargin": 0.03,
        "lines.linewidth": 1,  # Make the plot lines a little smaller
    }

    mpl.rcParams.update(style_options)

if __name__ == "__main__":
    main()

max_stars_repo_name: brownbaerchen/pySDC

max_issues_repo_name: brownbaerchen/pySDC

max_forks_repo_name: brownbaerchen/pySDC