text:
import serial
import threading
import numpy as np

from .message import Message

class Interface:
    '''
    This is a python replication of the dobot commands
    '''
    def __init__(self, port):
        threading.Thread.__init__(self)
        self.lock = threading.Lock()

        self.serial = serial.Serial(
            port=port,
            baudrate=115200,
            parity=serial.PARITY_NONE,
            stopbits=serial.STOPBITS_ONE,
            bytesize=serial.EIGHTBITS,
            timeout=2,
        )

    def send(self, message):
        self.lock.acquire()
        self.serial.write(message.package())
        self.serial.flush()
        response = Message.read(self.serial)
        self.lock.release()

        if response is None:
            return None

        return response.params

    @property
    def connected(self):
        '''
        returns if the serial port is open, not if it is connected to the dobot
        '''
        return self.serial.isOpen()

    def get_device_serial_number(self):
        request = Message([0xAA, 0xAA], 2, 0, False, False, [], direction='out')
        return self.send(request)

    def set_device_serial_number(self, serial_number):
        request = Message([0xAA, 0xAA], 2, 0, True, False, [serial_number], direction='out')
        return self.send(request)

    def get_device_name(self):
        request = Message([0xAA, 0xAA], 2, 1, False, False, [], direction='out')
        return self.send(request)

    def set_device_name(self, device_name):
        request = Message([0xAA, 0xAA], 2, 1, True, False, [device_name], direction='out')
        return self.send(request)

    def get_device_version(self):
        request = Message([0xAA, 0xAA], 2, 2, False, False, [], direction='out')
        return self.send(request)

    def set_sliding_rail_status(self, enable, version):
        request = Message([0xAA, 0xAA], 2, 3, True, False, [], direction='out')
        return self.send(request)

    # Time in milliseconds since start
    def get_device_time(self):
        request = Message([0xAA, 0xAA], 2, 4, False, False, [], direction='out')
        return self.send(request)

    def get_device_id(self):
        request = Message([0xAA, 0xAA], 2, 5, False, False, [], direction='out')
        return self.send(request)

    def get_pose(self):
        request = Message([0xAA, 0xAA], 2, 10, False, False, [], direction='out')
        return self.send(request)

    def reset_pose(self, manual, rear_arm_angle=None, front_arm_angle=None):
        '''

        Parameters
        ----------
        manual: int
            can be either 0 or 1. (sudo bool).
            0 is if you don't need to provide an angle measurement, 1 is you do
        rear_arm_angle: float (optional)
            if manual = 1, you need to provide the angle of the "rear arm"
        front_arm_angle: float (optional)
            if manual = 1, you need to provide the angle of the "front arm"

        '''
        if manual == 1:
            assert rear_arm_angle is not None
            assert front_arm_angle is not None
        else:
            rear_arm_angle=0
            front_arm_angle=0
        request = Message([0xAA, 0xAA], 2, 11, True, False, [manual, rear_arm_angle, front_arm_angle], direction='out')
        return self.send(request)

    def get_sliding_rail_pose(self):
        request = Message([0xAA, 0xAA], 2, 13, False, False, [], direction='out')
        return self.send(request)

    def get_alarms_state(self):
        request = Message([0xAA, 0xAA], 2, 20, False, False, [], direction='out')
        return self.send(request)

    def clear_alarms_state(self):
        request = Message([0xAA, 0xAA], 2, 21, True, False, [], direction='out')
        return self.send(request)

    def get_homing_paramaters(self):
        request = Message([0xAA, 0xAA], 2, 30, False, False, [], direction='out')
        return self.send(request)

    def set_homing_parameters(self, x, y, z, r, queue=True):
        request = Message([0xAA, 0xAA], 2, 30, True, queue, [x, y, z, r], direction='out')
        return self.send(request)

    def set_homing_command(self, command, queue=True):
        request = Message([0xAA, 0xAA], 2, 31, True, queue, [command], direction='out')
        return self.send(request)

    # TODO: Reference is wrong here, arm does not send the said value
    def get_auto_leveling(self):
        request = Message([0xAA, 0xAA], 2, 32, False, False, [], direction='out')
        return self.send(request)

    def set_auto_leveling(self, enable, accuracy, queue=True):
        request = Message([0xAA, 0xAA], 2, 32, True, queue, [enable, accuracy], direction='out')
        return self.send(request)

    def get_handheld_teaching_mode(self):
        request = Message([0xAA, 0xAA], 2, 40, False, False, [], direction='out')
        return self.send(request)

    def set_handheld_teaching_mode(self, mode):
        request = Message([0xAA, 0xAA], 2, 40, True, False, [mode], direction='out')
        return self.send(request)

    def get_handheld_teaching_state(self):
        request = Message([0xAA, 0xAA], 2, 41, False, False, [], direction='out')
        return self.send(request)

    def set_handheld_teaching_state(self, enable):
        request = Message([0xAA, 0xAA], 2, 41, True, False, [enable], direction='out')
        return self.send(request)

    def get_handheld_teaching_trigger(self):
        request = Message([0xAA, 0xAA], 2, 42, False, False, [], direction='out')
        return self.send(request)

    def get_end_effector_params(self):
        request = Message([0xAA, 0xAA], 2, 60, False, False, [], direction='out')
        return self.send(request)

    def set_end_effector_params(self, bias_x, bias_y, bias_z):
        request = Message([0xAA, 0xAA], 2, 60, True, False, [bias_x, bias_y, bias_z], direction='out')
        return self.send(request)

    def get_end_effector_laser(self):
        request = Message([0xAA, 0xAA], 2, 61, False, False, [], direction='out')
        return self.send(request)

    def set_end_effector_laser(self, enable_control, enable_laser, queue=True):
        request = Message([0xAA, 0xAA], 2, 61, True, queue, [enable_control, enable_laser], direction='out')
        return self.send(request)

    def get_end_effector_suction_cup(self):
        request = Message([0xAA, 0xAA], 2, 62, False, False, [], direction='out')
        return self.send(request)

    def set_end_effector_suction_cup(self, enable_control, enable_suction, queue=True):
        request = Message([0xAA, 0xAA], 2, 62, True, queue, [enable_control, enable_suction], direction='out')
        return self.send(request)

    def get_end_effector_gripper(self):
        request = Message([0xAA, 0xAA], 2, 63, False, False, [], direction='out')
        return self.send(request)

    def set_end_effector_gripper(self, enable_control, enable_grip, queue=True):
        request = Message([0xAA, 0xAA], 2, 63, True, queue, [enable_control, enable_grip], direction='out')
        return self.send(request)

    def get_jog_joint_params(self):
        request = Message([0xAA, 0xAA], 2, 70, False, False, [], direction='out')
        return self.send(request)

    # TODO: Does not work - but is implemented according to spec. Bad documentation?
    def set_jog_joint_params(self, velocity, acceleration, queue=True):
        request = Message([0xAA, 0xAA], 2, 70, True, queue, velocity + acceleration, direction='out')
        return self.send(request)

    def get_jog_coordinate_params(self):
        request = Message([0xAA, 0xAA], 2, 71, False, False, [], direction='out')
        return self.send(request)

    # TODO: Does not work - but is implemented according to spec. Bad documentation?
    def set_jog_coordinate_params(self, velocity, acceleration, queue=True):
        request = Message([0xAA, 0xAA], 2, 71, True, queue, velocity + acceleration, direction='out')
        return self.send(request)

    def get_jog_common_params(self):
        request = Message([0xAA, 0xAA], 2, 72, False, False, [], direction='out')
        return self.send(request)

    # TODO: Does not work - but is implemented according to spec. Bad documentation?
    def set_jog_common_params(self, velocity_ratio, acceleration_ratio, queue=True):
        request = Message([0xAA, 0xAA], 2, 72, True, queue, [velocity_ratio, acceleration_ratio], direction='out')
        return self.send(request)

    def set_jog_command(self, jog_type, command, queue=True):
        '''

        performs a jog command.

        Parameters
        ----------
        jog_type : int
            can be 0 or 1. 0 is coordinate jog, 1 is joint jog
        command : int
            can be 0-8. This is the coordinate or joint to move.
            zero can be used to stop the jog

            value   coordinate  joint
            0       null        null
            1       X+          Joint1+
            2       X-          Joint1-
            3       Y+          Joint2+
            4       Y-          Joint2-
            5       Z+          Joint3+
            6       Z-          Joint3-
            7       R+          Joint4+
            8       R-          Joint4-

        queue : bool
            If a queue is to be used.

        '''
        request = Message([0xAA, 0xAA], 2, 73, True, queue, [jog_type, command], direction='out')
        return self.send(request)

    def get_sliding_rail_jog_params(self):
        request = Message([0xAA, 0xAA], 2, 74, False, False, [], direction='out')
        return self.send(request)

    def set_sliding_rail_jog_params(self, velocity, acceleration, queue=True):
        request = Message([0xAA, 0xAA], 2, 74, True, queue, [velocity, acceleration], direction='out')
        return self.send(request)

    def get_point_to_point_joint_params(self):
        request = Message([0xAA, 0xAA], 2, 80, False, False, [], direction='out')
        return self.send(request)

    def set_point_to_point_joint_params(self, velocity, acceleration, queue=True):
        request = Message([0xAA, 0xAA], 2, 80, True, queue, velocity + acceleration, direction='out')
        return self.send(request)

    def get_point_to_point_coordinate_params(self):
        request = Message([0xAA, 0xAA], 2, 81, False, False, [], direction='out')
        return self.send(request)

    def set_point_to_point_coordinate_params(self, coordinate_velocity, effector_velocity, coordinate_acceleration, effector_acceleration, queue=True):
        request = Message([0xAA, 0xAA], 2, 81, True, queue, [coordinate_velocity, effector_velocity, coordinate_acceleration, effector_acceleration], direction='out')
        return self.send(request)

    def get_point_to_point_jump_params(self):
        request = Message([0xAA, 0xAA], 2, 82, False, False, [], direction='out')
        return self.send(request)

    def set_point_to_point_jump_params(self, jump_height, z_limit, queue=True):
        request = Message([0xAA, 0xAA], 2, 82, True, queue, [jump_height, z_limit], direction='out')
        return self.send(request)

    def get_point_to_point_common_params(self):
        request = Message([0xAA, 0xAA], 2, 83, False, False, [], direction='out')
        return self.send(request)

    def set_point_to_point_common_params(self, velocity_ratio, acceleration_ratio, queue=True):
        request = Message([0xAA, 0xAA], 2, 83, True, queue, [velocity_ratio, acceleration_ratio], direction='out')
        return self.send(request)

    def set_point_to_point_command(self, mode, x, y, z, r, queue=True):
        '''
        moves to dobot arm from one point to another.
        Taken from Dobot Communication Protocol V1.1.3

        Parameters
        ----------

        mode: int
            How to move
                int     meaning          coordinates
                0       jump             Cartesian
                1       joint movement   Cartesian
                2       linear movement  Cartesian
                3       jump             Joint - angles of the robot arms
                4       joint movement   Joint
                5       linear movement  Joint
                6       joint movment    Joint
                7       linear           Cartesian incriment in joint coordinate
                8       joint            Cartesian increment in Cartesian
                9       jump             Cartesian increment in Cartesian
        x: float
            the x dimension or angle of a joint i
        y: float
            the y dimension or angle of  joint j
        z: float
            the z dimension or the angle of joint k
        r: float
            the rotation of the hand -- if applicable ?

        '''

        request = Message([0xAA, 0xAA], 2, 84, True, queue, [mode, x, y, z, r], direction='out')
        return self.send(request)

    def get_point_to_point_sliding_rail_params(self):
        request = Message([0xAA, 0xAA], 2, 85, False, False, [], direction='out')
        return self.send(request)

    def set_point_to_point_sliding_rail_params(self, velocity, acceleration, queue=True):
        request = Message([0xAA, 0xAA], 2, 85, True, queue, [velocity, acceleration], direction='out')
        return self.send(request)

    def set_point_to_point_sliding_rail_command(self, mode, x, y, z, r, l, queue=True):
        request = Message([0xAA, 0xAA], 2, 86, True, queue, [mode, x, y, z, r, l], direction='out')
        return self.send(request)

    def get_point_to_point_jump2_params(self):
        request = Message([0xAA, 0xAA], 2, 87, False, False, [], direction='out')
        return self.send(request)

    def set_point_to_point_jump2_params(self, start_height, end_height, z_limit, queue=True):
        request = Message([0xAA, 0xAA], 2, 87, True, queue, [start_height, end_height, z_limit], direction='out')
        return self.send(request)

    # TODO: Reference is ambigious here - needs testing
    def set_point_to_point_po_command(self, mode, x, y, z, r, queue=True):
        request = Message([0xAA, 0xAA], 2, 88, True, queue, [mode, x, y, z, r], direction='out')
        return self.send(request)

    # TODO: Reference is ambigious here - needs testing
    def set_point_to_point_sliding_rail_po_command(self, ratio, address, level, queue=True):
        request = Message([0xAA, 0xAA], 2, 89, True, queue, [ratio, address, level], direction='out')
        return self.send(request)

    def get_continous_trajectory_params(self):
        request = Message([0xAA, 0xAA], 2, 90, False, False, [], direction='out')
        return self.send(request)

    def set_continous_trajectory_params(self, max_planned_acceleration, max_junction_velocity, acceleration, queue=True):
        request = Message([0xAA, 0xAA], 2, 90, True, queue, [max_planned_acceleration, max_junction_velocity, acceleration, 0], direction='out')
        return self.send(request)

    def set_continous_trajectory_real_time_params(self, max_planned_acceleration, max_junction_velocity, period, queue=True):
        request = Message([0xAA, 0xAA], 2, 90, True, queue, [max_planned_acceleration, max_junction_velocity, period, 1], direction='out')
        return self.send(request)

    def set_continous_trajectory_command(self, mode, x, y, z, velocity, queue=True):
        request = Message([0xAA, 0xAA], 2, 91, True, queue, [mode, x, y, z, velocity], direction='out')
        return self.send(request)

    def set_continous_trajectory_laser_engraver_command(self, mode, x, y, z, power, queue=True):
        request = Message([0xAA, 0xAA], 2, 92, True, queue, [mode, x, y, z, power], direction='out')
        return self.send(request)

    def get_arc_params(self):
        request = Message([0xAA, 0xAA], 2, 100, False, False, [], direction='out')
        return self.send(request)

    def set_arc_params(self, coordinate_velocity, effector_velocity, coordinate_acceleration, effector_acceleration, queue=True):
        request = Message([0xAA, 0xAA], 2, 100, True, queue, [coordinate_velocity, effector_velocity, coordinate_acceleration, effector_acceleration], direction='out')
        return self.send(request)

    def set_arc_command(self, circumference_point, ending_point, queue=True):
        request = Message([0xAA, 0xAA], 2, 101, True, queue, circumference_point + ending_point, direction='out')
        return self.send(request)

    def wait(self, milliseconds, queue=True):
        request = Message([0xAA, 0xAA], 2, 110, True, queue, [milliseconds], direction='out')
        return self.send(request)

    def set_trigger_command(self, address, mode, condition, threshold, queue=True):
        request = Message([0xAA, 0xAA], 2, 120, True, queue, [address, mode, condition, threshold], direction='out')
        return self.send(request)

    def get_io_multiplexing(self):
        request = Message([0xAA, 0xAA], 2, 130, False, False, [], direction='out')
        return self.send(request)

    def set_io_multiplexing(self, address, multiplex, queue=True):
        '''
        sets multiplexing for the extended items (things pluged into the arm)

        Parameters
        ----------
        address: 8bit int
            An EIO address. (Extensible I/O adrdress)
        multiplex: int
            A IOMultiplexing type
                value   meaning
                0       Not in use
                1       PWM Output
                2       Digital Output
                3       Digital Input
                4       Analgue input

            These appear in a different order in Dobit API V1.0.0
            to Dobot Communication Protocol V1.0.4
        '''
        request = Message([0xAA, 0xAA], 2, 130, True, queue, [address, multiplex], direction='out')
        return self.send(request)

    def get_io_do(self):
        request = Message([0xAA, 0xAA], 2, 131, False, False, [], direction='out')
        return self.send(request)

    def set_io_do(self, address, level, queue=True):
        '''
        sets an extensible I/O  to either high or low

        Parameters
        ----------
        address: 8bit int
            An EIO address. (Extensible I/O adrdress)
        level: int
            can be 0 or 1. (low or high)
        '''
        request = Message([0xAA, 0xAA], 2, 131, True, queue, [address, level], direction='out')
        return self.send(request)

    def get_io_pwm(self):
        request = Message([0xAA, 0xAA], 2, 132, False, False, [], direction='out')
        return self.send(request)

    def set_io_pwm(self, address, frequency, duty_cycle, queue=True):
        '''
        sets PWM to the extended I/O
        from to Dobot Communication Protocol V1.0.4

        Parameters
        ----------
        address: 8bit int
            An EIO address. (Extensible I/O adrdress)
        frequency: float
            the PWM frequency 10 - 10e6, in units of Hz
        duty_cycle: float
            is says 0-100.
        '''
        request = Message([0xAA, 0xAA], 2, 132, True, queue, [address, frequency, duty_cycle], direction='out')
        return self.send(request)

    def get_io_di(self):
        request = Message([0xAA, 0xAA], 2, 133, False, False, [], direction='out')
        return self.send(request)

    def get_io_adc(self):
        '''
        reads an analogue value?
        from to Dobot Communication Protocol V1.0.4
        '''
        request = Message([0xAA, 0xAA], 2, 134, False, False, [], direction='out')
        return self.send(request)

    def set_extended_motor_velocity(self, index, enable, speed, queue=True):
        '''
        sets the motot velocity of an extended motor

        Parameters
        ----------
        index : index
            can be either 0 or 1. 0 is stepper motor 1? and 1 is stepper motor 2?
        enable : bool
            Is motor control enabled
        speed : int -
            The number of the pulses per second. incorrectly labled in the offical
            documentation as float

        '''
        request = Message([0xAA, 0xAA], 2, 135, 1, queue, [index, enable, int(speed)], direction='out')
        return self.send(request)

    def get_color_sensor(self, index):
        request = Message([0xAA, 0xAA], 2, 137, False, False, [], direction='out')
        return self.send(request)

    def set_color_sensor(self, index, enable, port, version, queue=True):
        request = Message([0xAA, 0xAA], 2, 137, True, queue, [enable, port, version], direction='out')
        return self.send(request)

    def get_ir_switch(self, index):
        request = Message([0xAA, 0xAA], 2, 138, False, False, [], direction='out')
        return self.send(request)

    def set_ir_switch(self, index, enable, port, version, queue=True):
        request = Message([0xAA, 0xAA], 2, 138, True, queue, [enable, port, version], direction='out')
        return self.send(request)

    def get_angle_sensor_static_error(self, index):
        request = Message([0xAA, 0xAA], 2, 140, False, False, [], direction='out')
        return self.send(request)

    def set_angle_sensor_static_error(self, index, rear_arm_angle_error, front_arm_angle_error):
        request = Message([0xAA, 0xAA], 2, 140, True, False, [rear_arm_angle_error, front_arm_angle_error], direction='out')
        return self.send(request)

    def get_wifi_status(self):
        request = Message([0xAA, 0xAA], 2, 150, False, False, [], direction='out')
        return self.send(request)

    def set_wifi_status(self, index, enable):
        request = Message([0xAA, 0xAA], 2, 150, True, False, [enable], direction='out')
        return self.send(request)

    def get_wifi_ssid(self):
        request = Message([0xAA, 0xAA], 2, 151, False, False, [], direction='out')
        return self.send(request)

    def set_wifi_ssid(self, index, ssid):
        request = Message([0xAA, 0xAA], 2, 151, True, False, [ssid], direction='out')
        return self.send(request)

    def get_wifi_password(self):
        request = Message([0xAA, 0xAA], 2, 152, False, False, [], direction='out')
        return self.send(request)

    def set_wifi_password(self, index, ssid):
        request = Message([0xAA, 0xAA], 2, 152, True, False, [ssid], direction='out')
        return self.send(request)

    def get_wifi_address(self):
        request = Message([0xAA, 0xAA], 2, 153, False, False, [], direction='out')
        return self.send(request)

    # 192.168.1.1 = a.b.c.d
    def set_wifi_address(self, index, use_dhcp, a, b, c, d):
        request = Message([0xAA, 0xAA], 2, 153, True, False, [use_dhcp, a, b, c, d], direction='out')
        return self.send(request)

    def get_wifi_netmask(self):
        request = Message([0xAA, 0xAA], 2, 154, False, False, [], direction='out')
        return self.send(request)

    # 255.255.255.0 = a.b.c.d
    def set_wifi_netmask(self, index, a, b, c, d):
        request = Message([0xAA, 0xAA], 2, 154, True, False, [a, b, c, d], direction='out')
        return self.send(request)

    def get_wifi_gateway(self):
        request = Message([0xAA, 0xAA], 2, 155, False, False, [], direction='out')
        return self.send(request)

    # 192.168.1.1 = a.b.c.d
    def set_wifi_gateway(self, index, use_dhcp, a, b, c, d):
        request = Message([0xAA, 0xAA], 2, 155, True, False, [use_dhcp, a, b, c, d], direction='out')
        return self.send(request)

    def get_wifi_dns(self):
        request = Message([0xAA, 0xAA], 2, 156, False, False, [], direction='out')
        return self.send(request)

    # 192.168.1.1 = a.b.c.d
    def set_wifi_dns(self, index, use_dhcp, a, b, c, d):
        request = Message([0xAA, 0xAA], 2, 156, True, False, [use_dhcp, a, b, c, d], direction='out')
        return self.send(request)

    def get_wifi_connect_status(self):
        request = Message([0xAA, 0xAA], 2, 157, False, False, [], direction='out')
        return self.send(request)

    def set_lost_step_params(self, param):
        request = Message([0xAA, 0xAA], 2, 170, True, False, [param], direction='out')
        return self.send(request)

    def set_lost_step_command(self):
        request = Message([0xAA, 0xAA], 2, 171, True, False, [], direction='out')
        return self.send(request)

    def start_queue(self):
        request = Message([0xAA, 0xAA], 2, 240, True, False, [], direction='out')
        return self.send(request)

    def stop_queue(self, force=False):
        request = Message([0xAA, 0xAA], 2, 242 if force else 241, True, False, [], direction='out')
        return self.send(request)

    def start_queue_download(self, total_loop, line_per_loop):
        request = Message([0xAA, 0xAA], 2, 243, True, False, [total_loop, line_per_loop], direction='out')
        return self.send(request)

    def stop_queue_download(self):
        request = Message([0xAA, 0xAA], 2, 244, True, False, [], direction='out')
        return self.send(request)

    def clear_queue(self):
        request = Message([0xAA, 0xAA], 2, 245, True, False, [], direction='out')
        return self.send(request)

    def get_current_queue_index(self):
        request = Message([0xAA, 0xAA], 2, 246, True, False, [], direction='out')
        return self.send(request)

