<commit_before>"""Defines the :class:`SubRobotModel` class, which is :class:`RobotModel`-like
but only modifies selected degrees of freedom of the robot (e.g., an arm, a
leg).

Many, but not all, ``klampt`` module functions accept SubRobotModel in
the place of RobotModel.
"""

from ..robotsim import *
from .collide import self_collision_iter
from .trajectory import Trajectory,HermiteTrajectory
import warnings

class SubRobotModel:
    """A helper that lets you conveniently set/get quantities for a subset
    of moving links on a RobotModel.  This class has the same
    API as RobotModel, but everything is re-indexed so that configurations
    and link indices only modify the given subset of links.  As a reult,
    most methods applicable to a RobotModel can also be applied to
    a SubRobotModel.

    You provide the list of moving link indices or names in the constructor.

    The methods ``tofull`` and ``fromfull`` convert objects to and from the
    full robot.
    """
    def __init__(self,robot,links):
        """
        Args:
            robot (RobotModel or SubRobotModel): the robot to base this on.
            links (list of ints or strs): the links to use in this sub-robot.
        """
        assert isinstance(robot,(RobotModel,SubRobotModel)),"SubRobotModel constructor must be given a RobotModel or SubRobotModel as first argument"
        self._robot = robot
        self._links = links[:]
        self.index = robot.index
        self.world = robot.world
        if isinstance(robot,SubRobotModel):
            warnings.warn("Taking sub-robot of sub-robot... not tested yet")
        for i,l in enumerate(self._links):
            if isinstance(l,str):
                self._links[i] = robot.link(l).getIndex()

    def tofull(self,object,type=None):
        """Converts the given index, link, configuration, velocity, or trajectory of a sub robot
        to the corresponding object of the full robot.  Returns the object for the full robot.

        Args:
            object: an integer index, configuration, velocity, matrix, list of configurations,
                or Trajectory.
            type (str, optional): describes how to interpret object:

                * 'vector' or 'Velocity': object is treated like a velocity.
                * 'jacobian': object is treated like a jacobian matrix.
                * anything else: object is treated as a list of configurations.
        """
        if isinstance(object,int):
            return self._links[object]
        elif isinstance(object,SubRobotModelLink):
            return object._link
        elif isinstance(object,(list,tuple)):
            if hasattr(object[0],'__iter__'):
                if type == 'jacobian':
                    #treat this as a jacobian matrix
                    res = []
                    for row in object:
                        assert len(row) == len(self._links)
                        res.append(self.tofull(row,'velocity'))
                    return res
                else:
                    #treat this like a list of configurations
                    res = []
                    for row in object:
                        assert len(row) == len(self._links)
                        res.append(self.tofull(row))
                    return res
            else:
                assert len(object) == len(self._links)
                if type == 'velocity' or type == 'Vector':
                    #treat as a velocity
                    res = [0.0]*self._robot.numLinks()
                    for l,v in zip(self._links,object):
                        res[l] = v
                    return res
                else:
                    #treat as a configuration
                    res = self._robot.getConfig()
                    for l,v in zip(self._links,object):
                        res[l] = v
                    return res
        elif isinstance(object,Trajectory):
            if isinstance(object,HermiteTrajectory):
                raise NotImplementedError("Can't lift hermite trajectories to full robots yet")
            newmilestones = [self.tofull(v) for v in object.milestones]
            return object.constructor(object.times,newmilestones)
        else:
            raise ValueError("Invalid object type, not an integer, configuration, or Trajectory")

    def fromfull(self,object):
        """Converts the given index, configuration, velocity, or trajectory of a full robot
        to the corresponding object of the sub-robot.

        Args:
            object: an integer index, configuration, velocity, matrix, list of configurations,
                or Trajectory.

        Returns:
            : The corresponding object mapped to the sub-robot.

        Note:
            For indices, this is an O(n) operation where n is the size of the sub-robot.
            If the index doesn't belong to the sub-robot then None is returned.
        """
        if isinstance(object,int):
            for i,l in enumerate(self._links):
                if l == object:
                    return i
            return None
        elif isinstance(object,RobotModelLink):
            return SubRobotModelLink(object,self)
        elif isinstance(object,(list,tuple)):
            if hasattr(object[0],'__iter__'):
                #treat this like a list of configurations
                res = []
                for row in object:
                    assert len(row) == self._robot.numLinks()
                    res.append(self.fromfull(row))
                return res
            else:
                #treat as a configuration
                assert len(object) == self._robot.numLinks()
                return [object[i] for i in self._links]
        elif isinstance(object,Trajectory):
            if isinstance(object,HermiteTrajectory):
                raise NotImplementedError("Can't project hermite trajectories to sub-robots yet")
            newmilestones = [self.fromfull(v) for v in object.milestones]
            return object.constructor(object.times,newmilestones)
        else:
            raise ValueError("Invalid object type, not an integer, configuration, or Trajectory")

    def numLinks(self):
        return len(self._links)

    def link(self,index):
        if isinstance(index,str):
            return SubRobotModelLink(self._robot.link(index),self)
        else:
            return SubRobotModelLink(self._robot.link(self._links[index]),self)

    def numDrivers(self):
        raise NotImplementedError("TODO Accessing number of drivers in sub-robot")

    def driver(self,index):
        raise NotImplementedError("TODO Accessing drivers in sub-robot")

    def getConfig(self):
        q = self._robot.getConfig()
        return [q[i] for i in self._links]
    def getVelocity(self):
        q = self._robot.getVelocity()
        return [q[i] for i in self._links]
    def setConfig(self,q):
        assert len(q) == len(self._links)
        qfull = self._robot.getConfig()
        for i,v in zip(self._links,q):
            qfull[i] = v
        self._robot.setConfig(qfull)
    def setVelocity(self,q):
        assert len(q) == len(self._links)
        qfull = self._robot.getVelocity()
        for i,v in zip(self._links,q):
            qfull[i] = v
        self._robot.setVelocity(qfull)
    def getJointLimits(self):
        qmin,qmax = self._robot.getJointLimits()
        return [qmin[i] for i in self._links],[qmax[i] for i in self._links]
    def setJointLimits(self,qmin,qmax):
        assert len(qmin) == len(self._links)
        assert len(qmax) == len(self._links)
        qminfull,qmaxfull = self._robot.getJointLimits()
        for i,a,b in zip(self._links,qmin,qmax):
            qminfull[i] = a
            qmaxfull[i] = b
        self._robot.setJointLimits(qminfull,qmaxfull)
    def getVelocityLimits(self):
        q = self._robot.getVelocityLimits()
        return [q[i] for i in self._links]
    def setVelocityLimits(self,vmax):
        assert len(vmax) == len(self._links)
        qfull = self._robot.getVelocityLimits()
        for i,v in zip(self._links,vmax):
            qfull[i] = v
        self._robot.setVelocityLimits(qfull)
    def getAccelerationLimits(self):
        q = self._robot.getAccelerationLimits()
        return [q[i] for i in self._links]
    def setAccelerationLimits(self,amax):
        assert len(amax) == len(self._links)
        qfull = self._robot.getAccelerationLimits()
        for i,v in zip(self._links,amax):
            qfull[i] = v
        self._robot.setAccelerationLimits(qfull)
    def getTorqueLimits(self):
        q = self._robot.getTorqueLimits()
        return [q[i] for i in self._links]
    def setTorqueLimits(self,tmax):
        assert len(tmax) == len(self._links)
        qfull = self._robot.getTorqueLimits()
        for i,v in zip(self._links,tmax):
            qfull[i] = v
        self._robot.setTorqueLimits(qfull)

    def setDOFPosition(self,index,qi):
        if isinstance(index,str):
            self._robot.setDOFPosition(index,qi)
        else:
            self._robot.setDOFPosition(self._links[index],qi)
    def getDOFPosition(self,index):
        if isinstance(index,str):
            return self._robot.getDOFPosition(index)
        else:
            return self._robot.getDOFPosition(self._links[index])
    def getCom(self):
        raise NotImplementedError("TODO: getCom")
    def getComJacobian(self):
        raise NotImplementedError("TODO: getComJacobian")
    def getMassMatrix(self):
        raise NotImplementedError("TODO: getMassMatrix")
    def getMassMatrixInv(self):
        raise NotImplementedError("TODO: getMassMatrix")
    def getCoriolisForceMatrix(self):
        raise NotImplementedError("TODO: getCoriolisForceMatrix")
    def getCoriolisForces(self):
        raise NotImplementedError("TODO: getCoriolisForceMatrix")
    def getGravityForces(self,g):
        raise NotImplementedError("TODO:  getGravityForces")
    def torquesFromAccel(self,ddq):
        raise NotImplementedError("TODO: torquesFromAccel")
    def accelFromTorques(self,t):
        raise NotImplementedError("TODO: accelFromTorques")

    def interpolate(self,a,b,u):
        afull = self._robot.getConfig()
        bfull = afull[:]
        for i,ai,bi in zip(self._links,a,b):
            afull[i] =ai
            bfull[i] =bi
        q = self._robot.interpolate(afull,bfull,u)
        return [q[i] for i in self._links]
    def distance(self,a,b):
        afull = self._robot.getConfig()
        bfull = afull[:]
        for i,ai,bi in zip(self._links,a,b):
            afull[i] =ai
            bfull[i] =bi
        return self._robot.distance(afull,bfull)
    def interpolate_deriv(self,a,b):
        afull = self._robot.getConfig()
        bfull = afull[:]
        for i,ai,bi in zip(self._links,a,b):
            afull[i] =ai
            bfull[i] =bi
        q = self._robot.interpolate_deriv(afull,bfull)
        return [q[i] for i in self._links]
    def randomizeConfig(self,unboundedScale=1.0):
        oldfull = self._robot.getConfig()
        self._robot.randomizeConfig(unboundedScale)
        qfull = self._robot.getConfig()
        for i,l in enumerate(self._links):
            oldfull[l] = qfull[l]
        self._robot.setConfig(oldfull)
    def selfCollisionEnabled(self,link1,link2):
        if not isinstance(link1,str):
            link1 = self._links[link1]
        else:
            link1 = self._robot.link(link1).getIndex()
        if not isinstance(link2,str):
            link2 = self._links[link2]
        else:
            link2 = self._robot.link(link2).getIndex()
        return self._robot.selfCollisionEnabled(link1,link2)
    def enableSelfCollision(self,link1,link2,value=True):
        if not isinstance(link1,str):
            link1 = self._links[link1]
        else:
            link1 = self._robot.link(link1).getIndex()
        if not isinstance(link2,str):
            link2 = self._links[link2]
        else:
            link2 = self._robot.link(link2).getIndex()
        self._robot.enableSelfCollision(link1,link2,value)
    def selfCollides(self):
        geoms = [self._robot.link(i).geometry() for i in self._links]
        def dotest(i,j):
            return self._robot.selfCollisionEnabled(self._links[i],self._links[j])
        return any(self_collision_iter(geoms,dotest))
    def drawGL(self,keepAppearance=True):
        for i in self._links:
            self._robot.link(i).drawWorldGL(keepAppearance)

