drone-rigide/scripts/control_compute.py

#! /usr/bin/python

from scipy.signal import savgol_coeffs
import numpy as np
import sys
import threading
import functools

import rospy

from std_msgs.msg import Float64, Empty

import dynamic_reconfigure.server
from detect_targets.cfg import ProportionalNodeConfig, IntegralNodeConfig
from detect_targets.cfg import DerivativeNodeConfig, InputNodeConfig
from detect_targets.cfg import SaturateNodeConfig, RateNodeConfig


class ControlNode(object):
    def __init__(self, refresh_time=0.05, config_class=None):
        self.last_time = rospy.get_rostime()
        self.refresh_time = rospy.Duration.from_sec(refresh_time)
        self.output = 0.0

        self.output_topic = rospy.Publisher('output', Float64, queue_size=1)

        self.input_topics = {
            "reset": rospy.Subscriber("reset", Empty, self.on_reset),
        }


        if config_class is not None:
            self.server = dynamic_reconfigure.server.Server(config_class, self.on_reconf)

    def on_reset(self, value=None):
        self.output = 0.0

    def on_reconf(self, config, level):
        self.refresh_time = rospy.Duration.from_sec(config["refresh_time"])
        self.on_reset()
        return config

    def on_compute(self, value):
        self.output_topic.publish(data=self.output)
        self.last_time = rospy.get_rostime()

    def check_time(self, delta_time=0.0):
        t = rospy.get_rostime()
        if delta_time == 0.0:
            delta_time = t - self.last_time

        return delta_time >= self.refresh_time, delta_time.to_sec()

    def on_mainloop(self):
        rospy.spin()

class InputControlNode(ControlNode):
    def __init__(self, *args, **kwargs):
        super(InputControlNode, self).__init__(*args, **kwargs)
        self.input_topics["input"] = rospy.Subscriber("input", Float64, self.on_compute)

class RateNode(InputControlNode):
    def __init__(self, refresh_time=0.05):
        self.lock = threading.Lock()
        super(RateNode, self).__init__(refresh_time, RateNodeConfig)
        self.rate = rospy.Rate(1.0/self.refresh_time.to_sec())

    def on_reconf(self, config, level):
        c = super(RateNode, self).on_reconf(config, level)
        self.rate = rospy.Rate(self.refresh_time.to_sec())
        return c

    def on_compute(self, value):
        self.lock.acquire()
        self.output = value.data
        self.lock.release()

    def on_mainloop(self):
        while not rospy.is_shutdown():
            self.lock.acquire()
            self.output_topic.publish(self.output)
            self.lock.release()
            self.rate.sleep()

class ProportionalNode(InputControlNode):
    def __init__(self, k=1.0, refresh_time=0.05):
        super(ProportionalNode, self).__init__(refresh_time, ProportionalNodeConfig)
        self.k = k

    def on_reconf(self, config, level):
        self.k = config["k"]
        return super(ProportionalNode, self).on_reconf(config, level)

    def on_compute(self, value):
        time_ok,_ = self.check_time()
        if time_ok:
            self.output = self.k * value.data
            super(ProportionalNode, self).on_compute(value)

class SaturateNode(InputControlNode):
    def __init__(self, min=None, max=None, refresh_time=0.05, config=SaturateNodeConfig):
        super(SaturateNode, self).__init__(refresh_time, config)
        self.min = min
        self.max = max

    def on_reconf(self, config, level):
        self.min = config["min"] if config["use_min"] else None
        self.max = config["max"] if config["use_max"] else None
        return super(SaturateNode, self).on_reconf(config, level)

    def on_compute(self, value):
        time_ok,_ = self.check_time()
        if time_ok:
            self.output = min(self.max or value.data, value.data)
            self.output = max(self.min or self.output, self.output)
            super(SaturateNode, self).on_compute(value)

