def pupil_quant_cmh(avi_path):

    # -*- coding: utf-8 -*-
    #import libraries and define movie
    from CVROIGrabber_MRtest import CVROIGrabber
    from src.ext.PupilTracker_MRtest import PupilTracker
    import cv2, scipy.io, numpy as np

    #define pupil area
    cap = cv2.VideoCapture(avi_path) #avi_path is the video being called
    cap.set(1,1); #Take the first frame
    ret, frame = cap.read() # Read the frame
    cv = CVROIGrabber(frame)
    cv.grab() # you can drag the window. Q exits

    #Pupil size tracking parameters (why are these in a function)
    param = dict(
    relative_area_threshold=0.01,#.001
    ratio_threshold=2, #2
    error_threshold=1,
    min_contour_len=5,#5
    margin=0.02,
    contrast_threshold=0.2, #0.2
    speed_threshold=0.2,
    dr_threshold=0.2,#.2
    gaussian_blur=5,#5
    perc_high =90,
    perc_low = 10,
    perc_weight = 0.9,
    )

    tr = PupilTracker(param)
    Pupildata = tr.track(avi_path, cv.roi, display=True)

    #save
    np.save('Pupildata.npy',Pupildata)
    Pupildata = np.load('Pupildata.npy',allow_pickle=True)
    scipy.io.savemat(avi_path[:-4]+'_pupildata.mat', {"Pupildata": Pupildata})

    return


# coding=utf-8
from math import sqrt, floor, ceil   # Mathematische Hilfsfunktionen
from scipy.optimize import fmin      # Zum Ermitteln von Minima einer Funktion
from copy import deepcopy            # Zum Kopieren von Hashtables/Dictionaries
import os                            # Zum Lesen von Ordnern
import datetime                      # Zum Konvertieren von Sekunden in eine Uhrzeit

# Repraesentiert einen Punkt in Form einer Koordinate oder einen Vektor mit dessen x- und y-Wert
class Point:
    def __init__(self, x, y):
        self.x = x      # x-Koordinate
        self.y = y      # y-Koordinate

    # Hash-Funktion
    def __hash__(self):
        return hash((self.x, self.y))

    # Punkt als Zeichenkette
    def __repr__(self):
        return " (" + str(self.x) + " | " + str(self.y) + ") "

    # Vergleich-Funktion
    def __eq__(self, other):
        return (self.x, self.y) == (other.x, other.y)

    # Euklidische Distanz zweier Punkte berechnen (Anwendung des Satz des Pythagoras)
    def euclidean_distance(self, other_point):
        cathetus_a = other_point.x - self.x
        cathetus_b = other_point.y - self.y
        hypotenuse = sqrt(cathetus_a**2 + cathetus_b**2)
        return hypotenuse

# Repraesentiert eine Kante, bzw. eine Strecke mit zwei definierten Punkten A und B
class Edge:
    def __init__(self, point_a, point_b):
        self.point_a = point_a
        self.point_b = point_b

    # Vergleich-Funktion
    def __eq__(self, other):
        return (self.point_a, self.point_b) == (other.point_a, other.point_b) \
               or (self.point_a, self.point_b) == (other.point_b, other.point_a)

    # Erhalte Schnittpunkt mit einer anderen Kante, wobei beide Kanten als Geraden fortgefuehrt werden
    def get_intersection_with(self, other_edge):
        # Als allgm. Geradengleichung gilt: y=mx+b.
        # Hier: Zwei Geraden mit jeweils m1 bzwn. m2 und b1 bzw. b2
        # Es gilt: m1 = dy1/dx1 und m2 = dy2/dx2
        dy1 = self.point_b.y - self.point_a.y
        dy2 = other_edge.point_b.y - other_edge.point_a.y
        dx1 = float(self.point_b.x - self.point_a.x)
        dx2 = float(other_edge.point_b.x - other_edge.point_a.x)

        # Wenn Kante a parallel zur y-Achse ist
        if dx1 == 0.0:
            intersection_x = self.point_a.x
            m2 = dy2 / dx2
            b2 = other_edge.point_a.y - m2 * other_edge.point_a.x
            intersection_y = m2 * intersection_x + b2
            return Point(intersection_x, intersection_y)

        # Steigung m1 kann berechnet werden, da dx1 ungleich 0 ist
        m1 = dy1 / dx1
        b1 = self.point_a.y - m1 * self.point_a.x

        # Wenn Kante b parallel zur y-Achse ist
        if dx2 == 0.0:
            intersection_x = other_edge.point_a.x
            intersection_y = m1 * intersection_x + b1
            return Point(intersection_x, intersection_y)

        # Steigung m2 kann berechnet werden, da dx2 ungleich 0 ist
        m2 = dy2 / dx2
        b2 = other_edge.point_a.y - m2 * other_edge.point_a.x
        intersection_x = (b2-b1)/float(m1-m2)
        intersection_y = m1 * intersection_x + b1

        return Point(intersection_x, intersection_y)

    # Prueft, ob Kante sich ausserhalb der Flaeche von zwei aufgespannten Punkten befindet
    def is_out_of_area(self, point_a, point_b):
        # Pruefen, ob Kante im x-Bereich zwischen den Punkten A und B liegt
        if not point_b.x < point_a.x:
            # Punkt A ist links von B oder ueber bzw. unter B
            if self.point_a.x <= point_a.x and self.point_b.x <= point_a.x:
                # Kante liegt links ausserhalb von A
                return True
            if self.point_a.x >= point_b.x and self.point_b.x >= point_b.x:
                # Kante liegt rechts ausserhalb von B
                return True
        if not point_a.x < point_b.x:
            # Punkt A ist rechts von B oder ueber bzw. unter B
            if self.point_a.x >= point_a.x and self.point_b.x >= point_a.x:
                # Kante liegt rechts ausserhalb von A
                return True
            if self.point_a.x <= point_b.x and self.point_b.x <= point_b.x:
                # Kante liegt links ausserhalb von B
                return True

        # Pruefen, ob Kante im y-Bereich zwischen den Punkten A und B liegt
        if not point_b.y < point_a.y:
            # Punkt B ist ueber A oder auf gleicher hoehe
            if self.point_a.y <= point_a.y and self.point_b.y <= point_a.y:
                # Kante liegt unterhalb von A
                return True
            if self.point_b.y >= point_b.y and self.point_a.y >= point_b.y:
                # Kante liegt oberhalb von B
                return True
        if not point_a.y < point_b.y:
            # Punkt A ist ueber B oder auf gleicher hoehe
            if self.point_a.y >= point_a.y and self.point_b.y >= point_a.y:
                # Kante liegt oberhalb von A oder ist auf gleicher hoehe
                return True
            if self.point_a.y <= point_b.y and self.point_b.y <= point_b.y:
                # Kante liegt liegt unterhalb von B oder ist auf gleicher hoehe
                return True

        return False

    # Erhalte den Faktor mit dem die Kante erweitert werden muss, um auf einen bestimmten Punkt zu treffen
    def get_ratio(self, point):
        # Zu pruefende Kante ist parallel zur Y-Achse --> Verhaeltnis ueber y-Werte
        if self.point_b.x == self.point_a.x:
            ratio = (point.y - self.point_a.y) / float(self.point_b.y - self.point_a.y)
        # Sonst: Verhaeltnis ueber x-Werte
        else:
            ratio = (point.x - self.point_a.x) / float(self.point_b.x - self.point_a.x)

        return ratio

# Repraesentiert einen Knoten im Sichtbarkeitsgraphen
class Node:
    def __init__(self, point, id):
        self.point = point  # Jeder Knoten hat einen Punkt
        self.id = id        # Und eine ID, fuer Polygone gilt: ID > 0, Startknoten hat die ID -1, der Endknoten -2

        # Attribute, die fuer den A-Star-Algorithmus relevant sind
        self.g_cost = 0
        self.h_cost = 0
        self.f_cost = 0
        self.previous_node = 0

    # Methode zum Zuweisen der Heuristik (hier: euklidische Distanz)
    def assign_h_cost(self, end_node):
        self.h_cost = self.point.euclidean_distance(end_node.point)
        return self

    # Hash-Funktion, damit Knoten als Schluessel und Werte in einer Hash-Table sein koennen
    def __hash__(self):
        return hash((self.point, self.id))

    # Vergleich-Funktion (Knoten sind identisch, wenn sie den selben Punkt und die selbe ID haben)
    def __eq__(self, other):
        return (self.point, self.id) == (other.point, other.id)

    # Knoten als Zeichenkette
    def __repr__(self):
        if self.id == -2:   # Endknoten
            return " [Y, (" + str(self.point.x) + " | " + str(self.point.y) + ")]"
        if self.id == -1:
            return " [L, (" + str(self.point.x) + " | " + str(self.point.y) + ")]"
        return " [P" + str(self.id) + ", (" + str(self.point.x) + " | " + str(self.point.y) + ")]"

    # Prueft, ob aktueller Knoten A einen anderen Knoten B sehen kann
    def can_see(self, node_b):
        if node_b.id != -2:
            # Zugehoeriges Polygon des ersten Knoten
            polygon_of_first_node = Input.polygon_list[self.id-1]
        # Pruefen ob beide Knoten zu einem Polygon gehoeren
        if self.id == node_b.id:
            for edge in polygon_of_first_node.convert_to_edges():
                # Wenn Knoten sich eine Kante teilen, ...
                current_points = [edge.point_a, edge.point_b]
                if self.point in current_points and node_b.point in current_points:
                    # ... dann sehen sie sich in jedem Fall
                    return True

        # Richtungsvektor der Verbindung zwischen Knoten a und b
        vector_connection = Point(node_b.point.x - self.point.x, node_b.point.y - self.point.y)

        # Potentielle Hindernisse sind jene Kanten, die die Verbindung a zu b lediglich beruehren
        potential_obstacles = []

        # Pruefen, ob eine aller Kanten den Weg von node_a zu node_b verdeckt
        for edge in Input.list_of_all_edges:
            # Richtungsvektor der zu pruefenden Kante
            vector_edge = Point(edge.point_b.x - edge.point_a.x, edge.point_b.y - edge.point_a.y)
            # Pruefen, ob Verbindung a zu b und aktuelle Kante parallel sind
            if vector_connection.x*vector_edge.y == vector_connection.y * vector_edge.x:
                continue    # Kante kann Verbindung nicht schneiden

            # Falls sich die Kante ausserhalb des von a und b aufgespannten Rechtecks befindet
            if edge.is_out_of_area(self.point, node_b.point):
                continue    # Kante kann Verbindung nicht schneiden

            intersection_point = edge.get_intersection_with(Edge(self.point, node_b.point))
            ratio_edge = round(edge.get_ratio(intersection_point), 6)
            ratio_connection = round(Edge(self.point, node_b.point).get_ratio(intersection_point), 6)

            # Kante verdeckt beide Knoten
            if 0 < ratio_edge < 1 and 0 < ratio_connection < 1:
                return False

            # Verbindung koennte durch Eckpunkt eines Polygons gehen, Kanten beruehren die Verbindung a zu b
            elif (ratio_edge == 1 or ratio_edge == 0) and 0 < ratio_connection < 1:
                potential_obstacles.append(edge)

        # Knoten teilt Punkt mit anderen Knoten
        if node_b.point in Input.double_corners or \
                self.point in Input.double_corners:
            return False    # Doppelte Eckpunkte werden ausgelassen

        # Pruefen, ob Verbindung im Polygon liegt bei Knoten gleicher Polygone
        if self.id == node_b.id:
            # OA + AB/2 = OP, OP ist Ortsvektor des zu pruefenden Punktes
            local_vector_checkpoint = Point(self.point.x + vector_connection.x/2,
                                            self.point.y + vector_connection.y/2)
            # Beide Knoten koennen sich nicht sehen, da die Verbindung im Polygon liegt
            if polygon_of_first_node.contains(local_vector_checkpoint):
                return False
            else:
                return True

        # Pruefen, ob Verbindung durch eine Ecke geht, ohne Kanten zu schneiden
        # Dabei muss es benachbarte Kanten geben, die weder Knoten A noch B enthalten
        if len(potential_obstacles) > 1:
            for i in range(len(potential_obstacles)):
                for j in range(i+1, len(potential_obstacles)):
                    points_in_edges = [potential_obstacles[i].point_a, potential_obstacles[i].point_b,
                                       potential_obstacles[j].point_a, potential_obstacles[j].point_b]
                    # Kanten haben einen gemeinsamen Punkt, wenn es Duplikate in der
                    # Liste (bestehend aus jeweils beiden Punkten beider Kanten) gibt
                    if len(points_in_edges) != len(set(points_in_edges)):
                        # Wenn Knoten A und B nicht auf den Kanten liegen
                        if self.point not in points_in_edges and node_b.point not in points_in_edges:
                            return False
        return True

    # Prueft, ob Eckpunkt 2 im B-Bereich von Eckpunkt 1 liegt, wodurch die Verbindung von E2 zu E1 redundant waere
    def is_redundant_to(self, node_2):

        # Polygon des Eckpunkts 1
        polygon = Input.polygon_list[self.id-1]

        # Benachbarte Ecken von Eckpunkt 1
        adjacent_points = polygon.get_adjacent_points(self.point)

        # Richtungsvektoren der beiden fortgefuehrten Kanten an Eckpunkt 1
        r1 = Point(2*self.point.x-adjacent_points[0].x, 2*self.point.y-adjacent_points[0].y)
        r2 = Point(2*self.point.x-adjacent_points[1].x, 2*self.point.y-adjacent_points[1].y)

        # Richtungsvektor vom Eckpunkt 1 zum Eckpunkt 2
        vector_connection = Point(node_2.point.x - self.point.x, node_2.point.y - self.point.y)

        # Determinante der Orientierungs-Matrix berechnen
        def det(p, q, r):
            return (q.x * r.y + p.x * q.y + p.y * r.x) - (p.y * q.x + q.y * r.x + p.x * r.y)

        determinant_r1 = det(self.point, Point(r1.x + vector_connection.x, r1.y + vector_connection.y), node_2.point)
        determinant_r2 = det(self.point, Point(r2.x + vector_connection.x, r2.y + vector_connection.y), node_2.point)

        # Falls Eckpunkt 2 zwischen den beiden fortgefuehrten Kanten von Eckpunkt 1 liegt
        # Wenn also die Determinante der Orientierungsmatrix einmal negativ und einmal positiv ist
        if determinant_r1 < 0 < determinant_r2 or determinant_r2 < 0 < determinant_r1:
            return True

        return False