class Extented_interface(Interface):
    '''
    This includes extra helper functions to make things easier to use
    '''

    def set_conveyor_speed(self,index, speed, queue=True):
        '''
        Sets the conveyor to a speed in mm/s

        needs to be checked - info obtained from comments on agithub issue https://github.com/luismesas/pydobot/issues/22

        Parameters
        ----------
        index: int
            Index of the stepper motor. This is the number of the port the stepper motor
            is pluged into. It should be 0
        speed: float
            The speed in mm/s at which you want to run the speed
        queue: bool
            If to add to teh queue or just run immdeaitly
        '''
        MM_PER_REV = 33 * np.pi  # Seems to actually be closer to 36mm when measured but 34 works better
        STEP_ANGLE_DEG = 1.8
        STEPS_PER_REV = 360.0 / STEP_ANGLE_DEG * 10.0 * 16.0 / 2.0 # assume the 10*16/2 related pulses (clock) to steps
        pulses_per_second = int(speed/ MM_PER_REV * STEPS_PER_REV)
        print(index, pulses_per_second)
        self.set_extended_motor_velocity(index=index,enable=True, speed =pulses_per_second, queue=True)

    def set_conveyor_distance(self, index, distance, speed=100):
        '''
        moves the convery a specific distance in mm

        Parameters
        ----------
        index: int
            Index of the stepper motor. This is the number of the port the stepper motor
            is pluged into. It should be 0
        distance: float
            The distance in mm to move the conveyor
        speed: float
            The speed in mm/s at which you want to run the speed
        '''
        speed = abs(speed)

        time = abs(distance/speed)
        # get the time in seconds to wait

        if distance <0:
            speed *=-1

        self.set_conveyor_speed(index, speed, queue=True)

        self.wait(int(time*1000), queue=True)

        self.set_extended_motor_velocity(index=index,
                enable=True,
                speed =0,
                queue=True
                )


