<commit_before>
import tf
import rospy
from costar_robot_msgs.srv import *
from std_msgs.msg import String
from trajectory_msgs.msg import JointTrajectoryPoint
from std_srvs.srv import Empty as EmptyService
from sensor_msgs.msg import JointState
import tf_conversions.posemath as pm
import numpy as np

import PyKDL
import urdf_parser_py
from urdf_parser_py.urdf import URDF
from pykdl_utils.kdl_parser import kdl_tree_from_urdf_model
from pykdl_utils.kdl_kinematics import KDLKinematics

from costar_robot import SimplePlanning

from moveit_msgs.msg import *
from moveit_msgs.srv import *

from predicator_landmark import GetWaypointsService

class CostarArm(object):

    def __init__(self,
            base_link, end_link, planning_group,
            world="/world",
            listener=None,
            traj_step_t=0.1,
            max_acc=1,
            max_vel=1,
            max_goal_diff = 0.02,
            goal_rotation_weight = 0.01,
            max_q_diff = 1e-6,
            start_js_cb=True,
            base_steps=10,
            steps_per_meter=100):

        self.world = world
        self.base_link = base_link
        self.end_link = end_link
        self.planning_group = planning_group

        self.base_steps = base_steps
        self.steps_per_meter = steps_per_meter

        self.MAX_ACC = max_acc
        self.MAX_VEL = max_vel

        self.traj_step_t = traj_step_t

        self.max_goal_diff = max_goal_diff
        self.max_q_diff = max_q_diff
        self.goal_rotation_weight = goal_rotation_weight

        self.at_goal = True
        self.near_goal = True
        self.moving = False
        self.q0 = [0,0,0,0,0,0,0]
        self.old_q0 = [0,0,0,0,0,0,0]

        self.cur_stamp = 0

        self.teach_mode = rospy.Service('/costar/SetTeachMode',SetTeachMode,self.set_teach_mode_call)
        self.servo_mode = rospy.Service('/costar/SetServoMode',SetServoMode,self.set_servo_mode_call)
        self.shutdown = rospy.Service('/costar/ShutdownArm',EmptyService,self.shutdown_arm_call)
        self.servo = rospy.Service('/costar/ServoToPose',ServoToPose,self.servo_to_pose_call)
        self.plan = rospy.Service('/costar/PlanToPose',ServoToPose,self.plan_to_pose_call)
        self.smartmove = rospy.Service('/costar/SmartMove',SmartMove,self.smart_move_call)
        self.js_servo = rospy.Service('/costar/ServoToJointState',ServoToJointState,self.servo_to_joints_call)
        self.pt_publisher = rospy.Publisher('/joint_traj_pt_cmd',JointTrajectoryPoint,queue_size=1000)
        self.get_waypoints_srv = GetWaypointsService(world=world,service=False)
        self.driver_status = 'IDLE'
        self.status_publisher = rospy.Publisher('/costar/DriverStatus',String,queue_size=1000)
        self.robot = URDF.from_parameter_server()
        if start_js_cb:
            self.js_subscriber = rospy.Subscriber('joint_states',JointState,self.js_cb)
        self.tree = kdl_tree_from_urdf_model(self.robot)
        self.chain = self.tree.getChain(base_link, end_link)

        if listener is None:
            self.listener = tf.TransformListener()
        else:
            self.listener = listener

        #print self.tree.getNrOfSegments()
        #print self.chain.getNrOfJoints()
        self.kdl_kin = KDLKinematics(self.robot, base_link, end_link)

        self.display_pub = rospy.Publisher('costar/display_trajectory',DisplayTrajectory,queue_size=1000)

        #self.set_goal(self.q0)
        self.goal = None
        self.ee_pose = None

        self.planner = SimplePlanning(self.robot,base_link,end_link,self.planning_group)

    '''
    js_cb
    listen to robot joint state information
    '''
    def js_cb(self,msg):

        self.old_q0 = self.q0
        self.q0 = np.array(msg.position)
        self.update_position()

    '''
    update current position information
    '''
    def update_position(self):
        self.ee_pose = pm.fromMatrix(self.kdl_kin.forward(self.q0))

        if self.goal is not None:

            #goal_diff = np.abs(self.goal - self.q0).sum() / self.q0.shape[0]
            cart_diff = (self.ee_pose.p - self.goal.p).Norm()
            rot_diff = self.goal_rotation_weight * (pm.Vector(*self.ee_pose.M.GetRPY()) - pm.Vector(*self.goal.M.GetRPY())).Norm()
            goal_diff = cart_diff + rot_diff

            if goal_diff < self.max_goal_diff:
                self.at_goal = True

            if goal_diff < 10*self.max_goal_diff:
                self.near_goal = True

        q_diff = np.abs(self.old_q0 - self.q0).sum()

        if q_diff < self.max_q_diff:
            self.moving = False
        else:
            self.moving = True

    def check_req_speed_params(self,req):
        if req.accel > self.MAX_ACC:
            acceleration = self.MAX_ACC
        else:
            acceleration = req.accel
        if req.vel > self.MAX_VEL:
            velocity = self.MAX_VEL
        else:
            velocity = req.vel
        return (acceleration, velocity)

    '''
    Find any valid object that meets the requirements.
    Find a cartesian path or possibly longer path through joint space.
    '''
    def smart_move_call(self,req):

        if False and not self.driver_status == 'SERVO':
            rospy.logerr('DRIVER -- Not in servo mode!')
            return 'FAILURE - not in servo mode'

        (acceleration, velocity) = self.check_req_speed_params(req)
        (poses,names) = self.get_waypoints_srv.get_waypoints(
                req.obj_class, # object class to move to
                req.predicates, # predicates to match
                [req.pose], # offset/transform from each member of the class
                ["tmp"] # placeholder name
                )

        print req.obj_class
        print req.predicates
        print names
        print poses

        if not poses is None:
            msg = 'FAILURE - no valid objects found!'
            qs = []
            dists = []
            for (pose,name) in zip(poses,names):

                # figure out which tf frame we care about
                tf_path = name.split('/')
                tf_frame = ''
                for part in tf_path:
                    if len(part) > 0:
                        tf_frame = part
                        break

                if len(tf_frame) == 0:
                    continue

                # try to move to the pose until one succeeds
                T_base_world = pm.fromTf(self.listener.lookupTransform(self.world,self.base_link,rospy.Time(0)))
                T = T_base_world.Inverse()*pm.fromMsg(pose)

                # Check acceleration and velocity limits
                # Send command

                pt = JointTrajectoryPoint()

                #q = self.planner.ik(T, self.q0)
                traj = self.planner.getCartesianMove(T,self.q0,self.base_steps,self.steps_per_meter)
                #if len(traj.points) == 0:
                #    print T
                #    (code,res) = self.planner.getPlan(pm.toMsg(T),self.q0) # find a non-local movement
                #    traj = res.planned_trajectory.joint_trajectory

                print "Considering object with name = %s"%tf_frame

                if len(traj.points) > 0:
                    qs.append(traj)
                    dists.append((traj.points[-1].positions - self.q0).sum())
                else:
                    rospy.logwarn('SIMPLE DRIVER -- IK failed for %s'%name)

            if len(qs) == 0:
                msg = 'FAILURE - no joint configurations found!'

            possible_goals = zip(dists,qs)
            possible_goals.sort()

            print "POSSIBLE GOALS"
            print possible_goals

            (dist,traj) = possible_goals[0]
            rospy.logwarn("Trying to move to frame at distance %f"%(dist))

            msg = self.send_trajectory(traj,acceleration,velocity,cartesian=True)

            return msg

        else:
            msg = 'FAILURE - no match to predicate moves'
            return msg

    def set_goal(self,q):
        self.at_goal = False
        self.near_goal = False
        self.goal = pm.fromMatrix(self.kdl_kin.forward(q))

    '''
    Definitely do a planned motion.
    '''
    def plan_to_pose_call(self,req):
        #rospy.loginfo('Recieved servo to pose request')
        #print req
        if self.driver_status == 'SERVO':
            T = pm.fromMsg(req.target)

            # Check acceleration and velocity limits
            (acceleration, velocity) = self.check_req_speed_params(req)

            # Send command
            pt = JointTrajectoryPoint()
            traj = self.planner.getCartesianMove(T,self.q0,self.base_steps,self.steps_per_meter)
            if len(traj.points) > 0:
                (code,res) = self.planner.getPlan(req.target,traj.points[-1].positions)
            else:
                (code,res) = self.planner.getPlan(req.target,self.q0)

            #pt.positions = q
            if code > 0:
              self.at_goal = True
              return 'SUCCESS - at goal'
            if (not res is None) and len(res.planned_trajectory.joint_trajectory.points) > 0:

                disp = DisplayTrajectory()
                disp.trajectory.append(res.planned_trajectory)
                disp.trajectory_start = res.trajectory_start
                self.display_pub.publish(disp)

                traj = res.planned_trajectory.joint_trajectory

                return self.send_trajectory(traj,acceleration,velocity,cartesian=True)

            else:
                rospy.logerr(res)
                rospy.logerr('DRIVER -- PLANNING failed')
                return 'FAILURE - not in servo mode'
        else:
            rospy.logerr('DRIVER -- not in servo mode!')
            return 'FAILURE - not in servo mode'

    '''
    Send a whole joint trajectory message to a robot...
    that is listening to individual joint states.
    '''
    def send_trajectory(self,traj,acceleration=0.5,velocity=0.5,cartesian=False):
        rospy.logerr("Function 'send_trajectory' not implemented for base class!")
        return "FAILURE - running base class!"

    '''
    Send a whole sequence of points to a robot...
    that is listening to individual joint states.
    '''
    def send_sequence(self,traj,acceleration=0.5,velocity=0.5,cartesian=False):
        rospy.logerr("Function 'send_sequence' not implemented for base class!")
        return "FAILURE - running base class!"

    '''
    handle changes in driver mode
    '''
    def handle_mode_tick(self):
        rospy.logerr("Function 'handle_mode_tick' not implemented for base class!")
        return "FAILURE - running base class!"

    '''
    Standard movement call.
    Tries a cartesian move, then if that fails goes into a joint-space move.
    '''
    def servo_to_pose_call(self,req):
        #rospy.loginfo('Recieved servo to pose request')
        #print req
        if self.driver_status == 'SERVO':
            T = pm.fromMsg(req.target)

            # Check acceleration and velocity limits
            (acceleration, velocity) = self.check_req_speed_params(req)

            # inverse kinematics
            traj = self.planner.getCartesianMove(T,self.q0,self.base_steps,self.steps_per_meter)
            if len(traj.points) == 0:
                (code,res) = self.planner.getPlan(req.target,self.q0) # find a non-local movement
                if not res is None:
                    traj = res.planned_trajectory.joint_trajectory

            #if len(traj) == 0:
            #    traj.append(self.planner.ik(T,self.q0))

            # Send command
            if len(traj.points) > 0:
                rospy.logwarn("Robot moving to " + str(traj.points[-1].positions))
                return self.send_trajectory(traj,acceleration,velocity)
            else:
                rospy.logerr('SIMPLE DRIVER -- IK failed')
                return 'FAILURE - not in servo mode'
        else:
            rospy.logerr('SIMPLE DRIVER -- Not in servo mode')
            return 'FAILURE - not in servo mode'

    '''
    Standard move call.
    Make a joint space move to a destination.
    '''
    def servo_to_joints_call(self,req):

        if self.driver_status == 'SERVO':
            # Check acceleration and velocity limits
            (acceleration, velocity) = self.check_req_speed_params(req)
            return 'FAILURE - not yet implemented!'

        else:
            rospy.logerr('SIMPLE DRIVER -- Not in servo mode')
            return 'FAILURE - not in servo mode'

    '''
    set teach mode
    '''
    def set_teach_mode_call(self,req):
        if req.enable == True:

            # self.rob.set_freedrive(True)
            self.driver_status = 'TEACH'
            return 'SUCCESS - teach mode enabled'
        else:
            # self.rob.set_freedrive(False)
            self.driver_status = 'IDLE'
            return 'SUCCESS - teach mode disabled'

    '''
    send a single point
    '''
    def send_q(self,pt,acceleration,velocity):
        pt = JointTrajectoryPoint()
        pt.positions = self.q0

        self.pt_publisher.publish(pt)

    '''
    activate servo mode
    '''
    def set_servo_mode_call(self,req):
        if req.mode == 'SERVO':
            self.send_q(self.q0,0.1,0.1)

            self.driver_status = 'SERVO'
            return 'SUCCESS - servo mode enabled'
        elif req.mode == 'DISABLE':
            self.driver_status = 'IDLE'
            return 'SUCCESS - servo mode disabled'

    def shutdown_arm_call(self,req):
        self.driver_status = 'SHUTDOWN'
        pass

    '''
    robot-specific logic to update state every "tick"
    '''
    def handle_tick(self):
        rospy.logerr("Function 'handle_tick' not implemented for base class!")

    '''
    call this when "spinning" to keep updating things
    '''
    def tick(self):
        self.status_publisher.publish(self.driver_status)
        self.handle_tick()