class SubRobotModelLink:
    """A helper that lets you treat links of a subrobot just like a normal
    RobotModelLink. Correctly implements jacobians and indices with respect
    to the sub-robot.
    """
    def __init__(self,link,robot):
        self._link = link
        self._robot = robot
        self.robotIndex = link.robotIndex
        self.index = link.index
        self.getID = link.getID
        self.getName = link.getName
        self.setName = link.setName
        self.geometry = link.geometry
        self.appearance = link.appearance
        self.getMass = link.getMass
        self.setMass = link.setMass
        self.getParentTransform = link.getParentTransform
        self.setParentTransform = link.setParentTransform
        self.getAxis = link.getAxis
        self.setAxis = link.setAxis
        self.getWorldPosition = link.getWorldPosition
        self.getWorldDirection = link.getWorldDirection
        self.getLocalPosition = link.getLocalPosition
        self.getLocalDirection = link.getLocalDirection
        self.getTransform = link.getTransform
        self.setTransform = link.setTransform
        self.getVelocity = link.getVelocity
        self.getAngularVelocity = link.getAngularVelocity
        self.getPointVelocity = link.getPointVelocity
        self.drawLocalGL = link.drawLocalGL
        self.drawWorldGL = link.drawWorldGL
    def robot(self):
        return self._robot
    def getIndex(self):
        return self._robot.fromfull(self._link.getIndex())
    def getParent(self):
        p = self._robot.fromfull(self._link.getParent())
        if p == None: return -1
        return p
    def parent(self):
        p = self.getParent()
        return self._robot.link(p)
    def setParent(self,p):
        self._link.setParent(self._robot.tofull(p))
    def getJacobian(self,p):
        self._robot.tofull(self._link.getJacobian(p),'jacobian')
    def getPositionJacobian(self,p):
        self._robot.tofull(self._link.getPositionJacobian(p),'jacobian')
    def getOrientationJacobian(self):
        self._robot.tofull(self._link.getOrientationJacobian(),'jacobian')

<commit_msg>Fixed subrobotmodel for IK and other queries<commit_after>"""Defines the :class:`SubRobotModel` class, which is :class:`RobotModel`-like
but only modifies selected degrees of freedom of the robot (e.g., an arm, a
leg).

Many, but not all, ``klampt`` module functions accept SubRobotModel in
the place of RobotModel.
"""

from ..robotsim import *
from .collide import self_collision_iter
from .trajectory import Trajectory,HermiteTrajectory
import warnings

class SubRobotModel:
    """A helper that lets you conveniently set/get quantities for a subset
    of moving links on a RobotModel.  This class has the same
    API as RobotModel, but everything is re-indexed so that configurations
    and link indices only modify the given subset of links.  As a reult,
    most methods applicable to a RobotModel can also be applied to
    a SubRobotModel.

    You provide the list of moving link indices or names in the constructor.

    The methods ``tofull`` and ``fromfull`` convert objects to and from the
    full robot.
    """
    def __init__(self,robot,links):
        """
        Args:
            robot (RobotModel or SubRobotModel): the robot to base this on.
            links (list of ints or strs): the links to use in this sub-robot.
        """
        assert isinstance(robot,(RobotModel,SubRobotModel)),"SubRobotModel constructor must be given a RobotModel or SubRobotModel as first argument"
        self._robot = robot
        self._links = links[:]
        self.index = robot.index
        self.world = robot.world
        if isinstance(robot,SubRobotModel):
            warnings.warn("Taking sub-robot of sub-robot... not tested yet")
            self._robot = robot._robot
        for i,l in enumerate(self._links):
            if isinstance(l,str):
                self._links[i] = robot.link(l).getIndex()
        self._inv_links = dict((l,i) for (i,l) in enumerate(self._links))

    def tofull(self,object,reference=None):
        """Converts the given index, link, configuration, velocity, or
        trajectory of a sub robot to the corresponding object of the full
        robot.

        Args:
            object: an integer index, configuration, velocity, matrix, list of
                configurations, or Trajectory.
            reference (list, optional): describes the reference
                object that this should fill in for the indices not in this
                sub-robot. By default, uses the robot's current configuration.

        Returns:
            : The corresponding object mapped to the full robot.
        """
        if isinstance(object,int):
            return self._links[object]
        elif isinstance(object,SubRobotModelLink):
            return object._link
        elif isinstance(object,(list,tuple)):
            if hasattr(object[0],'__iter__'):
                #treat this as a list of configuration-like objects
                res = []
                if reference is not None:
                    if len(reference) != len(object):
                        if not hasattr(reference[0],'__iter__'):
                            reference = [reference]*len(object)
                        else:
                            raise ValueError("Invalid size of reference object")
                else:
                    reference = [None]*len(object)
                for i,row in enumerate(object):
                    assert len(row) == len(self._links)
                    res.append(self.tofull(row,reference=reference[i]))
                return res
            else:
                assert len(object) == len(self._links)
                if reference is None:
                    res = self._robot.getConfig()
                else:
                    res = [v for v in reference]
                for l,v in zip(self._links,object):
                    res[l] = v
                return res
        elif isinstance(object,Trajectory):
            if isinstance(object,HermiteTrajectory):
                raise NotImplementedError("Can't lift hermite trajectories to full robots yet")
            newmilestones = self.tofull(object.milestones,reference=reference)
            return object.constructor(object.times,newmilestones)
        else:
            raise ValueError("Invalid object type, not an integer, configuration, or Trajectory")

    def fromfull(self,object):
        """Converts the given index, configuration, velocity, or trajectory of
        a full robot to the corresponding object of the sub-robot.

        Args:
            object: an integer index, configuration, velocity, matrix, list of
                configurations, or Trajectory.

        Returns:
            : The corresponding object mapped to the sub-robot.

        Note:
            For indices, if the index doesn't belong to the sub-robot then None
            is returned.
        """
        if isinstance(object,int):
            return self._inv_links.get(object,None)
        elif isinstance(object,RobotModelLink):
            return SubRobotModelLink(object,self)
        elif isinstance(object,(list,tuple)):
            if hasattr(object[0],'__iter__'):
                #treat this like a list of configurations
                res = []
                for i,row in enumerate(object):
                    assert len(row) == self._robot.numLinks(),'Object {} needs to be a configuration of length {}'.format(i,self._robot.numLinks())
                    res.append(self.fromfull(row))
                return res
            else:
                #treat as a configuration
                assert len(object) == self._robot.numLinks(),'Object needs to be a configuration of length {}'.format(self._robot.numLinks())
                return [object[i] for i in self._links]
        elif isinstance(object,Trajectory):
            if isinstance(object,HermiteTrajectory):
                raise NotImplementedError("Can't project hermite trajectories to sub-robots yet")
            newmilestones = [self.fromfull(v) for v in object.milestones]
            return object.constructor(object.times,newmilestones)
        else:
            raise ValueError("Invalid object type, not an integer, configuration, or Trajectory")

    def numLinks(self):
        return len(self._links)

    def link(self,index):
        if isinstance(index,str):
            return SubRobotModelLink(self._robot.link(index),self)
        else:
            return SubRobotModelLink(self._robot.link(self._links[index]),self)

    def numDrivers(self):
        raise NotImplementedError("TODO Accessing number of drivers in sub-robot")

    def driver(self,index):
        raise NotImplementedError("TODO Accessing drivers in sub-robot")

    def getConfig(self):
        q = self._robot.getConfig()
        return [q[i] for i in self._links]
    def getVelocity(self):
        q = self._robot.getVelocity()
        return [q[i] for i in self._links]
    def setConfig(self,q):
        assert len(q) == len(self._links)
        qfull = self._robot.getConfig()
        for i,v in zip(self._links,q):
            qfull[i] = v
        self._robot.setConfig(qfull)
    def setVelocity(self,q):
        assert len(q) == len(self._links)
        qfull = self._robot.getVelocity()
        for i,v in zip(self._links,q):
            qfull[i] = v
        self._robot.setVelocity(qfull)
    def getJointLimits(self):
        qmin,qmax = self._robot.getJointLimits()
        return [qmin[i] for i in self._links],[qmax[i] for i in self._links]
    def setJointLimits(self,qmin,qmax):
        assert len(qmin) == len(self._links)
        assert len(qmax) == len(self._links)
        qminfull,qmaxfull = self._robot.getJointLimits()
        for i,a,b in zip(self._links,qmin,qmax):
            qminfull[i] = a
            qmaxfull[i] = b
        self._robot.setJointLimits(qminfull,qmaxfull)
    def getVelocityLimits(self):
        q = self._robot.getVelocityLimits()
        return [q[i] for i in self._links]
    def setVelocityLimits(self,vmax):
        assert len(vmax) == len(self._links)
        qfull = self._robot.getVelocityLimits()
        for i,v in zip(self._links,vmax):
            qfull[i] = v
        self._robot.setVelocityLimits(qfull)
    def getAccelerationLimits(self):
        q = self._robot.getAccelerationLimits()
        return [q[i] for i in self._links]
    def setAccelerationLimits(self,amax):
        assert len(amax) == len(self._links)
        qfull = self._robot.getAccelerationLimits()
        for i,v in zip(self._links,amax):
            qfull[i] = v
        self._robot.setAccelerationLimits(qfull)
    def getTorqueLimits(self):
        q = self._robot.getTorqueLimits()
        return [q[i] for i in self._links]
    def setTorqueLimits(self,tmax):
        assert len(tmax) == len(self._links)
        qfull = self._robot.getTorqueLimits()
        for i,v in zip(self._links,tmax):
            qfull[i] = v
        self._robot.setTorqueLimits(qfull)

    def setDOFPosition(self,index,qi):
        if isinstance(index,str):
            self._robot.setDOFPosition(index,qi)
        else:
            self._robot.setDOFPosition(self._links[index],qi)
    def getDOFPosition(self,index):
        if isinstance(index,str):
            return self._robot.getDOFPosition(index)
        else:
            return self._robot.getDOFPosition(self._links[index])
    def getCom(self):
        raise NotImplementedError("TODO: getCom")
    def getComJacobian(self):
        raise NotImplementedError("TODO: getComJacobian")
    def getLinearMomentum(self):
        vinit = self._robot.getVelocity()
        vtemp = self.tofull(self.getVelocity,[0]*self._robot.numLinks())
        self._robot.setVelocity(vtemp)
        res = self._robot.getLinearMomentum()
        self._robot.setVelocity(vinit)
        return res
    def getAngularMomentum(self):
        vinit = self._robot.getVelocity()
        vtemp = self.tofull(self.getVelocity,[0]*self._robot.numLinks())
        self._robot.setVelocity(vtemp)
        res = self._robot.getAngularMomentum()
        self._robot.setVelocity(vinit)
        return res
    def getKineticEnergy(self):
        vinit = self._robot.getVelocity()
        vtemp = self.tofull(self.getVelocity,[0]*self._robot.numLinks())
        self._robot.setVelocity(vtemp)
        res = self._robot.getKineticEnergy()
        self._robot.setVelocity(vinit)
        return res
    def getTotalInertia(self):
        raise NotImplementedError("TODO: getTotalInertia")
    def getMassMatrix(self):
        raise NotImplementedError("TODO: getMassMatrix")
    def getMassMatrixInv(self):
        raise NotImplementedError("TODO: getMassMatrix")
    def getCoriolisForceMatrix(self):
        raise NotImplementedError("TODO: getCoriolisForceMatrix")
    def getCoriolisForces(self):
        raise NotImplementedError("TODO: getCoriolisForceMatrix")
    def getGravityForces(self,g):
        raise NotImplementedError("TODO: getGravityForces")
    def torquesFromAccel(self,ddq):
        raise NotImplementedError("TODO: torquesFromAccel")
    def accelFromTorques(self,t):
        raise NotImplementedError("TODO: accelFromTorques")

    def interpolate(self,a,b,u):
        afull = self.tofull(a)
        bfull = self.tofull(b)
        q = self._robot.interpolate(afull,bfull,u)
        return [q[i] for i in self._links]
    def distance(self,a,b):
        afull = self.tofull(a)
        bfull = self.tofull(b)
        return self._robot.distance(afull,bfull)
    def interpolate_deriv(self,a,b):
        afull = self.tofull(a)
        bfull = self.tofull(b)
        q = self._robot.interpolate_deriv(afull,bfull)
        return [q[i] for i in self._links]
    def randomizeConfig(self,unboundedScale=1.0):
        oldfull = self._robot.getConfig()
        self._robot.randomizeConfig(unboundedScale)
        qrand = self._robot.getConfig()
        for i,l in enumerate(self._links):
            oldfull[l] = qrand[l]
        self._robot.setConfig(oldfull)
    def selfCollisionEnabled(self,link1,link2):
        if not isinstance(link1,str):
            link1 = self._links[link1]
        else:
            link1 = self._robot.link(link1).getIndex()
        if not isinstance(link2,str):
            link2 = self._links[link2]
        else:
            link2 = self._robot.link(link2).getIndex()
        return self._robot.selfCollisionEnabled(link1,link2)
    def enableSelfCollision(self,link1,link2,value=True):
        if not isinstance(link1,str):
            link1 = self._links[link1]
        else:
            link1 = self._robot.link(link1).getIndex()
        if not isinstance(link2,str):
            link2 = self._links[link2]
        else:
            link2 = self._robot.link(link2).getIndex()
        self._robot.enableSelfCollision(link1,link2,value)
    def selfCollides(self):
        geoms = [self._robot.link(i).geometry() for i in self._links]
        def dotest(i,j):
            return self._robot.selfCollisionEnabled(self._links[i],self._links[j])
        return any(self_collision_iter(geoms,dotest))
    def drawGL(self,keepAppearance=True):
        for i in self._links:
            self._robot.link(i).drawWorldGL(keepAppearance)
    def reduce(self):
        raise NotImplementedError("Can't reduce a sub-robot")
    def mount(self,link,subRobot,R,t):
        self._robot.mount(self.tofull(link),subRobot,R,t)
    def sensor(self,index):
        """Returns the SimSensorModel corresponding to index. Note however that
        you can't set the "link" setting according to this SubRobotModel.

        Args:
            index (int or str)
        """
        if isinstance(index,str):
            return self._robot.sensor(index)
        else:
            return self._robot.sensor(self.tofull(index))

