<commit_before># Copyright (c) 2009-2021 The Regents of the University of Michigan
# This file is part of the HOOMD-blue project, released under the BSD 3-Clause
# License.

"""Apply forces to particles."""

import hoomd
from hoomd import _hoomd
from hoomd.md import _md
from hoomd.operation import _HOOMDBaseObject
from hoomd.logging import log
from hoomd.data.typeparam import TypeParameter
from hoomd.data.typeconverter import OnlyTypes
from hoomd.data.parameterdicts import ParameterDict, TypeParameterDict
from hoomd.filter import ParticleFilter
from hoomd.md.constrain import ConstraintForce

def _ellip_preprocessing(constraint):
    if constraint is not None:
        if constraint.__class__.__name__ == "constraint_ellipsoid":
            return constraint
        else:
            raise RuntimeError(
                "Active force constraint is not accepted (currently only "
                "accepts ellipsoids)")
    else:
        return None

class _force:  # noqa - This will be removed eventually. Needed to build docs.
    pass

class Force(_HOOMDBaseObject):
    """Defines a force in HOOMD-blue.

    Pair, angle, bond, and other forces are subclasses of this class.

    Note:
        :py:class:`Force` is the base class for all loggable forces.
        Users should not instantiate this class directly.

    Initializes some loggable quantities.
    """

    @log
    def energy(self):
        """float: Sum of the energy of the whole system. :math:`[energy]`"""
        if self._attached:
            self._cpp_obj.compute(self._simulation.timestep)
            return self._cpp_obj.calcEnergySum()
        else:
            return None

    @log(category="particle")
    def energies(self):
        """(*N_particles*, ) `numpy.ndarray` of ``numpy.float64``: The \
        energies for all particles. :math:`[energy]`"""
        if self._attached:
            self._cpp_obj.compute(self._simulation.timestep)
            return self._cpp_obj.getEnergies()
        else:
            return None

    @log(category="particle")
    def forces(self):
        """(*N_particles*, 3) `numpy.ndarray` of ``numpy.float64``: The \
        forces for all particles. :math:`[force]`"""
        if self._attached:
            self._cpp_obj.compute(self._simulation.timestep)
            return self._cpp_obj.getForces()
        else:
            return None

    @log(category="particle")
    def torques(self):
        """(*N_particles*, 3) `numpy.ndarray` of ``numpy.float64``: The torque \
        for all particles. :math:`[force \cdot length]`"""
        if self._attached:
            self._cpp_obj.compute(self._simulation.timestep)
            return self._cpp_obj.getTorques()
        else:
            return None

    @log(category="particle")
    def virials(self):
        """(*N_particles*, ) `numpy.ndarray` of ``numpy.float64``: The virial \
        for all particles. :math:`[energy]`"""
        if self._attached:
            self._cpp_obj.compute(self._simulation.timestep)
            return self._cpp_obj.getVirials()
        else:
            return None

class constant(Force):  # noqa - this will be renamed when it is ported to v3
    R"""Constant force.

    Args:
        fvec (tuple): force vector :math:`[force]`
        tvec (tuple): torque vector :math:`[force \cdot length]`
        fx (float): x component of force, retained for backwards compatibility :math:`[force]`
        fy (float): y component of force, retained for backwards compatibility :math:`[force]`
        fz (float): z component of force, retained for backwards compatibility :math:`[force]`
        group (``hoomd.group``): Group for which the force will be set.
        callback (`callable`): A python callback invoked every time the forces
            are computed

    :py:class:`constant` specifies that a constant force should be added to
    every particle in the simulation or optionally to all particles in a group.

    Note:
        Forces are kept constant during the simulation. If a callback should
        re-compute particle forces every time step, it needs to overwrite the
        old forces of **all** particles with new values.

    Note:
        Per-particle forces take precedence over a particle group, which takes
        precedence over constant forces for all particles.

    Examples::

        force.constant(fx=1.0, fy=0.5, fz=0.25)
        const = force.constant(fvec=(0.4,1.0,0.5))
        const = force.constant(fvec=(0.4,1.0,0.5),group=fluid)
        const = force.constant(fvec=(0.4,1.0,0.5), tvec=(0,0,1) ,group=fluid)

        def updateForces(timestep):
            global const
            const.setForce(tag=1, fvec=(1.0*timestep,2.0*timestep,3.0*timestep))
        const = force.constant(callback=updateForces)
    """

    def __init__(
        self,
        fx=None,
        fy=None,
        fz=None,
        fvec=None,
        tvec=None,
        group=None,
        callback=None,
    ):

        if (fx is not None) and (fy is not None) and (fz is not None):
            self.fvec = (fx, fy, fz)
        elif fvec is not None:
            self.fvec = fvec
        else:
            self.fvec = (0, 0, 0)

        if tvec is not None:
            self.tvec = tvec
        else:
            self.tvec = (0, 0, 0)

        if (self.fvec == (0, 0, 0)) and (self.tvec == (0, 0, 0)
                                         and callback is None):
            hoomd.context.current.device.cpp_msg.warning(
                "The constant force specified has no non-zero components\n")

        # initialize the base class
        Force.__init__(self)

        # create the c++ mirror class
        if group is not None:
            self.cppForce = _hoomd.ConstForceCompute(
                hoomd.context.current.system_definition,
                group.cpp_group,
                self.fvec[0],
                self.fvec[1],
                self.fvec[2],
                self.tvec[0],
                self.tvec[1],
                self.tvec[2],
            )
        else:
            self.cppForce = _hoomd.ConstForceCompute(
                hoomd.context.current.system_definition,
                self.fvec[0],
                self.fvec[1],
                self.fvec[2],
                self.tvec[0],
                self.tvec[1],
                self.tvec[2],
            )

        if callback is not None:
            self.cppForce.setCallback(callback)

        hoomd.context.current.system.addCompute(self.cppForce, self.force_name)

    R""" Change the value of the constant force.

    Args:
        fx (float) New x-component of the force :math:`[force]`
        fy (float) New y-component of the force :math:`[force]`
        fz (float) New z-component of the force :math:`[force]`
        fvec (tuple) New force vector
        tvec (tuple) New torque vector
        group Group for which the force will be set
        tag (int) Particle tag for which the force will be set
            .. versionadded:: 2.3

     Using setForce() requires that you saved the created constant force in a
     variable. i.e.

     Examples:
        const = force.constant(fx=0.4, fy=1.0, fz=0.5)

        const.setForce(fx=0.2, fy=0.1, fz=-0.5)
        const.setForce(fx=0.2, fy=0.1, fz=-0.5, group=fluid)
        const.setForce(fvec=(0.2,0.1,-0.5), tvec=(0,0,1), group=fluid)
    """

    def setForce(  # noqa - this will be documented when it is ported to v3
        self,
        fx=None,
        fy=None,
        fz=None,
        fvec=None,
        tvec=None,
        group=None,
        tag=None,
    ):

        if (fx is not None) and (fy is not None) and (fx is not None):
            self.fvec = (fx, fy, fz)
        elif fvec is not None:
            self.fvec = fvec
        else:
            self.fvec = (0, 0, 0)

        if tvec is not None:
            self.tvec = tvec
        else:
            self.tvec = (0, 0, 0)

        if (fvec == (0, 0, 0)) and (tvec == (0, 0, 0)):
            hoomd.context.current.device.cpp_msg.warning(
                "You are setting the constant force to have no non-zero "
                "components\n")

        self.check_initialization()
        if group is not None:
            self.cppForce.setGroupForce(
                group.cpp_group,
                self.fvec[0],
                self.fvec[1],
                self.fvec[2],
                self.tvec[0],
                self.tvec[1],
                self.tvec[2],
            )
        elif tag is not None:
            self.cppForce.setParticleForce(
                tag,
                self.fvec[0],
                self.fvec[1],
                self.fvec[2],
                self.tvec[0],
                self.tvec[1],
                self.tvec[2],
            )
        else:
            self.cppForce.setForce(
                self.fvec[0],
                self.fvec[1],
                self.fvec[2],
                self.tvec[0],
                self.tvec[1],
                self.tvec[2],
            )

    R""" Set a python callback to be called before the force is evaluated

    Args:
        callback (`callable`) The callback function

     Examples:
        const = force.constant(fx=0.4, fy=1.0, fz=0.5)

        def updateForces(timestep):
            global const
            const.setForce(tag=1, fvec=(1.0*timestep,2.0*timestep,3.0*timestep))

        const.set_callback(updateForces)
        run(100)

        # Reset the callback
        const.set_callback(None)
    """

    def set_callback(self, callback=None):  # noqa - will be ported to v3
        self.cppForce.setCallback(callback)

    # there are no coeffs to update in the constant force compute
    def update_coeffs(self):  # noqa - will be ported to v3
        pass