text:
import numpy as np, torch, torch.nn as nn, torch.optim as optim, torch.nn.functional as F
import argparse, time, pickle, os
from tqdm import tqdm
from torch.utils.tensorboard import SummaryWriter
from torch.utils.data import DataLoader
from torch.utils.data.sampler import SubsetRandomSampler
from sklearn.metrics import f1_score, confusion_matrix, accuracy_score, classification_report, precision_recall_fscore_support
from model import CategoricalModel, MaskedNLLLoss, BiModalAttention
from dataloader import MOSEICategorical

np.random.seed(393)
torch.cuda.device([0])

def get_train_valid_sampler(trainset, valid=0.1):
    size = len(trainset)
    idx = list(range(size))
    split = int(valid*size)
    return SubsetRandomSampler(idx[split:]), SubsetRandomSampler(idx[:split])

def get_MOSEI_loaders(path, batch_size=128, valid=0.1, num_workers=0, pin_memory=False):
    trainset = MOSEICategorical(path=path)
    train_sampler, valid_sampler = get_train_valid_sampler(trainset, valid)
    train_loader = DataLoader(trainset, batch_size=batch_size, sampler=train_sampler, collate_fn=trainset.collate_fn, num_workers=num_workers, pin_memory=pin_memory)
    valid_loader = DataLoader(trainset, batch_size=batch_size, sampler=valid_sampler, collate_fn=trainset.collate_fn, num_workers=num_workers, pin_memory=pin_memory)
    testset = MOSEICategorical(path=path, train=False)
    test_loader = DataLoader(testset, batch_size=batch_size, collate_fn=testset.collate_fn, num_workers=num_workers, pin_memory=pin_memory)
    return train_loader, valid_loader, test_loader