class SubRobotModelLink:
    """A helper that lets you treat links of a subrobot just like a normal
    RobotModelLink. Correctly implements jacobians and indices with respect
    to the sub-robot.
    """
    def __init__(self,link,robot):
        self._link = link
        self._robot = robot
        self.robotIndex = link.robotIndex
        self.index = link.index
        self.getID = link.getID
        self.getName = link.getName
        self.setName = link.setName
        self.geometry = link.geometry
        self.appearance = link.appearance
        self.getMass = link.getMass
        self.setMass = link.setMass
        self.getParentTransform = link.getParentTransform
        self.setParentTransform = link.setParentTransform
        self.getAxis = link.getAxis
        self.setAxis = link.setAxis
        self.getWorldPosition = link.getWorldPosition
        self.getWorldDirection = link.getWorldDirection
        self.getLocalPosition = link.getLocalPosition
        self.getLocalDirection = link.getLocalDirection
        self.getTransform = link.getTransform
        self.setTransform = link.setTransform
        self.getVelocity = link.getVelocity
        self.getAngularVelocity = link.getAngularVelocity
        self.getPointVelocity = link.getPointVelocity
        self.drawLocalGL = link.drawLocalGL
        self.drawWorldGL = link.drawWorldGL
    def robot(self):
        return self._robot
    def getIndex(self):
        return self._robot.fromfull(self._link.getIndex())
    def getParent(self):
        p = self._robot.fromfull(self._link.getParent())
        if p == None: return -1
        return p
    def parent(self):
        p = self.getParent()
        return self._robot.link(p)
    def setParent(self,p):
        self._link.setParent(self._robot.tofull(p))
    def getJacobian(self,p):
        self._robot.fromfull(self._link.getJacobian(p))
    def getPositionJacobian(self,p):
        self._robot.fromfull(self._link.getPositionJacobian(p))
    def getOrientationJacobian(self):
        self._robot.fromfull(self._link.getOrientationJacobian())

<|endoftext|>


<commit_before>#!/usr/bin/env python2

###############################################################################
# ------------------------- Description ---------------------------------------
###############################################################################
# The purpose of this script is to regrid FINN or GFFED emissions data to ecmwf
# era-interim grid (or any lat lon passed) and save as a single daily nc file.

# TODO: Calculate and save the grid sizes m**2 to written nc files

import sys # for reading command line arguments
from netCDF4 import Dataset
import matplotlib.pyplot as plt
from mpl_toolkits.basemap import Basemap, cm, shiftgrid
import numpy as np
import scipy.interpolate
import time as timer

print 'Number of arguments:', len(sys.argv), 'arguments.'
print 'Argument List:', str(sys.argv)

if len(sys.argv) != 1:
	print 'Using arguments passed via command line.'
	year      =  str(sys.argv[1])
	inventory =  str(sys.argv[2])
	METGrid   =  str(sys.argv[3])
	species   =  str(sys.argv[4])

else:
	# Development environment. Set variables manually here.
	year      = str(2003)
	inventory = 'FINN'  # FINN | GFED4s
	METGrid   = 'ecmwf' #
	species   = 'CO2'   # Either CO2 or a specific vegetation type or C for GFED

# Consider making this directory a command line argument
dataDir  = '/barnes-scratch/sbrey/'
sanityCheck = False

# Get the era-interim (or other MET grid). They ALL live on the same grid, except
# time
# TODO: Handle time matching in this script as well!
metFile = dataDir + 'era_interim_nc_daily/sp_'+ year + '.nc'
met_nc = Dataset(metFile, 'r')
met_lat = met_nc.variables['latitude'][:]
met_lon = met_nc.variables['longitude'][:]
met_time= met_nc.variables['time']
sp_slice = met_nc.variables['sp'][180,:,:]

if sanityCheck:
	m = Basemap(projection='robin',lon_0=0,resolution='c')
	lons, lats = np.meshgrid(met_lon, met_lat)
	x,y=m(lons,lats)
	m.drawcoastlines()
	#m.fillcontinents(color='coral',lake_color='aqua')
	# draw parallels and meridians.
	#plt.pcolormesh(x,y, sp_slice)
	c = plt.contour(x, y, sp_slice )
	#m.drawmapboundary(fill_color='aqua')
	plt.title("ERA-Interim Grid")
	plt.show(block=False)
	#plt.close()

# Get emissions file (use yearly nc files)
if inventory == 'GDED4s':
	fire_file = dataDir + 'GFED4s/GFED4.1s_' + species + '_'+ year + '.nc'

elif inventory == 'FINN':
	fire_file = dataDir + 'FINN/FINN_' + species + '_'+ year + '.nc'

else:
	print 'Unknown emissions type'

fire_nc   = Dataset(fire_file, 'r')
fire_lat  = fire_nc.variables['latitude'][:]
fire_lon  = fire_nc.variables['longitude'][:]
fire_time = fire_nc.variables['time']

print '----------------------------------------------------------------------'
print 'Working on loading the really big emissions file'
print '----------------------------------------------------------------------'
timeStart = timer.time()
emissions = fire_nc.variables[species][:]
dt = (timer.time() - timeStart) / 60.
print '----------------------------------------------------------------------'
print 'It took ' + str(dt) + ' minutes to load emissions array into workspace'
print '----------------------------------------------------------------------'

print '----------------------------------------------------------------------'
print 'Working on shifting emissions from -180:180 grid to 0:360 grid.'
print '----------------------------------------------------------------------'
regridStart = timer.time()

emissions_new_grid = np.zeros(emissions.shape) # same size overall
fire_lon_new = np.zeros(fire_lon.shape)

eastOfGreenwich = np.where(fire_lon >= 0)[0]
westOfGreenwich = np.where(fire_lon < 0)[0]

# Assign the eastern hemisphere
fire_lon_new[range(len(eastOfGreenwich))] = fire_lon[eastOfGreenwich]
emissions_new_grid[:,:,range(len(eastOfGreenwich))] = emissions[:,:,eastOfGreenwich]
# assign the western hemisphere
assignment = np.where(fire_lon_new==0)[0]
fire_lon_new[assignment] = fire_lon[westOfGreenwich] + 360.
emissions_new_grid[:,:, assignment] = emissions[:,:,westOfGreenwich]

# Replace old grids and longitude with new
emissions = emissions_new_grid
fire_lon = fire_lon_new

dt = (timer.time() - regridStart) / 60.
print '-----------------------------------------------------------------------'
print 'It took ' + str(dt) + ' minutes to regrid (lons) of the emissions array'
print '-----------------------------------------------------------------------'

##############
# Sanity plot
##############
if sanityCheck:

	emissions_new_sum = np.sum(emissions_new_grid, axis=0)
	emissions_ma = np.ma.masked_where(emissions_new_sum==0,emissions_new_sum, copy=True)
	m = Basemap(projection='robin',lon_0=0,resolution='c')
	lons, lats = np.meshgrid(fire_lon_new, fire_lat)
	x,y=m(lons,lats)

	m = Basemap(projection='robin', lon_0=0,resolution='c')
	m.drawcoastlines()
	#m.fillcontinents(color='coral',lake_color='aqua')
	# draw parallels and meridians.
	m.drawparallels(np.arange(-90.,120.,30.))
	m.drawmeridians(np.arange(0.,360.,60.))
	c = m.pcolor(x, y, emissions_ma)
	#m.drawmapboundary(fill_color='aqua')
	plt.title("GFED GRID")
	plt.show()

# Get the total number of time steps
nTStep = len(fire_time)

# NOTE: float32 specified otherwise addition leads to small errors and mass
# NOTE: is not conserved in the regridding process.
fire_new_grid = np.zeros((nTStep, len(met_lat), len(met_lon)), dtype='float')
deltaMass = np.zeros(nTStep, dtype='float')