<commit_msg>fix up arm<commit_after>
import tf
import rospy
from costar_robot_msgs.srv import *
from std_msgs.msg import String
from trajectory_msgs.msg import JointTrajectoryPoint
from std_srvs.srv import Empty as EmptyService
from sensor_msgs.msg import JointState
import tf_conversions.posemath as pm
import numpy as np

import PyKDL
import urdf_parser_py
from urdf_parser_py.urdf import URDF
from pykdl_utils.kdl_parser import kdl_tree_from_urdf_model
from pykdl_utils.kdl_kinematics import KDLKinematics

from costar_robot import SimplePlanning

from moveit_msgs.msg import *
from moveit_msgs.srv import *

from predicator_landmark import GetWaypointsService

class CostarArm(object):

    def __init__(self,
            base_link, end_link, planning_group,
            world="/world",
            listener=None,
            traj_step_t=0.1,
            max_acc=1,
            max_vel=1,
            max_goal_diff = 0.02,
            goal_rotation_weight = 0.01,
            max_q_diff = 1e-6,
            start_js_cb=True,
            base_steps=10,
            steps_per_meter=100):

        self.world = world
        self.base_link = base_link
        self.end_link = end_link
        self.planning_group = planning_group

        self.base_steps = base_steps
        self.steps_per_meter = steps_per_meter

        self.MAX_ACC = max_acc
        self.MAX_VEL = max_vel

        self.traj_step_t = traj_step_t

        self.max_goal_diff = max_goal_diff
        self.max_q_diff = max_q_diff
        self.goal_rotation_weight = goal_rotation_weight

        self.at_goal = True
        self.near_goal = True
        self.moving = False
        self.q0 = [0,0,0,0,0,0,0]
        self.old_q0 = [0,0,0,0,0,0,0]

        self.cur_stamp = 0

        self.teach_mode = rospy.Service('/costar/SetTeachMode',SetTeachMode,self.set_teach_mode_call)
        self.servo_mode = rospy.Service('/costar/SetServoMode',SetServoMode,self.set_servo_mode_call)
        self.shutdown = rospy.Service('/costar/ShutdownArm',EmptyService,self.shutdown_arm_call)
        self.servo = rospy.Service('/costar/ServoToPose',ServoToPose,self.servo_to_pose_call)
        self.plan = rospy.Service('/costar/PlanToPose',ServoToPose,self.plan_to_pose_call)
        self.smartmove = rospy.Service('/costar/SmartMove',SmartMove,self.smart_move_call)
        self.js_servo = rospy.Service('/costar/ServoToJointState',ServoToJointState,self.servo_to_joints_call)
        self.save_frame = rospy.Service('/costar/SaveFrame',SaveFrame,self.save_frame_call)
        self.save_frame = rospy.Service('/costar/SaveJointPosition',SaveFrame,self.save_joints_call)
        self.pt_publisher = rospy.Publisher('/joint_traj_pt_cmd',JointTrajectoryPoint,queue_size=1000)
        self.get_waypoints_srv = GetWaypointsService(world=world,service=False)
        self.driver_status = 'IDLE'
        self.status_publisher = rospy.Publisher('/costar/DriverStatus',String,queue_size=1000)
        self.robot = URDF.from_parameter_server()
        if start_js_cb:
            self.js_subscriber = rospy.Subscriber('joint_states',JointState,self.js_cb)
        self.tree = kdl_tree_from_urdf_model(self.robot)
        self.chain = self.tree.getChain(base_link, end_link)

        if listener is None:
            self.listener = tf.TransformListener()
        else:
            self.listener = listener

        #print self.tree.getNrOfSegments()
        #print self.chain.getNrOfJoints()
        self.kdl_kin = KDLKinematics(self.robot, base_link, end_link)

        self.display_pub = rospy.Publisher('costar/display_trajectory',DisplayTrajectory,queue_size=1000)

        #self.set_goal(self.q0)
        self.goal = None
        self.ee_pose = None

        self.planner = SimplePlanning(self.robot,base_link,end_link,self.planning_group)

    '''
    js_cb
    listen to robot joint state information
    '''
    def js_cb(self,msg):

        self.old_q0 = self.q0
        self.q0 = np.array(msg.position)
        self.update_position()

    '''
    update current position information
    '''
    def update_position(self):
        self.ee_pose = pm.fromMatrix(self.kdl_kin.forward(self.q0))

        if self.goal is not None:

            #goal_diff = np.abs(self.goal - self.q0).sum() / self.q0.shape[0]
            cart_diff = (self.ee_pose.p - self.goal.p).Norm()
            rot_diff = self.goal_rotation_weight * (pm.Vector(*self.ee_pose.M.GetRPY()) - pm.Vector(*self.goal.M.GetRPY())).Norm()
            goal_diff = cart_diff + rot_diff

            if goal_diff < self.max_goal_diff:
                self.at_goal = True

            if goal_diff < 10*self.max_goal_diff:
                self.near_goal = True

        q_diff = np.abs(self.old_q0 - self.q0).sum()

        if q_diff < self.max_q_diff:
            self.moving = False
        else:
            self.moving = True

    def check_req_speed_params(self,req):
        if req.accel > self.MAX_ACC:
            acceleration = self.MAX_ACC
        else:
            acceleration = req.accel
        if req.vel > self.MAX_VEL:
            velocity = self.MAX_VEL
        else:
            velocity = req.vel
        return (acceleration, velocity)

    '''
    Save the current end effector pose as a frame that we can return to
    '''
    def save_frame_call(self,req):
      rospy.logwarn('Save frame does not check to see if your frame already exists!')

      return 'SUCCESS - '

    '''
    Save the current joint states as a frame that we can return to
    '''
    def save_joints_call(self,req):
      rospy.logwarn('Save frame does not check to see if your joint position already exists!')

      return 'SUCCESS - '

    '''
    Find any valid object that meets the requirements.
    Find a cartesian path or possibly longer path through joint space.
    '''
    def smart_move_call(self,req):

        if False and not self.driver_status == 'SERVO':
            rospy.logerr('DRIVER -- Not in servo mode!')
            return 'FAILURE - not in servo mode'

        (acceleration, velocity) = self.check_req_speed_params(req)
        (poses,names) = self.get_waypoints_srv.get_waypoints(
                req.obj_class, # object class to move to
                req.predicates, # predicates to match
                [req.pose], # offset/transform from each member of the class
                ["tmp"] # placeholder name
                )

        print req.obj_class
        print req.predicates
        print names
        print poses

        if not poses is None:
            msg = 'FAILURE - no valid objects found!'
            qs = []
            dists = []
            for (pose,name) in zip(poses,names):

                # figure out which tf frame we care about
                tf_path = name.split('/')
                tf_frame = ''
                for part in tf_path:
                    if len(part) > 0:
                        tf_frame = part
                        break

                if len(tf_frame) == 0:
                    continue

                # try to move to the pose until one succeeds
                T_base_world = pm.fromTf(self.listener.lookupTransform(self.world,self.base_link,rospy.Time(0)))
                T = T_base_world.Inverse()*pm.fromMsg(pose)

                # Check acceleration and velocity limits
                # Send command

                pt = JointTrajectoryPoint()

                #q = self.planner.ik(T, self.q0)
                traj = self.planner.getCartesianMove(T,self.q0,self.base_steps,self.steps_per_meter)
                #if len(traj.points) == 0:
                #    print T
                #    (code,res) = self.planner.getPlan(pm.toMsg(T),self.q0) # find a non-local movement
                #    traj = res.planned_trajectory.joint_trajectory

                print "Considering object with name = %s"%tf_frame

                if len(traj.points) > 0:
                    qs.append(traj)
                    dists.append((traj.points[-1].positions - self.q0).sum())
                else:
                    rospy.logwarn('SIMPLE DRIVER -- IK failed for %s'%name)

            if len(qs) == 0:
                msg = 'FAILURE - no joint configurations found!'

            possible_goals = zip(dists,qs)
            possible_goals.sort()

            print "POSSIBLE GOALS"
            print possible_goals

            (dist,traj) = possible_goals[0]
            rospy.logwarn("Trying to move to frame at distance %f"%(dist))

            msg = self.send_trajectory(traj,acceleration,velocity,cartesian=True)

            return msg

        else:
            msg = 'FAILURE - no match to predicate moves'
            return msg

    def set_goal(self,q):
        self.at_goal = False
        self.near_goal = False
        self.goal = pm.fromMatrix(self.kdl_kin.forward(q))

    '''
    Definitely do a planned motion.
    '''
    def plan_to_pose_call(self,req):
        #rospy.loginfo('Recieved servo to pose request')
        #print req
        if self.driver_status == 'SERVO':
            T = pm.fromMsg(req.target)

            # Check acceleration and velocity limits
            (acceleration, velocity) = self.check_req_speed_params(req)

            # Send command
            pt = JointTrajectoryPoint()
            traj = self.planner.getCartesianMove(T,self.q0,self.base_steps,self.steps_per_meter)
            if len(traj.points) > 0:
                (code,res) = self.planner.getPlan(req.target,traj.points[-1].positions)
            else:
                (code,res) = self.planner.getPlan(req.target,self.q0)

            #pt.positions = q
            if code > 0:
              self.at_goal = True
              return 'SUCCESS - at goal'
            if (not res is None) and len(res.planned_trajectory.joint_trajectory.points) > 0:

                disp = DisplayTrajectory()
                disp.trajectory.append(res.planned_trajectory)
                disp.trajectory_start = res.trajectory_start
                self.display_pub.publish(disp)

                traj = res.planned_trajectory.joint_trajectory

                return self.send_trajectory(traj,acceleration,velocity,cartesian=True)

            else:
                rospy.logerr(res)
                rospy.logerr('DRIVER -- PLANNING failed')
                return 'FAILURE - not in servo mode'
        else:
            rospy.logerr('DRIVER -- not in servo mode!')
            return 'FAILURE - not in servo mode'

    '''
    Send a whole joint trajectory message to a robot...
    that is listening to individual joint states.
    '''
    def send_trajectory(self,traj,acceleration=0.5,velocity=0.5,cartesian=False):
        rospy.logerr("Function 'send_trajectory' not implemented for base class!")
        return "FAILURE - running base class!"

    '''
    Send a whole sequence of points to a robot...
    that is listening to individual joint states.
    '''
    def send_sequence(self,traj,acceleration=0.5,velocity=0.5,cartesian=False):
        rospy.logerr("Function 'send_sequence' not implemented for base class!")
        return "FAILURE - running base class!"

    '''
    handle changes in driver mode
    '''
    def handle_mode_tick(self):
        rospy.logerr("Function 'handle_mode_tick' not implemented for base class!")
        return "FAILURE - running base class!"

    '''
    Standard movement call.
    Tries a cartesian move, then if that fails goes into a joint-space move.
    '''
    def servo_to_pose_call(self,req):
        #rospy.loginfo('Recieved servo to pose request')
        #print req
        if self.driver_status == 'SERVO':
            T = pm.fromMsg(req.target)

            # Check acceleration and velocity limits
            (acceleration, velocity) = self.check_req_speed_params(req)

            # inverse kinematics
            traj = self.planner.getCartesianMove(T,self.q0,self.base_steps,self.steps_per_meter)
            if len(traj.points) == 0:
                (code,res) = self.planner.getPlan(req.target,self.q0) # find a non-local movement
                if not res is None:
                    traj = res.planned_trajectory.joint_trajectory

            #if len(traj) == 0:
            #    traj.append(self.planner.ik(T,self.q0))

            # Send command
            if len(traj.points) > 0:
                rospy.logwarn("Robot moving to " + str(traj.points[-1].positions))
                return self.send_trajectory(traj,acceleration,velocity)
            else:
                rospy.logerr('SIMPLE DRIVER -- IK failed')
                return 'FAILURE - not in servo mode'
        else:
            rospy.logerr('SIMPLE DRIVER -- Not in servo mode')
            return 'FAILURE - not in servo mode'

    '''
    Standard move call.
    Make a joint space move to a destination.
    '''
    def servo_to_joints_call(self,req):

        if self.driver_status == 'SERVO':
            # Check acceleration and velocity limits
            (acceleration, velocity) = self.check_req_speed_params(req)
            return 'FAILURE - not yet implemented!'

        else:
            rospy.logerr('SIMPLE DRIVER -- Not in servo mode')
            return 'FAILURE - not in servo mode'

    '''
    set teach mode
    '''
    def set_teach_mode_call(self,req):
        if req.enable == True:

            # self.rob.set_freedrive(True)
            self.driver_status = 'TEACH'
            return 'SUCCESS - teach mode enabled'
        else:
            # self.rob.set_freedrive(False)
            self.driver_status = 'IDLE'
            return 'SUCCESS - teach mode disabled'

    '''
    send a single point
    '''
    def send_q(self,pt,acceleration,velocity):
        pt = JointTrajectoryPoint()
        pt.positions = self.q0

        self.pt_publisher.publish(pt)

    '''
    activate servo mode
    '''
    def set_servo_mode_call(self,req):
        if req.mode == 'SERVO':
            self.send_q(self.q0,0.1,0.1)

            self.driver_status = 'SERVO'
            return 'SUCCESS - servo mode enabled'
        elif req.mode == 'DISABLE':
            self.driver_status = 'IDLE'
            return 'SUCCESS - servo mode disabled'

    def shutdown_arm_call(self,req):
        self.driver_status = 'SHUTDOWN'
        pass

    '''
    robot-specific logic to update state every "tick"
    '''
    def handle_tick(self):
        rospy.logerr("Function 'handle_tick' not implemented for base class!")

    '''
    call this when "spinning" to keep updating things
    '''
    def tick(self):
        self.status_publisher.publish(self.driver_status)
        self.handle_tick()