class Active(Force):
    r"""Active force.

    Attributes:
        filter (:py:mod:`hoomd.filter`): Subset of particles on which to apply
            active forces.
        rotation_diff (float): rotational diffusion constant, :math:`D_r`, for
            all particles in the group. :math:`[time^{-1}]`
        active_force (tuple): active force vector in reference to the
            orientation of a particle. It is defined per particle type and stays
            constant during the simulation. :math:`[force]`
        active_torque (tuple): active torque vector in reference to the
            orientation of a particle. It is defined per particle type and stays
            constant during the simulation. :math:`[force \cdot length]`

    :py:class:`Active` specifies that an active force should be added to all
    particles.  Obeys :math:`\delta {\bf r}_i = \delta t v_0 \hat{p}_i`, where
    :math:`v_0` is the active velocity. In 2D :math:`\hat{p}_i = (\cos \theta_i,
    \sin \theta_i)` is the active force vector for particle :math:`i` and the
    diffusion of the active force vector follows :math:`\delta \theta / \delta t
    = \sqrt{2 D_r / \delta t} \Gamma`, where :math:`D_r` is the rotational
    diffusion constant, and the gamma function is a unit-variance random
    variable, whose components are uncorrelated in time, space, and between
    particles.  In 3D, :math:`\hat{p}_i` is a unit vector in 3D space, and
    diffusion follows :math:`\delta \hat{p}_i / \delta t = \sqrt{2 D_r / \delta
    t} \Gamma (\hat{p}_i (\cos \theta - 1) + \hat{p}_r \sin \theta)`, where
    :math:`\hat{p}_r` is an uncorrelated random unit vector. The persistence
    length of an active particle's path is :math:`v_0 / D_r`.  The rotational
    diffusion is applied to the orientation vector/quaternion of each particle.
    This implies that both the active force and the active torque vectors in the
    particle frame stay constant during the simulation. Hence, the active forces
    in the system frame are composed of the forces in particle frame and the
    current orientation of the particle.

    Examples::

        all = filter.All()
        active = hoomd.md.force.Active(
            filter=hoomd.filter.All(), rotation_diff=0.01
            )
        active.active_force['A','B'] = (1,0,0)
        active.active_torque['A','B'] = (0,0,0)
    """

    def __init__(self, filter, rotation_diff=0.1):
        # store metadata
        param_dict = ParameterDict(
            filter=ParticleFilter,
            rotation_diff=float(rotation_diff),
            constraint=OnlyTypes(ConstraintForce,
                                 allow_none=True,
                                 preprocess=_ellip_preprocessing),
        )
        param_dict.update(
            dict(
                constraint=None,
                rotation_diff=rotation_diff,
                filter=filter,
            ))
        # set defaults
        self._param_dict.update(param_dict)

        active_force = TypeParameter(
            "active_force",
            type_kind="particle_types",
            param_dict=TypeParameterDict((1.0, 0.0, 0.0), len_keys=1),
        )
        active_torque = TypeParameter(
            "active_torque",
            type_kind="particle_types",
            param_dict=TypeParameterDict((0.0, 0.0, 0.0), len_keys=1),
        )

        self._extend_typeparam([active_force, active_torque])

    def _add(self, simulation):
        """Add the operation to a simulation.

        Active forces use RNGs. Warn the user if they did not set the seed.
        """
        if simulation is not None:
            simulation._warn_if_seed_unset()

        super()._add(simulation)

    def _attach(self):
        # initialize the reflected c++ class
        if isinstance(self._simulation.device, hoomd.device.CPU):
            my_class = _md.ActiveForceCompute
        else:
            my_class = _md.ActiveForceComputeGPU

        self._cpp_obj = my_class(
            self._simulation.state._cpp_sys_def,
            self._simulation.state._get_group(self.filter),
            self.rotation_diff,
            _hoomd.make_scalar3(0, 0, 0),
            0,
            0,
            0,
        )

        # Attach param_dict and typeparam_dict
        super()._attach()
<commit_msg>reformatted units for md.force<commit_after># Copyright (c) 2009-2021 The Regents of the University of Michigan
# This file is part of the HOOMD-blue project, released under the BSD 3-Clause
# License.

"""Apply forces to particles."""

import hoomd
from hoomd import _hoomd
from hoomd.md import _md
from hoomd.operation import _HOOMDBaseObject
from hoomd.logging import log
from hoomd.data.typeparam import TypeParameter
from hoomd.data.typeconverter import OnlyTypes
from hoomd.data.parameterdicts import ParameterDict, TypeParameterDict
from hoomd.filter import ParticleFilter
from hoomd.md.constrain import ConstraintForce

def _ellip_preprocessing(constraint):
    if constraint is not None:
        if constraint.__class__.__name__ == "constraint_ellipsoid":
            return constraint
        else:
            raise RuntimeError(
                "Active force constraint is not accepted (currently only "
                "accepts ellipsoids)")
    else:
        return None

class _force:  # noqa - This will be removed eventually. Needed to build docs.
    pass

class Force(_HOOMDBaseObject):
    """Defines a force in HOOMD-blue.

    Pair, angle, bond, and other forces are subclasses of this class.

    Note:
        :py:class:`Force` is the base class for all loggable forces.
        Users should not instantiate this class directly.

    Initializes some loggable quantities.
    """

    @log
    def energy(self):
        """float: Sum of the energy of the whole system \
        :math:`[\\mathrm{energy}]`."""
        if self._attached:
            self._cpp_obj.compute(self._simulation.timestep)
            return self._cpp_obj.calcEnergySum()
        else:
            return None

    @log(category="particle")
    def energies(self):
        """(*N_particles*, ) `numpy.ndarray` of ``numpy.float64``: The \
        energies for all particles :math:`[\\mathrm{energy}]`."""
        if self._attached:
            self._cpp_obj.compute(self._simulation.timestep)
            return self._cpp_obj.getEnergies()
        else:
            return None

    @log(category="particle")
    def forces(self):
        """(*N_particles*, 3) `numpy.ndarray` of ``numpy.float64``: The \
        forces for all particles :math:`[\\mathrm{force}]`."""
        if self._attached:
            self._cpp_obj.compute(self._simulation.timestep)
            return self._cpp_obj.getForces()
        else:
            return None

    @log(category="particle")
    def torques(self):
        """(*N_particles*, 3) `numpy.ndarray` of ``numpy.float64``: The torque \
        for all \
        particles :math:`[\\mathrm{force} \\cdot \\mathrm{length}]`."""
        if self._attached:
            self._cpp_obj.compute(self._simulation.timestep)
            return self._cpp_obj.getTorques()
        else:
            return None

    @log(category="particle")
    def virials(self):
        """(*N_particles*, ) `numpy.ndarray` of ``numpy.float64``: The virial \
        for all particles :math:`[\\mathrm{energy}]`."""
        if self._attached:
            self._cpp_obj.compute(self._simulation.timestep)
            return self._cpp_obj.getVirials()
        else:
            return None

class constant(Force):  # noqa - this will be renamed when it is ported to v3
    R"""Constant force.

    Args:
        fvec (tuple): force vector :math:`[force]`
        tvec (tuple): torque vector :math:`[force \cdot length]`
        fx (float): x component of force, retained for backwards compatibility
          :math:`[\mathrm{force}]`
        fy (float): y component of force, retained for backwards compatibility
          :math:`[\mathrm{force}]`
        fz (float): z component of force, retained for backwards compatibility
          :math:`[\mathrm{force}]`
        group (``hoomd.group``): Group for which the force will be set.
        callback (`callable`): A python callback invoked every time the forces
            are computed

    :py:class:`constant` specifies that a constant force should be added to
    every particle in the simulation or optionally to all particles in a group.

    Note:
        Forces are kept constant during the simulation. If a callback should
        re-compute particle forces every time step, it needs to overwrite the
        old forces of **all** particles with new values.

    Note:
        Per-particle forces take precedence over a particle group, which takes
        precedence over constant forces for all particles.

    Examples::

        force.constant(fx=1.0, fy=0.5, fz=0.25)
        const = force.constant(fvec=(0.4,1.0,0.5))
        const = force.constant(fvec=(0.4,1.0,0.5),group=fluid)
        const = force.constant(fvec=(0.4,1.0,0.5), tvec=(0,0,1) ,group=fluid)

        def updateForces(timestep):
            global const
            const.setForce(tag=1, fvec=(1.0*timestep,2.0*timestep,3.0*timestep))
        const = force.constant(callback=updateForces)
    """

    def __init__(
        self,
        fx=None,
        fy=None,
        fz=None,
        fvec=None,
        tvec=None,
        group=None,
        callback=None,
    ):

        if (fx is not None) and (fy is not None) and (fz is not None):
            self.fvec = (fx, fy, fz)
        elif fvec is not None:
            self.fvec = fvec
        else:
            self.fvec = (0, 0, 0)

        if tvec is not None:
            self.tvec = tvec
        else:
            self.tvec = (0, 0, 0)

        if (self.fvec == (0, 0, 0)) and (self.tvec == (0, 0, 0)
                                         and callback is None):
            hoomd.context.current.device.cpp_msg.warning(
                "The constant force specified has no non-zero components\n")

        # initialize the base class
        Force.__init__(self)

        # create the c++ mirror class
        if group is not None:
            self.cppForce = _hoomd.ConstForceCompute(
                hoomd.context.current.system_definition,
                group.cpp_group,
                self.fvec[0],
                self.fvec[1],
                self.fvec[2],
                self.tvec[0],
                self.tvec[1],
                self.tvec[2],
            )
        else:
            self.cppForce = _hoomd.ConstForceCompute(
                hoomd.context.current.system_definition,
                self.fvec[0],
                self.fvec[1],
                self.fvec[2],
                self.tvec[0],
                self.tvec[1],
                self.tvec[2],
            )

        if callback is not None:
            self.cppForce.setCallback(callback)

        hoomd.context.current.system.addCompute(self.cppForce, self.force_name)

    R""" Change the value of the constant force.

    Args:
        fx (float) New x-component of the force :math:`[\mathrm{force}]`
        fy (float) New y-component of the force :math:`[\mathrm{force}]`
        fz (float) New z-component of the force :math:`[\mathrm{force}]`
        fvec (tuple) New force vector
        tvec (tuple) New torque vector
        group Group for which the force will be set
        tag (int) Particle tag for which the force will be set
            .. versionadded:: 2.3

     Using setForce() requires that you saved the created constant force in a
     variable. i.e.

     Examples:
        const = force.constant(fx=0.4, fy=1.0, fz=0.5)

        const.setForce(fx=0.2, fy=0.1, fz=-0.5)
        const.setForce(fx=0.2, fy=0.1, fz=-0.5, group=fluid)
        const.setForce(fvec=(0.2,0.1,-0.5), tvec=(0,0,1), group=fluid)
    """

    def setForce(  # noqa - this will be documented when it is ported to v3
        self,
        fx=None,
        fy=None,
        fz=None,
        fvec=None,
        tvec=None,
        group=None,
        tag=None,
    ):

        if (fx is not None) and (fy is not None) and (fx is not None):
            self.fvec = (fx, fy, fz)
        elif fvec is not None:
            self.fvec = fvec
        else:
            self.fvec = (0, 0, 0)

        if tvec is not None:
            self.tvec = tvec
        else:
            self.tvec = (0, 0, 0)

        if (fvec == (0, 0, 0)) and (tvec == (0, 0, 0)):
            hoomd.context.current.device.cpp_msg.warning(
                "You are setting the constant force to have no non-zero "
                "components\n")

        self.check_initialization()
        if group is not None:
            self.cppForce.setGroupForce(
                group.cpp_group,
                self.fvec[0],
                self.fvec[1],
                self.fvec[2],
                self.tvec[0],
                self.tvec[1],
                self.tvec[2],
            )
        elif tag is not None:
            self.cppForce.setParticleForce(
                tag,
                self.fvec[0],
                self.fvec[1],
                self.fvec[2],
                self.tvec[0],
                self.tvec[1],
                self.tvec[2],
            )
        else:
            self.cppForce.setForce(
                self.fvec[0],
                self.fvec[1],
                self.fvec[2],
                self.tvec[0],
                self.tvec[1],
                self.tvec[2],
            )

    R""" Set a python callback to be called before the force is evaluated

    Args:
        callback (`callable`) The callback function

     Examples:
        const = force.constant(fx=0.4, fy=1.0, fz=0.5)

        def updateForces(timestep):
            global const
            const.setForce(tag=1, fvec=(1.0*timestep,2.0*timestep,3.0*timestep))

        const.set_callback(updateForces)
        run(100)

        # Reset the callback
        const.set_callback(None)
    """

    def set_callback(self, callback=None):  # noqa - will be ported to v3
        self.cppForce.setCallback(callback)

    # there are no coeffs to update in the constant force compute
    def update_coeffs(self):  # noqa - will be ported to v3
        pass