def train_or_eval_model(model,loss_function, dataloader, epoch, optimizer=None, train=False):
    count = 0
    losses, preds, labels, masks, alphas_f, alphas_b, vids = [], [], [], [], [], [], []
    assert not train or optimizer!=None
    if train:
        model.train()
    else:
        model.eval()
    for data in dataloader:
        count+=1
        if train:
            optimizer.zero_grad()
        textf, visuf, acouf, qmask, umask, label =  [d.cuda() for d in data[:-1]] if cuda else data[:-1]
        log_prob ,alpha_f,alpha_b  = model(textf, acouf, visuf, textf, qmask, umask)
        lp_ = log_prob.transpose(0,1).contiguous().view(-1,log_prob.size()[2])
        labels_ = label.view(-1)
        loss = loss_function(lp_, labels_, umask)
        pred_ = torch.argmax(lp_,1)
        preds.append(pred_.data.cpu().numpy())
        labels.append(labels_.data.cpu().numpy())
        masks.append(umask.view(-1).cpu().numpy())
        losses.append(loss.item()*masks[-1].sum())
        if train:
            loss.backward()
            optimizer.step()
        else:
            alphas_f += alpha_f
            alphas_b += alpha_b
            vids += data[-1]
    if preds!=[]:
        preds  = np.concatenate(preds)
        labels = np.concatenate(labels)
        masks  = np.concatenate(masks)
    else:
        return float('nan'), float('nan'), [], [], [], float('nan'),[]
    avg_loss = round(np.sum(losses)/np.sum(masks),4)
    avg_accuracy = round(accuracy_score(labels,preds,sample_weight=masks)*100,2)
    avg_fscore = round(f1_score(labels,preds,sample_weight=masks,average='weighted')*100,2)
    return avg_loss, avg_accuracy, labels, preds, masks,avg_fscore, [alphas_f, alphas_b, vids]