<|endoftext|>


<commit_before>"""Module for preparing data from NCBI. Most low layer module for manipulating data."""

import os
import pickle
from collections import defaultdict

from Bio import Entrez
from Bio import SeqIO

# TODO: move to init
CACHE_DIR = "../../Diploma/cache"
if not os.path.isdir(CACHE_DIR):
    CACHE_DIR = "cache/"
    if not os.path.isdir(CACHE_DIR):
        os.makedirs(CACHE_DIR)

# ************  NCBI RECORD OPERATIONS  ************ #

def get_gids(term="Viruses[Organism] AND srcdb_refseq[PROP] AND complete_genome"):
    """
    Get genome IDs for given search term.

    :param term: search term for NCBI query
    :return: list of genome IDs for given term
    """
    # term = "Viruses[Organism] AND srcdb_refseq[PROP] AND complete_genome"
    handle = Entrez.esearch(db="nucleotide", term=term, retmax=100000)
    record = Entrez.read(handle)
    id_list = sorted(set(record["IdList"]))
    print((record["Count"], len(record["IdList"]), len(id_list)))
    return id_list

def get_rec(rec_id):
    """
    Get record for given genome id.

    :param rec_id: genome id
    :return: record
    """
    try:
        rec = pickle.load(open(os.path.join(CACHE_DIR, "%s.pkl.gz" % rec_id), "rb"))
    except IOError:  # , FileNotFoundError:
        print(("downloading sequence id:", rec_id))
        handle = Entrez.efetch(db="nucleotide", rettype="gb", id=rec_id)
        rec = SeqIO.read(handle, "gb")
        handle.close()
        pickle.dump(rec, open(os.path.join(CACHE_DIR, "%s.pkl.gz" % rec_id), "wb"), -1)
        print(("genome size:", len(rec.seq), rec.seq[:20] + "..."))
        print(("Taxonomy:", rec.annotations['taxonomy']))
        for a, t in list(rec.annotations.items()):
            print(("  %s: %s" % (a, str(t)[:15])))
        print()
    return rec