# Repraesentiert ein Polygon mit seiner ID und seinen Eckpunkten
class Polygon:
    def __init__(self, points, id):
        self.points = points
        self.id = id

    # Prueft, ob ein Punkt im Polygon befindet mithilfe der Strahl-Methode
    def contains(self, point):
        polygon_edges = self.convert_to_edges()  # Alle Kanten des Polygons
        count_to_left_of_point = 0               # Variable zum Zaehlen der Schnittpunkte des Strahls links vom Punkt
        for edge in polygon_edges:
            if edge.point_a.y <= point.y <= edge.point_b.y or edge.point_b.y <= point.y <= edge.point_a.y:
                # Punkt ist auf einer Y-Ebene mit der aktuellen Kante des Polygons
                if edge.point_a.y == edge.point_b.y == point.y:
                    # Kante liegt auf Strahl --> Wird aber als darueber gezaehlt
                    continue
                intersection_point = edge.get_intersection_with(Edge(Point(0, point.y), Point(1, point.y)))
                if intersection_point in [edge.point_a, edge.point_b]:
                    # Strahl geht durch Eckpunkt der Kante, Punkte auf dem Strahl werden als Punkte darueber gezaehlt
                    if edge.point_a.y == point.y:
                        # Punkt A der Kante liegt auf dem Strahl
                        if edge.point_b.y > point.y:
                            # Kante schneidet nicht Strahl, da Punkt B und Punkt A "ueber" dem Strahl
                            continue
                    else:
                        # Punkt B liegt auf dem Strahl
                        if edge.point_a.y > point.y:
                            # Kante schneidet nicht Strahl, da Punkt B und Punkt A "ueber" dem Strahl
                            continue

                if intersection_point.x < point.x:
                    # Schnittpunkt der Kante links vom zu pruefenden Punkt
                    count_to_left_of_point += 1

        # Wenn Anzahl der Schnittpunkte links vom Punkt ungerade ist --> Punkt liegt im Polygon
        if count_to_left_of_point % 2 != 0:
            return True
        else:
            return False

    # Erhalte benachbarten Punkte eines Punktes in einem Polygon
    def get_adjacent_points(self, point):
        adjacent_points = []
        for edge in self.convert_to_edges():
            if point == edge.point_a:
                adjacent_points.append(edge.point_b)
            if point == edge.point_b:
                adjacent_points.append(edge.point_a)

        return adjacent_points

    # Zusammenfuegen von Punkten eines Polygons zu Kanten
    def convert_to_edges(self):
        edges_list = []
        number_of_corners = len(self.points)
        for i in range(number_of_corners):
            if i == number_of_corners - 1:
                edges_list.append(Edge(self.points[i], self.points[0]))
            else:
                edges_list.append(Edge(self.points[i], self.points[i + 1]))
        return edges_list

    # Gibt nur konvexe Punkte eines Polygons zurueck
    def convex_vertices_only(self):
        points = self.points
        # Finde oberste linke Ecke, da diese in jedem Fall konvex ist
        top_left_corner_index = 0
        for i in range(1, len(points)):
            if points[i].x <= points[top_left_corner_index].x:
                if points[i].y >= points[top_left_corner_index].y:
                    top_left_corner_index = i

        # neue Liste der Eckpunkte mit der obersten linken Ecke als Startwert
        new_points_list = points[top_left_corner_index:]
        new_points_list.extend(points[0:top_left_corner_index])
        points = new_points_list

        # Alle Ecken, bei denen die Kante rechts von der vorherigen Kante angelegt wird, sind in list_right
        list_right = []
        # Die restlichen Ecken kommen in list_left
        list_left = []

        # Determinante der Orientierungs-Matrix berechnen
        def det(p, q, r):
            return (q.x * r.y + p.x * q.y + p.y * r.x) - (p.y * q.x + q.y * r.x + p.x * r.y)

        # Richtungen fuer jede Ecke bestimmen mithilfe von drei Punkten, der jeweils anliegenden zwei Kanten
        for i in range(len(points)):
            if i == 0:
                point_p = points[len(points)-1]
            else:
                point_p = points[i-1]

            point_q = points[i]

            if i == len(points)-1:
                point_r = points[0]
            else:
                point_r = points[i+1]

            determinant = det(point_p, point_q, point_r)

            # Orientierung ist im Uhrzeigersinn --> Rechtsrum
            if determinant < 0:
                list_right.append(point_q)
            # Orientierung gegen den Uhrzeigersinn --> Linksrum
            else:
                list_left.append(point_q)

        if points[0] in list_right:
            # Ecken, deren Kanten einen Rechts-Turn machen sind konvex

            return list_right
        else:
            # Ecken, deren Kanten einen Links-Turn machen sind konvex
            return list_left

class VisibilityGraph:

    # Gibt einen Sichtbarkeitsgraph als Hash-Table zurueck (Rekursive Methode)
    @staticmethod
    def get_graph(to_visit, current_graph):

        # Es muessen keine Knoten mehr besucht werden
        if len(to_visit) == 1:
            return current_graph

        # Knoten wird abgearbeitet
        current_node = to_visit[0]
        to_visit.pop(0)

        # Jeden zu besuchenden Knoten mit aktuellem Knoten auf Sichtbarkeit und Redundanz pruefen
        for node in to_visit:
            if node.can_see(current_node):
                if not node.is_redundant_to(current_node):
                    current_graph[current_node].append(node)
                if not current_node.is_redundant_to(node):
                    current_graph[node].append(current_node)

        # Naechsten zu besuchenden Knoten pruefen, Erweiterung des aktuellen Sichtbarkeitsgraphen
        return VisibilityGraph.get_graph(to_visit, current_graph)

    # Methode zum Hinzufuegen eines Start- oder Endknotens in einen bereits vorhandenen Sichtbarkeitsgraphen
    @staticmethod
    def add_node(to_add, visibility_graph):
        all_key_nodes = visibility_graph.keys()
        visibility_graph[to_add] = []

        # Jeden Knoten mit zu ergaenzendem Knoten auf Sichtbarkeit und Redundanz pruefen
        for node in all_key_nodes:
            if node.can_see(to_add):
                    if node.id >= 0:
                        if not node.is_redundant_to(to_add):
                            visibility_graph[node].append(to_add)
                            visibility_graph[to_add].append(node)
                    else:
                        visibility_graph[node].append(to_add)
                        visibility_graph[to_add].append(node)

        return visibility_graph

# Klasse, die alle wichtigen Methoden zum A* Pathfinding-Algorithmus enthaelt
class AStar:

    # Findet den kuerzesten Pfad in einem Sichtbarkeitsgraphen
    @staticmethod
    def get_shortest_path(visibility_graph):
        start_node = Input.start_node
        open_list = [start_node]
        current_node = start_node
        closed_list = []

        # Solange die Open-Liste nicht leer ist
        while len(open_list) != 0:

            # Alle sichtbaren Knoten des aktuellen Knotens
            visible_nodes = visibility_graph[current_node]

            # Knoten wird als abgeschlossen markiert
            closed_list.append(current_node)
            open_list.remove(current_node)

            # Iteriere durch sichtbare Knoten des aktuellen Knotens
            for i in range(len(visible_nodes)):
                visible_node = visible_nodes[i]
                # Berechne gCost (Abstand zum Startknoten)
                new_g_cost = current_node.g_cost + current_node.point.euclidean_distance(visible_node.point)

                # Wenn sichtbarer Knoten als abgeschlossen markiert und der neue Pfad zu ihm nicht guenstiger ist
                if visible_node in closed_list and new_g_cost >= visibility_graph[current_node][i].g_cost:
                    continue

                # Falls Knoten noch nicht entdeckt oder ein guenstigerer Pfad gefunden wurde
                if visible_node not in open_list or new_g_cost < visibility_graph[current_node][i].g_cost:
                    # Aktualisieren der gCost und Gesamtkosten fuer diesen Knoten
                    visibility_graph[current_node][i].g_cost = new_g_cost
                    visibility_graph[current_node][i].f_cost = new_g_cost + visible_node.h_cost

                    # Setzen des Zeigers auf den vorherigen Knoten
                    visibility_graph[current_node][i].previous_node = current_node

                    # Falls Knoten noch nicht entdeckt
                    if visible_node not in open_list:
                        # Knoten wird als offen markiert
                        open_list.append(visible_node)

            # Erhalte Knoten mit geringsten Gesamtkosten
            if len(open_list) > 0:
                best_next_node = open_list[0]
                for node in open_list:
                    if node.f_cost < best_next_node.f_cost:
                        best_next_node = node
                current_node = best_next_node

            else:
                return None     # Kein Pfad konnte gefunden werden

            # Wenn End-Knoten Knoten mit geringsten Gesamtkosten ist
            if current_node.id == -2:
                # Optimaler Pfad wurde gefunden, muss rekonstruiert werden
                return AStar.reconstruct_path(current_node)

        # Es konnte kein Pfad gefunden werden
        return None

    # Methode zur Rekonstruierung des optimalsten Wegs
    @staticmethod
    def reconstruct_path(end_node):
        shortest_path = [end_node]
        previous_node = end_node.previous_node
        total_length = end_node.point.euclidean_distance(previous_node.point)

        # Pfad erweitern mithilfe des Vorgaenger-Knotens, Erhoehung der Gesamtlaenge
        while previous_node.id != -1:                   # Solange der Startknoten nicht erreicht ist
            shortest_path.append(previous_node)
            total_length += previous_node.point.euclidean_distance(previous_node.previous_node.point)
            if previous_node.previous_node.id == -1:    # Falls der vorherige Knoten der Startknoten ist
                total_length += previous_node.point.euclidean_distance(previous_node.previous_node.point)
            previous_node = previous_node.previous_node

        shortest_path.append(previous_node)
        return shortest_path, total_length    # Rueckgabe des kuerzesten Wegs und der Gesamtlaenge als Tupel