class Active(Force):
    r"""Active force.

    Attributes:
        filter (:py:mod:`hoomd.filter`): Subset of particles on which to apply
            active forces.
        rotation_diff (float): rotational diffusion constant, :math:`D_r`, for
            all particles in the group. :math:`[\mathrm{time}^{-1}]`
        active_force (tuple): active force vector in reference to the
            orientation of a particle. It is defined per particle type and
            stays constant during the simulation. :math:`[\mathrm{force}]`
        active_torque (tuple): active torque vector in reference to the
            orientation of a particle. It is defined per particle type and
            stays constant during the simulation.
            :math:`[\mathrm{force} \cdot \mathrm{length}]`

    :py:class:`Active` specifies that an active force should be added to all
    particles.  Obeys :math:`\delta {\bf r}_i = \delta t v_0 \hat{p}_i`, where
    :math:`v_0` is the active velocity. In 2D :math:`\hat{p}_i = (\cos \theta_i,
    \sin \theta_i)` is the active force vector for particle :math:`i` and the
    diffusion of the active force vector follows :math:`\delta \theta / \delta t
    = \sqrt{2 D_r / \delta t} \Gamma`, where :math:`D_r` is the rotational
    diffusion constant, and the gamma function is a unit-variance random
    variable, whose components are uncorrelated in time, space, and between
    particles.  In 3D, :math:`\hat{p}_i` is a unit vector in 3D space, and
    diffusion follows :math:`\delta \hat{p}_i / \delta t = \sqrt{2 D_r / \delta
    t} \Gamma (\hat{p}_i (\cos \theta - 1) + \hat{p}_r \sin \theta)`, where
    :math:`\hat{p}_r` is an uncorrelated random unit vector. The persistence
    length of an active particle's path is :math:`v_0 / D_r`.  The rotational
    diffusion is applied to the orientation vector/quaternion of each particle.
    This implies that both the active force and the active torque vectors in the
    particle frame stay constant during the simulation. Hence, the active forces
    in the system frame are composed of the forces in particle frame and the
    current orientation of the particle.

    Examples::

        all = filter.All()
        active = hoomd.md.force.Active(
            filter=hoomd.filter.All(), rotation_diff=0.01
            )
        active.active_force['A','B'] = (1,0,0)
        active.active_torque['A','B'] = (0,0,0)
    """

    def __init__(self, filter, rotation_diff=0.1):
        # store metadata
        param_dict = ParameterDict(
            filter=ParticleFilter,
            rotation_diff=float(rotation_diff),
            constraint=OnlyTypes(ConstraintForce,
                                 allow_none=True,
                                 preprocess=_ellip_preprocessing),
        )
        param_dict.update(
            dict(
                constraint=None,
                rotation_diff=rotation_diff,
                filter=filter,
            ))
        # set defaults
        self._param_dict.update(param_dict)

        active_force = TypeParameter(
            "active_force",
            type_kind="particle_types",
            param_dict=TypeParameterDict((1.0, 0.0, 0.0), len_keys=1),
        )
        active_torque = TypeParameter(
            "active_torque",
            type_kind="particle_types",
            param_dict=TypeParameterDict((0.0, 0.0, 0.0), len_keys=1),
        )

        self._extend_typeparam([active_force, active_torque])

    def _add(self, simulation):
        """Add the operation to a simulation.

        Active forces use RNGs. Warn the user if they did not set the seed.
        """
        if simulation is not None:
            simulation._warn_if_seed_unset()

        super()._add(simulation)

    def _attach(self):
        # initialize the reflected c++ class
        if isinstance(self._simulation.device, hoomd.device.CPU):
            my_class = _md.ActiveForceCompute
        else:
            my_class = _md.ActiveForceComputeGPU

        self._cpp_obj = my_class(
            self._simulation.state._cpp_sys_def,
            self._simulation.state._get_group(self.filter),
            self.rotation_diff,
            _hoomd.make_scalar3(0, 0, 0),
            0,
            0,
            0,
        )

        # Attach param_dict and typeparam_dict
        super()._attach()
<|endoftext|>


<commit_before># -*- coding: utf8 -*-
#
# Module CACHE
#
# Part of Nutils: open source numerical utilities for Python. Jointly developed
# by HvZ Computational Engineering, TU/e Multiscale Engineering Fluid Dynamics,
# and others. More info at http://nutils.org <info@nutils.org>. (c) 2014

"""
The cache module.
"""

from . import core, log
import os, weakref

_property = property
def property( f ):
  def cache_property_wrapper( self, f=f ):
    try:
      value = self.__dict__[f.func_name]
    except KeyError:
      value = f( self )
      self.__dict__[f.func_name] = value
    return value
  assert not cache_property_wrapper.__closure__
  return _property(cache_property_wrapper)

class CallDict( dict ):
  'very simple cache object'

  hit = 0

  def __call__( self, *key ):
    'cache(func,*args): execute func(args) and cache result'
    value = self.get( key )
    if value is None:
      value = key[0]( *key[1:] )
      self[ key ] = value
    else:
      self.hit += 1
    return value

  def summary( self ):
    return 'not used' if not self \
      else 'effectivity %d%% (%d hits, %d misses)' % ( (100*self.hit)/(self.hit+len(self)), self.hit, len(self) )

class Immutable( object ):
  'weakly cache object instances based on init args'

  __slots__ = '__weakref__',

  class __metaclass__( type ):
    def __init__( cls, *args, **kwargs ):
      type.__init__( cls, *args, **kwargs )
      cls.cache = weakref.WeakValueDictionary()
    def __call__( cls, *args, **kwargs ):
      code = cls.__init__.func_code
      names = code.co_varnames[len(args)+1:code.co_argcount]
      for name in names:
        try:
          val = kwargs.pop(name)
        except KeyError:
          index = names.index(name)-len(names)
          try:
            val = cls.__init__.func_defaults[index]
          except Exception as e:
            raise TypeError, '%s.__init__ missing mandatory argument %r' % ( cls.__name__, name )
        args += val,
      assert not kwargs, '%s.__init__ got invalid arguments: %s' % ( cls.__name__, ', '.join(kwargs) )
      try:
        self = cls.cache[args]
      except KeyError:
        self = type.__call__( cls, *args )
        cls.cache[args] = self
      return self

class FileCache( object ):
  'cache'

  def __init__( self, *args ):
    'constructor'

    import hashlib
    strhash = ','.join( str(arg) for arg in args )
    md5hash = hashlib.md5( strhash ).hexdigest()
    log.info( 'using cache:', md5hash )
    cachedir = core.prop( 'cachedir', 'cache' )
    if not os.path.exists( cachedir ):
      os.makedirs( cachedir )
    path = os.path.join( cachedir, md5hash )
    self.data = file( path, 'ab+' if not core.prop( 'recache', False ) else 'wb+' )

  def __call__( self, func, *args, **kwargs ):
    'call'

    import cPickle
    name = func.__name__ + ''.join( ' %s' % arg for arg in args ) + ''.join( ' %s=%s' % item for item in kwargs.iteritems() )
    pos = self.data.tell()
    try:
      data = cPickle.load( self.data )
    except EOFError:
      data = func( *args, **kwargs)
      self.data.seek( pos )
      cPickle.dump( data, self.data, -1 )
      msg = 'written to'
    else:
      msg = 'loaded from'
    log.info( msg, 'cache:', name, '[%db]' % (self.data.tell()-pos) )
    return data

# vim:shiftwidth=2:foldmethod=indent:foldnestmax=2
<commit_msg>fix cache.FileCache<commit_after># -*- coding: utf8 -*-
#
# Module CACHE
#
# Part of Nutils: open source numerical utilities for Python. Jointly developed
# by HvZ Computational Engineering, TU/e Multiscale Engineering Fluid Dynamics,
# and others. More info at http://nutils.org <info@nutils.org>. (c) 2014

"""
The cache module.
"""

from . import core, log
import os, weakref

_property = property
def property( f ):
  def cache_property_wrapper( self, f=f ):
    try:
      value = self.__dict__[f.func_name]
    except KeyError:
      value = f( self )
      self.__dict__[f.func_name] = value
    return value
  assert not cache_property_wrapper.__closure__
  return _property(cache_property_wrapper)

class CallDict( dict ):
  'very simple cache object'

  hit = 0

  def __call__( self, *key ):
    'cache(func,*args): execute func(args) and cache result'
    value = self.get( key )
    if value is None:
      value = key[0]( *key[1:] )
      self[ key ] = value
    else:
      self.hit += 1
    return value

  def summary( self ):
    return 'not used' if not self \
      else 'effectivity %d%% (%d hits, %d misses)' % ( (100*self.hit)/(self.hit+len(self)), self.hit, len(self) )

class Immutable( object ):
  'weakly cache object instances based on init args'

  __slots__ = '__weakref__',

  class __metaclass__( type ):
    def __init__( cls, *args, **kwargs ):
      type.__init__( cls, *args, **kwargs )
      cls.cache = weakref.WeakValueDictionary()
    def __call__( cls, *args, **kwargs ):
      code = cls.__init__.func_code
      names = code.co_varnames[len(args)+1:code.co_argcount]
      for name in names:
        try:
          val = kwargs.pop(name)
        except KeyError:
          index = names.index(name)-len(names)
          try:
            val = cls.__init__.func_defaults[index]
          except Exception as e:
            raise TypeError, '%s.__init__ missing mandatory argument %r' % ( cls.__name__, name )
        args += val,
      assert not kwargs, '%s.__init__ got invalid arguments: %s' % ( cls.__name__, ', '.join(kwargs) )
      try:
        self = cls.cache[args]
      except KeyError:
        self = type.__call__( cls, *args )
        cls.cache[args] = self
      return self

class FileCache( object ):
  'cache'

  def __init__( self, *args ):
    'constructor'

    import hashlib
    strhash = ','.join( str(arg) for arg in args )
    md5hash = hashlib.md5( strhash ).hexdigest()
    log.info( 'using cache:', md5hash )
    cachedir = core.getprop( 'cachedir', 'cache' )
    if not os.path.exists( cachedir ):
      os.makedirs( cachedir )
    path = os.path.join( cachedir, md5hash )
    self.data = file( path, 'ab+' if not core.getprop( 'recache', False ) else 'wb+' )

  def __call__( self, func, *args, **kwargs ):
    'call'

    import cPickle
    name = func.__name__ + ''.join( ' %s' % arg for arg in args ) + ''.join( ' %s=%s' % item for item in kwargs.iteritems() )
    pos = self.data.tell()
    try:
      data = cPickle.load( self.data )
    except EOFError:
      data = func( *args, **kwargs)
      self.data.seek( pos )
      cPickle.dump( data, self.data, -1 )
      msg = 'written to'
    else:
      msg = 'loaded from'
    log.info( msg, 'cache:', name, '[%db]' % (self.data.tell()-pos) )
    return data