# For each time step, find the best match in terms of min lon lat
# distance
# TODO: Only loop over non zero locations? Over subregion?
timeStart = timer.time()

for t in range(nTStep): # nTStep
	#print str(t) + ' time step'
	#print str(float(t)/nTStep*100.) + " % complete"
	# Loop over fire lon values
	for x in range(len(fire_lon)):
		# to make it a point, loop over each lat for each lon
		for y in range(len(fire_lat)):

			# y and x represent the emissions data location we are
			# trying to assign!

			# Difference in regular grid arrays, units of deg
			dx = np.abs(fire_lon[x] - met_lon)
			dy = np.abs(fire_lat[y] - met_lat)

			# Index of best match
			xi = np.argmin(dx)
			yi = np.argmin(dy)

			# assign the data
			fire_new_grid[t, yi, xi] = (fire_new_grid[t, yi, xi] + emissions[t, y, x])

	# How much mass is lost or gained?
	deltaMass[t] = np.sum(fire_new_grid[t, :, :]) - np.sum(emissions[t, :, :])

dt = (timer.time() - timeStart) / 60.
print '----------------------------------------------------------------------'
print 'It took ' + str(dt) + ' minutes to run the re-grid loop'
print '----------------------------------------------------------------------'

######################################################################################
# Write the NETCDF data. Make sure to include all relevant units for a given species!
######################################################################################
print 'Working on writing the output as netCDF data'
nLat = len(met_lat)
nLon = len(met_lon)
nTime = len(fire_time)

# TODO: Make this save name (and whole section) dynamic based on the passed met grid
# TODO: to use

# Handle emissions specific labels
if inventory == 'GDED4s':
	outputFile = dataDir + 'GFED4s/' + 'GFED4.1s_' + METGrid + '_' + species + '_' +\
		         str(year) + '.nc'

elif inventory == 'FINN':
	outputFile = dataDir + 'FINN/' + 'FINN_' + METGrid + '_' + species + '_' +\
		         str(year) + '.nc'

	ncFile = Dataset(outputFile, 'w', format='NETCDF4')
	ncFile.description = 'Daily ' + inventory + ' regridded to era-interim grid.'
	ncFile.location = 'Global'
	ncFile.createDimension('time',  nTime )
	ncFile.createDimension('latitude', nLat )
	ncFile.createDimension('longitude', nLon )

	VAR_ = ncFile.createVariable(species, 'f4',('time','latitude','longitude'))
	VAR_.units = 'g ' + species + ' per grid cell per day'

# Create time variable
time_ = ncFile.createVariable('time', 'i4', ('time',))
time_.units = fire_time.units

# Keep track of the change in mass
deltaMass_ = ncFile.createVariable('deltaMass', 'f4', ('time',))
deltaMass_.units = 'change in grams whole grid each timestep due to regridding'

# create lat variable
latitude_ = ncFile.createVariable('latitude', 'f4', ('latitude',))
latitude_.units = 'degrees north'

# create longitude variable
longitude_ = ncFile.createVariable('longitude', 'f4', ('longitude',))
longitude_.units = 'degrees east'

# Write the actual data to these dimensions
VAR_[:]       = fire_new_grid[:]
latitude_[:]  = met_lat
longitude_[:] = met_lon
time_[:]      = fire_time[:]
deltaMass_[:] = deltaMass[:]

ncFile.close()

<commit_msg>ready to convert hysplit points spdf to ecmwf grid, and can handle the missing dates.<commit_after>#!/usr/bin/env python2

###############################################################################
# ------------------------- Description ---------------------------------------
###############################################################################
# The purpose of this script is to regrid FINN or GFFED emissions data to ecmwf
# era-interim grid (or any lat lon passed) and save as a single daily nc file.

# TODO: Calculate and save the grid sizes m**2 to written nc files

import sys # for reading command line arguments
from netCDF4 import Dataset
import matplotlib.pyplot as plt
from mpl_toolkits.basemap import Basemap, cm, shiftgrid
import numpy as np
import scipy.interpolate
import time as timer
import cesm_nc_manager as cnm

print 'Number of arguments:', len(sys.argv), 'arguments.'
print 'Argument List:', str(sys.argv)

if len(sys.argv) != 1:
	print 'Using arguments passed via command line.'
	year      =  str(sys.argv[1])
	inventory =  str(sys.argv[2])
	METGrid   =  str(sys.argv[3])
	species   =  str(sys.argv[4])

else:
	# Development environment. Set variables manually here.
	year      = str(2006)
	inventory = 'HMS'  # FINN | GFED4s | HMS
	METGrid   = 'ecmwf' #
	species   = 'SPDH'   # Either CO2 or a specific vegetation type or C for GFED, SPDH

# Consider making this directory a command line argument
dataDir  = '/barnes-scratch/sbrey/'
sanityCheck = False

# Get the era-interim (or other MET grid). They ALL live on the same grid, except
# time
# TODO: Handle time matching in this script as well!
metFile = dataDir + 'era_interim_nc_daily/sp_'+ year + '.nc'
met_nc = Dataset(metFile, 'r')
met_lat = met_nc.variables['latitude'][:]
met_lon = met_nc.variables['longitude'][:]
met_time= met_nc.variables['time']
sp_slice = met_nc.variables['sp'][180,:,:]

if sanityCheck:
	m = Basemap(projection='robin',lon_0=0,resolution='c')
	lons, lats = np.meshgrid(met_lon, met_lat)
	x,y=m(lons,lats)
	m.drawcoastlines()
	#m.fillcontinents(color='coral',lake_color='aqua')
	# draw parallels and meridians.
	#plt.pcolormesh(x,y, sp_slice)
	c = plt.contour(x, y, sp_slice )
	#m.drawmapboundary(fill_color='aqua')
	plt.title("ERA-Interim Grid")
	plt.show(block=False)
	#plt.close()

# Get emissions file (use yearly nc files)
if inventory == 'GDED4s':
	fire_file = dataDir + 'GFED4s/GFED4.1s_' + species + '_'+ year + '.nc'

elif inventory == 'FINN':
	fire_file = dataDir + 'FINN/FINN_' + species + '_'+ year + '.nc'
elif inventory == 'HMS':
	fire_file = dataDir + 'HMS/HYSPLITPoints_' + species  + '.nc'
else:
	print 'Unknown emissions type'

fire_nc   = Dataset(fire_file, 'r')
fire_lat  = fire_nc.variables['latitude'][:]
fire_lon  = fire_nc.variables['longitude'][:]
fire_time = fire_nc.variables['time']

print '----------------------------------------------------------------------'
print 'Working on loading the really big emissions file'
print '----------------------------------------------------------------------'
timeStart = timer.time()
emissions = fire_nc.variables[species][:]
dt = (timer.time() - timeStart) / 60.
print '----------------------------------------------------------------------'
print 'It took ' + str(dt) + ' minutes to load emissions array into workspace'
print '----------------------------------------------------------------------'

print '----------------------------------------------------------------------'
print 'Working on shifting emissions from -180:180 grid to 0:360 grid.'
print '----------------------------------------------------------------------'
regridStart = timer.time()

emissions_new_grid = np.zeros(emissions.shape) # same size overall
fire_lon_new = np.zeros(fire_lon.shape)

eastOfGreenwich = np.where(fire_lon >= 0)[0]
westOfGreenwich = np.where(fire_lon < 0)[0]

# Assign the eastern hemisphere
fire_lon_new[range(len(eastOfGreenwich))] = fire_lon[eastOfGreenwich]
emissions_new_grid[:,:,range(len(eastOfGreenwich))] = emissions[:,:,eastOfGreenwich]
# assign the western hemisphere
assignment = np.where(fire_lon_new==0)[0]
fire_lon_new[assignment] = fire_lon[westOfGreenwich] + 360.
emissions_new_grid[:,:, assignment] = emissions[:,:,westOfGreenwich]

# Replace old grids and longitude with new
emissions = emissions_new_grid
fire_lon = fire_lon_new

dt = (timer.time() - regridStart) / 60.
print '-----------------------------------------------------------------------'
print 'It took ' + str(dt) + ' minutes to regrid (lons) of the emissions array'
print '-----------------------------------------------------------------------'

##############
# Sanity plot
##############
if sanityCheck:

	emissions_new_sum = np.sum(emissions_new_grid, axis=0)
	emissions_ma = np.ma.masked_where(emissions_new_sum==0,emissions_new_sum, copy=True)
	m = Basemap(projection='robin',lon_0=0,resolution='c')
	lons, lats = np.meshgrid(fire_lon_new, fire_lat)
	x,y=m(lons,lats)

	m = Basemap(projection='robin', lon_0=0,resolution='c')
	m.drawcoastlines()
	#m.fillcontinents(color='coral',lake_color='aqua')
	# draw parallels and meridians.
	m.drawparallels(np.arange(-90.,120.,30.))
	m.drawmeridians(np.arange(0.,360.,60.))
	c = m.pcolor(x, y, emissions_ma)
	#m.drawmapboundary(fill_color='aqua')
	plt.title("REGRID GRID")
	plt.show()

fire_time_units = fire_time.units # preserve for writing nc data
# TODO: Subset by year range when HMS, since it is not in the file name
if inventory == 'HMS':
	t, month, years = cnm.get_era_interim_time(fire_time)
	yearMask = years == int(year)

	# apply the mask
	fire_time = fire_time[yearMask]
	emissions = emissions[yearMask,:,:]

# How make an array that represets all dates in this year to index the yearly data
# This is needed because HMS is missing dates
nTStep = len(met_time)

# NOTE: float32 specified otherwise addition leads to small errors and mass
# NOTE: is not conserved in the regridding process.
fire_new_grid = np.zeros((nTStep, len(met_lat), len(met_lon)), dtype='float')
deltaMass = np.zeros(nTStep, dtype='float')

# For each time step, find the best match in terms of min lon lat
# distance
# TODO: Only loop over non zero locations? Over subregion?
timeStart = timer.time()

for t in range(nTStep): # nTStep
	#print str(t) + ' time step'

	# Where does this date for this year fall in fire_time?
	tMask = met_time[t] == fire_time[:]

	# See if the date exists, if it does not that means there is no fire emission
	# data for this date and we can move forward with the time loop.
	if np.sum(tMask) > 0:

		print str(float(t)/nTStep*100.) + " % complete"
		# Loop over fire lon values
		for x in range(len(fire_lon)):
			# to make it a point, loop over each lat for each lon
			for y in range(len(fire_lat)):

				# y and x represent the emissions data location we are
				# trying to assign!

				# Difference in regular grid arrays, units of deg
				dx = np.abs(fire_lon[x] - met_lon)
				dy = np.abs(fire_lat[y] - met_lat)

				# Index of best match
				xi = np.argmin(dx)
				yi = np.argmin(dy)

				# assign the data, new grid is t referenced, emissions tMask
				fire_new_grid[t, yi, xi] = (fire_new_grid[t, yi, xi] + emissions[tMask, y, x])

		# How much mass is lost or gained?
		deltaMass[t] = np.sum(fire_new_grid[t, :, :]) - np.sum(emissions[tMask, :, :])
	else:
		print str(met_time[t]) + '  - missing in fire_emissions'