# Klasse, die allgemeine Methoden zur Loesung der Aufgabe enthaelt
class LisaRennt:

    # Gibt den optimalsten Pfad mit Sichtbarkeitsgraphen zurueck,
    # bei dem Lisa ihr Haus so spaet wie moeglich verlassen muss
    @staticmethod
    def get_optimal_result():
        to_visit = []       # Liste aller Knoten, die auf Sichtbarkeit ueberprueft werden muss
        start_graph = {}    # Anfangsgraph mit allen Polygonecken als Knoten

        # Konvexe Ecken aller Polygone als Ausgangsknoten hinzufuegen
        for polygon in Input.polygon_list:
            for convex_vertice in polygon.convex_vertices_only():
                # Erstellung eines Knotens aus einem konvexen Eckpunkt
                converted_node = Node(convex_vertice, polygon.id)

                # Heuristik fuer Knoten berechnen und hinzufuegen
                converted_node = converted_node.assign_h_cost(Input.end_node)

                # Knoten jeweils initialisieren
                to_visit.append(converted_node)
                start_graph[converted_node] = []

        # Liste aller Kanten generieren (Zur Erstellung des Sichtbarkeitsgraphen erforderlich)
        for polygon in Input.polygon_list:
            Input.list_of_all_edges.extend(polygon.convert_to_edges())

        # Sichtbarkeitsgraph mit Start- aber ohne Endknoten
        visibility_graph = VisibilityGraph.get_graph(to_visit, start_graph)
        visibility_graph = VisibilityGraph.add_node(Input.start_node, visibility_graph)

        # Untersten und obersten Knoten ermitteln (um Bereich zu definieren), Startknoten als Referenzwert
        smallest_y = Input.start_node.point.y
        biggest_y = Input.start_node.point.y
        for node in visibility_graph.keys():
            if node.point.y < smallest_y:
                smallest_y = node.point.y
            if node.point.y > biggest_y:
                biggest_y = node.point.y

        # Maximum von t(x), s. Dokumentation
        optimum_y = fmin(lambda x: -LisaRennt.t(x), 0, disp=False)

        # Bereich justieren, indem optimaler Weg gesucht wird, da das Optimum ein Teil dieses Bereichs sein muss
        if optimum_y < smallest_y:
            smallest_y = int(floor(optimum_y))
        if optimum_y > biggest_y:
            biggest_y = int(ceil(optimum_y))

        # Ermitteln des optimalsten Wegs durch Ausprobieren aller moeglichen Endknoten im vordefinierten Bereich
        full_visibility_graph = VisibilityGraph.add_node(Node(Point(0,smallest_y), -2), deepcopy(visibility_graph))
        best_visibility_graph = deepcopy(full_visibility_graph)
        calc_path = AStar.get_shortest_path(full_visibility_graph)
        if calc_path is None:
            return  # Kein Pfad
        best_path = calc_path
        Output.highest_time_after_departure = LisaRennt.time(smallest_y, calc_path[1])
        for i in range(smallest_y+1, biggest_y):
            full_visibility_graph = VisibilityGraph.add_node(Node(Point(0, i), -2), deepcopy(visibility_graph))
            calc_path = AStar.get_shortest_path(full_visibility_graph)
            time_after_departure = LisaRennt.time(i, calc_path[1])

            if time_after_departure > Output.highest_time_after_departure:
                Output.highest_time_after_departure = time_after_departure

                best_path = calc_path
                best_visibility_graph = deepcopy(full_visibility_graph)

        return best_path, best_visibility_graph  # Rueckgabe des optimalsten Wegs (inkl. Laenge) und Sichtbarkeitsgraph

    # Zeitdifferenz ohne Hindernisse
    @staticmethod
    def t(x):
        # Optimaler Punkt (0|x) auf der y-Achse, wenn es keine Hindernisse gibt
        s_x = Input.start_node.point.x  # Standort Lisa, x-Wert
        s_y = Input.start_node.point.y  # Standort Lisa, y-Wert
        return (3*x/25.0) - (6 * sqrt(s_x**2 + (s_y-x)**2))/25.0

    # Zeitdifferenz mit Hindernissen
    @staticmethod
    def time(x, path_length):
        return (3*x/25.0) - (6*path_length)/25.0

class Input:
    polygon_list = []
    list_of_all_edges = []
    start_node = Node(Point(0,0),-1)     # Platzhalter Start-Knoten
    end_node = Node(Point(0,0), -2)      # Platzhalter End-Knoten
    double_corners = []

    @staticmethod
    def polygon_liste_einlesen(file_name):
        lines = [line.rstrip('\n') for line in open('beispieldaten/' + file_name)]
        count_polygons = int(lines[0])
        all_vertices = []
        for i in range(count_polygons):
            corner_info = lines[i + 1].split(" ")
            count_corners = int(corner_info[0])
            corner_info.pop(0)
            corner_list = []
            x = 0
            for j in range(count_corners * 2):
                if j % 2 == 0:
                    x = int(corner_info[j])
                else:
                    y = int(corner_info[j])
                    corner_list.append(Point(x, y))
                    all_vertices.append(Point(x, y))

            Input.polygon_list.append(Polygon(corner_list, i + 1))

        start_information = lines[len(lines) - 1].split(" ")
        Input.start_node = Node(Point(int(start_information[0]), int(start_information[1])), -1)

        # Wenn Punkte doppelt vorkommen, haben Sie den Wert true, ansonsten false
        duplicates = {}
        for vertice in all_vertices:
            duplicates[vertice] = vertice in duplicates
        # Alle relevanten Eckpunkte, die nicht in mehr als einem Polygon vorkommen
        doulbe_corners = [i for i in duplicates if duplicates[i]]
        Input.double_corners.extend(doulbe_corners)

class Output:

    output_file = ""
    highest_time_after_departure = 0

    @staticmethod
    def display_result(result):
        final_output = "<svg xmlns='http://www.w3.org/2000/svg' version='1.1' viewBox='0 0 1024 768'> " \
                       "<g transform='scale(1 -1)'> <g transform='translate(0 -768)'>" \
                       "<line id='y' x1='0' x2='0' y1='0' y2='768' fill='none' stroke='#212121' stroke-width='3'/>"
        final_output += Output.draw_polygons()

        if result is None:
            print ("Es gibt keinen Pfad zum Ziel")
            return

        shortest_path = result[0][0]
        visibility_graph = result[1]

        bus_departure_time = 27000      # 7:30 in Sekunden
        lisas_leaving = int(round(Output.highest_time_after_departure))     # Lisas Verlassen des Hauses
        start_time = str(datetime.timedelta(seconds=bus_departure_time+lisas_leaving))
        print ("Startzeit: " + start_time)

        all_nodes = visibility_graph.keys()

        output_lines = ""
        for node in all_nodes:
            if node.id == -1:
                # Kreis zeichnen fuer Start-Knoten
                output_lines += "<circle cx='" + str(node.point.x) +"' cy='" + \
                                str(node.point.y) +"' r='3' stroke='black' stroke-width='3' fill='red'></circle>"

            if node.id == -2:
                end_time = int(bus_departure_time + node.point.y*(3.0/25.0))
                print ("Zielzeit: " + str(datetime.timedelta(seconds=end_time)))      # Zielzeit
                print ("y-Koordinate des Auftreffens: " + str(node.point.y))        # y-Koordinaten des Auftreffens
                # Kreis zeichnen fuer End-Knoten
                output_lines += "<circle cx='" + str(node.point.x) +"' cy='" + \
                                str(node.point.y) +"' r='3' stroke='black' stroke-width='3' fill='blue'></circle>"
            for visible_node in visibility_graph[node]:
                output_lines += "\n <line "
                output_lines += "x1='" + str(node.point.x) + "' y1='" + str(node.point.y) + "'"
                output_lines += " x2='" + str(visible_node.point.x) + "' y2='" + str(visible_node.point.y) + "' "
                output_lines += "stroke='green' stroke-width='1'/>"

        length_lisas_path = int(round(result[0][1]))  # Laenge von Lisas Route in Minuten
        print ("Länge von Lisas Route: " + str(length_lisas_path) + " Meter")

        duration_lisas_path = round((length_lisas_path * (6.0 / 25.0)) / 60.0, 2)  # Dauer von Lisas Route in Minuten
        print ("Dauer von Lisas Route: " + str(duration_lisas_path) + " Minuten")

        lisas_route = ""
        for i, first_node in reversed(list(enumerate(shortest_path))):
            lisas_route += str(first_node) + " -->"
            if i > 0:
                second_node = shortest_path[i-1]
                output_lines += "\n <line "
                output_lines += "x1='" + str(first_node.point.x) + "' y1='" + str(first_node.point.y) + "'"
                output_lines += " x2='" + str(second_node.point.x) + "' y2='" + str(second_node.point.y) + "' "
                output_lines += "stroke='red' stroke-width='3'/>"

        print("Lisas Route: " + lisas_route + " Ziel erreicht")

        final_output += output_lines
        final_output += "</g> </g> </svg>"

        if Output.output_file in os.listdir("output"):
            os.remove("output/" + Output.output_file)

        with open("output/" + Output.output_file, "a") as f:
            f.write(final_output)
            f.close()

        print ("\nDie generierte SVG-Datei befindet sich im Output-Ordner")

    @staticmethod
    def draw_polygons():
        output = ""
        # Eckpunkte die zu zwei Polygonen gehoeren werden ausgefiltert
        all_vertices = []
        for polygon in Input.polygon_list:
            for vertice in polygon.points:
                all_vertices.append(vertice)

        for polygon in Input.polygon_list:
            output += "<polygon id='P" + str(polygon.id) + "' points='"
            for point in polygon.points:
                output += str(point.x) + " " + str(point.y) + " "
            output += "' fill='#6B6B6B' stroke='#212121' stroke-width='2'/> \n"

        return output

class Interface:

    @staticmethod
    def start():
        print ("Welche Umgebung soll eingelesen werden?")
        all_text_files = []
        for file in os.listdir("beispieldaten"):
            if file.endswith(".txt"):
                all_text_files.append(file)

        for i in range(len(all_text_files)):
            print ("[" + str(i+1) + "] " + all_text_files[i])

        choice = int(input("\nAuswahl: "))-1
        print ("\nRoute wird berechnet... \n")
        Input.polygon_liste_einlesen(all_text_files[choice])
        Output.output_file = all_text_files[choice].replace(".txt", "") + ".svg"
        Output.display_result(LisaRennt.get_optimal_result())

Interface.start()


<gh_stars>0
import matplotlib.pyplot as plt
import numpy as np
import os
import torch
import torchvision.datasets
import torchvision.models
import torchvision.transforms

from AIS import get_obj_term_A, manual_AIS
from custom_rbm import RBM_custom, RBM_gaussian_custom
from data_process import image_data_collapse, binarize_image_data, data_mnist
from RBM_assess import plot_confusion_matrix, confusion_matrix_from_pred, get_X_y_dataset
from settings import MNIST_BINARIZATION_CUTOFF, DIR_OUTPUT, CLASSIFIER, BETA, DIR_MODELS

########## CONFIGURATION ##########
BATCH_SIZE = 100  # default 64
VISIBLE_UNITS = 784  # 28 x 28 images
HIDDEN_UNITS = 10  # was 128 but try 10
CD_K = 20
EPOCHS = 100  # was 10, or 51

LEARNING_RATE = 1*1e-4               # default was 1e-3, new base is 1e-4
learningrate_schedule = False         # swap from LEARNING_RATE to diff value at specified epoch
learningrate_schedule_value = 1*1e-4
learningrate_schedule_epoch = 25

AIS_STEPS = 0 #1000      # 0 or 1000 typically
AIS_CHAINS = 0 #100      # 100 or 500
USE_FIELDS = False
PLOT_WEIGHTS = False
POINTS_PER_EPOCH = 1

GAUSSIAN_RBM = True
if RBM_gaussian_custom:
    RBM = RBM_gaussian_custom
else:
    RBM = RBM_custom

def torch_data_loading(batch_size=BATCH_SIZE):
    print('Loading dataset...')
    DATA_FOLDER = 'data'
    train_dataset = torchvision.datasets.MNIST(root=DATA_FOLDER, train=True,
                                               transform=torchvision.transforms.ToTensor(), download=True)
    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size)
    test_dataset = torchvision.datasets.MNIST(root=DATA_FOLDER, train=False,
                                              transform=torchvision.transforms.ToTensor(), download=True)
    test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size)
    return train_dataset, train_loader, test_dataset, test_loader

def get_classloader(global_dataset, class_name):
    # see https://discuss.pytorch.org/t/how-to-use-one-class-of-number-in-mnist/26276/12
    def get_indices(dataset):
        indices = []
        for i in range(len(dataset.targets)):
            if dataset.targets[i] == class_name:
                indices.append(i)
        return indices

    idx = get_indices(global_dataset)
    num_class_samples = len(idx)
    loader = torch.utils.data.DataLoader(global_dataset, batch_size=BATCH_SIZE, sampler=torch.utils.data.sampler.SubsetRandomSampler(idx))

    return loader, num_class_samples