def get_gene(rec):
    """
    Get record and return gene sequence.

    :param rec: record
    :return: gene sequence
    """
    sequence = ""
    for f in rec.features:
        if f.type == "gene":
            start = f.location.nofuzzy_start
            end = f.location.nofuzzy_end
            if f.location.strand == 1:
                sequence += rec.seq[start:end]
            else:
                # ??
                sequence += rec.seq[start:end].complement()

    return str(sequence)

def load_oid_seq_classification(ids):
    """
    Build dictionary of sequences and taxonomies for every genome ID.

    :param ids: genome IDs
    :return: sequences and taxonomy annotations dictionaries for every genome ID
    """
    seq = defaultdict(list)
    tax = {}
    for oid in ids:
        rec = get_rec(oid)
        seq[oid] = str(rec.seq)
        tax[oid] = ';'.join(rec.annotations["taxonomy"])

    return seq, tax

# ************  TAXONOMY OPERATIONS  ************ #

def rec_dd():
    """
    Create dictionary of dictionaries to 'simulate' tree.

    :return: dictionary of dictionaries
    """
    return defaultdict(rec_dd)

def update_taxonomy(taxonomy, tax_path, genome_id):
    """
    Create dictionary with taxonomy name and IDs of sequences which belongs to specific taxonomy.

    :param taxonomy: current taxonomy
    :param tax_path: taxonomy path
    :param genome_id: genome_id
    :return: updated taxonomy
    """
    if not tax_path:
        return taxonomy

    tax = tax_path[0].lower()
    if tax in taxonomy:  # check if tax in taxonomy and update
        # temp_taxonomy[tax]["data"].append(seq_record.annotations["gi"])
        taxonomy[tax]["data"].append(genome_id)
        # taxonomy[tax]["data"].append(get_gene(rec))
        update_taxonomy(taxonomy[tax], tax_path[1:], genome_id)
    else:
        # temp_taxonomy[tax] = {"data": list({seq_record.annotations["gi"]})}
        taxonomy[tax] = {"data": list({genome_id})}
        # taxonomy[tax] = dict({"data": list({get_gene(rec)})})
        temp = update_taxonomy(taxonomy[tax], tax_path[1:], genome_id)
        if len(temp) > 1:  # 1 = data, 2 = data + key
            taxonomy = temp
    return taxonomy

def filter_classification(rec, to_filter):
    """
    Check if record is in filter list.

    :param rec: record
    :param to_filter: filter list
    :return: bool
    """
    in_to_filter = False
    for temp_tax in rec.annotations["taxonomy"]:
        temp_tax = temp_tax.lower().split()
        for temp_tax_el in temp_tax:
            if temp_tax_el in to_filter:
                in_to_filter = True
                print("filtered ", rec.annotations["taxonomy"])
    return in_to_filter

def print_nice(taxonomy, level=0):
    """
    Print taxonomy with tabs.

    :param taxonomy: taxonomy
    :param level: current level
    :return:
    """
    for i in sorted(taxonomy.keys()):
        if i == "data":
            if len(taxonomy) == 1:
                return
            else:
                continue
        else:
            print(level * "\t", i.replace("->", "", 1), len(taxonomy[i]["data"]))
            print_nice(taxonomy[i], level + 1)

def load_whole_taxonomy():
    """
    Build taxonomy and get list ids and labels.

    :return: data, label
    """
    taxonomy = get_taxonomy(get_gids())
    list_nodes = get_list_nodes_ids_labels(taxonomy)
    data, labels = list(zip(*list_nodes))
    for label in labels:
        print(label)
    label_number = -1
    temp_l = []
    label_n = []
    for l in labels:
        if l not in temp_l:
            temp_l.append(l)
            label_number += 1
        label_n.append(label_number)

    return data, label_n

def get_taxonomy(id_list, count=-1):
    # call: python get_viral_sequence.py>log.out 2>log.err

    # all virus sequences
    # term = "Viruses[Organism] NOT srcdb_refseq[PROP] NOT cellular organisms[ORGN] AND
    #           nuccore genome samespecies[Filter] NOT nuccore genome[filter] NOT gbdiv syn[prop]"
    # only reference (refSEQ) virues sequences
    # see distinction between the two, here:
    # http://www.ncbi.nlm.nih.gov/genomes/GenomesHome.cgi?taxid=10239&hopt=faq
    """
    Build taxonomy from Entrez search.

    :param id_list: list of genome ids we want to build taxonomy tree from
    :param count: how many elements we want in taxonomy; -1 means whole taxonomy
    :return: taxonomy
    """
    taxonomy = rec_dd()
    temp_count = 1
    for genome_id in id_list:
        try:
            rec = get_rec(genome_id)
            in_filter = filter_classification(rec, list({"bacteria", "unclassified", "unassigned"}))
            if not in_filter:
                update_taxonomy(taxonomy, rec.annotations["taxonomy"], genome_id)

                if count != -1:
                    if temp_count == count:
                        break
                    temp_count += 1
        except IOError as e:
            # efetch - Raises an IOError exception if there's a network error.
            # http://biopython.org/DIST/docs/api/Bio.Entrez-module.html
            print("IOError raised...")
            print(e)
        except ValueError as v:
            # http: // biopython.org / DIST / docs / api / Bio.SeqIO - module.html  # read
            print("problems with handling SeqIO...")
            print(v)
        except pickle.PicklingError as p:
            # https://docs.python.org/2/library/pickle.html#pickle.PicklingError
            print("problems with pickling object...")
            print(p)

    return taxonomy

def remove_small_nodes(taxonomy, threshold_size=100):
    """
    Remove small nodes from dataset.
    :param taxonomy: input taxonomy
    :param threshold_size: how many nodes do parent need to keep it
    :return: output taxonomy
    """
    if isinstance(taxonomy, (defaultdict, dict)):
        taxonomy_keys = [x for x in list(taxonomy.keys()) if x != "data"]
        for i in taxonomy_keys:
            print(i, len(taxonomy[i]['data']))
            if len(taxonomy[i]['data']) < threshold_size:
                taxonomy.pop(i)
            else:
                remove_small_nodes(taxonomy[i])
    else:
        return taxonomy

# ************  LIST OPERATIONS  ************ #

def remove_lists(taxonomy):
    """
    Remove all list nodes from taxonomy.

    :param taxonomy: taxonomy
    :return: taxonomy
    """
    # check for recurse exit
    if isinstance(taxonomy, (defaultdict, dict)):
        for i in [x for x in list(taxonomy.keys()) if x != "data"]:
            if set(taxonomy[i]) == set(list({"data"})):
                # if parent has only one list node, remove it
                # if len([x for x in taxonomy.keys() if x != "data"]) == 1:
                taxonomy.pop(i)
                continue
            else:
                remove_lists(taxonomy[i])
    else:
        return taxonomy

def get_list_nodes_unique(taxonomy, parent=""):
    """
    Get taxonomy and return unique list nodes.

    :param taxonomy: taxonomy
    :param parent: parent of current node
    :return: unique list nodes
    """
    # checked by hand and it works as expected
    list_nodes = list()
    keys = [x for x in list(taxonomy.keys()) if x != "data"]
    for i in keys:
        if set(taxonomy[i]) == set(list({"data"})):
            list_nodes.append(i)
        else:
            list_nodes += get_list_nodes_unique(taxonomy[i], parent + "->" + i)
    return list_nodes

def count_list_nodes(taxonomy):
    """
    Count list nodes and return sum.

    :param taxonomy: taxonomy
    :return: int
    """
    count = 0
    keys = [x for x in list(taxonomy.keys()) if x != "data"]
    for i in keys:
        if set(taxonomy[i]) == set(list({"data"})):
            if i == keys[-1]:
                count += 1
                return count
            else:
                count += 1
        else:
            count += count_list_nodes(taxonomy[i])
    return count

def get_list_nodes_ids_labels(taxonomy):
    """
    Get taxonomy and return tuples of all list nodes.

    :param taxonomy: taxonomy
    :return: list of tuples (id, class)
    """
    if len(list(taxonomy.keys())) > 1 or list(taxonomy.keys()) == ["viruses"]:
        temp = []
        for k in [x for x in list(taxonomy.keys()) if x != "data"]:
            temp += get_list_nodes_ids_labels(taxonomy[k])
        return temp
    # else:
    #     return [(x, parent) for x in taxonomy["data"]]

# ************  ALL NODES OPERATIONS  ************ #

def count_examples(taxonomy):
    """
    Get taxonomy, count examples in every node and return sum.

    :param taxonomy: taxonomy
    :return: sum of examples
    """
    count = 0
    keys = [x for x in list(taxonomy.keys()) if x != "data"]
    for i in keys:
        if set(taxonomy[i]) == set(list({"data"})):
            if i == keys[-1]:
                count += len(taxonomy[i]["data"])
                return count
            else:
                count += len(taxonomy[i]["data"])
        else:
            count += count_examples(taxonomy[i])
    return count