dt = (timer.time() - timeStart) / 60.
print '----------------------------------------------------------------------'
print 'It took ' + str(dt) + ' minutes to run the re-grid loop'
print '----------------------------------------------------------------------'

######################################################################################
# Write the NETCDF data. Make sure to include all relevant units for a given species!
######################################################################################
print 'Working on writing the output as netCDF data'
nLat = len(met_lat)
nLon = len(met_lon)
nTime = nTStep # should match number of days in calendar year of chosen year

# TODO: Make this save name (and whole section) dynamic based on the passed met grid
# TODO: to use

# Handle emissions specific labels
if inventory == 'GDED4s':
	outputFile = dataDir + 'GFED4s/' + 'GFED4.1s_' + METGrid + '_' + species + '_' +\
		         str(year) + '.nc'

elif inventory == 'FINN':
	outputFile = dataDir + 'FINN/' + 'FINN_' + METGrid + '_' + species + '_' +\
		         str(year) + '.nc'

elif inventory == 'HMS':
	outputFile = dataDir + 'HMS/' + 'HYSPLITPoints_' + METGrid + '_' + species +'_' +\
	             str(year) + '.nc'

ncFile = Dataset(outputFile, 'w', format='NETCDF4')
ncFile.description = 'Daily ' + inventory + ' regridded to era-interim grid.'
ncFile.location = 'Global'
ncFile.createDimension('time',  nTime )
ncFile.createDimension('latitude', nLat )
ncFile.createDimension('longitude', nLon )

VAR_ = ncFile.createVariable(species, 'f4',('time','latitude','longitude'))

if inventory != 'HMS':
	VAR_.units = 'g ' + species + ' per grid cell per day'
else:
	VAR_.units = 'smoke production duration hours'

# Create time variable
time_ = ncFile.createVariable('time', 'i4', ('time',))
time_.units = fire_time_units

# Keep track of the change in mass
deltaMass_ = ncFile.createVariable('deltaMass', 'f4', ('time',))
deltaMass_.units = 'change in grams whole grid each timestep due to regridding'

# create lat variable
latitude_ = ncFile.createVariable('latitude', 'f4', ('latitude',))
latitude_.units = 'degrees north'

# create longitude variable
longitude_ = ncFile.createVariable('longitude', 'f4', ('longitude',))
longitude_.units = 'degrees east'

# Write the actual data to these dimensions
VAR_[:]       = fire_new_grid[:]
latitude_[:]  = met_lat
longitude_[:] = met_lon
time_[:]      = met_time[:]
deltaMass_[:] = deltaMass[:]

ncFile.close()

<|endoftext|>


<commit_before>class Lexer:
    def __init__(self):
        self.__current_delim = None # This is bad. The lexer itself should not be stateful.

    def tokenize(self, characters):
        current_word = ''
        for c in characters:
            current_word += c
            if current_word in CHAR_TO_TYPE:
                yield CHAR_TO_TYPE[current_word](current_word)
                current_word = ''
            elif c in STRING_DELIMITERS:
                self.__handle_string_delimiter(c)
            elif not self.__lexing_string() and c == ' ':
                if current_word != ' ':
                    yield self.__get_token(current_word[:-1])
                current_word = ''
        if current_word != '':
            yield self.__get_token(current_word)

    def __get_token(self, word):
        if word[0] in STRING_DELIMITERS:
            return string_token(word)
        if word.isdigit():
            return integer_token(word)
        return identifier_token(word)

    def __handle_string_delimiter(self, delim):
        if self.__lexing_string() and self.__ends_string(delim):
            self.__current_delim = None
        elif not self.__lexing_string():
            self.__current_delim = delim

    def __ends_string(self, char):
        return self.__current_delim == char

    def __lexing_string(self):
        return self.__current_delim

class Token:
    def __init__(self, word, ttype):
        self.raw_value = word
        self.__ttype = ttype

    def is_a(self, token):
        return self.__ttype == token.__ttype

def integer_token(word=None):
    return Token(word, 0)
def string_token(word=None):
    return Token(word, 1)
def identifier_token(word=None):
    return Token(word, 2)
def left_paren_token(word=None):
    return Token(word, 3)
def right_paren_token(word=None):
    return Token(word, 4)
def right_square_bracket_token(word=None):
    return Token(word, 5)
def left_square_bracket_token(word=None):
    return Token(word, 6)

STRING_DELIMITERS = ['"', "'"]

CHAR_TO_TYPE = {
    '(': left_paren_token,
    ')': right_paren_token,
    ']': right_square_bracket_token,
    '[': left_square_bracket_token,
}

<commit_msg>Okay, there is more that could be done,<commit_after>class Lexer:
    def __init__(self):
        self.__current_delim = None # This is bad. The lexer itself should not be stateful.

    def tokenize(self, characters):
        current_word = ''
        for c in characters:
            current_word += c
            if current_word in CHAR_TO_TYPE:
                yield CHAR_TO_TYPE[current_word](current_word)
                current_word = ''
            elif c in STRING_DELIMITERS:
                self.__handle_string_delimiter(c)
            elif not self.__lexing_string() and c == ' ':
                if current_word != ' ':
                    yield self.__get_token(current_word[:-1])
                current_word = ''
        if current_word != '':
            yield self.__get_token(current_word)

    def __get_token(self, word):
        if word[0] in STRING_DELIMITERS:
            return string_token(word)
        if word.isdigit():
            return integer_token(word)
        return identifier_token(word)

    def __handle_string_delimiter(self, delim):
        if self.__lexing_string() and self.__ends_string(delim):
            self.__current_delim = None
        elif not self.__lexing_string():
            self.__current_delim = delim

    def __ends_string(self, char):
        return self.__current_delim == char

    def __lexing_string(self):
        return self.__current_delim

class Token:
    def __init__(self, word, ttype):
        self.raw_value = word
        self.__ttype = ttype

    def is_a(self, token):
        return self.__ttype == token.__ttype

def __make_token_types():
    token_type = -1
    while True:
        token_type += 1
        yield lambda word=None: Token(word, token_type)

token_types = __make_token_types()

integer_token = next(token_types)
string_token = next(token_types)
identifier_token = next(token_types)
left_paren_token = next(token_types)
right_paren_token = next(token_types)
right_square_bracket_token = next(token_types)
left_square_bracket_token = next(token_types)

STRING_DELIMITERS = ['"', "'"]

CHAR_TO_TYPE = {
    '(': left_paren_token,
    ')': right_paren_token,
    ']': right_square_bracket_token,
    '[': left_square_bracket_token,
}

<|endoftext|>


<commit_before>import argparse
import json
import os
import sys

import cleaner

parser = argparse.ArgumentParser(description='extract tweet texts from json')
parser.add_argument('-i', '--json-dir', type=str,
                    help='tweets json directory', required=True)
parser.add_argument('-o', '--out-dir', type=str,
                    help='the output directory.', required=True)
parser.add_argument('--exclude-redundant', type=str,
                    help='exclude redundant tweets', action='store_true')

def extract_tweets_from_json(json_reader, text_writer):
    json_tweets = json_reader.readlines()
    tweets_list = list()
    for json_tweet in json_tweets:
        try:
            if json_tweets:
                # load it as Python dict
                tweet = json.loads(json_tweet)
                text = tweet['text']
                text = cleaner.clean_tweet(text)
                if args.exclude_redundant:
                    if text not in tweets_list:
                        tweets_list.append(text)
                        text_writer.write(text)
                        text_writer.write("\n")
                else:
                    text_writer.write(text)
                    text_writer.write("\n")
        except json.decoder.JSONDecodeError as error:
            pass
        except UnicodeDecodeError as error:
            pass

def extract_tweets_from_json_files(json_dir, text_dir):
    for json_file in os.listdir(json_dir):
        if json_file.endswith('.json'):
            json_file = os.path.join(json_dir, json_file)
            filename, ext = os.path.splitext(json_file)
            text_file = os.path.join(text_dir, os.path.basename(filename) + '.txt')
            with open(json_file, mode='r', encoding='utf-8') as json_reader, open(text_file, mode='w',
                                                                                encoding='utf-8') as text_writer:
                print('extract tweets from {}'.format(json_file))
                extract_tweets_from_json(json_reader, text_writer)
                print('tweets extracted to {}'.format(text_file))
                print('--------------------')

if __name__ == '__main__':
    args = parser.parse_args()
    json_dir = args.json_dir
    out_dir = args.out_dir
    if json_dir == out_dir:
        print('error: input and output directories can not be the same.')
        sys.exit(-1)
    extract_tweets_from_json_files(json_dir, out_dir)
<commit_msg>add --exclude-redundant tweets option to json2text.py<commit_after>import argparse
import json
import os
import sys

import cleaner

parser = argparse.ArgumentParser(description='extract tweet texts from json')
parser.add_argument('-i', '--json-dir', type=str,
                    help='tweets json directory', required=True)
parser.add_argument('-o', '--out-dir', type=str,
                    help='the output directory.', required=True)
parser.add_argument('--exclude-redundant', type=str,
                    help='exclude redundant tweets', action='store_true')

def extract_tweets_from_json(json_reader, text_writer):
    json_tweets = json_reader.readlines()
    print('tweets in json file: {} tweets'.format(len(json_tweets)))
    tweets_list = list()
    for json_tweet in json_tweets:
        try:
            if json_tweets:
                # load it as Python dict
                tweet = json.loads(json_tweet)
                text = tweet['text']
                text = cleaner.clean_tweet(text)
                if args.exclude_redundant:
                    if text not in tweets_list:
                        tweets_list.append(text)
                        text_writer.write(text)
                        text_writer.write("\n")
                else:
                    text_writer.write(text)
                    text_writer.write("\n")
        except json.decoder.JSONDecodeError as error:
            pass
        except UnicodeDecodeError as error:
            pass

def extract_tweets_from_json_files(json_dir, text_dir):
    for json_file in os.listdir(json_dir):
        if json_file.endswith('.json'):
            json_file = os.path.join(json_dir, json_file)
            filename, ext = os.path.splitext(json_file)
            text_file = os.path.join(text_dir, os.path.basename(filename) + '.txt')
            with open(json_file, mode='r', encoding='utf-8') as json_reader, open(text_file, mode='w',
                                                                                encoding='utf-8') as text_writer:
                print('extract tweets from {}'.format(json_file))
                extract_tweets_from_json(json_reader, text_writer)
                print('tweets extracted to {}'.format(text_file))
                print('--------------------')

if __name__ == '__main__':
    args = parser.parse_args()
    json_dir = args.json_dir
    out_dir = args.out_dir
    if json_dir == out_dir:
        print('error: input and output directories can not be the same.')
        sys.exit(-1)
    extract_tweets_from_json_files(json_dir, out_dir)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