def custom_RBM_loop(train_loader, train_data_as_arr, hidden_units=HIDDEN_UNITS, init_weights=None,
                    use_fields=USE_FIELDS, beta=BETA, epochs=EPOCHS, cdk=CD_K,
                    outdir=None, classify=False, restart=False, points_per_epoch=POINTS_PER_EPOCH):
    assert beta == BETA         # TODO uncouple global STDEV in rbm class to make beta passable
    assert classify is False    # TODO need to add support for classify_with_rbm_hidden(...) at end

    if restart:
        fmod = ''  # was '_restart'; now use rundir name alone to store restart label
    else:
        fmod = ''

    if outdir is not None:
        if not os.path.exists(outdir):
            os.makedirs(outdir)
        trainingdir = outdir + os.sep + 'training' + fmod
        if not os.path.exists(trainingdir):
            os.makedirs(trainingdir)
    else:
        trainingdir = DIR_OUTPUT + os.sep + 'training'

    # initialize RBM class
    rbm = RBM(VISIBLE_UNITS, hidden_units, cdk, init_weights=init_weights, use_fields=use_fields,
              learning_rate=LEARNING_RATE)

    # prep object timeseries to store over training
    num_samples = train_data_as_arr.shape[0]
    batches_per_epoch = num_samples / BATCH_SIZE
    total_timepoints = points_per_epoch * epochs + 1
    iterations_per_pt = batches_per_epoch / points_per_epoch
    if POINTS_PER_EPOCH != 1:
        assert num_samples % BATCH_SIZE == 0
        assert batches_per_epoch % points_per_epoch == 0
    print("epochs:", epochs)
    print("num_samples:", num_samples)
    print("batches_per_epoch:", batches_per_epoch)
    print("iterations_per_pt:", iterations_per_pt)
    print("total_timepoints:", total_timepoints)
    iteration_arr = np.arange(total_timepoints) * iterations_per_pt  # iteration pt of the saved data
    print(iteration_arr)

    weights_timeseries = np.zeros((rbm.num_visible, rbm.num_hidden, total_timepoints))
    weights_timeseries[:, :, 0] = rbm.weights
    if use_fields:
        visible_bias_timeseries = np.zeros((rbm.num_visible, total_timepoints))
        visible_bias_timeseries[:, 0] = rbm.visible_bias.numpy()
        hidden_bias_timeseries = np.zeros((rbm.num_hidden, total_timepoints))
        hidden_bias_timeseries[:, 0] = rbm.hidden_bias.numpy()

    if PLOT_WEIGHTS:
        rbm.plot_model(title='epoch_0', outdir=trainingdir)

    obj_reconstruction = np.zeros(total_timepoints - 1)
    obj_logP_termA = np.zeros(total_timepoints)
    obj_logP_termB = np.zeros(total_timepoints)

    def loop_updates(iteration_idx, iteration_counter, epoch_recon_error=None):
        if PLOT_WEIGHTS:
            rbm.plot_model(title='iteration_%d' % (iteration_counter), outdir=trainingdir)

        if epoch_recon_error is not None:
            print('Epoch (Reconstruction) Error (epoch=%d) (iteration_idx:%d): %.4f' % (
                epoch + 1, iteration_idx, epoch_recon_error))
            obj_reconstruction[iteration_idx - 1] = epoch_recon_error

        if AIS_STEPS > 0:
            obj_logP_termA[iteration_idx] = get_obj_term_A(train_data_as_arr, rbm.weights, rbm.visible_bias,
                                                           rbm.hidden_bias, beta=beta)
            print('Estimating log Z...', )
            obj_logP_termB[iteration_idx], _ = manual_AIS(rbm, beta, nchains=AIS_CHAINS, nsteps=AIS_STEPS, CDK=1,
                                                          joint_mode=True)
        # save parameters each epoch
        weights_timeseries[:, :, iteration_idx] = rbm.weights.numpy()
        if use_fields:
            visible_bias_timeseries[:, iteration_idx] = rbm.visible_bias.numpy()
            hidden_bias_timeseries[:, iteration_idx] = rbm.hidden_bias.numpy()

        print('Term A:', obj_logP_termA[iteration_idx],
              '| Log Z:', obj_logP_termB[iteration_idx],
              '| Score:', obj_logP_termA[iteration_idx] - obj_logP_termB[iteration_idx])

    loop_updates(0, 0)

    print('Training RBM...')
    iteration_counter = 0
    for epoch in range(epochs):
        if epoch == learningrate_schedule_epoch and learningrate_schedule:
            rbm.learning_rate = learningrate_schedule_value
        print('LEARNING RATE: %.2E' % rbm.learning_rate)

        epoch_recon_error = 0.0
        for batch, _ in train_loader:
            iteration_counter += 1
            batch = batch.view(len(batch), VISIBLE_UNITS)  # flatten input data
            batch = (batch > MNIST_BINARIZATION_CUTOFF).float()  # convert to 0,1 form
            batch = -1 + batch * 2  # convert to -1,1 form
            batch_recon_error = rbm.contrastive_divergence(batch)
            epoch_recon_error += batch_recon_error

            if POINTS_PER_EPOCH != 1:
                if iteration_counter % iterations_per_pt == 0:
                    iteration_idx = int(iteration_counter / iterations_per_pt)
                    loop_updates(iteration_idx, iteration_counter, epoch_recon_error=epoch_recon_error)

        if POINTS_PER_EPOCH == 1:
            iteration_idx = epoch + 1
            loop_updates(iteration_idx, iteration_counter, epoch_recon_error=epoch_recon_error)

    score_arr = obj_logP_termA - obj_logP_termB

    if outdir is None:
        scoredir = DIR_OUTPUT + os.sep + 'logZ' + os.sep + 'rbm'
    else:
        scoredir = outdir

    # save parameters
    title_mod = '%dhidden_%dfields_%dcdk_%dstepsAIS_%.2fbeta%s' % (hidden_units, use_fields, cdk, AIS_STEPS, beta, fmod)
    fpath = scoredir + os.sep + 'objective_%s' % title_mod
    np.savez(fpath,
             epochs=range(epochs + 1),
             iterations=iteration_arr,
             termA=obj_logP_termA,
             logZ=obj_logP_termB,
             score=score_arr)

    fpath = scoredir + os.sep + 'weights_%s' % title_mod
    np.savez(fpath,
             epochs=range(epochs + 1),
             weights=weights_timeseries)  #TODO add iterations arr
    if use_fields:
        np.savez(scoredir + os.sep + 'visiblefield_%s' % title_mod,
                 epochs=range(epochs + 1),
                 iterations=iteration_arr,
                 visiblefield=visible_bias_timeseries)
        np.savez(scoredir + os.sep + 'hiddenfield_%s' % title_mod,
                 epochs=range(epochs + 1),
                 iterations=iteration_arr,
                 hiddenfield=hidden_bias_timeseries)

    plot_scores(iteration_arr, obj_logP_termA, obj_logP_termB, score_arr, scoredir, title_mod, 'iterations',
                obj_reconstruction=obj_reconstruction)

    if classify:
        classify_with_rbm_hidden(rbm, outdir=outdir)

    return rbm

def plot_scores(iteration_arr, obj_logP_termA, obj_logP_termB, score_arr, scoredir, title_mod, xlabel,
                obj_reconstruction=None):
    if obj_reconstruction is not None:
        plt.plot(iteration_arr[1:], obj_reconstruction)
        plt.xlabel(xlabel);
        plt.ylabel('reconstruction error')
        plt.savefig(scoredir + os.sep + 'rbm_recon_%s.pdf' % (title_mod)); plt.close()

    plt.plot(iteration_arr, obj_logP_termA)
    plt.xlabel(xlabel);
    plt.ylabel(r'$- \langle H(s) \rangle$')
    plt.savefig(scoredir + os.sep + 'rbm_termA_%s.pdf' % (title_mod)); plt.close()

    plt.plot(iteration_arr, obj_logP_termB)
    plt.xlabel(xlabel);
    plt.ylabel(r'$\ln \ Z$')
    plt.savefig(scoredir + os.sep + 'rbm_logZ_%s.pdf' % (title_mod)); plt.close()

    plt.plot(iteration_arr, score_arr)
    plt.xlabel(xlabel);
    plt.ylabel(r'$\langle\ln \ p(x)\rangle$')
    plt.savefig(scoredir + os.sep + 'rbm_score_%s.pdf' % (title_mod)); plt.close()
    return

def classify_with_rbm_hidden(rbm, train_dataset, train_loader, test_dataset, test_loader, outdir=None, beta=BETA):
    stdev = 1.0/np.sqrt(beta)

    print('Extracting features...')
    train_features = np.zeros((len(train_dataset), rbm.num_hidden))
    train_labels = np.zeros(len(train_dataset))
    test_features = np.zeros((len(test_dataset), rbm.num_hidden))
    test_labels = np.zeros(len(test_dataset))
    for i, (batch, labels) in enumerate(train_loader):
        batch = batch.view(len(batch), VISIBLE_UNITS)  # flatten input data
        batch = (batch > MNIST_BINARIZATION_CUTOFF).float()  # convert to 0,1 form
        batch = 2 * batch - 1  # convert to -1,1 form
        train_features[i * BATCH_SIZE:i * BATCH_SIZE + len(batch), :] = rbm.sample_hidden(batch, stdev=stdev)
        train_labels[i * BATCH_SIZE:i * BATCH_SIZE + len(batch)] = labels.numpy()

    for i, (batch, labels) in enumerate(test_loader):
        batch = batch.view(len(batch), VISIBLE_UNITS)  # flatten input data
        batch = (batch > MNIST_BINARIZATION_CUTOFF).float()  # convert to 0,1 form
        batch = 2 * batch - 1  # convert to -1,1 form
        test_features[i * BATCH_SIZE:i * BATCH_SIZE + len(batch), :] = rbm.sample_hidden(batch, stdev=stdev)
        test_labels[i * BATCH_SIZE:i * BATCH_SIZE + len(batch)] = labels.numpy()

    print('Training Classifier...')
    CLASSIFIER.fit(train_features, train_labels)
    print('Classifying...')
    predictions = CLASSIFIER.predict(test_features).astype(int)

    confusion_matrix = np.zeros((10, 10), dtype=int)
    matches = [False for _ in test_dataset]
    for idx, pair in enumerate(test_dataset):
        if pair[1] == predictions[idx]:
            matches[idx] = True
        confusion_matrix[pair[1], predictions[idx]] += 1
    title = "Successful test cases: %d/%d (%.3f)" % (
        matches.count(True), len(matches), float(matches.count(True) / len(matches)))
    if outdir is None:
        fpath = DIR_OUTPUT + os.sep + 'training' + os.sep + 'cm.jpg'
    else:
        fpath = outdir + os.sep + 'cm.jpg'
    cm = plot_confusion_matrix(confusion_matrix, title=title, save=fpath)
    print(title)
    return

def classifier_on_poe_scores(models, dataset_train, dataset_test, outpath, clfs=None, beta=2.0, plot=False):
    """
    models: dict of idx: rbm for idx in {0, ..., 9} i.e. the experts on each digit class
    clfs: list of classifiers
    """

    def score_1digit_model(model, img):
        """
        see hinton2002 ... .py for details
        """
        # should just be 0.5 * beta * Jij si sj
        beta = 1.0
        W = model.weights
        dotp = np.dot(W.T, img)
        # score = np.dot(dotp, dotp)  #
        score = beta * np.dot(dotp, dotp) / 2.0
        return score

    if clfs is None:
        clfs = [CLASSIFIER]

    features_order = list(models.keys())
    feature_dim = len(features_order)

    def get_X_y_features(dataset):
        X = np.zeros((len(dataset), feature_dim))
        y = np.zeros(len(dataset), dtype=int)
        for idx, pair in enumerate(dataset):
            elem_arr, elem_label = pair
            preprocessed_input = binarize_image_data(image_data_collapse(elem_arr), threshold=MNIST_BINARIZATION_CUTOFF)
            features = np.array([score_1digit_model(models[key], preprocessed_input) for key in features_order])
            #features = np.array([score_1digit_model(models[idx], preprocessed_input) for idx in range(10)])
            X[idx, :] = features
            y[idx] = elem_label
        return X, y

    print("[classifier_on_rbm_features] Step 1: get features for training")
    X_train_reduced, y_train = get_X_y_features(dataset_train)
    print("\tTraining data dimension", X_train_reduced.shape, y_train.shape)

    print("[classifier_on_rbm_features] Step 2: train classifier layer")
    for clf in clfs:
        print('fitting...')
        clf.fit(X_train_reduced, y_train)  # fit data

    print("[classifier_on_rbm_features] Step 3: get features for testing")
    X_test_reduced, y_test = get_X_y_features(dataset_test)

    print("[classifier_on_rbm_features] Step 4: classification metrics and confusion matrix")
    cms = [0] * len(clfs)
    accs = [0] * len(clfs)
    for idx, clf in enumerate(clfs):
        print('predicting...')
        predictions = clf.predict(X_test_reduced).astype(int)
        confusion_matrix, matches = confusion_matrix_from_pred(predictions, y_test)
        acc = float(matches.count(True) / len(matches))
        cms[idx] = confusion_matrix
        accs[idx] = acc
        title = "Successful test cases: %d/%d (%.3f)" % (matches.count(True), len(matches), acc)
        if plot:
            cm = plot_confusion_matrix(confusion_matrix, title=title, save=outpath)
            plt.close()
        print(title)
    return cms, accs