def get_all_nodes(taxonomy, parent=""):
    """
    Get taxonomy and return all nodes (including list nodes).

    :param parent: parent of current node - default ""
    :param taxonomy: taxonomy
    :return: all nodes
    """
    all_nodes = list()
    keys = [x for x in list(taxonomy.keys()) if x != "data"]
    for i in keys:
        # if we want all non-list nodes, than this stays, otherwise comment this
        # if len([x for x in taxonomy[i].keys() if x != "data"]) == 0:
        # continue
        if i == "rest":
            all_nodes.append(parent + "->" + i)
        else:
            all_nodes.append(i)
        all_nodes += get_all_nodes(taxonomy[i], i)
    return all_nodes

# ************  OTHER  ************ #

def load_seqs_from_ncbi(seq_len=100, skip_read=0, overlap=50, taxonomy_el_count=-1):
    """
    Load sequences from NCBI database.

    Prepare sequences, sliced to seq_len length. Skip every skip_read and overlap two reads with overlap nucleotides.
    Overlap 50 means that half of the read is going to be overlapped with next read.
    If seq_len is -1, load whole sequences (do not strip them) - usually using with fasta format as we slice sequences
    later.
    :param seq_len: read length
    :param skip_read: number of skipped reads
    :param overlap: overlapping nucleotides count
    :param taxonomy_el_count: how many elements we want in taxonomy; -1 means whole taxonomy
    :return:    dictionary reads - each genome ID key contains list of reads for specific genome,
                dictionary taxonomy - each genome ID key contains taxonomy for specific genome
    """
    data, _ = run(taxonomy_el_count)
    print("getting sequences...")
    seqs, tax = load_oid_seq_classification(data)

    reads = defaultdict(list)

    if seq_len != -1:
        for oid, seq in seqs.items():
            while seq:
                if len(seq) < seq_len:
                    # we don't want shorter sequences than seq_len (shorter than 100)
                    break
                reads[oid].append(seq[:seq_len])
                seq = seq[seq_len - overlap + ((seq_len - overlap) * skip_read):]
    else:
        reads = seqs

    return reads, tax

def run(taxonomy_el_count=-1):
    """
    Build taxonomy and get list ids and labels.

    :param taxonomy_el_count: how many elements we want in taxonomy; -1 means whole taxonomy
    :return: data, label
    """
    taxonomy = get_taxonomy(get_gids(), count=taxonomy_el_count)
    # remove_lists(taxonomy)
    print_nice(taxonomy)
    remove_small_nodes(taxonomy, 100)
    print_nice(taxonomy)
    list_nodes = get_list_nodes_ids_labels(taxonomy)
    data, labels = list(zip(*list_nodes))
    # for label in labels:
    #     print label
    label_number = -1
    temp_l = []
    label_n = []
    for l in labels:
        if l not in temp_l:
            temp_l.append(l)
            label_number += 1
        label_n.append(label_number)

    return data, label_n

if __name__ == "__main__":
    # a = load_seqs_from_ncbi(taxonomy_el_count=20)
    temp_taxonomy = get_taxonomy(get_gids())
    print("no of examples after taxonomy was built: %d" % count_examples(temp_taxonomy))
    print("no of list nodes after taxonomy was built: %d" % count_list_nodes(temp_taxonomy))
    print_nice(temp_taxonomy)
    remove_small_nodes(temp_taxonomy, 100)
    # remove_lists(temp_taxonomy)
    print_nice(temp_taxonomy)
    run()
<commit_msg>another pylint fix<commit_after>"""Module for preparing data from NCBI. Most low layer module for manipulating data."""

import os
import pickle
from collections import defaultdict

from Bio import Entrez
from Bio import SeqIO

# TODO: move to init
CACHE_DIR = "../../Diploma/cache"
if not os.path.isdir(CACHE_DIR):
    CACHE_DIR = "cache/"
    if not os.path.isdir(CACHE_DIR):
        os.makedirs(CACHE_DIR)

# ************  NCBI RECORD OPERATIONS  ************ #

def get_gids(term="Viruses[Organism] AND srcdb_refseq[PROP] AND complete_genome"):
    """
    Get genome IDs for given search term.

    :param term: search term for NCBI query
    :return: list of genome IDs for given term
    """
    # term = "Viruses[Organism] AND srcdb_refseq[PROP] AND complete_genome"
    handle = Entrez.esearch(db="nucleotide", term=term, retmax=100000)
    record = Entrez.read(handle)
    id_list = sorted(set(record["IdList"]))
    print((record["Count"], len(record["IdList"]), len(id_list)))
    return id_list

def get_rec(rec_id):
    """
    Get record for given genome id.

    :param rec_id: genome id
    :return: record
    """
    try:
        rec = pickle.load(open(os.path.join(CACHE_DIR, "%s.pkl.gz" % rec_id), "rb"))
    except IOError:  # , FileNotFoundError:
        print(("downloading sequence id:", rec_id))
        handle = Entrez.efetch(db="nucleotide", rettype="gb", id=rec_id)
        rec = SeqIO.read(handle, "gb")
        handle.close()
        pickle.dump(rec, open(os.path.join(CACHE_DIR, "%s.pkl.gz" % rec_id), "wb"), -1)
        print(("genome size:", len(rec.seq), rec.seq[:20] + "..."))
        print(("Taxonomy:", rec.annotations['taxonomy']))
        for a, t in list(rec.annotations.items()):
            print(("  %s: %s" % (a, str(t)[:15])))
        print()
    return rec

def get_gene(rec):
    """
    Get record and return gene sequence.

    :param rec: record
    :return: gene sequence
    """
    sequence = ""
    for f in rec.features:
        if f.type == "gene":
            start = f.location.nofuzzy_start
            end = f.location.nofuzzy_end
            if f.location.strand == 1:
                sequence += rec.seq[start:end]
            else:
                # ??
                sequence += rec.seq[start:end].complement()

    return str(sequence)

def load_oid_seq_classification(ids):
    """
    Build dictionary of sequences and taxonomies for every genome ID.

    :param ids: genome IDs
    :return: sequences and taxonomy annotations dictionaries for every genome ID
    """
    seq = defaultdict(list)
    tax = {}
    for oid in ids:
        rec = get_rec(oid)
        seq[oid] = str(rec.seq)
        tax[oid] = ';'.join(rec.annotations["taxonomy"])

    return seq, tax

# ************  TAXONOMY OPERATIONS  ************ #

def rec_dd():
    """
    Create dictionary of dictionaries to 'simulate' tree.

    :return: dictionary of dictionaries
    """
    return defaultdict(rec_dd)

def update_taxonomy(taxonomy, tax_path, genome_id):
    """
    Create dictionary with taxonomy name and IDs of sequences which belongs to specific taxonomy.

    :param taxonomy: current taxonomy
    :param tax_path: taxonomy path
    :param genome_id: genome_id
    :return: updated taxonomy
    """
    if not tax_path:
        return taxonomy

    tax = tax_path[0].lower()
    if tax in taxonomy:  # check if tax in taxonomy and update
        # temp_taxonomy[tax]["data"].append(seq_record.annotations["gi"])
        taxonomy[tax]["data"].append(genome_id)
        # taxonomy[tax]["data"].append(get_gene(rec))
        update_taxonomy(taxonomy[tax], tax_path[1:], genome_id)
    else:
        # temp_taxonomy[tax] = {"data": list({seq_record.annotations["gi"]})}
        taxonomy[tax] = {"data": list({genome_id})}
        # taxonomy[tax] = dict({"data": list({get_gene(rec)})})
        temp = update_taxonomy(taxonomy[tax], tax_path[1:], genome_id)
        if len(temp) > 1:  # 1 = data, 2 = data + key
            taxonomy = temp
    return taxonomy

def filter_classification(rec, to_filter):
    """
    Check if record is in filter list.

    :param rec: record
    :param to_filter: filter list
    :return: bool
    """
    in_to_filter = False
    for temp_tax in rec.annotations["taxonomy"]:
        temp_tax = temp_tax.lower().split()
        for temp_tax_el in temp_tax:
            if temp_tax_el in to_filter:
                in_to_filter = True
                print("filtered ", rec.annotations["taxonomy"])
    return in_to_filter

def print_nice(taxonomy, level=0):
    """
    Print taxonomy with tabs.

    :param taxonomy: taxonomy
    :param level: current level
    :return:
    """
    for i in sorted(taxonomy.keys()):
        if i == "data":
            if len(taxonomy) == 1:
                return
            else:
                continue
        else:
            print(level * "\t", i.replace("->", "", 1), len(taxonomy[i]["data"]))
            print_nice(taxonomy[i], level + 1)

def load_whole_taxonomy():
    """
    Build taxonomy and get list ids and labels.

    :return: data, label
    """
    taxonomy = get_taxonomy(get_gids())
    list_nodes = get_list_nodes_ids_labels(taxonomy)
    data, labels = list(zip(*list_nodes))
    for label in labels:
        print(label)
    label_number = -1
    temp_l = []
    label_n = []
    for l in labels:
        if l not in temp_l:
            temp_l.append(l)
            label_number += 1
        label_n.append(label_number)

    return data, label_n