if __name__ == '__main__':
    parser = argparse.ArgumentParser("Trains a categorical model for the emotion label given by --emotion")
    parser.add_argument('--no-cuda', action='store_true', default=False, help='does not use GPU')
    parser.add_argument('--lr', type=float, default=0.0001, metavar='LR', help='learning rate')
    parser.add_argument('--l2', type=float, default=0.00001, metavar='L2', help='L2 regularization weight')
    parser.add_argument('--rec-dropout', type=float, default=0.1, metavar='rec_dropout', help='rec_dropout rate')
    parser.add_argument('--dropout', type=float, default=0.5, metavar='dropout', help='dropout rate')
    parser.add_argument('--batch-size', type=int, default=128, metavar='BS', help='batch size')
    parser.add_argument('--epochs', type=int, default=50, metavar='E', help='number of epochs')
    parser.add_argument('--class-weight', action='store_true', default=True, help='class weight')
    parser.add_argument('--emotion', type=str, default='happiness', metavar='emotion', help='emotion to train')
    parser.add_argument('--log_dir', type=str, default='logs/mosei_emotion', help='Directory for tensorboard logs')
    args = parser.parse_args()
    os.makedirs(args.log_dir, exist_ok = True)
    writer = SummaryWriter(args.log_dir)
    print(args)

    # Run on either GPU or  CPU
    args.cuda = torch.cuda.is_available() and not args.no_cuda
    if args.cuda:
        print('Running on GPU')
    else:
        print('Running on CPU')
    print("Tensorboard logs in " + args.log_dir)

    batch_size = args.batch_size
    n_classes  = 2
    cuda       = args.cuda
    n_epochs   = args.epochs
    emotion = str(args.emotion)
    D_m_text, D_m_audio, D_m_video, D_m_context = 300, 384, 35, 300
    D_g, D_p, D_e, D_h, D_a = 150, 150, 100, 100, 100

    # Instantiate model
    model = CategoricalModel(D_m_text,D_m_audio,D_m_video, D_m_context, D_g, D_p, D_e, D_h, n_classes=n_classes, dropout_rec=args.rec_dropout, dropout=args.dropout)

    if cuda:
        model.cuda()
    if emotion=='anger': loss_weights = torch.FloatTensor([1/0.2128, 1/0.7872])
    elif emotion=='happiness': loss_weights = torch.FloatTensor([1/0.5381, 1/0.4619])
    elif emotion=='sadness': loss_weights = torch.FloatTensor([1/0.2546, 1/0.7454])
    elif emotion=='disgust': loss_weights = torch.FloatTensor([1/0.1699, 1/0.8301])
    elif emotion=='fear': loss_weights = torch.FloatTensor([1/0.0816, 1/0.9184])
    elif emotion=='surprise': loss_weights = torch.FloatTensor([1/0.0986, 1/0.9014])
    if args.class_weight:
        loss_function  = MaskedNLLLoss(loss_weights.cuda() if cuda else loss_weights)
    else:
        loss_function = MaskedNLLLoss()

    # Get optimizers and relevanat dataloaders
    optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.l2)
    train_loader, valid_loader, test_loader = get_MOSEI_loaders('./data/' + str(emotion) + '.pkl', valid=0.0, batch_size=batch_size, num_workers=0)
    best_loss, best_label, best_pred, best_mask = None, None, None, None

    # Training loop
    for e in tqdm(range(n_epochs), desc = 'MOSEI Emotion (' + str(emotion) + ')'):
        train_loss, train_acc, _,_,_,train_fscore,_ = train_or_eval_model(model, loss_function, train_loader, e, optimizer, True)
        test_loss, test_acc, test_label, test_pred, test_mask, test_fscore, attentions = train_or_eval_model(model,loss_function, test_loader, e)
        writer.add_scalar("Train Loss - MOSEI Emotion (' + str(emotion) + ')", train_loss, e)
        writer.add_scalar("Test Loss - MOSEI Emotion (' + str(emotion) + ')", test_loss, e)
        if best_loss == None or best_loss > test_loss:
            best_loss, best_label, best_pred, best_mask, best_attn = test_loss, test_label, test_pred, test_mask, attentions

    print('Test performance..')
    print('Loss {} accuracy {}'.format(best_loss, round(accuracy_score(best_label,best_pred,sample_weight=best_mask)*100,2)))
    print(classification_report(best_label,best_pred,sample_weight=best_mask,digits=4))
    print(confusion_matrix(best_label,best_pred,sample_weight=best_mask))

ext: py

max_stars_repo_name: Vakihito/multilogue-net

max_issues_repo_name: Vakihito/multilogue-net

max_forks_repo_name: Vakihito/multilogue-net

reason: import numpy


text:
import cv2
from pathlib import Path
import numpy as np

class AsciiImage:

    ASCII_CHARS = r"B8&WM#YXQO{}[]()I1i!pao;:,.    "
    N_CHARS = len(ASCII_CHARS)

    def __init__(self, filename=None, neighborhood=1,
                 aspect_ratio=None, padding=True, mode="gray"):
        self.neighborhood = neighborhood
        self.aspect_ratio = aspect_ratio
        self.padding = padding
        self.mode = mode
        if filename is not None:
            self.load_image(filename=filename)
        self.convert()

    def load_image(self, filename=None, image_rgb=None):
        if image_rgb is not None:
            self.original_image = image_rgb
        elif Path(filename).exists():
            filename = Path(filename)
            self.original_image = cv2.imread(filename.as_posix())[::-1]
        else:
            raise ValueError(f"{filename} doesn't exists.")
        self.original_dim = self.original_image.shape
        if self.mode == "gray":
            self.src_image = cv2.cvtColor(self.original_image, cv2.COLOR_RGB2GRAY)
        else:
            #TODO: add knn color
            #self.src_image = cv2.cvtColor(self.original_image, cv2.COLOR_RGB2GRAY)
            raise NotImplementedError()

        if self.neighborhood != 1:
            f_ = 1./self.neighborhood
            self.src_image = cv2.resize(self.src_image, None, fx=f_, fy=f_)

    def get_ascii(self, pixels_value):
        return AsciiImage.ASCII_CHARS[int(((AsciiImage.N_CHARS-1)*pixels_value)/256)]

    def convert(self):
        self.ascii = []
        rows, cols = self.src_image.shape
        for i in range(rows):
            self.ascii.append("")
            for j in range(cols):
                pixel = self.src_image[i,j]
                char = self.get_ascii(pixel)
                self.ascii[i] += char
        self.ascii = self.ascii[::-1]

    def save(self, output_path):
        output_filepath = Path(output_path)
        if output_filepath.suffix == ".txt":
            self.save_to_textfile(output_path)
        elif output_filepath.suffix in (".jpg", ".jpeg"):
            self.save_to_image(output_path)
        else:
            raise ValueError(f"Invalid extension: '{output_filepath.suffix}'")

    def save_to_textfile(self, output_path):
        with open(output_path, "w") as f:
            f.write(str(self) + "\n")

    def save_to_image(self, output_path):
        pass

    def __str__(self):
        return "\n".join(self.ascii)

max_stars_repo_name: yast-ia/ascii-art

max_issues_repo_name: yast-ia/ascii-art

max_forks_repo_name: yast-ia/ascii-art


text:
import unittest
import heat
import numpy as np