#
# python-json-patch - An implementation of the JSON Patch format
# https://github.com/stefankoegl/python-json-patch
#
# Copyright (c) 2011 Stefan Kögl <stefan@skoegl.net>
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#
# 1. Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
# 3. The name of the author may not be used to endorse or promote products
#    derived from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
# IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
# OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
# IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT,
# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
# NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
# THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#

"""Apply JSON-Patches according to http://tools.ietf.org/html/draft-pbryan-json-patch-01"""

# Will be parsed by setup.py to determine package metadata
__author__ = 'Stefan Kögl <stefan@skoegl.net>'
__version__ = '0.1'
__website__ = 'https://github.com/stefankoegl/python-json-patch'
__license__ = 'Modified BSD License'

import copy
import json

class JsonPatchException(Exception):
    pass

class JsonPatchConflict(JsonPatchException):
    pass

def apply_patch(doc, patch):
    """
    >>> obj = { 'baz': 'qux', 'foo': 'bar' }
    >>> patch = [ { 'remove': '/baz' } ]
    >>> apply_patch(obj, patch)
    {'foo': 'bar'}
    """

    p = JsonPatch(patch)
    return p.apply(doc)

class JsonPatch(object):
    """ A JSON Patch is a list of Patch Operations """

    def __init__(self, patch):
        self.patch = patch

        self.OPERATIONS = {
            'remove': RemoveOperation,
            'add': AddOperation,
            'replace': ReplaceOperation,
            'move': MoveOperation
        }

    @classmethod
    def from_string(cls, patch_str):
        patch = json.loads(patch_str)
        return cls(patch)

    def apply(self, obj):
        """ Applies the patch to a copy of the given object """

        obj = copy.deepcopy(obj)

        for operation in self.patch:
            op = self._get_operation(operation)
            op.apply(obj)

        return obj

    def _get_operation(self, operation):
        for action, op_cls in self.OPERATIONS.items():
            if action in operation:
                location = operation[action]
                op = op_cls(location, operation)
                return op

        raise JsonPatchException("invalid operation '%s'" % operation)

class PatchOperation(object):
    """ A single operation inside a JSON Patch """

    def __init__(self, location, operation):
        self.location = location
        self.operation = operation

    def locate(self, obj, location, last_must_exist=True):
        """ Walks through the object according to location

        Returns the last step as (sub-object, last location-step) """

        parts = location.split('/')
        if parts.pop(0) != '':
            raise JsonPatchException('location must starts with /')

        for part in parts[:-1]:
            obj, loc_part = self._step(obj, part)

        _, last_loc = self._step(obj, parts[-1], must_exist=last_must_exist)
        return obj, last_loc

    def _step(self, obj, loc_part, must_exist=True):
        """ Goes one step in a locate() call """

        # Its not clear if a location "1" should be considered as 1 or "1"
        # We prefer the integer-variant if possible
        part_variants = self._try_parse(loc_part) + [loc_part]

        for variant in part_variants:
            try:
                return obj[variant], variant
            except:
                continue

        if must_exist:
            raise JsonPatchConflict('key %s not found' % loc_part)
        else:
            return obj, part_variants[0]

    @staticmethod
    def _try_parse(val, cls=int):
        try:
            return [cls(val)]
        except:
            return []

class RemoveOperation(PatchOperation):
    """ Removes an object property or an array element

    >>> obj = { 'baz': 'qux', 'foo': 'bar' }
    >>> patch = JsonPatch( [ { 'remove': '/baz' } ] )
    >>> patch.apply(obj)
    {'foo': 'bar'}

    >>> obj = { 'foo': [ 'bar', 'qux', 'baz' ] }
    >>> patch = JsonPatch( [ { "remove": "/foo/1" } ] )
    >>> patch.apply(obj)
    {'foo': ['bar', 'baz']}
    """

    def apply(self, obj):
        subobj, part = self.locate(obj, self.location)
        del subobj[part]

class AddOperation(PatchOperation):
    """ Adds an object property or an array element

    >>> obj = { "foo": "bar" }
    >>> patch = JsonPatch([ { "add": "/baz", "value": "qux" } ])
    >>> patch.apply(obj)
    {'foo': 'bar', 'baz': 'qux'}

    >>> obj = { "foo": [ "bar", "baz" ] }
    >>> patch = JsonPatch([ { "add": "/foo/1", "value": "qux" } ])
    >>> patch.apply(obj)
    {'foo': ['bar', 'qux', 'baz']}
    """

    def apply(self, obj):
        value = self.operation["value"]
        subobj, part = self.locate(obj, self.location, last_must_exist=False)

        if isinstance(subobj, list):
            if part > len(subobj) or part < 0:
                raise JsonPatchConflict("can't insert outside of list")

            subobj.insert(part, value)

        elif isinstance(subobj, dict):
            if part in subobj:
                raise JsonPatchConflict("object '%s' already exists" % part)

            subobj[part] = value

        else:
            raise JsonPatchConflict("can't add to type '%s'" % subobj.__class__.__name__)

class ReplaceOperation(PatchOperation):
    """ Replaces a value

    >>> obj = { "baz": "qux", "foo": "bar" }
    >>> patch = JsonPatch([ { "replace": "/baz", "value": "boo" } ])
    >>> patch.apply(obj)
    {'foo': 'bar', 'baz': 'boo'}
    """

    def apply(self, obj):
        value = self.operation["value"]
        subobj, part = self.locate(obj, self.location)

        if isinstance(subobj, list):
            if part > len(subobj) or part < 0:
                raise JsonPatchConflict("can't replace outside of list")

        elif isinstance(subobj, dict):
            if not part in subobj:
                raise JsonPatchConflict("can't replace non-existant object '%s'" % part)

        else:
            raise JsonPatchConflict("can't replace in type '%s'" % subobj.__class__.__name__)

        subobj[part] = value

class MoveOperation(PatchOperation):
    """ Moves a value

    >>> obj = {'foo': {'bar': 'baz', 'waldo': 'fred'}, 'qux': {'corge': 'grault'}}
    >>> patch = JsonPatch([{'move': '/foo/waldo', 'to': '/qux/thud'}])
    >>> patch.apply(obj)
    {'qux': {'thud': 'fred', 'corge': 'grault'}, 'foo': {'bar': 'baz'}}

    >>> obj =  {'foo': ['all', 'grass', 'cows', 'eat']}
    >>> patch = JsonPatch([{'move': '/foo/1', 'to': '/foo/3'}])
    >>> patch.apply(obj)
    {'foo': ['all', 'cows', 'eat', 'grass']}
    """

    def apply(self, obj):
        subobj, part = self.locate(obj, self.location)
        value = subobj[part]
        RemoveOperation(self.location, self.operation).apply(obj)
        AddOperation(self.operation['to'], {'value': value}).apply(obj)
<commit_msg>Add support for test operation<commit_after># -*- coding: utf-8 -*-
#
# python-json-patch - An implementation of the JSON Patch format
# https://github.com/stefankoegl/python-json-patch
#
# Copyright (c) 2011 Stefan Kögl <stefan@skoegl.net>
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#
# 1. Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
# 3. The name of the author may not be used to endorse or promote products
#    derived from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
# IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
# OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
# IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT,
# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
# NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
# THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#

"""Apply JSON-Patches according to http://tools.ietf.org/html/draft-pbryan-json-patch-01"""

# Will be parsed by setup.py to determine package metadata
__author__ = 'Stefan Kögl <stefan@skoegl.net>'
__version__ = '0.1'
__website__ = 'https://github.com/stefankoegl/python-json-patch'
__license__ = 'Modified BSD License'

import copy
import json

class JsonPatchException(Exception):
    pass

class JsonPatchConflict(JsonPatchException):
    pass

def apply_patch(doc, patch):
    """
    >>> obj = { 'baz': 'qux', 'foo': 'bar' }
    >>> patch = [ { 'remove': '/baz' } ]
    >>> apply_patch(obj, patch)
    {'foo': 'bar'}
    """

    p = JsonPatch(patch)
    return p.apply(doc)

class JsonPatch(object):
    """ A JSON Patch is a list of Patch Operations """

    def __init__(self, patch):
        self.patch = patch

        self.OPERATIONS = {
            'remove': RemoveOperation,
            'add': AddOperation,
            'replace': ReplaceOperation,
            'move': MoveOperation,
            'test': TestOperation
        }

    @classmethod
    def from_string(cls, patch_str):
        patch = json.loads(patch_str)
        return cls(patch)

    def apply(self, obj):
        """ Applies the patch to a copy of the given object """

        obj = copy.deepcopy(obj)

        for operation in self.patch:
            op = self._get_operation(operation)
            op.apply(obj)

        return obj

    def _get_operation(self, operation):
        for action, op_cls in self.OPERATIONS.items():
            if action in operation:
                location = operation[action]
                op = op_cls(location, operation)
                return op

        raise JsonPatchException("invalid operation '%s'" % operation)

class PatchOperation(object):
    """ A single operation inside a JSON Patch """

    def __init__(self, location, operation):
        self.location = location
        self.operation = operation

    def locate(self, obj, location, last_must_exist=True):
        """ Walks through the object according to location

        Returns the last step as (sub-object, last location-step) """

        parts = location.split('/')
        if parts.pop(0) != '':
            raise JsonPatchException('location must starts with /')

        for part in parts[:-1]:
            obj, loc_part = self._step(obj, part)

        _, last_loc = self._step(obj, parts[-1], must_exist=last_must_exist)
        return obj, last_loc

    def _step(self, obj, loc_part, must_exist=True):
        """ Goes one step in a locate() call """

        # Its not clear if a location "1" should be considered as 1 or "1"
        # We prefer the integer-variant if possible
        part_variants = self._try_parse(loc_part) + [loc_part]

        for variant in part_variants:
            try:
                return obj[variant], variant
            except:
                continue

        if must_exist:
            raise JsonPatchConflict('key %s not found' % loc_part)
        else:
            return obj, part_variants[0]

    @staticmethod
    def _try_parse(val, cls=int):
        try:
            return [cls(val)]
        except:
            return []

class RemoveOperation(PatchOperation):
    """ Removes an object property or an array element

    >>> obj = { 'baz': 'qux', 'foo': 'bar' }
    >>> patch = JsonPatch( [ { 'remove': '/baz' } ] )
    >>> patch.apply(obj)
    {'foo': 'bar'}

    >>> obj = { 'foo': [ 'bar', 'qux', 'baz' ] }
    >>> patch = JsonPatch( [ { "remove": "/foo/1" } ] )
    >>> patch.apply(obj)
    {'foo': ['bar', 'baz']}
    """

    def apply(self, obj):
        subobj, part = self.locate(obj, self.location)
        del subobj[part]