# vim:shiftwidth=2:foldmethod=indent:foldnestmax=2
<|endoftext|>


<commit_before>import forecastio
import numpy as np
from api_key import key

from flask import Flask, render_template

app = Flask(__name__)

api_key = key()

lat = 51.2330
lng = 0.2859

@app.route('/')
def index():

    forecast = forecastio.load_forecast(api_key, lat, lng)
    current = forecast.currently()

    return render_template('index.html', text1=lat, text2=lng, text3=current.summary)

if __name__ == '__main__':
    app.run(debug=True)
<commit_msg>adding weather index and comparison<commit_after>import forecastio
import numpy as np
from api_key import key

from flask import Flask, render_template

app = Flask(__name__)

api_key = key()

lat = 51.2330
lng = 0.2859

# Some locations:
here = (lat, lng)
sydney = (-33.86, -151.2094)
new_york = (40.7127, 74.0059)
tokyo = (35.6895, -139.6917)
riodj = (22.9083, 43.1964)

@app.route('/')
def index():
    forecast = forecastio.load_forecast(api_key, lat, lng)
    current = forecast.currently()
    current1, current2, index_ratio = compare_locations(here, new_york)
    return render_template('index.html', text1=lat, text2=lng, text3=current.summary)

# compares the weather index of two locations
# if > 1, loc1 is better than loc2.
def compare_locations(loc1, loc2):
    lat1, lng1 = loc1
    lat2, lng2 = loc2

    forecast1 = forecastio.load_forecast(api_key, lat1, lng1)
    current1 = forecast1.currently()
    index1 = weather_index(current1)

    forecast2 = forecastio.load_forecast(api_key, lat2, lng2)
    current2 = forecast2.currently()
    index2 = weather_index(current2)
    return current1, current2, index1/index2

# calculates a weather index for a location, based on its current forecast
# currently uses temperature only
def weather_index(current):
    temp = current.temperature
    feels_like = current.apparentTemperature
    humidity = current.humidity*100
    precip = current.precipIntensity
    precipProb = current.precipProbability
    if current.precipIntensity != 0:
        precipType = current.precipType
    return temp

if __name__ == '__main__':
    app.run(debug=True)

# Can't get a google maps API key until we have a URL!
# from googlemaps import GoogleMaps
# gmaps = GoogleMaps(api_key)
# address = 'Constitution Ave NW & 10th St NW, Washington, DC'
# lat, lng = gmaps.address_to_latlng(address)
# print lat, lng
<|endoftext|>


<commit_before>"""
Backend for SAML 2.0 support

Terminology:

"Service Provider" (SP): Your web app
"Identity Provider" (IdP): The third-party site that is authenticating users via SAML
"""
from onelogin.saml2.auth import OneLogin_Saml2_Auth
from onelogin.saml2.settings import OneLogin_Saml2_Settings
from social.backends.base import BaseAuth
from social.exceptions import AuthFailed

# Helpful constants:
OID_COMMON_NAME = "urn:oid:2.5.4.3"
OID_EDU_PERSON_PRINCIPAL_NAME = "urn:oid:1.3.6.1.4.1.5923.1.1.1.6"
OID_EDU_PERSON_ENTITLEMENT = "urn:oid:1.3.6.1.4.1.5923.1.1.1.7"
OID_GIVEN_NAME = "urn:oid:2.5.4.42"
OID_MAIL = "urn:oid:0.9.2342.19200300.100.1.3"
OID_SURNAME = "urn:oid:2.5.4.4"
OID_USERID = "urn:oid:0.9.2342.19200300.100.1.1"

class SAMLIdentityProvider(object):
    """
    Wrapper around configuration for a SAML Identity provider
    """

    def __init__(self, name, **kwargs):
        """ Load and parse configuration """
        self.name = name
        # name should be a slug and must not contain a colon, which could conflict with uid prefixing:
        assert ':' not in self.name and ' ' not in self.name, "IdP 'name' should be a slug (short, no spaces)"
        self.conf = kwargs

    def get_user_permanent_id(self, attributes):
        """
        The most important method: Get a permanent, unique identifier for this user from the
        attributes supplied by the IdP.

        If you want to use the NameID, it's available via attributes['name_id']
        """
        return attributes[self.conf.get('user_permanent_id', OID_USERID)][0]

    # Attributes processing:
    def get_user_details(self, attributes):
        """
        Given the SAML attributes extracted from the SSO response, get the user data like name.
        """
        return {
            'fullname': self.get_attr(attributes, 'attr_full_name', OID_COMMON_NAME),
            'first_name': self.get_attr(attributes, 'attr_first_name', OID_GIVEN_NAME),
            'last_name': self.get_attr(attributes, 'attr_last_name', OID_SURNAME),
            'username': self.get_attr(attributes, 'attr_username', OID_USERID),
            'email': self.get_attr(attributes, 'attr_email', OID_MAIL),
        }

    def get_attr(self, attributes, conf_key, default_attribute):
        """
        Internal helper method.
        Get the attribute 'default_attribute' out of the attributes, unless self.conf[conf_key]
        overrides the default by specifying another attribute to use.
        """
        key = self.conf.get(conf_key, default_attribute)
        return attributes[key][0] if key in attributes else None

    @property
    def entity_id(self):
        """ Get the entity ID for this IdP """
        return self.conf['entity_id']  # Required. e.g. "https://idp.testshib.org/idp/shibboleth"

    @property
    def sso_url(self):
        """ Get the SSO URL for this IdP """
        return self.conf['url']  # Required. e.g. "https://idp.testshib.org/idp/profile/SAML2/Redirect/SSO"

    @property
    def sso_binding(self):
        """ Get the method used to submit our request to the SSO URL """
        return self.conf.get('binding', 'urn:oasis:names:tc:SAML:2.0:bindings:HTTP-Redirect')

    @property
    def x509cert(self):
        """ X.509 Public Key Certificate for this IdP """
        return self.conf['x509cert']

    @property
    def saml_config_dict(self):
        """ Get the IdP configuration dict in the format required by python-saml """
        return {
            "entityId": self.entity_id,
            "singleSignOnService": {
                "url": self.sso_url,
                "binding": self.sso_binding,
            },
            "x509cert": self.x509cert,
        }

class DummySAMLIdentityProvider(SAMLIdentityProvider):
    """
    A placeholder IdP used when we must specify something, e.g. when generating SP metadata.

    If OneLogin_Saml2_Auth is modified to not always require IdP config, this can be removed.
    """
    def __init__(self):
        super(DummySAMLIdentityProvider, self).__init__(
            "dummy",
            entity_id="https://dummy.none/saml2",
            url="https://dummy.none/SSO",
            x509cert='',
        )