if __name__ == '__main__':

    train_dataset, train_loader, test_dataset, test_loader = torch_data_loading()
    TRAINING, TESTING = data_mnist(binarize=True)
    X, _ = get_X_y_dataset(TRAINING, dim_visible=VISIBLE_UNITS, binarize=True)

    test_data_loader = False

    num_runs = 5
    hopfield_runs = False
    random_runs = False
    restart_random_runs = False

    load_scores = False
    load_weights = False

    poe_mode_train = False
    poe_mode_classify = True
    # TODO print settings file for each run

    rescore_ais_trained_rbms = False

    if test_data_loader:
        seven_loader = get_classloader(train_dataset, 7)

    bigruns = DIR_OUTPUT + os.sep + 'archive' + os.sep + 'big_runs'
    if hopfield_runs:

        HEBBIAN = False
        PCA = False
        MOD_SVD = False
        MOD_SQRT = False
        if HEBBIAN:
            hmod = '_hebbian'
            assert not PCA
        elif PCA:
            hmod = '_PCA'
        else:
            hmod = ''

        # load hopfield weights
        npzpath = DIR_MODELS + os.sep + 'saved' + os.sep + 'hopfield_mnist_%d%s.npz' % (HIDDEN_UNITS, hmod)
        print("Loading weights from %s" % npzpath)
        arr = np.load(npzpath)['Q']

        # modify initial weight matrix
        if MOD_SVD or MOD_SQRT:
            assert HEBBIAN
            import scipy as sp
            A = np.dot(arr.T, arr)
            A_inv = np.linalg.inv(A)
            A_sqrt = sp.linalg.sqrtm(A)  # A_sqrt = mtrx_sqrt(A)
            A_sqrt_inv = np.linalg.inv(A_sqrt)

            if MOD_SVD:
                # svd of XI
                U, Sigma, V = np.linalg.svd(arr, full_matrices=False)

                init_weights = torch.from_numpy(U).float()
            else:
                assert MOD_SQRT
                K = np.dot(arr, A_sqrt_inv)
                init_weights = torch.from_numpy(K).float()
        else:
            init_weights = torch.from_numpy(arr).float()

        outdir = bigruns + os.sep + 'rbm' + os.sep + 'hopfield%s_%dhidden_%dfields_%.2fbeta_%dbatch_%depochs_%dcdk_%.2Eeta_%dais_%dppEpoch' % \
                 (hmod, HIDDEN_UNITS, USE_FIELDS, BETA, BATCH_SIZE, EPOCHS, CD_K, LEARNING_RATE, AIS_STEPS, POINTS_PER_EPOCH)

        for idx in range(num_runs):
            rundir = outdir + os.sep + 'run%d' % idx
            custom_RBM_loop(train_loader, X, init_weights=init_weights.clone(), outdir=rundir, classify=False)

    if random_runs:
        for idx in range(num_runs):
            outdir = bigruns + os.sep + 'rbm' + os.sep + 'normal_%dhidden_%dfields_%.2fbeta_%dbatch_%depochs_%dcdk_%.2Eeta_%dais_%dppEpoch' % \
                     (HIDDEN_UNITS, USE_FIELDS, BETA, BATCH_SIZE, EPOCHS, CD_K, LEARNING_RATE, AIS_STEPS, POINTS_PER_EPOCH)
            rundir = outdir + os.sep + 'run%d' % idx
            custom_RBM_loop(train_loader, X, init_weights=None, outdir=rundir, classify=False)

    if restart_random_runs:
        EPOCHS_RESTART = 100
        HIDDEN_RESTART = 50
        ETA_RESTART = 1e-4
        AIS_STEPS_RESTART = 200
        EPOCHS_TO_EXTEND = 100

        for idx in range(num_runs):
            # load pre-trained weights
            indir = bigruns + os.sep + 'rbm' + os.sep + 'normal_%dhidden_%dfields_%.2fbeta_%dbatch_%depochs_%dcdk_%.2Eeta_%dais' % \
                     (HIDDEN_RESTART, USE_FIELDS, BETA, BATCH_SIZE, EPOCHS_RESTART, CD_K, ETA_RESTART, AIS_STEPS_RESTART)
            rundir = indir + os.sep + 'run%d' % idx
            print("Loading PRE_TRAINED weights from %s" % rundir)
            npzpath = rundir + os.sep + 'weights_%dhidden_%dfields_%dcdk_%dstepsAIS_%.2fbeta.npz' % (HIDDEN_RESTART, USE_FIELDS, CD_K, AIS_STEPS_RESTART, BETA)
            loaded_weights_timeseries_np = np.load(npzpath)['weights']
            init_weights_np = loaded_weights_timeseries_np[:, :, -1]
            init_weights = torch.from_numpy(init_weights_np).float()
            # specify new filename modifier
            outdir = rundir + '_restart'
            custom_RBM_loop(train_loader, X, init_weights=init_weights, outdir=outdir, classify=False, restart=True, epochs=EPOCHS_TO_EXTEND)

    if load_scores:
        outdir = bigruns + os.sep + 'rbm' + os.sep + 'C_beta2duringTraining_%dbatch_%depochs_%dcdk_%.2Eeta_%dais' % (BATCH_SIZE, EPOCHS, CD_K, LEARNING_RATE, AIS_STEPS)
        fname = 'objective_10hidden_0fields_20cdk_200stepsAIS_2.00beta.npz'
        dataobj = np.load(outdir + os.sep + fname)

        obj_logP_termA = dataobj['termA']
        obj_logP_termB = dataobj['logZ']
        score_arr = dataobj['score']
        epochs = dataobj['epochs']
        timesteps = len(epochs) - 1

        title_mod = '10hidden_0fields_20cdk_200stepsAIS_2.00beta'
        plot_scores(timesteps, obj_logP_termA, obj_logP_termB, score_arr, outdir, title_mod,
                    'epochs', obj_reconstruction=None)

    if load_weights:
        # Note: looks like generative training does not help with classification at first glance
        local_beta = 200
        hidden = 1000
        epoch_idx = 199

        outdir = bigruns + os.sep + 'rbm' + os.sep + 'normal_%dhidden_0fields_2.00beta_100batch_200epochs_1cdk_1.00E-04eta_0ais' % hidden \
                 + os.sep + 'run0'
        fname = 'weights_%dhidden_0fields_1cdk_0stepsAIS_2.00beta.npz' % hidden
        dataobj = np.load(outdir + os.sep + fname)
        arr = dataobj['weights'][:, :, epoch_idx]

        rbm = RBM(VISIBLE_UNITS, hidden, 0, init_weights=None, use_fields=False, learning_rate=0)
        rbm.weights = torch.from_numpy(arr).float()

        classify_with_rbm_hidden(rbm, train_dataset, train_loader, test_dataset, test_loader, outdir=outdir, beta=local_beta)

    if poe_mode_train:
        # TODO X_for_digit make (will be called if AIS steps > 0 ) -- second arg to custom_RBM_loop()
        # TODO beta in scoring
        use_hopfield = True
        k_range = [10, 20, 100] #, 200, 500, 250, 300]
        runs = 3

        HEBBIAN = False
        PCA = True
        MOD_SVD = False
        MOD_SQRT = False
        if HEBBIAN:
            hmod = '_hebbian'
            assert not PCA
        elif PCA:
            hmod = '_PCA'
        else:
            hmod = ''

        loader_dict = {idx: get_classloader(train_dataset, idx)[0] for idx in range(10)}
        loader_num_samples_dict = {idx: get_classloader(train_dataset, idx)[1] for idx in range(10)}

        for idx in range(runs):
            for k in k_range:
                print("Training POE for k=%d" % k)
                for digit in range(10):
                    print("Training POE for k=%d (digit: %d)" % (k, digit))
                    dir_modifier = '%ddigit_%dhidden_%dfields_%.2fbeta_%dbatch_%depochs_%dcdk_%.2Eeta_%dais' % \
                                   (digit, k, USE_FIELDS, BETA, BATCH_SIZE, EPOCHS, CD_K, LEARNING_RATE, AIS_STEPS)

                    fake_empty_data = np.zeros((loader_num_samples_dict[digit], 28**2))

                    if use_hopfield:
                        outdir = bigruns + os.sep + 'poe' + os.sep + 'run%d' % idx + os.sep + \
                                 'poe_hopfield%s_%s' % (hmod, dir_modifier)
                        rundir = outdir
                        # load hopfield weights
                        npzpath = DIR_MODELS + os.sep + 'poe' + os.sep + 'hopfield_digit%d_p%d%s.npz' % (digit, 10*k, hmod)
                        print("Loading weights from %s" % npzpath)
                        arr = np.load(npzpath)['Q']
                        init_weights = torch.from_numpy(arr).float()

                        custom_RBM_loop(loader_dict[digit], fake_empty_data, hidden_units=k, init_weights=init_weights.clone(), outdir=rundir, classify=False)
                    else:
                        outdir = bigruns + os.sep + 'poe' + os.sep + 'run%d' % idx + os.sep + \
                                 'poe_normal_%s' % dir_modifier
                        rundir = outdir
                        custom_RBM_loop(loader_dict[digit], fake_empty_data, hidden_units=k, init_weights=None, outdir=rundir, classify=False)

    if poe_mode_classify:
        init_weights_type = 'hopfield'  # hopfield or normal

        HEBBIAN = False
        PCA = False
        if HEBBIAN:
            assert init_weights_type == 'hopfield'
            assert not PCA
            init_weights_type += '_hebbian'
        elif PCA:
            assert init_weights_type == 'hopfield'
            init_weights_type += '_PCA'
        else:
            hmod = ''

        k_range = [10, 20, 100]    #, 200, 250, 300, 500]  # range(1, 110)  # 110
        epochs = [0, 1, 2, 3, 4] + list(range(5, 101, 5))
        runs = [2] # [0,1,2]          # TODO care
        accs = np.zeros((len(epochs), len(k_range)))

        for run in runs:
            for epoch_idx, epoch in enumerate(epochs):
                for k_idx, k in enumerate(k_range):
                    print('epoch, k:', epoch, k)
                    hidden_units = k
                    dir_modifier = '0fields_2.00beta_100batch_100epochs_%dcdk_1.00E-04eta_0ais' % (CD_K)
                    models10 = {}
                    for digit in range(10):
                        rbm_digit = RBM(VISIBLE_UNITS, hidden_units, 0, init_weights=None, use_fields=False, learning_rate=0)
                        # load weights for given epoch
                        run_dir = bigruns + os.sep + 'poe' + os.sep + '%s' % init_weights_type + os.sep + 'run%d' % run
                        npzpdir = run_dir + os.sep + \
                                  'poe_%s_%ddigit_%dhidden_%s' % (init_weights_type, digit, hidden_units, dir_modifier)
                        weights_path = npzpdir + os.sep + 'weights_%dhidden_0fields_%dcdk_0stepsAIS_2.00beta.npz' % (hidden_units, CD_K)
                        arr = np.load(weights_path)['weights'][:, :, epoch]
                        rbm_digit.weights = torch.from_numpy(arr).float()
                        # set as model for that digit
                        models10[digit] = rbm_digit
                    fpath = DIR_OUTPUT + os.sep + 'training' + os.sep + 'cm_%s_k%d_epoch%d.jpg' % (init_weights_type, k, epoch)
                    cm_list, acc_list = classifier_on_poe_scores(models10, TRAINING, TESTING, fpath, clfs=None, beta=2.0, plot=False)
                    accs[epoch_idx, k_idx] = acc_list[0]

            # save data
            fpath = DIR_OUTPUT + os.sep + 'training' + os.sep + 'poe_scores_kNum%d_epochsNum%d_%s%d.npz' % \
                    (len(k_range), len(epochs), init_weights_type, run)
            np.savez(fpath, accs_epoch_by_k=accs, epochs=epochs, k_range=k_range)
            print(accs)
            # plot data
            plt.figure()
            error_pct = 100 * (1 - accs)
            for idx, epoch in enumerate(epochs):
                plt.plot(k_range, error_pct[idx, :], label='Epoch: %d' % epoch)
            plt.xlabel(r'$k$')
            plt.ylabel('Error')
            plt.legend()
            plt.savefig(DIR_OUTPUT + os.sep + 'training' + os.sep + 'score_vs_k%d_vs_epoch%d_%s%d.jpg' %
                        (len(k_range), len(epochs), init_weights_type, run))
            plt.show()
            # TODO check if hinton use logreg or SVM
            # TODO increase max iter to remove warning: ConvergenceWarning: The max_iter was reached which means the coef_ did not converge
            #   "the coef_ did not converge", ConvergenceWarning)
            # TODO more epochs on whichever k worked best? depends on how the error scales with epoch at large k
            # TODO compare vs normal dist (note this run of 20epochs x 110 k values was ~11GB)

    if rescore_ais_trained_rbms:

        # AIS settings
        beta = 2.0
        nsteps = 1000
        nchains = 100
        ntest = 1
        nsteps_rev = 0
        nchains_rev = 0
        runs = 1
        assert runs == 1
        # hebbian = False
        """
        if hebbian:
            strmod = '_hebbian'
        else:
            strmod = '' """

        # prep dataset
        training_subsample = TRAINING[:]
        X, _ = get_X_y_dataset(training_subsample, dim_visible=28 ** 2, binarize=True)

        # prep models to load
        epoch_list = list(range(0, 51))  #71
        VISIBLE_UNITS = 28 ** 2
        hidden_units = 50
        CD_K = 20
        bigruns = DIR_OUTPUT + os.sep + 'archive' + os.sep + 'big_runs' + os.sep + 'rbm'
        model_dir = 'hopfield_PCA_%dhidden_0fields_2.00beta_100batch_50epochs_20cdk_1.00E-04eta_0ais_1ppEpoch' % hidden_units
        model_runs = [2,3,4] #[0,1,2,3,4]  #[0, 1, 2, 3, 4]
        model_paths = [
            bigruns + os.sep + model_dir + os.sep + 'run%d' % a + os.sep +
            'weights_%dhidden_0fields_20cdk_0stepsAIS_2.00beta.npz' % hidden_units
            for a in model_runs
        ]

        for training_run_idx, weights_path in enumerate(model_paths):

            # load weights for given epoch
            weights_timeseries_np = np.load(weights_path)['weights']

            rbm = RBM(VISIBLE_UNITS, hidden_units, 0, init_weights=None, use_fields=False, learning_rate=0)
            N = rbm.num_visible
            p = rbm.num_hidden
            zero_field_visible = np.zeros(N)
            zero_field_hidden = np.zeros(p)

            termA_arr = np.zeros(len(epoch_list))
            logZ_arr = np.zeros(len(epoch_list))
            score_arr = np.zeros(len(epoch_list))

            for idx in range(len(epoch_list)):
                epoch = epoch_list[idx]
                # specify new local rbm class for this epoch only
                weights_epoch_np = weights_timeseries_np[:, :, epoch]
                rbm.weights = torch.from_numpy(weights_epoch_np).float()

                obj_term_A = get_obj_term_A(X, weights_epoch_np, zero_field_visible, zero_field_hidden, beta=beta)
                termA_arr[idx] = obj_term_A

                # Forward AIS
                logZ_fwd, _ = manual_AIS(rbm, beta, nchains=nchains, nsteps=nsteps)
                score_fwd = obj_term_A - logZ_fwd
                print('training_run_idx, epoch:', training_run_idx, epoch)
                print('mean log p(data):', score_fwd,
                      '(run 1 only, beta=%.2f, A=%.2f, B=%.2f)' % (beta, obj_term_A, logZ_fwd))
                logZ_arr[idx] = logZ_fwd
                score_arr[idx] = score_fwd

            # save updated AIS data
            out_dir = os.path.dirname(weights_path)
            title_mod = '%dhidden_%dfields_%dcdk_%dstepsAIS_%.2fbeta' % (hidden_units, False, CD_K, nsteps, beta)
            fpath = out_dir + os.sep + 'objective_%s' % title_mod
            np.savez(fpath,
                     epochs=epoch_list,
                     termA=termA_arr,
                     logZ=logZ_arr,
                     score=score_arr)