def get_taxonomy(id_list, count=-1):
    # call: python get_viral_sequence.py>log.out 2>log.err

    # all virus sequences
    # term = "Viruses[Organism] NOT srcdb_refseq[PROP] NOT cellular organisms[ORGN] AND
    #           nuccore genome samespecies[Filter] NOT nuccore genome[filter] NOT gbdiv syn[prop]"
    # only reference (refSEQ) virues sequences
    # see distinction between the two, here:
    # http://www.ncbi.nlm.nih.gov/genomes/GenomesHome.cgi?taxid=10239&hopt=faq
    """
    Build taxonomy from Entrez search.

    :param id_list: list of genome ids we want to build taxonomy tree from
    :param count: how many elements we want in taxonomy; -1 means whole taxonomy
    :return: taxonomy
    """
    taxonomy = rec_dd()
    temp_count = 1
    for genome_id in id_list:
        try:
            rec = get_rec(genome_id)
            in_filter = filter_classification(rec, list({"bacteria", "unclassified", "unassigned"}))
            if not in_filter:
                update_taxonomy(taxonomy, rec.annotations["taxonomy"], genome_id)

                if count != -1:
                    if temp_count == count:
                        break
                    temp_count += 1
        except IOError as e:
            # efetch - Raises an IOError exception if there's a network error.
            # http://biopython.org/DIST/docs/api/Bio.Entrez-module.html
            print("IOError raised...")
            print(e)
        except ValueError as v:
            # http: // biopython.org / DIST / docs / api / Bio.SeqIO - module.html  # read
            print("problems with handling SeqIO...")
            print(v)
        except pickle.PicklingError as p:
            # https://docs.python.org/2/library/pickle.html#pickle.PicklingError
            print("problems with pickling object...")
            print(p)

    return taxonomy

def remove_small_nodes(taxonomy, threshold_size=100):
    """
    Remove small nodes from dataset.

    :param taxonomy: input taxonomy
    :param threshold_size: how many nodes do parent need to keep it
    :return: output taxonomy
    """
    if isinstance(taxonomy, (defaultdict, dict)):
        taxonomy_keys = [x for x in list(taxonomy.keys()) if x != "data"]
        for i in taxonomy_keys:
            print(i, len(taxonomy[i]['data']))
            if len(taxonomy[i]['data']) < threshold_size:
                taxonomy.pop(i)
            else:
                remove_small_nodes(taxonomy[i])
    else:
        return taxonomy

# ************  LIST OPERATIONS  ************ #

def remove_lists(taxonomy):
    """
    Remove all list nodes from taxonomy.

    :param taxonomy: taxonomy
    :return: taxonomy
    """
    # check for recurse exit
    if isinstance(taxonomy, (defaultdict, dict)):
        for i in [x for x in list(taxonomy.keys()) if x != "data"]:
            if set(taxonomy[i]) == set(list({"data"})):
                # if parent has only one list node, remove it
                # if len([x for x in taxonomy.keys() if x != "data"]) == 1:
                taxonomy.pop(i)
                continue
            else:
                remove_lists(taxonomy[i])
    else:
        return taxonomy

def get_list_nodes_unique(taxonomy, parent=""):
    """
    Get taxonomy and return unique list nodes.

    :param taxonomy: taxonomy
    :param parent: parent of current node
    :return: unique list nodes
    """
    # checked by hand and it works as expected
    list_nodes = list()
    keys = [x for x in list(taxonomy.keys()) if x != "data"]
    for i in keys:
        if set(taxonomy[i]) == set(list({"data"})):
            list_nodes.append(i)
        else:
            list_nodes += get_list_nodes_unique(taxonomy[i], parent + "->" + i)
    return list_nodes

def count_list_nodes(taxonomy):
    """
    Count list nodes and return sum.

    :param taxonomy: taxonomy
    :return: int
    """
    count = 0
    keys = [x for x in list(taxonomy.keys()) if x != "data"]
    for i in keys:
        if set(taxonomy[i]) == set(list({"data"})):
            if i == keys[-1]:
                count += 1
                return count
            else:
                count += 1
        else:
            count += count_list_nodes(taxonomy[i])
    return count

def get_list_nodes_ids_labels(taxonomy):
    """
    Get taxonomy and return tuples of all list nodes.

    :param taxonomy: taxonomy
    :return: list of tuples (id, class)
    """
    if len(list(taxonomy.keys())) > 1 or list(taxonomy.keys()) == ["viruses"]:
        temp = []
        for k in [x for x in list(taxonomy.keys()) if x != "data"]:
            temp += get_list_nodes_ids_labels(taxonomy[k])
        return temp
    # else:
    #     return [(x, parent) for x in taxonomy["data"]]

# ************  ALL NODES OPERATIONS  ************ #

def count_examples(taxonomy):
    """
    Get taxonomy, count examples in every node and return sum.

    :param taxonomy: taxonomy
    :return: sum of examples
    """
    count = 0
    keys = [x for x in list(taxonomy.keys()) if x != "data"]
    for i in keys:
        if set(taxonomy[i]) == set(list({"data"})):
            if i == keys[-1]:
                count += len(taxonomy[i]["data"])
                return count
            else:
                count += len(taxonomy[i]["data"])
        else:
            count += count_examples(taxonomy[i])
    return count

def get_all_nodes(taxonomy, parent=""):
    """
    Get taxonomy and return all nodes (including list nodes).

    :param parent: parent of current node - default ""
    :param taxonomy: taxonomy
    :return: all nodes
    """
    all_nodes = list()
    keys = [x for x in list(taxonomy.keys()) if x != "data"]
    for i in keys:
        # if we want all non-list nodes, than this stays, otherwise comment this
        # if len([x for x in taxonomy[i].keys() if x != "data"]) == 0:
        # continue
        if i == "rest":
            all_nodes.append(parent + "->" + i)
        else:
            all_nodes.append(i)
        all_nodes += get_all_nodes(taxonomy[i], i)
    return all_nodes

# ************  OTHER  ************ #

def load_seqs_from_ncbi(seq_len=100, skip_read=0, overlap=50, taxonomy_el_count=-1):
    """
    Load sequences from NCBI database.

    Prepare sequences, sliced to seq_len length. Skip every skip_read and overlap two reads with overlap nucleotides.
    Overlap 50 means that half of the read is going to be overlapped with next read.
    If seq_len is -1, load whole sequences (do not strip them) - usually using with fasta format as we slice sequences
    later.
    :param seq_len: read length
    :param skip_read: number of skipped reads
    :param overlap: overlapping nucleotides count
    :param taxonomy_el_count: how many elements we want in taxonomy; -1 means whole taxonomy
    :return:    dictionary reads - each genome ID key contains list of reads for specific genome,
                dictionary taxonomy - each genome ID key contains taxonomy for specific genome
    """
    data, _ = run(taxonomy_el_count)
    print("getting sequences...")
    seqs, tax = load_oid_seq_classification(data)

    reads = defaultdict(list)

    if seq_len != -1:
        for oid, seq in seqs.items():
            while seq:
                if len(seq) < seq_len:
                    # we don't want shorter sequences than seq_len (shorter than 100)
                    break
                reads[oid].append(seq[:seq_len])
                seq = seq[seq_len - overlap + ((seq_len - overlap) * skip_read):]
    else:
        reads = seqs

    return reads, tax

def run(taxonomy_el_count=-1):
    """
    Build taxonomy and get list ids and labels.

    :param taxonomy_el_count: how many elements we want in taxonomy; -1 means whole taxonomy
    :return: data, label
    """
    taxonomy = get_taxonomy(get_gids(), count=taxonomy_el_count)
    # remove_lists(taxonomy)
    print_nice(taxonomy)
    remove_small_nodes(taxonomy, 100)
    print_nice(taxonomy)
    list_nodes = get_list_nodes_ids_labels(taxonomy)
    data, labels = list(zip(*list_nodes))
    # for label in labels:
    #     print label
    label_number = -1
    temp_l = []
    label_n = []
    for l in labels:
        if l not in temp_l:
            temp_l.append(l)
            label_number += 1
        label_n.append(label_number)

    return data, label_n

if __name__ == "__main__":
    # a = load_seqs_from_ncbi(taxonomy_el_count=20)
    temp_taxonomy = get_taxonomy(get_gids())
    print("no of examples after taxonomy was built: %d" % count_examples(temp_taxonomy))
    print("no of list nodes after taxonomy was built: %d" % count_list_nodes(temp_taxonomy))
    print_nice(temp_taxonomy)
    remove_small_nodes(temp_taxonomy, 100)
    # remove_lists(temp_taxonomy)
    print_nice(temp_taxonomy)
    run()
<|endoftext|>


<commit_before>#!/usr/bin/env python
#
# Copyright (C) 2011, 2012, 2013  Strahinja Val Markovic  <val@markovic.io>
#
# This file is part of YouCompleteMe.
#
# YouCompleteMe is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# YouCompleteMe is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with YouCompleteMe.  If not, see <http://www.gnu.org/licenses/>.

import vim
import vimsupport
from completers.completer import Completer

class OmniCompleter( Completer ):
  def __init__( self ):
    super( OmniCompleter, self ).__init__()
    self.omnifunc = None
    self.stored_candidates = None

  def SupportedFiletypes( self ):
    return []

  def ShouldUseNowInner( self, start_column ):
    if not self.omnifunc:
      return False
    return super( OmniCompleter, self ).ShouldUseNowInner( start_column )

  def CandidatesForQueryAsyncInner( self, query ):
    if not self.omnifunc:
      self.stored_candidates = None
      return

    return_value = vim.eval( self.omnifunc + '(1,"")' )
    if return_value < 0:
      self.stored_candidates = None
      return

    omnifunc_call = [ self.omnifunc,
                      "(0,'",
                      vimsupport.EscapeForVim( query ),
                      "')" ]

    self.stored_candidates = vim.eval( ''.join( omnifunc_call ) )

  def AsyncCandidateRequestReadyInner( self ):
    return True

  def OnFileReadyToParse( self ):
    self.omnifunc = vim.eval( '&omnifunc' )

  def CandidatesFromStoredRequestInner( self ):
    return self.stored_candidates