class SAMLAuth(BaseAuth):
    """
    PSA Backend that implements SAML 2.0 Service Provider (SP) functionality.

    Unlike all of the other backends, this one can be configured to work with
    many identity providers (IdPs). For example, a University that belongs to a
    Shibboleth federation may support authentication via ~100 partner
    universities. Also, the IdP configuration can be changed at runtime if you
    require that functionality - just subclass this and override `get_idp()`.

    Several settings are required. Here's an example:

    SOCIAL_AUTH_SAML_SP_ENTITY_ID = "https://saml.example.com/"
    SOCIAL_AUTH_SAML_SP_PUBLIC_CERT = "... X.509 certificate string ..."
    SOCIAL_AUTH_SAML_SP_PRIVATE_KEY = "... private key ..."
    SOCIAL_AUTH_SAML_ORG_INFO = {
        "en-US": {"name": "example", "displayname": "Example Inc.", "url": "http://example.com", },
    }
    SOCIAL_AUTH_SAML_TECHNICAL_CONTACT = {"givenName": "Tech Gal", "emailAddress": "technical@example.com", }
    SOCIAL_AUTH_SAML_SUPPORT_CONTACT = {"givenName": "Support Guy", "emailAddress": "support@example.com", }
    SOCIAL_AUTH_SAML_ENABLED_IDPS = {
        "testshib": {
            "entity_id": "https://idp.testshib.org/idp/shibboleth",
            "url": "https://idp.testshib.org/idp/profile/SAML2/Redirect/SSO",
            "x509cert": "MIIEDjCCAvagAwIBAgIBADANBgkqhkiG9w0B ... 8Bbnl+ev0peYzxFyF5sQA==",
        }
    }

    Optional settings:
    SOCIAL_AUTH_SAML_SP_EXTRA = {}
    SOCIAL_AUTH_SAML_SECURITY_CONFIG = {}
    SOCIAL_AUTH_SAML_SP_NAMEID_FORMATS = []
    """
    name = "saml"

    def get_idp(self, idp_name):
        """ Given the name of an IdP, get a SAMLIdentityProvider instance """
        idp_config = self.setting("ENABLED_IDPS")[idp_name]
        return SAMLIdentityProvider(idp_name, **idp_config)

    def generate_saml_config(self, idp):
        """
        Generate the configuration required to instantiate OneLogin_Saml2_Auth
        """
        # The shared absolute URL that all IdPs redirect back to - this is specified in our metadata.xml:
        abs_completion_url = self.redirect_uri

        config = {
            "contactPerson": {
                "technical": self.setting("TECHNICAL_CONTACT"),
                "support": self.setting("SUPPORT_CONTACT"),
            },
            "debug": True,
            "idp": idp.saml_config_dict,
            "organization": self.setting("ORG_INFO"),
            "security": {
                'metadataValidUntil': '',
                'metadataCacheDuration': 'P10D',  # metadata valid for ten days
            },
            "sp": {
                "assertionConsumerService": {
                    "url": abs_completion_url,
                    "binding": "urn:oasis:names:tc:SAML:2.0:bindings:HTTP-POST",
                },
                "entityId": self.setting("SP_ENTITY_ID"),
                "NameIDFormats": self.setting("SP_NAMEID_FORMATS", []),
                "x509cert": self.setting("SP_PUBLIC_CERT"),
                "privateKey": self.setting("SP_PRIVATE_KEY"),
            },
            "strict": True,  # We must force strict mode - for security
        }
        config["security"].update(self.setting("SECURITY_CONFIG", {}))
        config["sp"].update(self.setting("SP_EXTRA", {}))
        return config

    def generate_metadata_xml(self):
        """
        Helper method that can be used from your web app to generate the XML metadata required
        to link your web app as a Service Provider with each IdP you wish to use.

        Returns (metadata XML string, list of errors)

        Example usage (Django):
            from social.apps.django_app.utils import load_strategy, load_backend
            def saml_metadata_view(request):
                complete_url = reverse('social:complete', args=("saml", ))
                saml_backend = load_backend(load_strategy(request), "saml", complete_url)
                metadata, errors = saml_backend.generate_metadata_xml()
                if not errors:
                    return HttpResponse(content=metadata, content_type='text/xml')
                return HttpResponseServerError(content=', '.join(errors))
        """
        idp = DummySAMLIdentityProvider()  # python-saml requires us to specify something here even though it's not used
        config = self.generate_saml_config(idp)
        saml_settings = OneLogin_Saml2_Settings(config)
        metadata = saml_settings.get_sp_metadata()
        errors = saml_settings.validate_metadata(metadata)
        return metadata, errors

    def _create_saml_auth(self, idp):
        """
        Get an instance of OneLogin_Saml2_Auth
        """
        config = self.generate_saml_config(idp)
        request_info = {
            'https': 'on' if self.strategy.request_is_secure() else 'off',
            'http_host': self.strategy.request_host(),
            'script_name': self.strategy.request_path(),
            'server_port': self.strategy.request_port(),
            'get_data': self.strategy.request_get(),
            'post_data': self.strategy.request_post(),
        }
        return OneLogin_Saml2_Auth(request_info, config)

    def auth_url(self):
        """ Get the URL to which we must redirect in order to authenticate the user """
        idp_name = self.strategy.request_data()['idp']
        auth = self._create_saml_auth(idp=self.get_idp(idp_name))
        # Below, return_to sets the RelayState, which can contain arbitrary data.
        # We use it to store the specific SAML IdP backend name, since we combine
        # many backends to a single URL.
        return auth.login(return_to=idp_name)

    def get_user_details(self, response):
        """
        Get user details like full name, email, etc. from the response - see auth_complete
        """
        idp = self.get_idp(response['idp_name'])
        return idp.get_user_details(response['attributes'])

    def get_user_id(self, details, response):
        """
        Get the permanent ID for this user from the response.
        We prefix each ID with the name of the IdP so that we can connect multiple IdPs to this
        user.
        """
        idp = self.get_idp(response['idp_name'])
        uid = idp.get_user_permanent_id(response['attributes'])
        return '{}:{}'.format(idp.name, uid)

    def auth_complete(self, *args, **kwargs):
        """
        The user has been redirected back from the IdP and we should now log them in, if
        everything checks out.
        """
        idp_name = self.strategy.request_data()['RelayState']
        idp = self.get_idp(idp_name)
        auth = self._create_saml_auth(idp)
        auth.process_response()
        errors = auth.get_errors()
        if errors or not auth.is_authenticated():
            reason = auth.get_last_error_reason()
            raise AuthFailed(self, 'SAML login failed: {} ({})'.format(errors, reason))

        attributes = auth.get_attributes()
        attributes['name_id'] = auth.get_nameid()

        self._check_entitlements(idp, attributes)

        response = {
            'idp_name': idp_name,
            'attributes': attributes,
            'session_index': auth.get_session_index(),
        }

        kwargs.update({'response': response, 'backend': self})

        return self.strategy.authenticate(*args, **kwargs)

    def _check_entitlements(self, idp, attributes):
        """
        Additional verification of a SAML response before authenticating the user.

        Subclasses can override this method if they need custom validation code,
        such as requiring the presence of an eduPersonEntitlement.

        raise social.exceptions.AuthForbidden if the user should not be authenticated,
        or do nothing to allow the login pipeline to continue.
        """
        pass
<commit_msg>Minor consistency fix<commit_after>"""
Backend for SAML 2.0 support

Terminology:

"Service Provider" (SP): Your web app
"Identity Provider" (IdP): The third-party site that is authenticating users via SAML
"""
from onelogin.saml2.auth import OneLogin_Saml2_Auth
from onelogin.saml2.settings import OneLogin_Saml2_Settings
from social.backends.base import BaseAuth
from social.exceptions import AuthFailed

# Helpful constants:
OID_COMMON_NAME = "urn:oid:2.5.4.3"
OID_EDU_PERSON_PRINCIPAL_NAME = "urn:oid:1.3.6.1.4.1.5923.1.1.1.6"
OID_EDU_PERSON_ENTITLEMENT = "urn:oid:1.3.6.1.4.1.5923.1.1.1.7"
OID_GIVEN_NAME = "urn:oid:2.5.4.42"
OID_MAIL = "urn:oid:0.9.2342.19200300.100.1.3"
OID_SURNAME = "urn:oid:2.5.4.4"
OID_USERID = "urn:oid:0.9.2342.19200300.100.1.1"

class SAMLIdentityProvider(object):
    """
    Wrapper around configuration for a SAML Identity provider
    """

    def __init__(self, name, **kwargs):
        """ Load and parse configuration """
        self.name = name
        # name should be a slug and must not contain a colon, which could conflict with uid prefixing:
        assert ':' not in self.name and ' ' not in self.name, "IdP 'name' should be a slug (short, no spaces)"
        self.conf = kwargs

    def get_user_permanent_id(self, attributes):
        """
        The most important method: Get a permanent, unique identifier for this user from the
        attributes supplied by the IdP.

        If you want to use the NameID, it's available via attributes['name_id']
        """
        return attributes[self.conf.get('attr_user_permanent_id', OID_USERID)][0]

    # Attributes processing:
    def get_user_details(self, attributes):
        """
        Given the SAML attributes extracted from the SSO response, get the user data like name.
        """
        return {
            'fullname': self.get_attr(attributes, 'attr_full_name', OID_COMMON_NAME),
            'first_name': self.get_attr(attributes, 'attr_first_name', OID_GIVEN_NAME),
            'last_name': self.get_attr(attributes, 'attr_last_name', OID_SURNAME),
            'username': self.get_attr(attributes, 'attr_username', OID_USERID),
            'email': self.get_attr(attributes, 'attr_email', OID_MAIL),
        }

    def get_attr(self, attributes, conf_key, default_attribute):
        """
        Internal helper method.
        Get the attribute 'default_attribute' out of the attributes, unless self.conf[conf_key]
        overrides the default by specifying another attribute to use.
        """
        key = self.conf.get(conf_key, default_attribute)
        return attributes[key][0] if key in attributes else None

    @property
    def entity_id(self):
        """ Get the entity ID for this IdP """
        return self.conf['entity_id']  # Required. e.g. "https://idp.testshib.org/idp/shibboleth"

    @property
    def sso_url(self):
        """ Get the SSO URL for this IdP """
        return self.conf['url']  # Required. e.g. "https://idp.testshib.org/idp/profile/SAML2/Redirect/SSO"

    @property
    def sso_binding(self):
        """ Get the method used to submit our request to the SSO URL """
        return self.conf.get('binding', 'urn:oasis:names:tc:SAML:2.0:bindings:HTTP-Redirect')

    @property
    def x509cert(self):
        """ X.509 Public Key Certificate for this IdP """
        return self.conf['x509cert']

    @property
    def saml_config_dict(self):
        """ Get the IdP configuration dict in the format required by python-saml """
        return {
            "entityId": self.entity_id,
            "singleSignOnService": {
                "url": self.sso_url,
                "binding": self.sso_binding,
            },
            "x509cert": self.x509cert,
        }

class DummySAMLIdentityProvider(SAMLIdentityProvider):
    """
    A placeholder IdP used when we must specify something, e.g. when generating SP metadata.

    If OneLogin_Saml2_Auth is modified to not always require IdP config, this can be removed.
    """
    def __init__(self):
        super(DummySAMLIdentityProvider, self).__init__(
            "dummy",
            entity_id="https://dummy.none/saml2",
            url="https://dummy.none/SSO",
            x509cert='',
        )