import numpy as np
from scipy.special import roots_legendre

class Contour():
    def __init__(self, emin, emax=0.0):
        self.emin=emin
        self.emax=emax

    def build_path_semicircle(self, npoints, endpoint=True):
        R= (self.emax-self.emin)/2.0
        R0= (self.emin+self.emax)/2.0
        phi=np.linspace(np.pi, 0, num=npoints+1, endpoint=endpoint)
        p=R0+R*np.exp(1.0j * phi)
        if endpoint:
            self.path=p
            self.de=np.diff(p)
        else:
            self.path=(p[:-1]+p[1:])/2
            self.de=p[1:]-p[:-1]

    def build_path_legendre(self, npoints, endpoint=True):
        p=13
        x,w= roots_legendre(npoints)
        R= (self.emax-self.emin)/2.0
        R0= (self.emin+self.emax)/2.0
        y1 = -np.log(1+np.pi*p)
        y2 = 0
        y   = (y2-y1)/2*x+(y2+y1)/2
        phi = (np.exp(-y)-1)/p
        path  = R0+R*np.exp(1.0j*phi)
        #weight= -(y2-y1)/2*np.exp(-y)/p*1j*(path-R0)*w
        if endpoint:
            self.path=path
            self.de=np.diff(path)
        else:
            self.path=(path[:-1]+path[1:])/2
            self.de=path[1:]-path[:-1]

    def build_path_rectangle(self, height=0.1, nz1=50, nz2=200, nz3=50):
        """
        prepare list of energy for integration.
        The path has three segments:
         emin --1-> emin + 1j*height --2-> emax+1j*height --3-> emax
        """
        nz1, nz2, nz3 = nz1, nz2, nz3
        nz = nz1+nz2+nz3
        p= np.zeros(nz + 1, dtype='complex128')
        p[:nz1] = self.emin + np.linspace(
            0, height, nz1, endpoint=False) * 1j
        p[nz1:nz1 + nz2] = self.emin + height * 1j + np.linspace(
            0, self.emax - self.emin, nz2, endpoint=False)
        p[nz1 + nz2:nz] = self.emax + height * 1j + np.linspace(
            0, -height, nz3, endpoint=False) * 1j
        p[-1] = self.emax  # emax
        self.path=(p[:-1]+p[1:])/2
        self.de=p[1:]-p[:-1]

    @property
    def npoints(self):
        return len(self.path)

    def de(self):
        return self.de

    def plot(self, ax=None):
        import matplotlib.pyplot as plt
        if ax is None:
            fig,ax=plt.subplots()
        plt.plot(self.path.real, self.path.imag, marker='.')
        plt.show()

def test():
    ct=Contour(emin=-16, emax=0)
    #ct.build_path_semicircle(npoints=100)
    #ct.build_path_rectangle()
    ct.build_path_legendre(npoints=50)
    print(ct.npoints)
    ct.plot()
if __name__ == '__main__':
    test()


<reponame>Noone65536/CORSMAL-Challenge-2022-Squids
from logging.config import valid_ident
import numpy as np
import cv2
import torch
import os
import scipy.stats as stats
from torch.utils.data import Dataset
from tqdm import tqdm
from PIL import Image
from torchvision import datasets, transforms

def randomlyAug(crop, depth, label, max_val=640, square=False, normal=False, depth_aug=False):
  h, w, c = crop.shape

  if h >= w:
    max_dim = h
  else:
    max_dim = w
  max_rand = max_val / max_dim

  lower, upper = 0.5, 1.5
  if upper > max_rand:
      upper = max_rand
  mu, sigma = 1, 0.5
  X = stats.truncnorm(
    (lower - mu) / sigma, (upper - mu) / sigma, loc=mu, scale=sigma)

  if not normal:
      rand_num = np.random.uniform(0.5,max_rand,1).item()
  else:
      rand_num = X.rvs(1).item()

      while rand_num < lower or rand_num > upper:
          rand_num = X.rvs(1).item()

  width = int(w * rand_num)
  height = int(h * rand_num)
  dim = (width, height)

  # resize image
  crop = cv2.resize(crop, dim, interpolation = cv2.INTER_AREA)
  depth = cv2.resize(depth, dim, interpolation = cv2.INTER_NEAREST)[:, :, np.newaxis]

  if square:
      label *= (height / h)**2
  else:
      label *= (height / h)

  # depth aug

  lower, upper = 0.8, 1.2

  mu, sigma = 1, 0.5
  X = stats.truncnorm(
    (lower - mu) / sigma, (upper - mu) / sigma, loc=mu, scale=sigma)

  if depth_aug:
      rand_num = X.rvs(1).item()

      while rand_num < lower or rand_num > upper:
              rand_num = X.rvs(1).item()

      depth = (depth * rand_num).astype(np.uint8)

      if square:
          label *= rand_num**2
      else:
          label *= rand_num

  return crop, depth, label

def get_annotation(id,input,anno_path='/content/labels'):
    anno = np.load(os.path.join(anno_path,'{:06d}.npy'.format(id)),allow_pickle=True).item()
    return anno.get(input)

class MiniDataset(Dataset):
    def __init__(self, base_p, label_f, depth, crops_rgb_f, aug=False, square=False, normal=False, depth_aug=False, label_name=['container capacity']):
      self.label_f = label_f #_f = folder
      self.depth = depth
      self.base = base_p
      self.label_name = label_name
      self.crops_rgb_f = crops_rgb_f
      self.samples = os.listdir(crops_rgb_f)
      self.ids = [ int(x.split('.')[0]) for x in self.samples]
      self.transform = transforms.Compose([
                                             transforms.Resize((320, 320)),
                                             transforms.ToTensor(),
                                             transforms.ConvertImageDtype(torch.float),
                                             ])
      self.aug = aug
      self.square = square
      self.normal = normal
      self.depth_aug = depth_aug
    def __len__(self):
      return len(self.ids)

    def __getitem__(self, idx):
      id_ = self.ids[idx]

      # depth
      depth = np.asarray(Image.open(os.path.join(self.depth,'{:06d}.png'.format(id_))))[:,:,np.newaxis]

      # rgb_cropped
      crop = np.asarray(Image.open(os.path.join(self.crops_rgb_f,'{:06d}.png'.format(id_))))
      # label
      label = np.array([get_annotation(id_,name,os.path.join(self.base, 'labels')) for name in self.label_name]).astype(np.float)

      if self.aug:
          crop, depth, label = randomlyAug(crop, depth, label, max_val=640, square=self.square, normal=self.normal, depth_aug=self.depth_aug)

      h, w, c = crop.shape

      resX = 640 - h
      resY = 640 - w

      up = resX // 2
      down = up
      if resX % 2 != 0:
        down +=1

      left = resY // 2
      right = left

      if resY % 2 != 0:
        left += 1

      padding = transforms.Pad((left, up, right, down))

      image = Image.fromarray(np.concatenate((crop, depth), axis=2))
      image = padding(image)
      image = self.transform(image)

      return image, label

class audioDataSet(Dataset):
  def __init__(self,root_pth,test=False,transform = None):
    print("Dataset initializing...")
    class_num=4
    self.audio_pth = os.path.join(root_pth, 'audios', 'mfcc')
    filling_type = np.load(os.path.join(root_pth, 'audios', 'filling_type.npy'))
    pouring_or_shaking = np.load(os.path.join(root_pth,  'audios', 'pouring_or_shaking.npy'))
    self.label = filling_type * pouring_or_shaking
    self.is_test=test
    self.each_class_size = []
    for i in range(class_num):
        self.each_class_size.append(np.count_nonzero(self.label==i))
    mx=0
    mn=1000
    for idx in tqdm(range(self.label.shape[0])):
        data=np.load(os.path.join(self.audio_pth, "{0:06d}".format(idx+1) + '.npy'), allow_pickle=True)
        tmp_max=np.max(data)
        tmp_min=np.min(data)
        if mx<tmp_max:
            mx=tmp_max
        if mn>tmp_min:
            mn=tmp_min
    self.mn=mn
    self.mx=mx
  def __len__(self):
    return self.label.shape[0]
  def __getitem__(self, idx):
        if torch.is_tensor(idx):
            idx = idx.tolist()

        lbl = -1

        if self.is_test is False:
            lbl = self.label[idx]
        data=np.load(os.path.join(self.audio_pth, "{0:06d}".format(idx+1) + '.npy'), allow_pickle=True)
        data= (data-self.mn)/(self.mx-self.mn)
        data=data.transpose(2,0,1)
        data=torch.from_numpy(data.astype(np.float32))
        return data , lbl

  def get_each_class_size(self):
    return np.array(self.each_class_size)

class Padding(object):
    def __init__(self, seq_len):
        self.seq_len = seq_len

    def __call__(self, sample, pred):
        #np.clip(pred, 0,1,out=pred)
        sample_len, input_dim = sample.shape
        #for i in range(sample_len):
        #    sample[i, :] *= pred[i]

        if sample_len >= self.seq_len:
            features = sample[:self.seq_len, :]
            return features
        else:
            start_seq = np.random.randint(0, self.seq_len - sample_len+1)
            #ini=[1]+[0]*(input_dim-1)
            ini=[0]*(input_dim)
            features = np.full((self.seq_len, input_dim),ini, dtype = float)
            features[start_seq:start_seq+sample_len, :] = sample
            return features