<commit_msg>Handle omnifunc returning dict with 'words'<commit_after>#!/usr/bin/env python
#
# Copyright (C) 2011, 2012, 2013  Strahinja Val Markovic  <val@markovic.io>
#
# This file is part of YouCompleteMe.
#
# YouCompleteMe is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# YouCompleteMe is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with YouCompleteMe.  If not, see <http://www.gnu.org/licenses/>.

import vim
import vimsupport
from completers.completer import Completer

class OmniCompleter( Completer ):
  def __init__( self ):
    super( OmniCompleter, self ).__init__()
    self.omnifunc = None
    self.stored_candidates = None

  def SupportedFiletypes( self ):
    return []

  def ShouldUseNowInner( self, start_column ):
    if not self.omnifunc:
      return False
    return super( OmniCompleter, self ).ShouldUseNowInner( start_column )

  def CandidatesForQueryAsyncInner( self, query ):
    if not self.omnifunc:
      self.stored_candidates = None
      return

    return_value = vim.eval( self.omnifunc + '(1,"")' )
    if return_value < 0:
      self.stored_candidates = None
      return

    omnifunc_call = [ self.omnifunc,
                      "(0,'",
                      vimsupport.EscapeForVim( query ),
                      "')" ]

    items = vim.eval( ''.join( omnifunc_call ) )
    self.stored_candidates = items.words if hasattr( items, 'words' ) else items

  def AsyncCandidateRequestReadyInner( self ):
    return True

  def OnFileReadyToParse( self ):
    self.omnifunc = vim.eval( '&omnifunc' )

  def CandidatesFromStoredRequestInner( self ):
    return self.stored_candidates

<|endoftext|>


<commit_before># overall python packages
import glob
import astropy.io.fits as fits
import os
import time
import numpy as n
import sys

# specific functions
from scipy.stats import norm
from scipy.integrate import quad
from scipy.interpolate import interp1d

# dedicated packages
import ClusterScalingRelations
cl = ClusterScalingRelations.ClusterScalingRelations_Mantz2016()
import StellarMass
import XrayLuminosity
xr = XrayLuminosity.XrayLuminosity()
dV = -9999

from numpy import exp, random

logmbh = lambda logm200c : 8.18 + 1.55 * ( logm200c - 13. )

LX_Z_lo, FRACTION_Z_lo = n.loadtxt(os.path.join(os.environ['GIT_NBODY_DIR'],'data','fraction_zgt0_zlt10.txt' ), unpack=True)
LX_Z_me, FRACTION_Z_me = n.loadtxt(os.path.join(os.environ['GIT_NBODY_DIR'],'data','fraction_zgt10_zlt25.txt'), unpack=True)
LX_Z_hi, FRACTION_Z_hi = n.loadtxt(os.path.join(os.environ['GIT_NBODY_DIR'],'data','fraction_zgt25_zlt50.txt'), unpack=True)
zmin_obscur = n.array([-0.01,1.0,2.5])
zmax_obscur = n.array([1.0,2.5,10.0])

# read the Xray AGN luminosity function and add a condition to reproduce it

def create_catalogs_out(fileList, z, snap_name):
	"""
	Adds Xray emission mass using the Bongiorno et al. 2016 model to the rockstar outputs.
	"""
	#
	out_duty_cycle = os.path.join(os.environ['MD10'],"duty_cycle", snap_name + "_duty_cycle.txt")
	log_stellar_mass, duty_cycle = n.loadtxt(out_duty_cycle, unpack="True")
	percentage_active = interp1d(n.hstack((-200., 0,n.min(log_stellar_mass)-0.01,log_stellar_mass,n.max(log_stellar_mass)+0.01,15)), n.hstack(( 0., 0., 0., duty_cycle, 0., 0.)))

	# set up the x ray lambda SAR
	logMs = n.arange(4.5,14.5,0.01)
	cdfs_interpolations = []
	XXS = n.arange(32,36.1,0.1)
	for mass in logMs:
		norming = xr.Phi_stellar_mass(mass, z)
		cdfs_interpolations.append( interp1d(n.hstack((n.array([xr.Phi_stellar_mass_to_X(X, mass, z) for X in XXS ])/norming, 1.)), n.hstack((XXS, XXS[-1]+0.1))) )

	cdfs_interpolations = n.array(cdfs_interpolations)

	if z > zmin_obscur[0] and z<zmax_obscur[0]:
		obscured_fraction_interpolated = interp1d(LX_Z_lo, FRACTION_Z_lo)
	elif z > zmin_obscur[1] and z<zmax_obscur[1]:
		obscured_fraction_interpolated = interp1d(LX_Z_me, FRACTION_Z_me)
	else:
		obscured_fraction_interpolated = interp1d(LX_Z_hi, FRACTION_Z_hi)

	# loops over files
	for fileName in fileList:
		t0=time.time()
		outFile = fileName[:-5]+"_Xray.fits"
		# opens all relevant files
		msFile = fileName[:-5]+"_Ms.fits"
		hd = fits.open(fileName)
		hm = fits.open(msFile)

		stellar_mass = hm[1].data['stellar_mass_Mo13_mvir']
		selection = hm[1].data['stellar_mass_reliable']
		Nhalo=len(stellar_mass)
		randomX = n.random.rand(Nhalo)
		active_gn = ( percentage_active(stellar_mass) > randomX )
		# lambda SAR addition
		indexes = n.searchsorted(logMs,stellar_mass)
		indexes[selection] = n.zeros_like(indexes[selection])
		lambda_sar_Bo16 = n.array([ cdfs_interpolations[indexes[ii]](randomX[ii]) for ii in range(Nhalo) ])
		# obscuration Merloni et al. 2015 as a function of luminosity
		# randomObscuration = n.random.rand(Nhalo)
		# obscured = (xr.obscured_fraction_optical_Merloni2015(lambda_sar_Bo16 + stellar_mass) < randomObscuration )

		# obscuration Buchner et al. 2015 + 2017
		# add the log NH of the logNH_host
		# 35 % have a thick obscuration 24 - 26
		# 65 % have a thin obscuration that depends on stellar mass and Xray luminosity
		logNH = random.uniform(20, 22,len(Nhalo))
		obs_type = n.zeros(len(Nhalo))
		# 35% of thick, 24-26
		randomNH = n.random.rand(Nhalo)
		thick_obscuration = (randomNH < 0.35)
		thin_obscuration = (randomNH >= 0.35)
		logNH[thick_obscuration] = random.uniform(24, 26, len(logNH[thick_obscuration]))
		obs_type[thick_obscuration] = n.ones_like(logNH[thick_obscuration]))*2
		# the thin : about 40 % are thin whatever happens: 22-24
		logNH_host_mean = 21.7 + (stellar_mass - 9.5)*0.38
		logNH_host = random.gauss(logNH_host_mean, 0.5)
		logNH[(thin_obscuration)&(logNH_host>=22)] = random.uniform(22, 24, len(logNH[(thin_obscuration)&(logNH_host>=22)]))
		obs_type[(thin_obscuration)&(logNH_host>=22)] =  n.ones_like(logNH[(thin_obscuration)&(logNH_host>=22)])
		# a few more are thin depending on their Xray luminosity: 22-24
		rest = (thin_obscuration)&(logNH_host<22)
		randomNH2 = n.random.rand(Nhalo)
		rest_obscured = (rest)&(randomNH2 < obscured_fraction_interpolated(lambda_sar_Bo16 + stellar_mass))
		logNH[(rest_obscured)] = random.uniform(22, 24, len(logNH[(rest_obscured)]))
		obs_type[(rest_obscured)] =  n.ones_like(logNH[(rest_obscured)])
		# the rest has: 20-22 by default
		print logNH, obs_type

		# columns related to Xray AGN
		col1 = fits.Column(name='lambda_sar_Bo16',format='D', array = lambda_sar_Bo16 )
		col2 = fits.Column(name='activity',format='L', array = active_gn )
		col3 = fits.Column(name='obscuration_type_Buchner2017',format='K', array = obs_type.astype('int') )
		#col4 = fits.Column(name='M_BH'  ,format='D', array =logmbh(hd[1].data['M200c'])))
		col4 = fits.Column(name='log_NH_Buchner2017'  ,format='D', array = logNH)

		# columns related to clusters
		col5 = fits.Column(name='Mgas_cluster'  ,format='D', array =n.log10(cl.logM500_to_logMgas(hd[1].data['M500c'], z)))
		col6 = fits.Column(name='kT_cluster'    ,format='D', unit='keV', array =cl.logM500_to_kT(hd[1].data['M500c'], z))
		col7 = fits.Column(name='Lx_bol_cluster',format='D', array =n.log10(cl.logM500_to_L(hd[1].data['M500c'], z)))
		col8 = fits.Column(name='Lx_ce_cluster' ,format='D', array =n.log10(cl.logM500_to_Lce(hd[1].data['M500c'], z)))

		#define the table hdu
		colArray = [col1]
		#for col in hd[1].columns :
			#colArray.append(col)

		# AGN Mvir cols
		colArray.append(col2)
		colArray.append(col3)
		# Clusters columns
		colArray.append(col4)
		colArray.append(col5)
		colArray.append(col6)
		colArray.append(col7)
		colArray.append(col8)

		hdu_cols  = fits.ColDefs(colArray)
		tb_hdu = fits.BinTableHDU.from_columns( hdu_cols )
		#define the header
		prihdr = fits.Header()
		prihdr['author'] = 'JC'
		prihdu = fits.PrimaryHDU(header=prihdr)
		#writes the file
		thdulist = fits.HDUList([prihdu, tb_hdu])
		if os.path.isfile(outFile):
			os.system("rm "+outFile)

		thdulist.writeto(outFile)
		print time.time()-t0