class SAMLAuth(BaseAuth):
    """
    PSA Backend that implements SAML 2.0 Service Provider (SP) functionality.

    Unlike all of the other backends, this one can be configured to work with
    many identity providers (IdPs). For example, a University that belongs to a
    Shibboleth federation may support authentication via ~100 partner
    universities. Also, the IdP configuration can be changed at runtime if you
    require that functionality - just subclass this and override `get_idp()`.

    Several settings are required. Here's an example:

    SOCIAL_AUTH_SAML_SP_ENTITY_ID = "https://saml.example.com/"
    SOCIAL_AUTH_SAML_SP_PUBLIC_CERT = "... X.509 certificate string ..."
    SOCIAL_AUTH_SAML_SP_PRIVATE_KEY = "... private key ..."
    SOCIAL_AUTH_SAML_ORG_INFO = {
        "en-US": {"name": "example", "displayname": "Example Inc.", "url": "http://example.com", },
    }
    SOCIAL_AUTH_SAML_TECHNICAL_CONTACT = {"givenName": "Tech Gal", "emailAddress": "technical@example.com", }
    SOCIAL_AUTH_SAML_SUPPORT_CONTACT = {"givenName": "Support Guy", "emailAddress": "support@example.com", }
    SOCIAL_AUTH_SAML_ENABLED_IDPS = {
        "testshib": {
            "entity_id": "https://idp.testshib.org/idp/shibboleth",
            "url": "https://idp.testshib.org/idp/profile/SAML2/Redirect/SSO",
            "x509cert": "MIIEDjCCAvagAwIBAgIBADANBgkqhkiG9w0B ... 8Bbnl+ev0peYzxFyF5sQA==",
        }
    }

    Optional settings:
    SOCIAL_AUTH_SAML_SP_EXTRA = {}
    SOCIAL_AUTH_SAML_SECURITY_CONFIG = {}
    SOCIAL_AUTH_SAML_SP_NAMEID_FORMATS = []
    """
    name = "saml"

    def get_idp(self, idp_name):
        """ Given the name of an IdP, get a SAMLIdentityProvider instance """
        idp_config = self.setting("ENABLED_IDPS")[idp_name]
        return SAMLIdentityProvider(idp_name, **idp_config)

    def generate_saml_config(self, idp):
        """
        Generate the configuration required to instantiate OneLogin_Saml2_Auth
        """
        # The shared absolute URL that all IdPs redirect back to - this is specified in our metadata.xml:
        abs_completion_url = self.redirect_uri

        config = {
            "contactPerson": {
                "technical": self.setting("TECHNICAL_CONTACT"),
                "support": self.setting("SUPPORT_CONTACT"),
            },
            "debug": True,
            "idp": idp.saml_config_dict,
            "organization": self.setting("ORG_INFO"),
            "security": {
                'metadataValidUntil': '',
                'metadataCacheDuration': 'P10D',  # metadata valid for ten days
            },
            "sp": {
                "assertionConsumerService": {
                    "url": abs_completion_url,
                    "binding": "urn:oasis:names:tc:SAML:2.0:bindings:HTTP-POST",
                },
                "entityId": self.setting("SP_ENTITY_ID"),
                "NameIDFormats": self.setting("SP_NAMEID_FORMATS", []),
                "x509cert": self.setting("SP_PUBLIC_CERT"),
                "privateKey": self.setting("SP_PRIVATE_KEY"),
            },
            "strict": True,  # We must force strict mode - for security
        }
        config["security"].update(self.setting("SECURITY_CONFIG", {}))
        config["sp"].update(self.setting("SP_EXTRA", {}))
        return config

    def generate_metadata_xml(self):
        """
        Helper method that can be used from your web app to generate the XML metadata required
        to link your web app as a Service Provider with each IdP you wish to use.

        Returns (metadata XML string, list of errors)

        Example usage (Django):
            from social.apps.django_app.utils import load_strategy, load_backend
            def saml_metadata_view(request):
                complete_url = reverse('social:complete', args=("saml", ))
                saml_backend = load_backend(load_strategy(request), "saml", complete_url)
                metadata, errors = saml_backend.generate_metadata_xml()
                if not errors:
                    return HttpResponse(content=metadata, content_type='text/xml')
                return HttpResponseServerError(content=', '.join(errors))
        """
        idp = DummySAMLIdentityProvider()  # python-saml requires us to specify something here even though it's not used
        config = self.generate_saml_config(idp)
        saml_settings = OneLogin_Saml2_Settings(config)
        metadata = saml_settings.get_sp_metadata()
        errors = saml_settings.validate_metadata(metadata)
        return metadata, errors

    def _create_saml_auth(self, idp):
        """
        Get an instance of OneLogin_Saml2_Auth
        """
        config = self.generate_saml_config(idp)
        request_info = {
            'https': 'on' if self.strategy.request_is_secure() else 'off',
            'http_host': self.strategy.request_host(),
            'script_name': self.strategy.request_path(),
            'server_port': self.strategy.request_port(),
            'get_data': self.strategy.request_get(),
            'post_data': self.strategy.request_post(),
        }
        return OneLogin_Saml2_Auth(request_info, config)

    def auth_url(self):
        """ Get the URL to which we must redirect in order to authenticate the user """
        idp_name = self.strategy.request_data()['idp']
        auth = self._create_saml_auth(idp=self.get_idp(idp_name))
        # Below, return_to sets the RelayState, which can contain arbitrary data.
        # We use it to store the specific SAML IdP backend name, since we combine
        # many backends to a single URL.
        return auth.login(return_to=idp_name)

    def get_user_details(self, response):
        """
        Get user details like full name, email, etc. from the response - see auth_complete
        """
        idp = self.get_idp(response['idp_name'])
        return idp.get_user_details(response['attributes'])

    def get_user_id(self, details, response):
        """
        Get the permanent ID for this user from the response.
        We prefix each ID with the name of the IdP so that we can connect multiple IdPs to this
        user.
        """
        idp = self.get_idp(response['idp_name'])
        uid = idp.get_user_permanent_id(response['attributes'])
        return '{}:{}'.format(idp.name, uid)

    def auth_complete(self, *args, **kwargs):
        """
        The user has been redirected back from the IdP and we should now log them in, if
        everything checks out.
        """
        idp_name = self.strategy.request_data()['RelayState']
        idp = self.get_idp(idp_name)
        auth = self._create_saml_auth(idp)
        auth.process_response()
        errors = auth.get_errors()
        if errors or not auth.is_authenticated():
            reason = auth.get_last_error_reason()
            raise AuthFailed(self, 'SAML login failed: {} ({})'.format(errors, reason))

        attributes = auth.get_attributes()
        attributes['name_id'] = auth.get_nameid()

        self._check_entitlements(idp, attributes)

        response = {
            'idp_name': idp_name,
            'attributes': attributes,
            'session_index': auth.get_session_index(),
        }

        kwargs.update({'response': response, 'backend': self})

        return self.strategy.authenticate(*args, **kwargs)

    def _check_entitlements(self, idp, attributes):
        """
        Additional verification of a SAML response before authenticating the user.

        Subclasses can override this method if they need custom validation code,
        such as requiring the presence of an eduPersonEntitlement.

        raise social.exceptions.AuthForbidden if the user should not be authenticated,
        or do nothing to allow the login pipeline to continue.
        """
        pass
<|endoftext|>


<commit_before># Copyright (c) 2009-2019 The Regents of the University of Michigan
# This file is part of the HOOMD-blue project, released under the BSD 3-Clause
# License.

# Maintainer: joaander / All Developers are free to add commands for new
# features

R""" Apply forces to particles.
"""

import hoomd
from hoomd import _hoomd
from hoomd.md import _md
from hoomd.operation import _HOOMDBaseObject
from hoomd.logging import log
from hoomd.data.typeparam import TypeParameter
from hoomd.data.typeconverter import OnlyType
from hoomd.data.parameterdicts import ParameterDict, TypeParameterDict
from hoomd.filter import ParticleFilter
from hoomd.md.constrain import ConstraintForce

def ellip_preprocessing(constraint):
    if constraint is not None:
        if (constraint.__class__.__name__ == "constraint_ellipsoid") :
            return act_force
        else:
            raise RuntimeError("Active force constraint is not accepted (currently only accepts ellipsoids)")
    else:
        return None

class _force():
    pass

class Force(_HOOMDBaseObject):
    '''Defines a force in HOOMD-blue.

    Pair, angle, bond, and other forces are subclasses of this class.

    Note:
        :py:class:`Force` is the base class for all loggable forces.
        Users should not instantiate this class directly.

    Initializes some loggable quantities.
    '''

    def _attach(self):
        super()._attach()

    @log
    def energy(self):
        """float: Sum of the energy of the whole system."""
        if self._attached:
            self._cpp_obj.compute(self._simulation.timestep)
            return self._cpp_obj.calcEnergySum()
        else:
            return None

    @log(flag='particle')
    def energies(self):
        """(*N_particles*, ) `numpy.ndarray` of ``numpy.float64``: The energies for all particles."""
        if self._attached:
            self._cpp_obj.compute(self._simulation.timestep)
            return self._cpp_obj.getEnergies()
        else:
            return None

    @log(flag='particle')
    def forces(self):
        """(*N_particles*, 3) `numpy.ndarray` of ``numpy.float64``: The forces for all particles."""
        if self._attached:
            self._cpp_obj.compute(self._simulation.timestep)
            return self._cpp_obj.getForces()
        else:
            return None

    @log(flag='particle')
    def torques(self):
        """(*N_particles*, 3) `numpy.ndarray` of ``numpy.float64``: The torque for all particles."""
        if self._attached:
            self._cpp_obj.compute(self._simulation.timestep)
            return self._cpp_obj.getTorques()
        else:
            return None

    @log(flag='particle')
    def virials(self):
        """(*N_particles*, ) `numpy.ndarray` of ``numpy.float64``: The virial for all particles."""
        if self._attached:
            self._cpp_obj.compute(self._simulation.timestep)
            return self._cpp_obj.getVirials()
        else:
            return None

class constant(Force):
    R""" Constant force.

    Args:
        fvec (tuple): force vector (in force units)
        tvec (tuple): torque vector (in torque units)
        fx (float): x component of force, retained for backwards compatibility
        fy (float): y component of force, retained for backwards compatibility
        fz (float): z component of force, retained for backwards compatibility
        group (``hoomd.group``): Group for which the force will be set.
        callback (`callable`): A python callback invoked every time the forces are computed

    :py:class:`constant` specifies that a constant force should be added to every
    particle in the simulation or optionally to all particles in a group.

    Note:
        Forces are kept constant during the simulation. If a callback should re-compute
        particle forces every time step, it needs to overwrite the old forces of **all**
        particles with new values.

    Note:
        Per-particle forces take precedence over a particle group, which takes precedence over constant forces for all particles.

    Examples::

        force.constant(fx=1.0, fy=0.5, fz=0.25)
        const = force.constant(fvec=(0.4,1.0,0.5))
        const = force.constant(fvec=(0.4,1.0,0.5),group=fluid)
        const = force.constant(fvec=(0.4,1.0,0.5), tvec=(0,0,1) ,group=fluid)

        def updateForces(timestep):
            global const
            const.setForce(tag=1, fvec=(1.0*timestep,2.0*timestep,3.0*timestep))
        const = force.constant(callback=updateForces)
    """
    def __init__(self, fx=None, fy=None, fz=None, fvec=None, tvec=None, group=None, callback=None):

        if (fx is not None) and (fy is not None) and (fz is not None):
            self.fvec = (fx,fy,fz)
        elif (fvec is not None):
            self.fvec = fvec
        else:
            self.fvec = (0,0,0)

        if (tvec is not None):
            self.tvec = tvec
        else:
            self.tvec = (0,0,0)

        if (self.fvec == (0,0,0)) and (self.tvec == (0,0,0) and callback is None):
            hoomd.context.current.device.cpp_msg.warning("The constant force specified has no non-zero components\n")

        # initialize the base class
        Force.__init__(self)

        # create the c++ mirror class
        if (group is not None):
            self.cppForce = _hoomd.ConstForceCompute(hoomd.context.current.system_definition,
                group.cpp_group,
                self.fvec[0],
                self.fvec[1],
                self.fvec[2],
                self.tvec[0],
                self.tvec[1],
                self.tvec[2])
        else:
            self.cppForce = _hoomd.ConstForceCompute(hoomd.context.current.system_definition,
                self.fvec[0],
                self.fvec[1],
                self.fvec[2],
                self.tvec[0],
                self.tvec[1],
                self.tvec[2])

        if callback is not None:
            self.cppForce.setCallback(callback)

        hoomd.context.current.system.addCompute(self.cppForce, self.force_name)

    R""" Change the value of the constant force.

    Args:
        fx (float) New x-component of the force (in force units)
        fy (float) New y-component of the force (in force units)
        fz (float) New z-component of the force (in force units)
        fvec (tuple) New force vector
        tvec (tuple) New torque vector
        group Group for which the force will be set
        tag (int) Particle tag for which the force will be set
            .. versionadded:: 2.3

     Using setForce() requires that you saved the created constant force in a variable. i.e.

     Examples:
        const = force.constant(fx=0.4, fy=1.0, fz=0.5)

        const.setForce(fx=0.2, fy=0.1, fz=-0.5)
        const.setForce(fx=0.2, fy=0.1, fz=-0.5, group=fluid)
        const.setForce(fvec=(0.2,0.1,-0.5), tvec=(0,0,1), group=fluid)
    """
    def setForce(self, fx=None, fy=None, fz=None, fvec=None, tvec=None, group=None, tag=None):

        if (fx is not None) and (fy is not None) and (fx is not None):
            self.fvec = (fx,fy,fz)
        elif fvec is not None:
            self.fvec = fvec
        else:
            self.fvec = (0,0,0)

        if tvec is not None:
            self.tvec = tvec
        else:
            self.tvec = (0,0,0)

        if (fvec==(0,0,0)) and (tvec==(0,0,0)):
            hoomd.context.current.device.cpp_msg.warning("You are setting the constant force to have no non-zero components\n")

        self.check_initialization()
        if (group is not None):
            self.cppForce.setGroupForce(group.cpp_group, self.fvec[0], self.fvec[1], self.fvec[2],
                                                          self.tvec[0], self.tvec[1], self.tvec[2])
        elif (tag is not None):
            self.cppForce.setParticleForce(tag, self.fvec[0], self.fvec[1], self.fvec[2],
                                                 self.tvec[0], self.tvec[1], self.tvec[2])
        else:
            self.cppForce.setForce(self.fvec[0], self.fvec[1], self.fvec[2], self.tvec[0], self.tvec[1], self.tvec[2])

    R""" Set a python callback to be called before the force is evaluated

    Args:
        callback (`callable`) The callback function

     Examples:
        const = force.constant(fx=0.4, fy=1.0, fz=0.5)

        def updateForces(timestep):
            global const
            const.setForce(tag=1, fvec=(1.0*timestep,2.0*timestep,3.0*timestep))

        const.set_callback(updateForces)
        run(100)

        # Reset the callback
        const.set_callback(None)
    """
    def set_callback(self, callback=None):
        self.cppForce.setCallback(callback)

    # there are no coeffs to update in the constant force compute
    def update_coeffs(self):
        pass