class TestHeatEquation(unittest.TestCase):
    def test_zero(self):
        initial_data = lambda x: 0
        dt = 1 / 1024.
        dx = dt

        end_time = 1.25

        solution_to_heat_equation = heat.solve_heat_equation(initial_data, dt, dx, end_time)

        self.assertEqual(int(1 / dt), solution_to_heat_equation.shape[0])
        self.assertTrue(np.all(solution_to_heat_equation == np.zeros_like(solution_to_heat_equation)))

    def test_sine_single(self):
        # we do a quick convergence test to make sure it is indeed second order
        initial_data = lambda x: np.sin(np.pi * x)

        resolutions = 2.0 ** np.arange(-5, -12, -1)
        errors = []

        end_time = 1.25

        for dx in resolutions:
            dt = dx

            solution_to_heat_equation = heat.solve_heat_equation(initial_data, dt, dx, end_time)

            self.assertEqual(int(1 / dx), solution_to_heat_equation.shape[0])

            x = np.arange(0, 1, dx)

            exact_solution = np.exp(-np.pi ** 2 * end_time) * initial_data(x)

            difference_in_l2_norm = np.linalg.norm((exact_solution - solution_to_heat_equation) * dx, ord=2)

            errors.append(difference_in_l2_norm)

        convergence_rate = np.polyfit(np.log(resolutions), np.log(errors), 1)[0]

        self.assertGreaterEqual(convergence_rate, 2)

    def test_sine_three_modes(self):
        # we do a quick convergence test to make sure it is indeed second order
        resolutions = 2.0 ** np.arange(-5, -12, -1)
        errors = []

        end_time = 1.25

        coefficients = [0.4, 0.2, 0.7]
        for dx in resolutions:
            dt = dx

            initial_data = heat.InitialDataControlSine(coefficients)

            solution_to_heat_equation = heat.solve_heat_equation(initial_data, dt, dx, end_time)

            self.assertEqual(int(1 / dx), solution_to_heat_equation.shape[0])

            x = np.arange(0, 1, dx)

            exact_solution = initial_data.exact_solution(x, end_time)

            difference_in_l2_norm = np.linalg.norm((exact_solution - solution_to_heat_equation) * dx, ord=2)

            errors.append(difference_in_l2_norm)

        convergence_rate = np.polyfit(np.log(resolutions), np.log(errors), 1)[0]

        self.assertGreaterEqual(convergence_rate, 2)

    def test_sine_single_different_coefficient(self):
        # we do a quick convergence test to make sure it is indeed second order
        initial_data = lambda x: np.sin(np.pi * x)

        resolutions = 2.0 ** np.arange(-5, -12, -1)
        errors = []

        end_time = 1.25

        q = 0.8
        for dx in resolutions:
            dt = dx

            solution_to_heat_equation = heat.solve_heat_equation(initial_data, dt, dx, end_time, q=q)

            self.assertEqual(int(1 / dx), solution_to_heat_equation.shape[0])

            x = np.arange(0, 1, dx)

            exact_solution = np.exp(-q * np.pi ** 2 * end_time) * initial_data(x)

            difference_in_l2_norm = np.linalg.norm((exact_solution - solution_to_heat_equation) * dx, ord=2)

            errors.append(difference_in_l2_norm)

        convergence_rate = np.polyfit(np.log(resolutions), np.log(errors), 1)[0]

        self.assertGreaterEqual(convergence_rate, 1.9)

    def test_sine_three_modes_different_coefficient(self):
        # we do a quick convergence test to make sure it is indeed second order
        resolutions = 2.0 ** np.arange(-5, -12, -1)
        errors = []

        end_time = 1.25

        coefficients = [0.4, 0.2, 0.7]

        q = 1.3
        for dx in resolutions:
            dt = dx

            initial_data = heat.InitialDataControlSine(coefficients)

            solution_to_heat_equation = heat.solve_heat_equation(initial_data, dt, dx, end_time, q=q)

            self.assertEqual(int(1 / dx), solution_to_heat_equation.shape[0])

            x = np.arange(0, 1, dx)

            exact_solution = initial_data.exact_solution(x, end_time, q)

            difference_in_l2_norm = np.linalg.norm((exact_solution - solution_to_heat_equation) * dx, ord=2)

            errors.append(difference_in_l2_norm)

        convergence_rate = np.polyfit(np.log(resolutions), np.log(errors), 1)[0]

        self.assertGreaterEqual(convergence_rate, 1.9)

if __name__ == '__main__':
    unittest.main()


text:
import numpy as np

#Making an empty list with 16 elements
n = 16
x = np.zeros(n)

#Setting x[0] and x[1] to their given values
x[0] = 1
x[1] = 1

#Printing out x[0] and x[1]
print(x[0])
print(x[1])

#Calculating the rest of the fiboncacci string using recursion
for i in range(2,n):
    x[i]= float(x[i-1]) + float(x[i-2])
    print(f"{x[i]}")


text:
"""
This file is part of midas-cpp

MIT License

Copyright (c) 2020 Shing Yan Loo (lsyan@ualberta.ca)

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
"""

import torch
import midas

import glob
import os

import cv2
import numpy as np

def visualize_minmax(arr, cmap=None, min_val=None, max_val=None):
    """
    Given a 2D NumPy array (h x w), visualize the value (pixel)
    based on the min-max range

    :param arr: 2D NumPy array
    :param cmap: (Optional) One of the OpenCV colormaps (https://docs.opencv.org/master/d3/d50/group__imgproc__colormap.html)
    :param min_val: (Optional) clip the min range
    :param max_val: (Optional) clip the max range
    :return: heatmap
    """
    # Clip min-max range
    if min_val is not None:
        arr[arr < min_val] = min_val
    if max_val is not None:
        arr[arr > max_val] = max_val

    # Normalize the range
    arr_norm = cv2.normalize(arr, None, norm_type=cv2.NORM_MINMAX)

    # Apply OpenCV colormap
    if cmap is not None:
        arr_norm = np.expand_dims(arr_norm, axis=-1) * 255
        arr_norm = arr_norm.astype(np.uint8)
        arr_norm = cv2.applyColorMap(arr_norm, cmap)

    return arr_norm

if __name__ == '__main__':
    # Load an image
    # image = cv2.imread('../sample_image/bicycle.jpg')
    image = cv2.imread('../sample_image/horses.jpg')

    # Load MiDaS
    height, width = image.shape[:2]
    midas = midas.MiDas(width, height, '../traced_model.pt')

    # Predict depth
    depth = midas.inference(image)

    # Visualize depth
    depth_visual = visualize_minmax(depth, cmap=cv2.COLORMAP_JET)
    cv2.imshow('depth prediction', np.vstack((image, depth_visual)))
    cv2.waitKey(0)


text:
# MIT License: Copyright (c) 2021 Lorenzo Loconte, Gennaro Gala

from __future__ import annotations
import abc
from typing import Optional, Union, List, Iterator
from collections import deque, defaultdict

import numpy as np
from scipy.special import logsumexp

class Node(abc.ABC):
    def __init__(self, scope: List[int], children: Optional[List[Node]] = None):
        """
        Initialize a SPN node given the children list and its scope.

        :param scope: The scope.
        :param children: A list of nodes. If None, children are initialized as an empty list.
        :raises ValueError: If the scope is empty.
        :raises ValueError: If the scope contains duplicates.
        """
        if not scope:
            raise ValueError("The scope must not be empty")
        if len(scope) != len(set(scope)):
            raise ValueError("The scope must not contain duplicates")
        if children is None:
            children = list()

        self.id = 0
        self.scope = scope
        self.children = children

    @abc.abstractmethod
    def likelihood(self, x: np.ndarray) -> np.ndarray:
        """
        Compute the likelihood of the node given some input.

        :param x: The inputs.
        :return: The resulting likelihoods.
        """

    @abc.abstractmethod
    def log_likelihood(self, x: np.ndarray) -> np.ndarray:
        """
        Compute the logarithmic likelihood of the node given some input.

        :param x: The inputs.
        :return: The resulting log-likelihoods.
        """