# open the output file_type
summ = fits.open(os.path.join(os.environ["MD10"], 'output_MD_1.0Gpc.fits'))[1].data

for el in summ:
	print el
	fileList_snap = n.array(glob.glob(os.path.join(os.environ["MD10"], 'work_agn', 'out_'+el['snap_name']+'_SAM_Nb_?.fits')))
	fileList_snap.sort()
	print fileList_snap
	create_catalogs_out(fileList_snap, el['redshift'], el['snap_name'])
<commit_msg>message<commit_after># overall python packages
import glob
import astropy.io.fits as fits
import os
import time
import numpy as n
import sys

# specific functions
from scipy.stats import norm
from scipy.integrate import quad
from scipy.interpolate import interp1d

# dedicated packages
import ClusterScalingRelations
cl = ClusterScalingRelations.ClusterScalingRelations_Mantz2016()
import StellarMass
import XrayLuminosity
xr = XrayLuminosity.XrayLuminosity()
dV = -9999

from numpy import exp, random

logmbh = lambda logm200c : 8.18 + 1.55 * ( logm200c - 13. )

LX_Z_lo, FRACTION_Z_lo = n.loadtxt(os.path.join(os.environ['GIT_NBODY_NPT'],'data','fraction_zgt0_zlt10.txt' ), unpack=True)
LX_Z_me, FRACTION_Z_me = n.loadtxt(os.path.join(os.environ['GIT_NBODY_NPT'],'data','fraction_zgt10_zlt25.txt'), unpack=True)
LX_Z_hi, FRACTION_Z_hi = n.loadtxt(os.path.join(os.environ['GIT_NBODY_NPT'],'data','fraction_zgt25_zlt50.txt'), unpack=True)
zmin_obscur = n.array([-0.01,1.0,2.5])
zmax_obscur = n.array([1.0,2.5,10.0])

# read the Xray AGN luminosity function and add a condition to reproduce it

def create_catalogs_out(fileList, z, snap_name):
	"""
	Adds Xray emission mass using the Bongiorno et al. 2016 model to the rockstar outputs.
	"""
	#
	out_duty_cycle = os.path.join(os.environ['MD10'],"duty_cycle", snap_name + "_duty_cycle.txt")
	log_stellar_mass, duty_cycle = n.loadtxt(out_duty_cycle, unpack="True")
	percentage_active = interp1d(n.hstack((-200., 0,n.min(log_stellar_mass)-0.01,log_stellar_mass,n.max(log_stellar_mass)+0.01,15)), n.hstack(( 0., 0., 0., duty_cycle, 0., 0.)))

	# set up the x ray lambda SAR
	logMs = n.arange(4.5,14.5,0.01)
	cdfs_interpolations = []
	XXS = n.arange(32,36.1,0.1)
	for mass in logMs:
		norming = xr.Phi_stellar_mass(mass, z)
		cdfs_interpolations.append( interp1d(n.hstack((n.array([xr.Phi_stellar_mass_to_X(X, mass, z) for X in XXS ])/norming, 1.)), n.hstack((XXS, XXS[-1]+0.1))) )

	cdfs_interpolations = n.array(cdfs_interpolations)

	if z > zmin_obscur[0] and z<zmax_obscur[0]:
		obscured_fraction_interpolated = interp1d(LX_Z_lo, FRACTION_Z_lo)
	elif z > zmin_obscur[1] and z<zmax_obscur[1]:
		obscured_fraction_interpolated = interp1d(LX_Z_me, FRACTION_Z_me)
	else:
		obscured_fraction_interpolated = interp1d(LX_Z_hi, FRACTION_Z_hi)

	# loops over files
	for fileName in fileList:
		t0=time.time()
		outFile = fileName[:-5]+"_Xray.fits"
		# opens all relevant files
		msFile = fileName[:-5]+"_Ms.fits"
		hd = fits.open(fileName)
		hm = fits.open(msFile)

		stellar_mass = hm[1].data['stellar_mass_Mo13_mvir']
		selection = hm[1].data['stellar_mass_reliable']
		Nhalo=len(stellar_mass)
		randomX = n.random.rand(Nhalo)
		active_gn = ( percentage_active(stellar_mass) > randomX )
		# lambda SAR addition
		indexes = n.searchsorted(logMs,stellar_mass)
		indexes[selection] = n.zeros_like(indexes[selection])
		lambda_sar_Bo16 = n.array([ cdfs_interpolations[indexes[ii]](randomX[ii]) for ii in range(Nhalo) ])
		# obscuration Merloni et al. 2015 as a function of luminosity
		# randomObscuration = n.random.rand(Nhalo)
		# obscured = (xr.obscured_fraction_optical_Merloni2015(lambda_sar_Bo16 + stellar_mass) < randomObscuration )

		# obscuration Buchner et al. 2015 + 2017
		# add the log NH of the logNH_host
		# 35 % have a thick obscuration 24 - 26
		# 65 % have a thin obscuration that depends on stellar mass and Xray luminosity
		logNH = random.uniform(20, 22,Nhalo)
		obs_type = n.zeros(Nhalo)
		# 35% of thick, 24-26
		randomNH = n.random.rand(Nhalo)
		thick_obscuration = (randomNH < 0.35)
		thin_obscuration = (randomNH >= 0.35)
		logNH[thick_obscuration] = random.uniform(24, 26, len(logNH[thick_obscuration]))
		obs_type[thick_obscuration] = n.ones_like(logNH[thick_obscuration])*2
		# the thin : about 40 % are thin whatever happens: 22-24
		logNH_host_mean = 21.7 + (stellar_mass - 9.5)*0.38
		logNH_host = random.normal(logNH_host_mean, 0.5)
		logNH[(thin_obscuration)&(logNH_host>=22)] = random.uniform(22, 24, len(logNH[(thin_obscuration)&(logNH_host>=22)]))
		obs_type[(thin_obscuration)&(logNH_host>=22)] =  n.ones_like(logNH[(thin_obscuration)&(logNH_host>=22)])
		# a few more are thin depending on their Xray luminosity: 22-24
		rest = (thin_obscuration)&(logNH_host<22)
		randomNH2 = n.random.rand(Nhalo)
		rest_obscured = (rest)&(randomNH2 < obscured_fraction_interpolated(lambda_sar_Bo16 + stellar_mass))
		logNH[(rest_obscured)] = random.uniform(22, 24, len(logNH[(rest_obscured)]))
		obs_type[(rest_obscured)] =  n.ones_like(logNH[(rest_obscured)])
		# the rest has: 20-22 by default
		print logNH, obs_type

		# columns related to Xray AGN
		col1 = fits.Column(name='lambda_sar_Bo16',format='D', array = lambda_sar_Bo16 )
		col2 = fits.Column(name='activity',format='L', array = active_gn )
		col3 = fits.Column(name='obscuration_type_Buchner2017',format='K', array = obs_type.astype('int') )
		#col4 = fits.Column(name='M_BH'  ,format='D', array =logmbh(hd[1].data['M200c'])))
		col4 = fits.Column(name='log_NH_Buchner2017'  ,format='D', array = logNH)

		# columns related to clusters
		col5 = fits.Column(name='Mgas_cluster'  ,format='D', array =n.log10(cl.logM500_to_logMgas(hd[1].data['M500c'], z)))
		col6 = fits.Column(name='kT_cluster'    ,format='D', unit='keV', array =cl.logM500_to_kT(hd[1].data['M500c'], z))
		col7 = fits.Column(name='Lx_bol_cluster',format='D', array =n.log10(cl.logM500_to_L(hd[1].data['M500c'], z)))
		col8 = fits.Column(name='Lx_ce_cluster' ,format='D', array =n.log10(cl.logM500_to_Lce(hd[1].data['M500c'], z)))

		#define the table hdu
		colArray = [col1]
		#for col in hd[1].columns :
			#colArray.append(col)

		# AGN Mvir cols
		colArray.append(col2)
		colArray.append(col3)
		# Clusters columns
		colArray.append(col4)
		colArray.append(col5)
		colArray.append(col6)
		colArray.append(col7)
		colArray.append(col8)

		hdu_cols  = fits.ColDefs(colArray)
		tb_hdu = fits.BinTableHDU.from_columns( hdu_cols )
		#define the header
		prihdr = fits.Header()
		prihdr['author'] = 'JC'
		prihdu = fits.PrimaryHDU(header=prihdr)
		#writes the file
		thdulist = fits.HDUList([prihdu, tb_hdu])
		if os.path.isfile(outFile):
			os.system("rm "+outFile)

		thdulist.writeto(outFile)
		print time.time()-t0

# open the output file_type
summ = fits.open(os.path.join(os.environ["MD10"], 'output_MD_1.0Gpc.fits'))[1].data

for el in summ:
	print el
	fileList_snap = n.array(glob.glob(os.path.join(os.environ["MD10"], 'work_agn', 'out_'+el['snap_name']+'_SAM_Nb_?.fits')))
	fileList_snap.sort()
	print fileList_snap
	create_catalogs_out(fileList_snap, el['redshift'], el['snap_name'])
<|endoftext|>