class Active(Force):
    R""" Active force.

    Attributes:
        filter (:py:mod:`hoomd.filter`): Subset of particles on which to apply
            active forces.
        seed (int): required user-specified seed number for random number
            generator.
        rotation_diff (float): rotational diffusion constant, :math:`D_r`, for
            all particles in the group.
        active_force (tuple): active force vector in reference to the
            orientation of a particle. It is defined per particle type and stays
            constant during the simulation.
        active_torque (tuple): active torque vector in reference to the
            orientation of a particle. It is defined per particle type and stays
            constant during the simulation.

    :py:class:`Active` specifies that an active force should be added to all
    particles.  Obeys :math:`\delta {\bf r}_i = \delta t v_0 \hat{p}_i`, where
    :math:`v_0` is the active velocity. In 2D :math:`\hat{p}_i = (\cos \theta_i,
    \sin \theta_i)` is the active force vector for particle :math:`i` and the
    diffusion of the active force vector follows :math:`\delta \theta / \delta t
    = \sqrt{2 D_r / \delta t} \Gamma`, where :math:`D_r` is the rotational
    diffusion constant, and the gamma function is a unit-variance random
    variable, whose components are uncorrelated in time, space, and between
    particles.  In 3D, :math:`\hat{p}_i` is a unit vector in 3D space, and
    diffusion follows :math:`\delta \hat{p}_i / \delta t = \sqrt{2 D_r / \delta
    t} \Gamma (\hat{p}_i (\cos \theta - 1) + \hat{p}_r \sin \theta)`, where
    :math:`\hat{p}_r` is an uncorrelated random unit vector. The persistence
    length of an active particle's path is :math:`v_0 / D_r`.  The rotational
    diffusion is applied to the orientation vector/quaternion of each particle.
    This implies that both the active force and the active torque vectors in the
    particle frame stay constant during the simulation. Hence, the active forces
    in the system frame are composed of the forces in particle frame and the
    current orientation of the particle.

    Examples::

        all = filter.All()
        active = hoomd.md.force.Active(filter=hoomd.filter.All(), seed=1,rotation_diff=0.01)
        active.active_force['A','B'] = (1,0,0)
        active.active_torque['A','B'] = (0,0,0)
    """

    def __init__(self, filter, seed, rotation_diff=0.1):
        # store metadata
        param_dict = ParameterDict(
            filter=ParticleFilter,
            seed=int(seed),
            rotation_diff=float(rotation_diff),
            constraint=OnlyType(ConstraintForce, allow_none=True,
                                preprocess=ellip_preprocessing),
            )
        param_dict.update(dict(constraint=None,
                               rotation_diff=rotation_diff, seed=seed, filter=filter))
        # set defaults
        self._param_dict.update(param_dict)

        active_force = TypeParameter(
            'active_force', type_kind='particle_types', param_dict=TypeParameterDict((1, 0, 0), len_keys=1))
        active_torque = TypeParameter(
            'active_torque', type_kind='particle_types', param_dict=TypeParameterDict((0, 0, 0), len_keys=1))

        self._extend_typeparam([active_force, active_torque])

    def _attach(self):
        # initialize the reflected c++ class
        if isinstance(self._simulation.device, hoomd.device.CPU):
            my_class = _md.ActiveForceCompute
        else:
            my_class = _md.ActiveForceComputeGPU

        self._cpp_obj = my_class(
            self._simulation.state._cpp_sys_def,
            self._simulation.state._get_group(self.filter),
            self.seed, self.rotation_diff,
            _hoomd.make_scalar3(0, 0, 0), 0, 0, 0)

        # Attach param_dict and typeparam_dict
        super()._attach()

class dipole(Force):
    R""" Treat particles as dipoles in an electric field.

    Args:
        field_x (float): x-component of the field (units?)
        field_y (float): y-component of the field (units?)
        field_z (float): z-component of the field (units?)
        p (float): magnitude of the particles' dipole moment in the local z direction

    Examples::

        force.external_field_dipole(field_x=0.0, field_y=1.0 ,field_z=0.5, p=1.0)
        const_ext_f_dipole = force.external_field_dipole(field_x=0.0, field_y=1.0 ,field_z=0.5, p=1.0)
    """
    def __init__(self, field_x,field_y,field_z,p):

        # initialize the base class
        Force.__init__(self)

        # create the c++ mirror class
        self.cppForce = _md.ConstExternalFieldDipoleForceCompute(hoomd.context.current.system_definition, field_x, field_y, field_z, p)

        hoomd.context.current.system.addCompute(self.cppForce, self.force_name)

        # store metadata
        self.field_x = field_x
        self.field_y = field_y
        self.field_z = field_z

    def set_params(field_x, field_y,field_z,p):
        R""" Change the constant field and dipole moment.

        Args:
            field_x (float): x-component of the field (units?)
            field_y (float): y-component of the field (units?)
            field_z (float): z-component of the field (units?)
            p (float): magnitude of the particles' dipole moment in the local z direction

        Examples::

            const_ext_f_dipole = force.external_field_dipole(field_x=0.0, field_y=1.0 ,field_z=0.5, p=1.0)
            const_ext_f_dipole.setParams(field_x=0.1, field_y=0.1, field_z=0.0, p=1.0))

        """
        self.check_initialization()

        self.cppForce.setParams(field_x,field_y,field_z,p)

    # there are no coeffs to update in the constant ExternalFieldDipoleForceCompute
    def update_coeffs(self):
        pass
<commit_msg>blacked<commit_after># Copyright (c) 2009-2019 The Regents of the University of Michigan
# This file is part of the HOOMD-blue project, released under the BSD 3-Clause
# License.

# Maintainer: joaander / All Developers are free to add commands for new
# features

R""" Apply forces to particles.
"""

import hoomd
from hoomd import _hoomd
from hoomd.md import _md
from hoomd.operation import _HOOMDBaseObject
from hoomd.logging import log
from hoomd.data.typeparam import TypeParameter
from hoomd.data.typeconverter import OnlyType
from hoomd.data.parameterdicts import ParameterDict, TypeParameterDict
from hoomd.filter import ParticleFilter
from hoomd.md.constrain import ConstraintForce

def ellip_preprocessing(constraint):
    if constraint is not None:
        if constraint.__class__.__name__ == "constraint_ellipsoid":
            return act_force
        else:
            raise RuntimeError(
                "Active force constraint is not accepted (currently only "
                "accepts ellipsoids)"
            )
    else:
        return None

class _force:
    pass

class Force(_HOOMDBaseObject):
    """Defines a force in HOOMD-blue.

    Pair, angle, bond, and other forces are subclasses of this class.

    Note:
        :py:class:`Force` is the base class for all loggable forces.
        Users should not instantiate this class directly.

    Initializes some loggable quantities.
    """

    def _attach(self):
        super()._attach()

    @log
    def energy(self):
        """float: Sum of the energy of the whole system."""
        if self._attached:
            self._cpp_obj.compute(self._simulation.timestep)
            return self._cpp_obj.calcEnergySum()
        else:
            return None

    @log(flag="particle")
    def energies(self):
        """(*N_particles*, ) `numpy.ndarray` of ``numpy.float64``: The energies
        for all particles."""
        if self._attached:
            self._cpp_obj.compute(self._simulation.timestep)
            return self._cpp_obj.getEnergies()
        else:
            return None

    @log(flag="particle")
    def forces(self):
        """(*N_particles*, 3) `numpy.ndarray` of ``numpy.float64``: The forces
        for all particles."""
        if self._attached:
            self._cpp_obj.compute(self._simulation.timestep)
            return self._cpp_obj.getForces()
        else:
            return None

    @log(flag="particle")
    def torques(self):
        """(*N_particles*, 3) `numpy.ndarray` of ``numpy.float64``: The torque
        for all particles."""
        if self._attached:
            self._cpp_obj.compute(self._simulation.timestep)
            return self._cpp_obj.getTorques()
        else:
            return None

    @log(flag="particle")
    def virials(self):
        """(*N_particles*, ) `numpy.ndarray` of ``numpy.float64``: The virial
        for all particles."""
        if self._attached:
            self._cpp_obj.compute(self._simulation.timestep)
            return self._cpp_obj.getVirials()
        else:
            return None