class MyLSTMDataset(torch.utils.data.Dataset):
    def __init__(self,root_pth,label=None, test=False,transform = None, padding_size = 100):
        class_num=3
        self.mid_pth = os.path.join(root_pth,'features', 'T2_mid_test')
        self.pred_pth = os.path.join(root_pth,'features', 'T2_pred_test')
        self.label = label  # gt['filling_level'].to_numpy()
        self.is_test=test
        self.each_class_size = []
        self.each_class_sum = [0]*class_num
        for i in range(class_num):
            self.each_class_size.append(np.count_nonzero(self.label==i))
        mx=0
        mn=1000
        len_mx = 0

        if label is None:
            self.label = np.zeros((len(os.listdir(self.mid_pth))))

        for idx in range(len(os.listdir(self.mid_pth))):
            data=np.load(os.path.join(self.mid_pth, "{0:06d}".format(idx) + '.npy'), allow_pickle=True)
            #self.each_class_sum[self.label[idx]]+=data.shape[0]
            if data.shape[0] > len_mx:
                len_mx=data.shape[0]
            tmp_max=np.max(data)
            tmp_min=np.min(data)
            if mx<tmp_max:
                mx=tmp_max
            if mn>tmp_min:
                mn=tmp_min
        self.mn=mn
        self.mx=mx
        self.pad = Padding(padding_size)
        print(len_mx)

    def __len__(self):
        return self.label.shape[0]

    def __getitem__(self, idx):
        if torch.is_tensor(idx):
            idx = idx.tolist()

        lbl = -1

        if self.is_test is False:
            lbl = self.label[idx]

        data=np.load(os.path.join(self.mid_pth, "{0:06d}".format(idx) + '.npy'), allow_pickle=True)
        pred=np.load(os.path.join(self.pred_pth, "{0:06d}".format(idx) + '.npy'), allow_pickle=True)
        data = (data-self.mn)/(self.mx-self.mn)
        data = self.pad(data, pred)

        #np.clip(data, 0,1,out=data)
        data=torch.from_numpy(data.astype(np.float32))
        return data , lbl

    def get_each_class_size(self):
        return np.array(self.each_class_size)

    def get_each_class_avg_len(self):
        each_class_avg_len =  np.array(self.each_class_sum)/np.array(self.each_class_size)
        all_class_avg_len = np.sum(np.array(self.each_class_sum))/np.sum(np.array(self.each_class_size))
        return each_class_avg_len, all_class_avg_len

class BatchProcess(Dataset):
    """
    type  = train or test
    """
    def __init__(self,rgbs,depths):

        self.rgbs = rgbs
        self.depths = depths
        self.transform = transforms.Compose([transforms.Resize((320, 320)),
                                             transforms.ToTensor(),
                                             transforms.ConvertImageDtype(torch.float),
                                             ])

    def __len__(self):
        return len(self.rgbs)

    def __getitem__(self, idx):

        depth = np.asarray(self.depths[idx])[:,:,np.newaxis]

        # rgb_cropped
        crop = np.asarray(self.rgbs[idx])
        h, w, c = crop.shape
        resX = 640 - h
        resY = 640 - w
        up = resX // 2
        down = up
        if resX % 2 != 0:
            down +=1
        left = resY // 2
        right = left
        if resY % 2 != 0:
            left += 1
        padding = transforms.Pad((left, up, right, down))
        image = Image.fromarray(np.concatenate((crop, depth), axis=2))
        image = padding(image)
        image = self.transform(image)

        return image

class MyLSTMDataset_combine(torch.utils.data.Dataset):
    def __init__(self,root_pth,label=None, test=False,transform = None, padding_size = 100):
        class_num=3
        self.mid_pth = os.path.join(root_pth,'features', 'T2_mid_test')
        self.pred_pth = os.path.join(root_pth,'features', 'T2_pred_test')
        self.video_pth = os.path.join(root_pth,'features_video_test')
        self.label = label  # gt['filling_level'].to_numpy()
        self.is_test=test
        self.each_class_size = []
        self.each_class_sum = [0]*class_num
        for i in range(class_num):
            self.each_class_size.append(np.count_nonzero(self.label==i))
        mx=0
        mn=1000
        len_mx = 0

        if label is None:
            self.label = np.zeros((len(os.listdir(self.mid_pth))))

        for idx in range(len(os.listdir(self.mid_pth))):
            data=np.load(os.path.join(self.mid_pth, "{0:06d}".format(idx) + '.npy'), allow_pickle=True)
            #self.each_class_sum[self.label[idx]]+=data.shape[0]
            if data.shape[0] > len_mx:
                len_mx=data.shape[0]
            tmp_max=np.max(data)
            tmp_min=np.min(data)
            if mx<tmp_max:
                mx=tmp_max
            if mn>tmp_min:
                mn=tmp_min
        self.mn=mn
        self.mx=mx
        self.pad = Padding(padding_size)

    def __len__(self):
        return self.label.shape[0]

    def __getitem__(self, idx):
        if torch.is_tensor(idx):
            idx = idx.tolist()

        lbl = -1

        if self.is_test is False:
            lbl = self.label[idx]

        data=np.load(os.path.join(self.mid_pth, "{0:06d}".format(idx) + '.npy'), allow_pickle=True)
        data_video=np.load(os.path.join(self.video_pth, "{0:06d}".format(idx) + '.npy'), allow_pickle=True)
        pred=np.load(os.path.join(self.pred_pth, "{0:06d}".format(idx) + '.npy'), allow_pickle=True)
        data = (data-self.mn)/(self.mx-self.mn)
        data = self.pad(data, pred)
        data_video =self.pad(self.pad(data, pred), pred)
        data_combine=np.concatenate((data, data_video),axis=0)
        data_combine=torch.from_numpy(data_combine.astype(np.float32))
        return data_combine , lbl

    def get_each_class_size(self):
        return np.array(self.each_class_size)

    def get_each_class_avg_len(self):
        each_class_avg_len =  np.array(self.each_class_sum)/np.array(self.each_class_size)
        all_class_avg_len = np.sum(np.array(self.each_class_sum))/np.sum(np.array(self.each_class_size))
        return each_class_avg_len, all_class_avg_len


<reponame>msgoff/sympy<filename>sympy/polys/tests/test_orthopolys.py
"""Tests for efficient functions for generating orthogonal polynomials. """

from sympy import Poly, S, Rational as Q
from sympy.testing.pytest import raises

from sympy.polys.orthopolys import (
    jacobi_poly,
    gegenbauer_poly,
    chebyshevt_poly,
    chebyshevu_poly,
    hermite_poly,
    legendre_poly,
    laguerre_poly,
)

from sympy.abc import x, a, b

def test_jacobi_poly():
    raises(ValueError, lambda: jacobi_poly(-1, a, b, x))

    assert jacobi_poly(1, a, b, x, polys=True) == Poly(
        (a / 2 + b / 2 + 1) * x + a / 2 - b / 2, x, domain="ZZ(a,b)"
    )

    assert jacobi_poly(0, a, b, x) == 1
    assert jacobi_poly(1, a, b, x) == a / 2 - b / 2 + x * (a / 2 + b / 2 + 1)
    assert jacobi_poly(2, a, b, x) == (
        a ** 2 / 8
        - a * b / 4
        - a / 8
        + b ** 2 / 8
        - b / 8
        + x ** 2
        * (a ** 2 / 8 + a * b / 4 + a * Q(7, 8) + b ** 2 / 8 + b * Q(7, 8) + Q(3, 2))
        + x * (a ** 2 / 4 + a * Q(3, 4) - b ** 2 / 4 - b * Q(3, 4))
        - S.Half
    )

    assert jacobi_poly(1, a, b, polys=True) == Poly(
        (a / 2 + b / 2 + 1) * x + a / 2 - b / 2, x, domain="ZZ(a,b)"
    )

def test_gegenbauer_poly():
    raises(ValueError, lambda: gegenbauer_poly(-1, a, x))

    assert gegenbauer_poly(1, a, x, polys=True) == Poly(2 * a * x, x, domain="ZZ(a)")

    assert gegenbauer_poly(0, a, x) == 1
    assert gegenbauer_poly(1, a, x) == 2 * a * x
    assert gegenbauer_poly(2, a, x) == -a + x ** 2 * (2 * a ** 2 + 2 * a)
    assert gegenbauer_poly(3, a, x) == x ** 3 * (
        4 * a ** 3 / 3 + 4 * a ** 2 + a * Q(8, 3)
    ) + x * (-2 * a ** 2 - 2 * a)

    assert gegenbauer_poly(1, S.Half).dummy_eq(x)
    assert gegenbauer_poly(1, a, polys=True) == Poly(2 * a * x, x, domain="ZZ(a)")

def test_chebyshevt_poly():
    raises(ValueError, lambda: chebyshevt_poly(-1, x))

    assert chebyshevt_poly(1, x, polys=True) == Poly(x)

    assert chebyshevt_poly(0, x) == 1
    assert chebyshevt_poly(1, x) == x
    assert chebyshevt_poly(2, x) == 2 * x ** 2 - 1
    assert chebyshevt_poly(3, x) == 4 * x ** 3 - 3 * x
    assert chebyshevt_poly(4, x) == 8 * x ** 4 - 8 * x ** 2 + 1
    assert chebyshevt_poly(5, x) == 16 * x ** 5 - 20 * x ** 3 + 5 * x
    assert chebyshevt_poly(6, x) == 32 * x ** 6 - 48 * x ** 4 + 18 * x ** 2 - 1

    assert chebyshevt_poly(1).dummy_eq(x)
    assert chebyshevt_poly(1, polys=True) == Poly(x)

def test_chebyshevu_poly():
    raises(ValueError, lambda: chebyshevu_poly(-1, x))

    assert chebyshevu_poly(1, x, polys=True) == Poly(2 * x)

    assert chebyshevu_poly(0, x) == 1
    assert chebyshevu_poly(1, x) == 2 * x
    assert chebyshevu_poly(2, x) == 4 * x ** 2 - 1
    assert chebyshevu_poly(3, x) == 8 * x ** 3 - 4 * x
    assert chebyshevu_poly(4, x) == 16 * x ** 4 - 12 * x ** 2 + 1
    assert chebyshevu_poly(5, x) == 32 * x ** 5 - 32 * x ** 3 + 6 * x
    assert chebyshevu_poly(6, x) == 64 * x ** 6 - 80 * x ** 4 + 24 * x ** 2 - 1

    assert chebyshevu_poly(1).dummy_eq(2 * x)
    assert chebyshevu_poly(1, polys=True) == Poly(2 * x)

def test_hermite_poly():
    raises(ValueError, lambda: hermite_poly(-1, x))

    assert hermite_poly(1, x, polys=True) == Poly(2 * x)

    assert hermite_poly(0, x) == 1
    assert hermite_poly(1, x) == 2 * x
    assert hermite_poly(2, x) == 4 * x ** 2 - 2
    assert hermite_poly(3, x) == 8 * x ** 3 - 12 * x
    assert hermite_poly(4, x) == 16 * x ** 4 - 48 * x ** 2 + 12
    assert hermite_poly(5, x) == 32 * x ** 5 - 160 * x ** 3 + 120 * x
    assert hermite_poly(6, x) == 64 * x ** 6 - 480 * x ** 4 + 720 * x ** 2 - 120

    assert hermite_poly(1).dummy_eq(2 * x)
    assert hermite_poly(1, polys=True) == Poly(2 * x)

def test_legendre_poly():
    raises(ValueError, lambda: legendre_poly(-1, x))

    assert legendre_poly(1, x, polys=True) == Poly(x)

    assert legendre_poly(0, x) == 1
    assert legendre_poly(1, x) == x
    assert legendre_poly(2, x) == Q(3, 2) * x ** 2 - Q(1, 2)
    assert legendre_poly(3, x) == Q(5, 2) * x ** 3 - Q(3, 2) * x
    assert legendre_poly(4, x) == Q(35, 8) * x ** 4 - Q(30, 8) * x ** 2 + Q(3, 8)
    assert legendre_poly(5, x) == Q(63, 8) * x ** 5 - Q(70, 8) * x ** 3 + Q(15, 8) * x
    assert legendre_poly(6, x) == Q(231, 16) * x ** 6 - Q(315, 16) * x ** 4 + Q(
        105, 16
    ) * x ** 2 - Q(5, 16)

    assert legendre_poly(1).dummy_eq(x)
    assert legendre_poly(1, polys=True) == Poly(x)