class IntegralNode(SaturateNode):
    def __init__(self, k=0.0, min=None, max=None, refresh_time=0.05):
        super(IntegralNode, self).__init__(min, max, refresh_time, IntegralNodeConfig)
        self.k = k

    def on_reconf(self, config, level):
        self.k = config["k"]
        return super(IntegralNode, self).on_reconf(config, level)

    def on_compute(self, value):
        time_ok, delta = self.check_time()
        if time_ok:
            old_output = self.output
            self.output += self.k * value.data * delta
            return super(IntegralNode, self).on_compute(Float64(self.output))

class DerivativeNode(InputControlNode):
    def __init__(self, k=0.05, size=9, polyorder=3, deriv=1, refresh_time=0.05):
        super(DerivativeNode, self).__init__(refresh_time, DerivativeNodeConfig)
        self.k = k
        self.filter = savgol_coeffs(size, polyorder, deriv, delta=refresh_time, use='dot')
        self.last_points = np.zeros(size)

    def on_reconf(self, config, level):
        self.k = config["k"]
        try:
            self.filter = savgol_coeffs(
                config["size"],
                config["poly_order"],
                config["deriv"],
                delta=config["refresh_time"],
                use='dot'
            )
        except ValueError as e:
            rospy.logerr(e)
        self.last_points = np.zeros(config["size"])
        return super(DerivativeNode, self).on_reconf(config, level)

    def on_compute(self, value):
        time_ok, delta = self.check_time()
        if time_ok:
            self.last_points = np.concatenate([self.last_points[1:], [value.data]])
            self.output = self.k * self.last_points.dot(self.filter)
            return super(DerivativeNode, self).on_compute(value)

    def on_reset(self, value=None):
        self.last_points = np.zeros(len(self.filter))
        return super(DerivativeNode, self).on_reset(value)

class DifferenciateNode(InputControlNode):
    def __init__(self):
        self.measure = 0.0
        self.input = 0.0
        super(DifferenciateNode, self).__init__()
        self.input_topics["measure"] = rospy.Subscriber("measure", Float64, self.on_compute_measure)

    def on_compute(self, value):
        self.input = value.data
        self.output = self.input - self.measure
        super(DifferenciateNode, self).on_compute(value)

    def on_compute_measure(self, value):
        self.measure = value.data
        self.output = self.input - self.measure
        super(DifferenciateNode, self).on_compute(value)

class InputNode(InputControlNode):
    def __init__(self):
        super(InputNode, self).__init__(refresh_time=0.05, config_class=InputNodeConfig)

    def on_reconf(self, config, level):
        self.output = config["value"]
        super(InputNode, self).on_compute(Float64(config["value"]))
        return config

    def on_compute(self, value):
        self.output = value.data
        super(InputNode, self).on_compute(value)

class SumNode(ControlNode):
    def __init__(self, nb_topics):
        super(SumNode, self).__init__(refresh_time=0.05)
        self.nb_topics = int(nb_topics)
        self.inputs = dict()
        for i in range(self.nb_topics):
            topic_name = 'input_'+str(i)
            self.input_topics[topic_name] = rospy.Subscriber(topic_name, Float64, functools.partial(self.update_value, topic_name))

    def update_value(self, key, value):
        self.inputs.update({key:value.data})
        self.on_compute(value)

    def on_compute(self, value):
        self.output = sum(self.inputs.values())
        super(SumNode, self).on_compute(value)

if __name__ == '__main__':
    nodes = {
        'sum' : (SumNode, True),
        'differenciate' : (DifferenciateNode, False),
        'input' : (InputNode, False),
        'saturate' : (SaturateNode, False),
        'derivative' : (DerivativeNode, False),
        'integral' : (IntegralNode, False),
        'proportional' : (ProportionalNode, False),
        'rate': (RateNode, False),
    }
    rospy.init_node('control_node', anonymous=True)
    rospy.loginfo("Running node " + sys.argv[1])
    try:
        node_class, need_param = nodes[sys.argv[1]]
    except KeyError as e:
        raise ValueError(sys.argv[1] + " is not a valid node name.")

    if need_param:
        param = sys.argv[2]
        node = node_class(param)
    else:
        node = node_class()

    node.on_mainloop()