class constant(Force):
    R"""Constant force.

    Args:
        fvec (tuple): force vector (in force units)
        tvec (tuple): torque vector (in torque units)
        fx (float): x component of force, retained for backwards compatibility
        fy (float): y component of force, retained for backwards compatibility
        fz (float): z component of force, retained for backwards compatibility
        group (``hoomd.group``): Group for which the force will be set.
        callback (`callable`): A python callback invoked every time the forces
            are computed

    :py:class:`constant` specifies that a constant force should be added to
    every particle in the simulation or optionally to all particles in a group.

    Note:
        Forces are kept constant during the simulation. If a callback should
        re-compute particle forces every time step, it needs to overwrite the
        old forces of **all** particles with new values.

    Note:
        Per-particle forces take precedence over a particle group, which takes
        precedence over constant forces for all particles.

    Examples::

        force.constant(fx=1.0, fy=0.5, fz=0.25)
        const = force.constant(fvec=(0.4,1.0,0.5))
        const = force.constant(fvec=(0.4,1.0,0.5),group=fluid)
        const = force.constant(fvec=(0.4,1.0,0.5), tvec=(0,0,1) ,group=fluid)

        def updateForces(timestep):
            global const
            const.setForce(tag=1, fvec=(1.0*timestep,2.0*timestep,3.0*timestep))
        const = force.constant(callback=updateForces)
    """

    def __init__(
        self,
        fx=None,
        fy=None,
        fz=None,
        fvec=None,
        tvec=None,
        group=None,
        callback=None,
    ):

        if (fx is not None) and (fy is not None) and (fz is not None):
            self.fvec = (fx, fy, fz)
        elif fvec is not None:
            self.fvec = fvec
        else:
            self.fvec = (0, 0, 0)

        if tvec is not None:
            self.tvec = tvec
        else:
            self.tvec = (0, 0, 0)

        if (self.fvec == (0, 0, 0)) and (
            self.tvec == (0, 0, 0) and callback is None
        ):
            hoomd.context.current.device.cpp_msg.warning(
                "The constant force specified has no non-zero components\n"
            )

        # initialize the base class
        Force.__init__(self)

        # create the c++ mirror class
        if group is not None:
            self.cppForce = _hoomd.ConstForceCompute(
                hoomd.context.current.system_definition,
                group.cpp_group,
                self.fvec[0],
                self.fvec[1],
                self.fvec[2],
                self.tvec[0],
                self.tvec[1],
                self.tvec[2],
            )
        else:
            self.cppForce = _hoomd.ConstForceCompute(
                hoomd.context.current.system_definition,
                self.fvec[0],
                self.fvec[1],
                self.fvec[2],
                self.tvec[0],
                self.tvec[1],
                self.tvec[2],
            )

        if callback is not None:
            self.cppForce.setCallback(callback)

        hoomd.context.current.system.addCompute(self.cppForce, self.force_name)

    R""" Change the value of the constant force.

    Args:
        fx (float) New x-component of the force (in force units)
        fy (float) New y-component of the force (in force units)
        fz (float) New z-component of the force (in force units)
        fvec (tuple) New force vector
        tvec (tuple) New torque vector
        group Group for which the force will be set
        tag (int) Particle tag for which the force will be set
            .. versionadded:: 2.3

     Using setForce() requires that you saved the created constant force in a
     variable. i.e.

     Examples:
        const = force.constant(fx=0.4, fy=1.0, fz=0.5)

        const.setForce(fx=0.2, fy=0.1, fz=-0.5)
        const.setForce(fx=0.2, fy=0.1, fz=-0.5, group=fluid)
        const.setForce(fvec=(0.2,0.1,-0.5), tvec=(0,0,1), group=fluid)
    """

    def setForce(
        self,
        fx=None,
        fy=None,
        fz=None,
        fvec=None,
        tvec=None,
        group=None,
        tag=None,
    ):

        if (fx is not None) and (fy is not None) and (fx is not None):
            self.fvec = (fx, fy, fz)
        elif fvec is not None:
            self.fvec = fvec
        else:
            self.fvec = (0, 0, 0)

        if tvec is not None:
            self.tvec = tvec
        else:
            self.tvec = (0, 0, 0)

        if (fvec == (0, 0, 0)) and (tvec == (0, 0, 0)):
            hoomd.context.current.device.cpp_msg.warning(
                "You are setting the constant force to have no non-zero "
                "components\n"
            )

        self.check_initialization()
        if group is not None:
            self.cppForce.setGroupForce(
                group.cpp_group,
                self.fvec[0],
                self.fvec[1],
                self.fvec[2],
                self.tvec[0],
                self.tvec[1],
                self.tvec[2],
            )
        elif tag is not None:
            self.cppForce.setParticleForce(
                tag,
                self.fvec[0],
                self.fvec[1],
                self.fvec[2],
                self.tvec[0],
                self.tvec[1],
                self.tvec[2],
            )
        else:
            self.cppForce.setForce(
                self.fvec[0],
                self.fvec[1],
                self.fvec[2],
                self.tvec[0],
                self.tvec[1],
                self.tvec[2],
            )

    R""" Set a python callback to be called before the force is evaluated

    Args:
        callback (`callable`) The callback function

     Examples:
        const = force.constant(fx=0.4, fy=1.0, fz=0.5)

        def updateForces(timestep):
            global const
            const.setForce(tag=1, fvec=(1.0*timestep,2.0*timestep,3.0*timestep))

        const.set_callback(updateForces)
        run(100)

        # Reset the callback
        const.set_callback(None)
    """

    def set_callback(self, callback=None):
        self.cppForce.setCallback(callback)

    # there are no coeffs to update in the constant force compute
    def update_coeffs(self):
        pass

class Active(Force):
    R"""Active force.

    Attributes:
        filter (:py:mod:`hoomd.filter`): Subset of particles on which to apply
            active forces.
        seed (int): required user-specified seed number for random number
            generator.
        rotation_diff (float): rotational diffusion constant, :math:`D_r`, for
            all particles in the group.
        active_force (tuple): active force vector in reference to the
            orientation of a particle. It is defined per particle type and stays
            constant during the simulation.
        active_torque (tuple): active torque vector in reference to the
            orientation of a particle. It is defined per particle type and stays
            constant during the simulation.

    :py:class:`Active` specifies that an active force should be added to all
    particles.  Obeys :math:`\delta {\bf r}_i = \delta t v_0 \hat{p}_i`, where
    :math:`v_0` is the active velocity. In 2D :math:`\hat{p}_i = (\cos \theta_i,
    \sin \theta_i)` is the active force vector for particle :math:`i` and the
    diffusion of the active force vector follows :math:`\delta \theta / \delta t
    = \sqrt{2 D_r / \delta t} \Gamma`, where :math:`D_r` is the rotational
    diffusion constant, and the gamma function is a unit-variance random
    variable, whose components are uncorrelated in time, space, and between
    particles.  In 3D, :math:`\hat{p}_i` is a unit vector in 3D space, and
    diffusion follows :math:`\delta \hat{p}_i / \delta t = \sqrt{2 D_r / \delta
    t} \Gamma (\hat{p}_i (\cos \theta - 1) + \hat{p}_r \sin \theta)`, where
    :math:`\hat{p}_r` is an uncorrelated random unit vector. The persistence
    length of an active particle's path is :math:`v_0 / D_r`.  The rotational
    diffusion is applied to the orientation vector/quaternion of each particle.
    This implies that both the active force and the active torque vectors in the
    particle frame stay constant during the simulation. Hence, the active forces
    in the system frame are composed of the forces in particle frame and the
    current orientation of the particle.

    Examples::

        all = filter.All()
        active = hoomd.md.force.Active(
            filter=hoomd.filter.All(), seed=1, rotation_diff=0.01
            )
        active.active_force['A','B'] = (1,0,0)
        active.active_torque['A','B'] = (0,0,0)
    """

    def __init__(self, filter, seed, rotation_diff=0.1):
        # store metadata
        param_dict = ParameterDict(
            filter=ParticleFilter,
            seed=int(seed),
            rotation_diff=float(rotation_diff),
            constraint=OnlyType(
                ConstraintForce, allow_none=True, preprocess=ellip_preprocessing
            ),
        )
        param_dict.update(
            dict(
                constraint=None,
                rotation_diff=rotation_diff,
                seed=seed,
                filter=filter,
            )
        )
        # set defaults
        self._param_dict.update(param_dict)

        active_force = TypeParameter(
            "active_force",
            type_kind="particle_types",
            param_dict=TypeParameterDict((1, 0, 0), len_keys=1),
        )
        active_torque = TypeParameter(
            "active_torque",
            type_kind="particle_types",
            param_dict=TypeParameterDict((0, 0, 0), len_keys=1),
        )

        self._extend_typeparam([active_force, active_torque])

    def _attach(self):
        # initialize the reflected c++ class
        if isinstance(self._simulation.device, hoomd.device.CPU):
            my_class = _md.ActiveForceCompute
        else:
            my_class = _md.ActiveForceComputeGPU

        self._cpp_obj = my_class(
            self._simulation.state._cpp_sys_def,
            self._simulation.state._get_group(self.filter),
            self.seed,
            self.rotation_diff,
            _hoomd.make_scalar3(0, 0, 0),
            0,
            0,
            0,
        )

        # Attach param_dict and typeparam_dict
        super()._attach()

class dipole(Force):
    R"""Treat particles as dipoles in an electric field.

    Args:
        field_x (float): x-component of the field (units?)
        field_y (float): y-component of the field (units?)
        field_z (float): z-component of the field (units?)
        p (float): magnitude of the particles' dipole moment in the local z
            direction

    Examples::

        force.external_field_dipole(
            field_x=0.0, field_y=1.0 ,field_z=0.5, p=1.0
            )
        const_ext_f_dipole = force.external_field_dipole(
            field_x=0.0, field_y=1.0 ,field_z=0.5, p=1.0
            )
    """

    def __init__(self, field_x, field_y, field_z, p):

        # initialize the base class
        Force.__init__(self)

        # create the c++ mirror class
        self.cppForce = _md.ConstExternalFieldDipoleForceCompute(
            hoomd.context.current.system_definition,
            field_x,
            field_y,
            field_z,
            p,
        )

        hoomd.context.current.system.addCompute(self.cppForce, self.force_name)

        # store metadata
        self.field_x = field_x
        self.field_y = field_y
        self.field_z = field_z

    def set_params(field_x, field_y, field_z, p):
        R"""Change the constant field and dipole moment.

        Args:
            field_x (float): x-component of the field (units?)
            field_y (float): y-component of the field (units?)
            field_z (float): z-component of the field (units?)
            p (float): magnitude of the particles' dipole moment in the local z
                direction

        Examples::

            const_ext_f_dipole = force.external_field_dipole(
                field_x=0.0, field_y=1.0 ,field_z=0.5, p=1.0
                )
            const_ext_f_dipole.setParams(
                field_x=0.1, field_y=0.1, field_z=0.0, p=1.0
                )

        """
        self.check_initialization()

        self.cppForce.setParams(field_x, field_y, field_z, p)

    # there are no coeffs to update in the constant
    # ExternalFieldDipoleForceCompute
    def update_coeffs(self):
        pass
<|endoftext|>