def test_laguerre_poly():
    raises(ValueError, lambda: laguerre_poly(-1, x))

    assert laguerre_poly(1, x, polys=True) == Poly(-x + 1)

    assert laguerre_poly(0, x) == 1
    assert laguerre_poly(1, x) == -x + 1
    assert laguerre_poly(2, x) == Q(1, 2) * x ** 2 - Q(4, 2) * x + 1
    assert (
        laguerre_poly(3, x) == -Q(1, 6) * x ** 3 + Q(9, 6) * x ** 2 - Q(18, 6) * x + 1
    )
    assert (
        laguerre_poly(4, x)
        == Q(1, 24) * x ** 4
        - Q(16, 24) * x ** 3
        + Q(72, 24) * x ** 2
        - Q(96, 24) * x
        + 1
    )
    assert (
        laguerre_poly(5, x)
        == -Q(1, 120) * x ** 5
        + Q(25, 120) * x ** 4
        - Q(200, 120) * x ** 3
        + Q(600, 120) * x ** 2
        - Q(600, 120) * x
        + 1
    )
    assert (
        laguerre_poly(6, x)
        == Q(1, 720) * x ** 6
        - Q(36, 720) * x ** 5
        + Q(450, 720) * x ** 4
        - Q(2400, 720) * x ** 3
        + Q(5400, 720) * x ** 2
        - Q(4320, 720) * x
        + 1
    )

    assert laguerre_poly(0, x, a) == 1
    assert laguerre_poly(1, x, a) == -x + a + 1
    assert (
        laguerre_poly(2, x, a)
        == x ** 2 / 2 + (-a - 2) * x + a ** 2 / 2 + a * Q(3, 2) + 1
    )
    assert (
        laguerre_poly(3, x, a)
        == -(x ** 3) / 6
        + (a / 2 + Q(3) / 2) * x ** 2
        + (-(a ** 2) / 2 - a * Q(5, 2) - 3) * x
        + a ** 3 / 6
        + a ** 2
        + a * Q(11, 6)
        + 1
    )

    assert laguerre_poly(1).dummy_eq(-x + 1)
    assert laguerre_poly(1, polys=True) == Poly(-x + 1)


#!/usr/bin/env python
"""Collection of simple functions useful in computational chemistry scripting.

Many of the following functions are used to make operations on xyz coordinates
of molecular structure. When refering to ``xyz_data`` bellow, the following
structures (also used in :py:mod:`~comp_chem_utils.molecule_data`) is assumed::

    atom 1 label  and corresponding xyz coordinate
    atom 2 label  and corresponding xyz coordinate
    : : :
    atom N label  and corresponding xyz coordinate

For example the ``xyz_data`` of a Hydrogen molecule along the z-axis
should be passed as::

    >>> xyz_data
    [['H', 0.0, 0.0, 0.0], ['H', 0.0, 0.0, 1.0]]

"""

__author__="<NAME>"
__email__="<EMAIL>"

import os
import shutil
import numpy as np
from scipy.spatial.distance import pdist

from comp_chem_utils.periodic import element

def vel_auto_corr(vel, max_corr_index, tstep):
    """Calculate velocity autocorrelation function.

    Args:
        vel (list): The velocities along the trajectory are given as
            a list of ``np.array()`` of dimensions N_atoms . 3.

        max_corr_index (int): Maximum number of steps to consider
            for the auto correlation function. In other words, it
            corresponds to the number of points in the output function.
    """

    max_step_index = len(vel) - max_corr_index
    natoms = vel[0].shape[0]

    G = np.zeros((max_corr_index), dtype='float64')

    for itau in range(max_corr_index):
        for it in range(max_step_index):
            #for i in range(natoms):
            #    G[itau] += np.dot(vel[it][i,:], vel[it+itau][i,:])

            G[itau] += np.trace(np.dot(vel[it],np.transpose(vel[it+itau])))

    G = G / (natoms * max_corr_index * tstep)
    xpts = np.arange(max_corr_index)
    xpts = xpts * tstep

    return xpts, G

def get_lmax_from_atomic_charge(charge):
    """Return the maximum angular momentum based on the input atomic charge.

    This function is designed to return LMAX in a CPMD input file.

    Args:
        charge (int): Atomic charge.

    Returns:
        'S' for H or He; 'P' for second row elements; and 'D' for heavier elements.
    """

    if charge <= 2:
        # H or He
        lmax = 'S'
    elif charge <= 10:
        lmax = 'P'
    else:
        lmax = 'D'

    return lmax

def get_file_as_list(filename, raw=False):
    """Read a file and return it as a list of lines (str).

    By default comments (i.e. lines starting with #)
    and empty lines are ommited. This can be changed
    by setting ``raw=True``

    Args:
        filename (str): Name of the file to read.
        raw (bool, optional): To return the file as it is,
            i.e. with comments and blank lines. Default
            is ``raw=False``.

    Returns:
        A list of lines (str)."""

    lines = []
    with open(filename,'r') as myfile:
        for line in myfile:
            if raw:
                lines.append(line)
            else:
                # remove empty lines
                if line.strip():
                    # remove comments
                    if line.strip()[0] != '#':
                        lines.append(line)

    return lines

def make_new_dir(dirn):
    """Make new empty directory.

    If the directory already exists it is erased and replaced.

    Args:
        dirn (str): Name for the new directory (can include path).
    """

    if not os.path.exists(dirn):
        os.makedirs(dirn)
    else:
        try:
            os.removedirs(dirn)
        except(OSError):
            print("WARNING: erasing (not empty) directory! {}".format(dirn))
            shutil.rmtree(dirn)

        os.makedirs(dirn)

def center_of_mass(xyz_data):
    """Calculate center of mass of a molecular structure based on xyz coordinates.

    Args:
        xyz_data (list): xyz atomic coordinates arranged as described above.

    Returns:
        3-dimensional ``np.array()`` containing the xyz coordinates of the
        center of mass of the molecular structure. The unit of the center of mass
        matches the xyz input units (usually Angstroms).
    """

    totM = 0.0
    COM = np.zeros((3))
    for line in xyz_data:
        symbol = line[0]
        coord = np.array(line[1:])

        mass = element(symbol).mass
        totM += mass
        COM += coord*mass

    COM = COM/totM
    return COM

def change_vector_norm(fix, mob, R):
    """Scale a 3-D vector defined by two points in space to have a new norm R.

    The input vector is defined by a fix position in 3-D space ``fix``,
    and a mobile position ``mob``. The function returns a new mobile
    position such that the new vector has the norm R.

    Args:
        fix (np.array): xyz coordinates of the fix point.
        mob (np.array): Original xyz coordinates of the mobile point.
        R (float): Desired norm for the new vector.

    Returns:
        The new mobile position as an ``np.array()`` of dimenssion 3.
    """

    unit = mob - fix
    unit = unit/np.linalg.norm(unit)

    # return new position
    return fix + R * unit

def change_vector_norm_sym(pos1, pos2, R):
    """Symmetric version of change_vector_norm function.

    In other word both positions are modified symmetrically.
    """

    unit = pos2 - pos1
    norm = np.linalg.norm(unit)
    unit = unit/norm

    shift = (R-norm)/2.0

    # return new position
    new_pos1 = pos1 - unit * shift
    new_pos2 = pos2 + unit * shift
    return new_pos1, new_pos2

def get_rmsd(xyz_data1, xyz_data2):
    """Calculate RMSD between two sets of coordinates.

    The Root-mean-square deviation of atomic positions is calculated as

    .. math::
        RMSD = \\sqrt{ \\frac{1}{N} \\sum_{i=1}^N \\delta_{i}^{2} }

    Where ``\delta_i`` is the distance between atom i in ``xyz_data1`` and in
    ``xyz_data2``.

    Args:
        xyz_data1 (list): List of atomic coordinates for the first structure
            arranged as described above for xyz_data.
        xyz_data2 (list): Like ``xyz_data1`` but for the second structure.

    Returns:
        The RMSD (float).
    """

    rmsd = 0
    for c1, c2 in zip(xyz_data1, xyz_data2):

        d1 = np.array([c1[x] for x in range(1,4)])
        d2 = np.array([c2[x] for x in range(1,4)])
        vector = d2 - d1
        rmsd += np.dot(vector, vector)

    rmsd = rmsd/(len(xyz_data1))
    return np.sqrt(rmsd)

def get_distance(xyz_data, atoms, box_size=None):
    """Calculate distance between two atoms in xyz_data.

    Args:
        xyz_data (list): xyz atomic coordinates arranged as described above.
        atoms (list): list of two indices matching the two rows of the
            xyz_data for wich the distance should be calculated.

    Returns:
        Distance between the two atoms in the list as a ``float``, the
        unit will match the input unit in the ``xyz_data``.

    """
    coord1 = np.array([xyz_data[atoms[0]][x] for x in range(1,4)])
    coord2 = np.array([xyz_data[atoms[1]][x] for x in range(1,4)])
    vector = coord2 - coord1

    if box_size:
       for i,x in enumerate(vector):
          if abs(x) > box_size/2.0:
             vector[i] = box_size - abs(x)

    return np.linalg.norm(vector)

def get_distance_matrix(xyz_data, box_size=None):

   # make np.array
   natoms = len(xyz_data)
   coord = np.zeros((natoms,3), dtype='float')
   for i,line in enumerate(xyz_data):
      coord[i,:] = line[1:]

   if box_size:
      npairs = natoms * (natoms - 1)
      matrix = np.zeros((npairs,3), dtype='float')

      for i in range(natoms):
         for j in range(i+1,natoms):
            ij = i + (j-1)*natoms

            matrix[ij,:] = coord[i,:] - coord[j,:]

      # find out which element to shift:
      # basically shift is an array of same shape as matrix
      # with zeros every where except where the elements of
      # matrix are larger than box_size/2.0
      # in that case shift as the value box_size
      shift = box_size * (matrix > box_size/2.0).astype(int)

      # we can now shift the matrix as follows:
      matrix = abs(shift - matrix)

      # and get the distances...
      matrix = np.linalg.norm(matrix, axis=1)

   else:
      matrix = pdist(coord)

   return matrix

def get_distance_matrix_2(xyz_data1, xyz_data2, box_size=None):
   # repeat as above for 2 different sets of coordinates
   nat1 = len(xyz_data1)
   coord1 = np.zeros((nat1,3), dtype='float')
   for i,line in enumerate(xyz_data1):
      coord1[i,:] = line[1:]

   nat2 = len(xyz_data2)
   coord2 = np.zeros((nat2,3), dtype='float')
   for i,line in enumerate(xyz_data2):
      coord2[i,:] = line[1:]

   if box_size:
      matrix = np.zeros((nat1,nat2,3), dtype='float')

      for i in range(nat1):
         for j in range(nat2):
            matrix[i,j,:] = coord[i,:] - coord[j,:]

      # find out which element to shift:
      # basically shift is an array of same shape as matrix
      # with zeros every where except where the elements of
      # matrix are larger than box_size/2.0
      # in that case shift as the value box_size
      shift = box_size * (matrix > box_size/2.0).astype(int)

      # we can now shift the matrix as follows:
      matrix = abs(shift - matrix)

      # and get the distances...
      matrix = np.linalg.norm(matrix, axis=2)

   else:
      matrix = cdist(coord1, coord2)

   return matrix

def get_angle(xyz_data, atoms):
    """Calculate angle between three atoms in xyz_data.

    Args:
        xyz_data (list): xyz atomic coordinates arranged as described above.
        atoms (list): list of three indices matching the rows of the
            xyz_data for wich the angle should be calculated.

    Returns:
        Angle between the three atoms in the list as a ``float`` in degrees.

    """

    coord1 = np.array([xyz_data[atoms[0]][x] for x in range(1,4)])
    coord2 = np.array([xyz_data[atoms[1]][x] for x in range(1,4)])
    coord3 = np.array([xyz_data[atoms[2]][x] for x in range(1,4)])
    vec1 = coord1-coord2
    vec2 = coord3-coord2

    return np.degrees( np.arccos( np.dot(vec1,vec2) / (np.linalg.norm(vec1) * np.linalg.norm(vec2)) ) )

def get_dihedral_angle(table, atoms):
    """Calculate dihedral angle defined by 4 atoms in xyz_data.

    It relies on the praxeolitic formula (1 sqrt, 1 cross product).

    Args:
        xyz_data (list): xyz atomic coordinates arranged as described above.
        atoms (list): list of 4 indices matching the rows of the
            xyz_data for wich the dihedral angle should be calculated.

    Returns:
        Dihedral angle defined by the 4 atoms in the list as a ``float`` in degrees.

    """

    p0 = np.array([table[atoms[0]][x] for x in range(1,4)])
    p1 = np.array([table[atoms[1]][x] for x in range(1,4)])
    p2 = np.array([table[atoms[2]][x] for x in range(1,4)])
    p3 = np.array([table[atoms[3]][x] for x in range(1,4)])

    b0 = -1.0*(p1 - p0)
    b1 = p2 - p1
    b2 = p3 - p2

    # normalize b1 so that it does not influence magnitude of vector
    # rejections that come next
    b1 /= np.linalg.norm(b1)

    # vector rejections
    # v = projection of b0 onto plane perpendicular to b1
    #   = b0 minus component that aligns with b1
    # w = projection of b2 onto plane perpendicular to b1
    #   = b2 minus component that aligns with b1
    v = b0 - np.dot(b0, b1)*b1
    w = b2 - np.dot(b2, b1)*b1

    # angle between v and w in a plane is the torsion angle
    # v and w may not be normalized but that's fine since tan is y/x
    x = np.dot(v, w)
    y = np.dot(np.cross(b1, v), w)

    return np.degrees(np.arctan2(y, x))