class AddOperation(PatchOperation):
    """ Adds an object property or an array element

    >>> obj = { "foo": "bar" }
    >>> patch = JsonPatch([ { "add": "/baz", "value": "qux" } ])
    >>> patch.apply(obj)
    {'foo': 'bar', 'baz': 'qux'}

    >>> obj = { "foo": [ "bar", "baz" ] }
    >>> patch = JsonPatch([ { "add": "/foo/1", "value": "qux" } ])
    >>> patch.apply(obj)
    {'foo': ['bar', 'qux', 'baz']}
    """

    def apply(self, obj):
        value = self.operation["value"]
        subobj, part = self.locate(obj, self.location, last_must_exist=False)

        if isinstance(subobj, list):
            if part > len(subobj) or part < 0:
                raise JsonPatchConflict("can't insert outside of list")

            subobj.insert(part, value)

        elif isinstance(subobj, dict):
            if part in subobj:
                raise JsonPatchConflict("object '%s' already exists" % part)

            subobj[part] = value

        else:
            raise JsonPatchConflict("can't add to type '%s'" % subobj.__class__.__name__)

class ReplaceOperation(PatchOperation):
    """ Replaces a value

    >>> obj = { "baz": "qux", "foo": "bar" }
    >>> patch = JsonPatch([ { "replace": "/baz", "value": "boo" } ])
    >>> patch.apply(obj)
    {'foo': 'bar', 'baz': 'boo'}
    """

    def apply(self, obj):
        value = self.operation["value"]
        subobj, part = self.locate(obj, self.location)

        if isinstance(subobj, list):
            if part > len(subobj) or part < 0:
                raise JsonPatchConflict("can't replace outside of list")

        elif isinstance(subobj, dict):
            if not part in subobj:
                raise JsonPatchConflict("can't replace non-existant object '%s'" % part)

        else:
            raise JsonPatchConflict("can't replace in type '%s'" % subobj.__class__.__name__)

        subobj[part] = value

class MoveOperation(PatchOperation):
    """ Moves a value

    >>> obj = {'foo': {'bar': 'baz', 'waldo': 'fred'}, 'qux': {'corge': 'grault'}}
    >>> patch = JsonPatch([{'move': '/foo/waldo', 'to': '/qux/thud'}])
    >>> patch.apply(obj)
    {'qux': {'thud': 'fred', 'corge': 'grault'}, 'foo': {'bar': 'baz'}}

    >>> obj =  {'foo': ['all', 'grass', 'cows', 'eat']}
    >>> patch = JsonPatch([{'move': '/foo/1', 'to': '/foo/3'}])
    >>> patch.apply(obj)
    {'foo': ['all', 'cows', 'eat', 'grass']}
    """

    def apply(self, obj):
        subobj, part = self.locate(obj, self.location)
        value = subobj[part]
        RemoveOperation(self.location, self.operation).apply(obj)
        AddOperation(self.operation['to'], {'value': value}).apply(obj)

class TestOperation(PatchOperation):
    """ Test value by specified location

    >>> obj = {'baz': 'qux', 'foo': ['a', 2, 'c']}
    >>> patch = JsonPatch([
    ...    {'test': '/baz', 'value': 'qux'},
    ...    {'test': '/foo/1', 'value': 2}
    ... ])
    >>> patch.apply(obj)
    {'foo': ['a', 2, 'c'], 'baz': 'qux'}

    >>> patch = JsonPatch([
    ...    {'test': '/foo/1', 'value': 'BOOM!'}
    ... ])
    >>> try:
    ...     patch.apply(obj)
    ... except AssertionError:
    ...     pass
    ... else:
    ...     assert False, 'test should fall'
    """

    def apply(self, obj):
        value = self.operation['value']
        subobj, part = self.locate(obj, self.location)
        assert subobj[part] == value
<|endoftext|>


<commit_before>import numpy as np
import math
import random
import itertools
import collections

class NeuralNet:
    def __init__(self,n_inputs, n_outputs, n_hiddens, n_hidden_layers, activation_functions):
        self.n_inputs = n_inputs                # Number of network input signals
        self.n_outputs = n_outputs              # Number of desired outputs from the network
        self.n_hiddens = n_hiddens              # Number of nodes in each hidden layer
        self.n_hidden_layers = n_hidden_layers  # Number of hidden layers in the network
        self.activation_functions = activation_functions

        assert len(activation_functions)==(n_hidden_layers+1), "Requires "+(n_hidden_layers+1)+" activation functions, got: "+len(activation_functions)+"."

        if n_hidden_layers == 0:
            # Count the necessary number of weights for the input->output connection.
            # input -> [] -> output
            self.n_weights = ((n_inputs+1)*n_outputs)
        else:
            # Count the necessary number of weights summed over all the layers.
            # input -> [n_hiddens -> n_hiddens] -> output
            self.n_weights = (n_inputs+1)*n_hiddens+\
                             (n_hiddens**2+n_hiddens)*(n_hidden_layers-1)+\
                             n_hiddens*n_outputs+n_outputs

        # Initialize the network with new randomized weights
        self.set_weights( self.generate_weights() )
    #end

    def generate_weights(self, low=-0.1, high=0.1):
        # Generate new random weights for all the connections in the network
        if not False:
            # Support NumPy
            return [random.uniform(low,high) for _ in xrange(self.n_weights)]
        else:
            return np.random.uniform(low, high, size=(1,self.n_weights)).tolist()[0]
    #end

    def unpack(self, weight_list ):
        # This method will create a list of weight matrices. Each list element
        # corresponds to the connection between two layers.
        if self.n_hidden_layers == 0:
            return [ np.array(weight_list).reshape(self.n_inputs+1,self.n_outputs) ]
        else:
            weight_layers = [ np.array(weight_list[:(self.n_inputs+1)*self.n_hiddens]).reshape(self.n_inputs+1,self.n_hiddens) ]
            weight_layers += [ np.array(weight_list[(self.n_inputs+1)*self.n_hiddens+(i*(self.n_hiddens**2+self.n_hiddens)):(self.n_inputs+1)*self.n_hiddens+((i+1)*(self.n_hiddens**2+self.n_hiddens))]).reshape(self.n_hiddens+1,self.n_hiddens) for i in xrange(self.n_hidden_layers-1) ]
            weight_layers += [ np.array(weight_list[(self.n_inputs+1)*self.n_hiddens+((self.n_hidden_layers-1)*(self.n_hiddens**2+self.n_hiddens)):]).reshape(self.n_hiddens+1,self.n_outputs) ]

        return weight_layers
    #end

    def set_weights(self, weight_list ):
        # This is a helper method for setting the network weights to a previously defined list.
        # This is useful for utilizing a previously optimized neural network weight set.
        self.weights = self.unpack( weight_list )
    #end

    def get_weights(self, ):
        # This will stack all the weights in the network on a list, which may be saved to the disk.
        return [w for l in self.weights for w in l.flat]
    #end

    def backpropagation(self, trainingset, ERROR_LIMIT=1e-3, learning_rate=0.3, momentum_factor=0.9  ):
        def addBias(A):
            # Add 1 as bias.
            return np.hstack((np.ones((A.shape[0],1)),A))
        #end addBias

        assert trainingset[0].features.shape[0] == self.n_inputs, "ERROR: input size varies from the defined input setting"
        assert trainingset[0].targets.shape[0] == self.n_outputs, "ERROR: output size varies from the defined output setting"

        training_data = np.array( [instance.features for instance in trainingset ] )
        training_targets = np.array( [instance.targets for instance in trainingset ] )

        MSE      = ( ) # inf
        neterror = None
        momentum = collections.defaultdict( int )

        epoch = 0
        while MSE>ERROR_LIMIT:
            epoch += 1

            input_layers      = self.update( training_data, trace=True )
            out               = input_layers[-1]

            error             = training_targets - out
            delta             = error
            MSE               = np.mean( np.power(error,2) )

            # Loop over the weight layers in reversed order to calculate the deltas
            loop  = itertools.izip(
                            xrange(len(self.weights)-1, -1, -1),
                            reversed(self.weights),
                            reversed(input_layers[:-1]),
                        )

            for i, weight_layer, input_signals in loop:
                # Calculate weight change
                dW = learning_rate * np.dot( addBias(input_signals).T, delta )  + momentum_factor * momentum[i]

                if i!= 0:
                    """Do not calculate the delta unnecessarily."""
                    # Skipping the bias weight during calculation.
                    weight_delta = np.dot( delta, weight_layer[1:,:].T )

                    # Calculate the delta for the subsequent layer
                    delta = np.multiply(  weight_delta, self.activation_functions[i-1]( input_signals, derivative=True) )

                # Store momentum
                momentum[i] = dW

                # Update weights
                self.weights[ i ] += dW

            if epoch%1000==0: print "* current network error (MSE):", MSE

        print "* Converged to error bound (%.3g) with MSE = %.3g." % ( ERROR_LIMIT, MSE )
        print "* Trained for %d epochs." % epoch
    # end backprop

    def update(self, input_values, trace=False ):
        # This is a forward operation in the network. This is how we calculate the network output
        # from a set of input signals.

        output = input_values
        if trace: tracelist = [ output ]

        for i, weight_layer in enumerate(self.weights):
            # Loop over the network layers and calculate the output
            output = np.dot( output, weight_layer[1:,:] ) + weight_layer[0:1,:] # implicit bias
            output = self.activation_functions[i]( output )
            if trace: tracelist.append( output )

        if trace: return tracelist

        return output
    #end
#end class<commit_msg>Feature: added a saving and loading<commit_after>import numpy as np
import math
import random
import itertools
import collections