class Sum(Node):
    def __init__(
        self,
        scope: Optional[List[int]] = None,
        children: Optional[List[Node]] = None,
        weights: Optional[Union[List[float], np.ndarray]] = None,
    ):
        """
        Initialize a SPN sum node given a list of children and their weights and a scope.

        :param scope: The scope. If None, the scope is initialized based on children scopes.
        :param children: A list of nodes. If None, children are initialized as an empty list.
        :param weights: The weights associated to each children node. It can be None.
        :raises ValueError: If both scope and children are None.
        :raises ValueError: If children nodes have different scopes.
        :raises ValueError: If the length of weights and children are different.
        :raises ValueError: If weights don't sum up to 1.
        """
        if children is None:
            if scope is None:
                raise ValueError("Cannot infer Sum node's scope without children")
        else:
            if scope is None:
                scope = children[0].scope
            s_scope = set(scope)
            if any(map(lambda c: set(c.scope) != s_scope, children[1:])):
                raise ValueError("Children of Sum node have different scopes")
            if weights is not None and len(weights) != len(children):
                raise ValueError("Weights and children length mismatch")

        if weights is not None:
            if isinstance(weights, list):
                weights = np.array(weights, dtype=np.float32)
            if not np.isclose(np.sum(weights), 1.0):
                raise ValueError("Weights don't sum up to 1")
        self.weights = weights

        super().__init__(scope, children)

    def em_init(self, random_state: np.random.RandomState):
        """
        Random initialize the node's parameters for Expectation-Maximization (EM).

        :param random_state: The random state.
        """
        weights = random_state.dirichlet(np.ones(len(self.children)))
        self.weights = weights.astype(np.float32)

    def em_step(self, stats: np.ndarray, step_size: float):
        """
        Compute a batch Expectation-Maximization (EM) step.

        :param stats: The sufficient statistics of each sample.
        :param step_size: The step size of update.
        """
        unnorm_weights = self.weights * np.sum(stats, axis=1) + np.finfo(np.float32).eps
        weights = unnorm_weights / np.sum(unnorm_weights)

        # Update the parameters
        self.weights = (1.0 - step_size) * self.weights + step_size * weights

    def likelihood(self, x: np.ndarray) -> np.ndarray:
        return np.expand_dims(np.dot(x, self.weights), axis=1)

    def log_likelihood(self, x: np.ndarray) -> np.ndarray:
        return logsumexp(x, b=self.weights, axis=1, keepdims=True)

class Product(Node):
    def __init__(
        self,
        scope: Optional[List[int]] = None,
        children: Optional[List[Node]] = None
    ):
        """
        Initialize a product node given a list of children and its scope.

        :param scope: The scope. If None, the scope is initialized based on children scopes.
        :param children: A list of nodes. If None, children are initialized as an empty list.
        :raises ValueError: If both scope and children are None.
        :raises ValueError: If children nodes don't have disjointed scopes.
        """
        if children is None:
            if scope is None:
                raise ValueError("Cannot infer Product node's scope without children")
        else:
            c_scope = list(sum([c.scope for c in children], []))
            s_scope = set(c_scope)
            if scope is None:
                if len(c_scope) != len(s_scope):
                    raise ValueError("Children of Product node don't have disjointed scopes")
                scope = c_scope
            elif set(scope) != s_scope:
                raise ValueError("Children of Product node don't have disjointed scopes")

        super().__init__(scope, children)

    def likelihood(self, x: np.ndarray) -> np.ndarray:
        return np.prod(x, axis=1, keepdims=True)

    def log_likelihood(self, x: np.append) -> np.ndarray:
        return np.sum(x, axis=1, keepdims=True)

def assign_ids(root: Node) -> Node:
    """
    Assign the ids to the nodes of a SPN.

    :param root: The root of the SPN.
    :return: The same SPN with each node having modified ids.
    :raises ValueError: If the SPN structure is not a DAG.
    """
    nodes = topological_order(root)
    if nodes is None:
        raise ValueError("SPN structure is not a directed acyclic graph (DAG)")

    next_id = 0
    for node in nodes:
        node.id = next_id
        next_id += 1
    return root

def bfs(root: Node) -> Iterator[Node]:
    """
    Compute the Breadth First Search (BFS) ordering for a SPN.

    :param root: The root of the SPN.
    :return: The BFS nodes iterator.
    """
    seen, queue = {root}, deque([root])
    while queue:
        node = queue.popleft()
        yield node
        for c in node.children:
            if c not in seen:
                seen.add(c)
                queue.append(c)

def dfs_post_order(root: Node) -> Iterator[Node]:
    """
    Compute Depth First Search (DFS) Post-Order ordering for a SPN.

    :param root: The root of the SPN.
    :return: The DFS Post-Order nodes iterator.
    """
    seen, stack = {root}, [root]
    while stack:
        node = stack[-1]
        if set(node.children).issubset(seen):
            stack.pop()
            yield node
            continue
        for c in node.children:
            if c not in seen:
                seen.add(c)
                stack.append(c)

def topological_order(root: Node) -> Optional[List[Node]]:
    """
    Compute the Topological Ordering for a SPN, using the Kahn's Algorithm.

    :param root: The root of the SPN.
    :return: A list of nodes that form a topological ordering.
             If the SPN graph is not acyclic, it returns None.
    """
    ordering = list()
    num_outgoings = defaultdict(int)
    num_outgoings[root] = 0

    # Initialize the number of outgoings edges for each node
    for node in bfs(root):
        for c in node.children:
            num_outgoings[c] += 1

    # Check the unusual case where the root node have outgoings edges, i.e. a trivial cycle has been found
    if num_outgoings[root] != 0:
        return None

    # Non-layered topological ordering implementation
    queue = deque([root])
    while queue:
        node = queue.popleft()
        ordering.append(node)
        for c in node.children:
            num_outgoings[c] -= 1
            if num_outgoings[c] == 0:
                queue.append(c)

    # Check if a cycle has been found
    if sum(num_outgoings.values()) != 0:
        return None
    return ordering

def topological_order_layered(root: Node) -> Optional[List[List[Node]]]:
    """
    Compute the Topological Ordering Layered for a SPN, using the Kahn's Algorithm.

    :param root: The root of the SPN.
    :return: A list of layers that form a topological ordering.
             If the SPN graph is not acyclic, it returns None.
    """
    ordering = list()
    num_outgoings = defaultdict(int)
    num_outgoings[root] = 0

    # Initialize the number of outgoings edges for each node
    for node in bfs(root):
        for c in node.children:
            num_outgoings[c] += 1

    # Check the unusual case where the root node have outgoings edges, i.e. a trivial cycle has been found
    if num_outgoings[root] != 0:
        return None

    # Layered topological ordering implementation
    ordering.append([root])
    while True:
        layer = list()
        for node in ordering[-1]:
            for c in node.children:
                num_outgoings[c] -= 1
                if num_outgoings[c] == 0:
                    layer.append(c)
        if not layer:
            break
        ordering.append(layer)

    # Check if a cycle has been found
    if sum(num_outgoings.values()) != 0:
        return None
    return ordering

max_stars_repo_name: deeprob-org/deeprob-kit

max_issues_repo_name: deeprob-org/deeprob-kit

max_forks_repo_name: deeprob-org/deeprob-kit

reason: import numpy,from scipy


text:
from math import sqrt, exp, pi
from itertools import zip_longest
import numpy as np
import matplotlib.pyplot as plt