class NeuralNet:
    def __init__(self, n_inputs, n_outputs, n_hiddens, n_hidden_layers, activation_functions ):
        self.n_inputs = n_inputs                # Number of network input signals
        self.n_outputs = n_outputs              # Number of desired outputs from the network
        self.n_hiddens = n_hiddens              # Number of nodes in each hidden layer
        self.n_hidden_layers = n_hidden_layers  # Number of hidden layers in the network
        self.activation_functions = activation_functions

        assert len(activation_functions)==(n_hidden_layers+1), "Requires "+(n_hidden_layers+1)+" activation functions, got: "+len(activation_functions)+"."

        if n_hidden_layers == 0:
            # Count the necessary number of weights for the input->output connection.
            # input -> [] -> output
            self.n_weights = ((n_inputs+1)*n_outputs)
        else:
            # Count the necessary number of weights summed over all the layers.
            # input -> [n_hiddens -> n_hiddens] -> output
            self.n_weights = (n_inputs+1)*n_hiddens+\
                             (n_hiddens**2+n_hiddens)*(n_hidden_layers-1)+\
                             n_hiddens*n_outputs+n_outputs

        # Initialize the network with new randomized weights
        self.set_weights( self.generate_weights() )
    #end

    def generate_weights(self, low=-0.1, high=0.1):
        # Generate new random weights for all the connections in the network
        if not False:
            # Support NumPy
            return [random.uniform(low,high) for _ in xrange(self.n_weights)]
        else:
            return np.random.uniform(low, high, size=(1,self.n_weights)).tolist()[0]
    #end

    def unpack(self, weight_list ):
        # This method will create a list of weight matrices. Each list element
        # corresponds to the connection between two layers.
        if self.n_hidden_layers == 0:
            return [ np.array(weight_list).reshape(self.n_inputs+1,self.n_outputs) ]
        else:
            weight_layers = [ np.array(weight_list[:(self.n_inputs+1)*self.n_hiddens]).reshape(self.n_inputs+1,self.n_hiddens) ]
            weight_layers += [ np.array(weight_list[(self.n_inputs+1)*self.n_hiddens+(i*(self.n_hiddens**2+self.n_hiddens)):(self.n_inputs+1)*self.n_hiddens+((i+1)*(self.n_hiddens**2+self.n_hiddens))]).reshape(self.n_hiddens+1,self.n_hiddens) for i in xrange(self.n_hidden_layers-1) ]
            weight_layers += [ np.array(weight_list[(self.n_inputs+1)*self.n_hiddens+((self.n_hidden_layers-1)*(self.n_hiddens**2+self.n_hiddens)):]).reshape(self.n_hiddens+1,self.n_outputs) ]

        return weight_layers
    #end

    def set_weights(self, weight_list ):
        # This is a helper method for setting the network weights to a previously defined list.
        # This is useful for utilizing a previously optimized neural network weight set.
        self.weights = self.unpack( weight_list )
    #end

    def get_weights(self, ):
        # This will stack all the weights in the network on a list, which may be saved to the disk.
        return [w for l in self.weights for w in l.flat]
    #end

    def backpropagation(self, trainingset, ERROR_LIMIT=1e-3, learning_rate=0.3, momentum_factor=0.9  ):
        def addBias(A):
            # Add 1 as bias.
            return np.hstack(( np.ones((A.shape[0],1)), A ))
        #end addBias

        assert trainingset[0].features.shape[0] == self.n_inputs, "ERROR: input size varies from the defined input setting"
        assert trainingset[0].targets.shape[0] == self.n_outputs, "ERROR: output size varies from the defined output setting"

        training_data = np.array( [instance.features for instance in trainingset ] )
        training_targets = np.array( [instance.targets for instance in trainingset ] )

        MSE      = ( ) # inf
        neterror = None
        momentum = collections.defaultdict( int )

        epoch = 0
        while MSE > ERROR_LIMIT:
            epoch += 1

            input_layers      = self.update( training_data, trace=True )
            out               = input_layers[-1]

            error             = training_targets - out
            delta             = error
            MSE               = np.mean( np.power(error,2) )

            loop  = itertools.izip(
                            xrange(len(self.weights)-1, -1, -1),
                            reversed(self.weights),
                            reversed(input_layers[:-1]),
                        )

            for i, weight_layer, input_signals in loop:
                # Loop over the weight layers in reversed order to calculate the deltas

                # Calculate weight change
                dW = learning_rate * np.dot( addBias(input_signals).T, delta ) + momentum_factor * momentum[i]

                if i!= 0:
                    """Do not calculate the delta unnecessarily."""
                    # Skipping the bias weight during calculation.
                    weight_delta = np.dot( delta, weight_layer[1:,:].T )

                    # Calculate the delta for the subsequent layer
                    delta = np.multiply(  weight_delta, self.activation_functions[i-1]( input_signals, derivative=True) )

                # Store the momentum
                momentum[i] = dW

                # Update the weights
                self.weights[ i ] += dW

            if epoch%1000==0:
                # Show the current training status
                print "* current network error (MSE):", MSE

        print "* Converged to error bound (%.4g) with MSE = %.4g." % ( ERROR_LIMIT, MSE )
        print "* Trained for %d epochs." % epoch
    # end backprop

    def update(self, input_values, trace=False ):
        # This is a forward operation in the network. This is how we calculate the network output
        # from a set of input signals.

        output = input_values
        if trace: tracelist = [ output ]

        for i, weight_layer in enumerate(self.weights):
            # Loop over the network layers and calculate the output
            output = np.dot( output, weight_layer[1:,:] ) + weight_layer[0:1,:] # implicit bias
            output = self.activation_functions[i]( output )
            if trace: tracelist.append( output )

        if trace: return tracelist

        return output
    #end

    def save_to_file(self, filename = "network.pkl" ):
        import cPickle
        """
        This save method pickles the parameters of the current network into a
        binary file for persistant storage.
        """
        with open( filename , 'wb') as file:
            store_dict = {
                "n_inputs"             : self.n_inputs,
                "n_outputs"            : self.n_outputs,
                "n_hiddens"            : self.n_hiddens,
                "n_hidden_layers"      : self.n_hidden_layers,
                "activation_functions" : self.activation_functions,
                "n_weights"            : self.n_weights,
                "weights"              : self.weights

            }
            cPickle.dump( store_dict, file, 2 )
    #end

    @staticmethod
    def load_from_file( filename = "network.pkl" ):
        """
        Load the complete configuration of a previously stored network.
        """
        network = NeuralNet( 0, 0, 0, 0, [0] )

        with open( filename , 'rb') as file:
            import cPickle
            store_dict = cPickle.load(file)

            network.n_inputs             = store_dict["n_inputs"]
            network.n_outputs            = store_dict["n_outputs"]
            network.n_hiddens            = store_dict["n_hiddens"]
            network.n_hidden_layers      = store_dict["n_hidden_layers"]
            network.n_weights            = store_dict["n_weights"]
            network.weights              = store_dict["weights"]
            network.activation_functions = store_dict["activation_functions"]

        return network
    #end
#end class<|endoftext|>


<commit_before>import json

def dispose(json_str):
    """Clear all comments in json_str.

    Clear JS-style comments like // and /**/ in json_str.
    Accept a str or unicode as input.

    Args:
        json_str: A json string of str or unicode to clean up comment

    Returns:
        str: The str without comments (or unicode if you pass in unicode)
    """
    result_str = list(json_str)
    escaped = False
    normal = True
    sl_comment = False
    ml_comment = False
    quoted = False

    a_step_from_comment = False
    a_step_from_comment_away = False

    former_index = None

    for index, char in enumerate(json_str):

        if escaped:  # We have just met a '\'
            escaped = False
            continue

        if a_step_from_comment:  # We have just met a '/'
            if char != '/' and char != '*':
                a_step_from_comment = False
                normal = True
                continue

        if char == '"':
            if normal and not escaped:
                # We are now in a string
                quoted = True
                normal = False
            elif quoted and not escaped:
                # We are now out of a string
                quoted = False
                normal = True

        elif char == '\\':
            # '\' should not take effect in comment
            if normal or quoted:
                escaped = True

        elif char == '/':
            if a_step_from_comment:
                # Now we are in single line comment
                a_step_from_comment = False
                sl_comment = True
                normal = False
                former_index = index - 1
            elif a_step_from_comment_away:
                # Now we are out of comment
                a_step_from_comment_away = False
                normal = True
                ml_comment = False
                for i in range(former_index, index + 1):
                    result_str[i] = ""

            elif normal:
                # Now we are just one step away from comment
                a_step_from_comment = True
                normal = False

        elif char == '*':
            if a_step_from_comment:
                # We are now in multi-line comment
                a_step_from_comment = False
                ml_comment = True
                normal = False
                former_index = index - 1
            elif ml_comment:
                a_step_from_comment_away = True
        elif char == '\n':
            if sl_comment:
                sl_comment = False
                normal = True
                for i in range(former_index, index + 1):
                    result_str[i] = ""
        elif char == ']' or char == '}':
            if normal:
                _remove_last_comma(result_str, index)

    # Show respect to original input if we are in python2
    return ("" if isinstance(json_str, str) else u"").join(result_str)

# There may be performance suffer backtracking the last comma
def _remove_last_comma(str_list, before_index):
    i = before_index - 1
    while str_list[i].isspace() or not str_list[i]:
        i -= 1

    # This is the first none space char before before_index
    if str_list[i] == ',':
        str_list[i] = ''

# Below are just some wrapper function around the standard json module.

def loads(text, **kwargs):
    return json.loads(dispose(text), **kwargs)

def load(fp, **kwargs):
    return loads(fp.read(), **kwargs)

def dumps(obj, **kwargs):
    return json.dumps(obj, **kwargs)

def dump(obj, fp, **kwargs):
    json.dump(obj, fp, **kwargs)
<commit_msg>Update jstyleson.py<commit_after>import json

def dispose(json_str):
    """Clear all comments in json_str.

    Clear JS-style comments like // and /**/ in json_str.
    Accept a str or unicode as input.

    Args:
        json_str: A json string of str or unicode to clean up comment

    Returns:
        str: The str without comments (or unicode if you pass in unicode)
    """
    result_str = list(json_str)
    escaped = False
    normal = True
    sl_comment = False
    ml_comment = False
    quoted = False

    a_step_from_comment = False
    a_step_from_comment_away = False

    former_index = None

    for index, char in enumerate(json_str):

        if escaped:  # We have just met a '\'
            escaped = False
            continue

        if a_step_from_comment:  # We have just met a '/'
            if char != '/' and char != '*':
                a_step_from_comment = False
                normal = True
                continue

        if a_step_from_comment_away:  # We have just met a '*'
            if char != '/':
                a_step_from_comment_away = False

        if char == '"':
            if normal and not escaped:
                # We are now in a string
                quoted = True
                normal = False
            elif quoted and not escaped:
                # We are now out of a string
                quoted = False
                normal = True

        elif char == '\\':
            # '\' should not take effect in comment
            if normal or quoted:
                escaped = True

        elif char == '/':
            if a_step_from_comment:
                # Now we are in single line comment
                a_step_from_comment = False
                sl_comment = True
                normal = False
                former_index = index - 1
            elif a_step_from_comment_away:
                # Now we are out of comment
                a_step_from_comment_away = False
                normal = True
                ml_comment = False
                for i in range(former_index, index + 1):
                    result_str[i] = ""

            elif normal:
                # Now we are just one step away from comment
                a_step_from_comment = True
                normal = False

        elif char == '*':
            if a_step_from_comment:
                # We are now in multi-line comment
                a_step_from_comment = False
                ml_comment = True
                normal = False
                former_index = index - 1
            elif ml_comment:
                a_step_from_comment_away = True
        elif char == '\n':
            if sl_comment:
                sl_comment = False
                normal = True
                for i in range(former_index, index + 1):
                    result_str[i] = ""
        elif char == ']' or char == '}':
            if normal:
                _remove_last_comma(result_str, index)

    # Show respect to original input if we are in python2
    return ("" if isinstance(json_str, str) else u"").join(result_str)

# There may be performance suffer backtracking the last comma
def _remove_last_comma(str_list, before_index):
    i = before_index - 1
    while str_list[i].isspace() or not str_list[i]:
        i -= 1

    # This is the first none space char before before_index
    if str_list[i] == ',':
        str_list[i] = ''

# Below are just some wrapper function around the standard json module.

def loads(text, **kwargs):
    return json.loads(dispose(text), **kwargs)

def load(fp, **kwargs):
    return loads(fp.read(), **kwargs)

def dumps(obj, **kwargs):
    return json.dumps(obj, **kwargs)

def dump(obj, fp, **kwargs):
    json.dump(obj, fp, **kwargs)
<|endoftext|>