class Kalman1D(object):
    def __init__(self, init_mean, init_var, measure_var, motion_var):
        """ Initialize prior mean and variance """
        self.mean = init_mean
        self.var = init_var
        self.measure_var = measure_var
        self.motion_var = motion_var

    def gaussian(self, x):
        """ Gaussian Function, for a given mean, variance and x return a Gaussian value """
        coefficient = 1.0 / sqrt(2.0 * pi * self.var)
        exponential = exp(-0.5 * (x - self.mean)**2 / self.var)
        return coefficient * exponential

    def update(self, measurement_mean, measurement_var):
        """ Updates the Gaussian parameters based on a measurements mean and variance after measurement"""
        self.mean = (measurement_var * self.mean + self.var * measurement_mean) / (measurement_var + self.var)
        # Variance is smaller than previous, which means that after measurement we are more certain of location
        self.var = 1 / (1/self.var + 1/measurement_var)

    def predict(self, motion_mean, motion_var):
        """ Updates the Gaussian parameters based on motion mean and variance after motion """
        self.mean += motion_mean
        # Variance is summed, which means that after motion we are more uncertain of location
        self.var += motion_var

    def cycle(self, motions, measurements):
        """ Cycles the update and predict operations of object in environment """
        if len(motions) == len(measurements):
            for motion, measurement in zip(motions, measurements):
                self.update(measurement, self.measure_var)
                self.predict(motion, self.motion_var)
        # If there is more motions than measurements
        elif len(motions) > len(measurements):
            for motion, measurement in zip_longest(motions, measurements):
                if measurement is not None:
                    self.update(measurement, self.measure_var)
                    self.predict(motion, self.motion_var)
                else:
                    # Do remaining motion
                    self.predict(motion, self.motion_var)
        # If there is more measurements than motions
        elif len(motions) < len(measurements):
            for motion, measurement in zip_longest(motions, measurements):
                if motion is not None:
                    self.update(measurement, self.measure_var)
                    self.predict(motion, self.motion_var)
                else:
                    # Do remaining measurement
                    self.update(measurement, self.measure_var)

    def display_distribution(self):
        """ Plots the probability distribution of object location """
        x_axis = np.arange(-30, 30, 0.1)
        distribution = []
        for x in x_axis:
            distribution.append(self.gaussian(x))
        print(f'Gaussian Distribution: \nMean: {round(self.mean, 2)} Variance: {round(self.var, 2)}')
        plt.plot(x_axis, distribution, color='brown')
        plt.title('Probability distribution of object location')
        plt.xlabel('Location')
        plt.ylabel('Probability')
        plt.grid()
        plt.show()

max_stars_repo_name: TomislavZupanovic/CVND-Localization

max_issues_repo_name: TomislavZupanovic/CVND-Localization

max_forks_repo_name: TomislavZupanovic/CVND-Localization


text:
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.4'
#       jupytext_version: 1.1.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # S_SizeSignals [<img src="https://www.arpm.co/lab/icons/icon_permalink.png" width=30 height=30 style="display: inline;">](https://www.arpm.co/lab/redirect.php?code=S_SizeSignals&codeLang=Python)
# For details, see [here](https://www.arpm.co/lab/redirect.php?permalink=eb-size-signal).

# ## Prepare the environment

# +
import os
import os.path as path
import sys

sys.path.append(path.abspath('../../functions-legacy'))

import numpy as np
from numpy import arange, array, argsort, linspace, diag, round, log, exp, sqrt, zeros, sum as npsum

from scipy.io import loadmat

import matplotlib.pyplot as plt
from matplotlib.pyplot import plot, xlim, ylim, subplots, title

plt.style.use('seaborn')
np.seterr(invalid='ignore')

from CONFIG import GLOBAL_DB, TEMPORARY_DB
from ARPM_utils import save_plot, date_mtop
from FPmeancov import FPmeancov
from EwmaFP import EwmaFP
# -

# ## Upload the database

# +
try:
    db = loadmat(os.path.join(GLOBAL_DB, 'db_strategies'), squeeze_me=True)
except FileNotFoundError:
    db = loadmat(os.path.join(TEMPORARY_DB, 'db_strategies'), squeeze_me=True)

last_price = db['last_price']
dates = db['dates']
s_siz = db['s_siz']

# settings
v = last_price
[n_, t_] = v.shape
t_start = 252*2  # starting point of the strategy
# -

# ## Compute the smoothed signals

# +
tauHL_smoo = log(2) / 10
t_smoo = 180
s_siz_smoo = zeros((n_,t_-t_smoo+1))

for t in range(t_smoo,s_siz.shape[1]+1):
    s_siz_smoo[:, [t - t_smoo]] = EwmaFP(s_siz[:, t - t_smoo :t], tauHL_smoo)[0]
# -

# ## Use the smoothed signals to compute the scored signal

t_scor = 252
s_siz_scor = zeros((n_,s_siz_smoo.shape[1]-t_scor+1))
tauHL_scor = log(2) / 120
p_scor = exp(-tauHL_scor*arange(t_scor - 1, 0 + -1, -1)).reshape(1,-1) / npsum(exp(-tauHL_scor*arange(t_scor - 1, 0 + -1, -1)))
for t in arange(t_scor,s_siz_smoo.shape[1]+1):
    mu_siz, cov_siz = FPmeancov(s_siz_smoo[:, t - (t_scor):t], p_scor)
    s_siz_scor[:, t - t_scor] = (s_siz_smoo[:,t-1] - mu_siz.flatten()) / sqrt(diag(cov_siz))

# ## Use the scored signals to compute the ranked signals

s_siz_rk = zeros((n_,s_siz_scor.shape[1]))
for t in range(s_siz_scor.shape[1]):
    rk = argsort(s_siz_scor[:,t])
    rk_signal = argsort(rk)+1
    s_siz_rk[:,t] = (rk_signal - 0.5*n_)*(2 / n_)

# ## Compare the plots of one signal, one smoothed signal and one scored signal

dates = dates[t_start-1:]
grid_dates = linspace(0, len(dates)-1, 5)
grid_dates = list(map(int,round(grid_dates)))  # integer numbers
index = argsort(s_siz_rk[:,-1])

# ## Compare the plots of a cluster of 4 scored signals with their ranked counterparts

dates_dt = array([date_mtop(i) for i in dates])
date_tick = dates_dt[grid_dates]
f, ax = subplots(2,1)
plt.sca(ax[0])
xx = t_start-1
plot(dates_dt, s_siz[index[int(round(n_*0.2))-1], xx:])
plt.xticks(dates_dt[grid_dates])
xlim([dates_dt[0], dates_dt[-1]])
title('Size versus smoothed size signal')
xx = t_start - t_smoo
plot(dates_dt, s_siz_smoo[index[int(round(n_*0.2))-1], xx:], 'r')
plt.xticks(dates_dt[grid_dates])
xlim([dates_dt[0], dates_dt[-1]])
plt.sca(ax[1])
xx = t_start - t_smoo - t_scor+1
plot(dates_dt, s_siz_scor[index[int(round(n_*0.2))-1], xx:])
plt.xticks(dates_dt[grid_dates])
xlim([dates_dt[0], dates_dt[-1]])
title('Scored size signal')
plt.tight_layout();
# save_plot(ax=plt.gca(), extension='png', scriptname=os.path.basename('.')[:-3], count=plt.get_fignums()[-1])
f, ax = subplots(2,1)
plt.sca(ax[0])
plot(dates_dt, s_siz_scor[[index[int(round(n_*0.2))-1], index[int(round(n_*0.4))-1], index[int(round(n_*0.6))-1],index[int(round(n_*0.8))-1]],xx:].T)
plt.xticks(dates_dt[grid_dates])
xlim([dates_dt[0], dates_dt[-1]])
title('Scored size signal cluster')
plt.sca(ax[1])
plot(dates_dt,s_siz_rk[[index[int(round(n_*0.2))-1], index[int(round(n_*0.4))-1], index[int(round(n_*0.6))-1], index[int(round(n_*0.8))-1]], xx:].T)
plt.xticks(dates_dt[grid_dates])
xlim([dates_dt[0], dates_dt[-1]])
ylim([-1.05, 1.05])
title('Ranked size signal cluster')
plt.tight_layout();
# save_plot(ax=plt.gca(), extension='png', scriptname=os.path.basename('.')[:-3], count=plt.get_fignums()[-1])

max_stars_repo_name: dpopadic/arpmRes

max_issues_repo_name: dpopadic/arpmRes

max_forks_repo_name: dpopadic/arpmRes

reason: import numpy,from numpy,from scipy