text:
# @Email: rah3156@rit.edu
# @Website: https://ryanhartzell.github.io/scanimate
# @License: https://github.com/RyanHartzell/scanimate/blob/master/LICENSE
# @github: https://github.com/RyanHartzell/scanimate
#
# Copyright (c) 2020 Ryan Hartzell
#
"""
Purpose of this module should be two major things:

    1) Moire class for generating obscuration stripes/grids/patterns
        Constant Patterns (don't change frame to frame)
            * horizontal and vertical striped, uniform stride and width
            * horizontal and vertical striped, variable/logarithmic/sigmoid stride and width (to play with animation keyframes)
            - diagonal striped (anti-aliased?)
            - SINUSOIDAL PATTERNS (shift by phase)
            - blurry stipes or slight overlap for motion blurring?
            - radial "sun" striped, uniform stride and width
            - radial "pinwheel" striped, uniform stride and width
            - concentric bands, play with "zoom" style clipping for frequency content (and distance viewing) based animation? Might be pretty cluttered
            - hybrid image method? clip at different frequencies for different images, and then see different animation depending on grid you choose
        Variable Patterns (different spatial mask per frame)
        - topographical / isomap where each "altitude" is clipped and rotated through (digital only!!!)
        - random -> take a unique (or slightly overlapping...) random but uniformly distributed selection of pixels and clip to that, then do a different selection for each frame. might result in a more cohesive looking image than the consistent banding, or might break the illusion

    2) Functions for shifting the patterns:
        - laterally
        - radially
        - potentially zooming them in or out

    3) Clipper/cropper functions for taking a stack of images/frames and
        aggregating the clipped versions into a single static animation frame

    4) Sizing/scaling/re-sampling utilities for preping images for different types of display. Computer screen is easy.
        If you wanted to print a tattoo and grid onto a transparency or into an app though, that needs to be done correctly,
        or the illusion will break.

"""
import numpy as np
import matplotlib.pyplot as plt
import imageio

class Moire(object):
    def __init__(self, pattern, stripe_width=1, num_frames=6):
        self._pattern = pattern
        self._stripe_width = stripe_width # in pixels
        self._num_frames = num_frames

        self._replications = 24 # Will be set by target image dimensions eventually
        self.make_grid()

    def __call__(self, arr, schema="BFCHW", mode="static"):
        """Applies grid to stack of images. takes arr (default shape of [batch, frames, channel, height, width] / BFCHW
                (where height is rows, width is columns), but also takes FCWH, WHCF, HWCF (for color==True),
                FWH, FHW, WHF, HWF (for color==False aka grayscale),
                or any of the previous preceded or followed by a dimension B for batch."""
        # 1) check arr for len == num_frames
        if not len(arr) == num_frames:
            raise ValueError(f"Current implementation requires first axis of array to be equal to {self._num_frames}")
        # 2) apply grid to image stack along axis who's length equals num_frames
        #       this will effectively be performed as follows:
        #           a) Multiply grid with arr
        #           b) Sum arr along pixels to combine
        # 3) return single image of scanimation *OR* multiplied stack, which will show illusion when played as gif (really just for testing purposes)
        pass

    @property
    def replications(self):
        return self._replications

    @replications.setter
    def replications(self, val):
        self._replications = val

    def scale(self, arr):
        """Scale input image to size of grid"""
        pass

    def crop(self, arr):
        """Crop input image to size of grid if larger than grid"""
        pass

    def pad(self, arr, mode="zeros"):
        """Pad input image to size of grid if smaller"""
        pass

    def make_grid(self):
        size = self._stripe_width * self._num_frames * self._replications
        self.grid = np.zeros((size, size)) # assuming square for now, not worrying about an input image stack's dimensions
        self.grid[0::self._num_frames] = 1

    def apply_grid(self, arr):
        # This will be expanded so we know how to deal with all different modes, or we'll make a utility for applying a function efficiently against the various shapes. Also may add complexity for low frequency and high frequency
        return self.grid * arr

if __name__ == "__main__":
    import numpy as np
    import matplotlib.pyplot as plt
    import imageio

    # RUN THIS SCRIPT AS ```python core/moire.py```

    # 1) Example of circle moving, monochrome (144, 144)
    # Define control parameters
    stripe_width = 1 # size of each stripe in pixels, stripe per frame in animation. For now, only single pixel width works without fancy indexing and or fancy windowed striding
    num_frames = 5 # number of frames in animation (image stack size)
    replications = 64 # number of times the pattern defined by (stripe_width * num_frames) should be repeated

    # Next, setup your moire patterns
    size = stripe_width * num_frames * replications

    # For horizontal motion, a square wave in the x direction works well
    hgrid = np.zeros((size, size))
    hgrid[:, 0::num_frames] = 1

    # Likewise, a square wave in the y direction is good for vertical motion
    vgrid = np.zeros((size, size))
    vgrid[0::num_frames] = 1

    # We'll use the horizontal one from here on out
    # Let's make a stack of frames showing a circle moving
    X,Y = np.meshgrid(np.linspace(-5,5,size), np.linspace(-5,5,size))
    r = np.sqrt(X**2+Y**2)

    # Now lets build our circle
    circle = np.zeros_like(r)
    circle[r < 2.5] = 1

    # and let's assemble the masked frames
    masked_circle = np.asarray([np.zeros_like(circle)]*5)
    for i in range(5):
        masked_circle[i] = np.roll(circle, (i*30,i*20), axis=(0,1)) * np.roll(hgrid, stripe_width*i, axis=1)

    # At this point these are the pattern-clipped frames (which can be visualized), but must be summed for printable static animation (scanimation)
    masked_circle = np.sum(masked_circle, axis=0).astype(np.uint8)

    # show the result
    plt.imshow(masked_circle)
    plt.title("Static")
    plt.show()

    # 2) Example of michael.gif into static animation, RGBA
    # Load a gif of Michael from The Office
    img = np.asarray(imageio.mimread("images/michael.gif"))
    print("GIF SHAPE [BEFORE SUBSAMP] = ", img.shape)

    # Subsample gif frames (axis=0)
    img = img[0::2]
    print("GIF SHAPE [AFTER SUBSAMP]  = ", img.shape)

    grid = np.zeros_like(img)

    # this step should likely be done with stride tricks sliding window view. Could also do that to build static animation image from source automatically,
    #       but likely only for the vertical and horizontal striped cases. It's *possible* that I could work out a tiled  pattern and use that as the "grid"
    #       using this method too, but I think that's outside of the scope here. Would like to come up with an efficient mask construction case regardless.
    #       One other method would be to grab a view of a complex source mask, then "smear" it by interleaving how ever many copy columns from a view into a new array
    #       then cropping to the size of the image to be masked. Something like that could work nicely. This may not be a one size fits all approach either though

    for i in range(stripe_width):
        grid[:,:,i::img.shape[0]*stripe_width,:] = 1

    # for i in range(stripe_width):
    #     grid[:,i::img.shape[0]*stripe_width,:,:] = 1

    grid[:,:,:,-1] = 1 # takes care of alpha channel being always 1

    # Now we have two options, either store the rolled frames of the grid, or roll during multiplication
    for i in range(1, len(grid)):
        grid[i] = np.roll(grid[i], i*stripe_width, axis=1) # axis = 1 for columnwise/horizontal shift, axis = 0 for rowwise/vertical shift

    # Multiply through, sum all frames together (taking care to replace alpha channel with 255.) and converting all to np.uint8
    res = img.astype(np.float32) * grid.astype(np.float32)

    # At this point these are the pattern-clipped frames (which can be visualized), but must be summed for printable static animation (scanimation)
    res = np.sum(res, axis=0)
    res[:,:,-1] = 255.
    res = res.astype(np.uint8)

    # show the result
    plt.imshow(res)
    plt.title("Static")
    plt.show()

    # show the result animated by changing the grid's phase for some number of steps
    plt.ion()

    fig, ax = plt.subplots(1, figsize=(15,15))

    for n in range(5):
        for i in range(len(grid)):
            ax.clear()
            ax.set_title("Dynamic")
            ax.imshow(res*grid[i,:,:])
            plt.pause(1/(len(grid)))

    plt.ioff()

    # Save out both static image and grid (first frame only, since you'll print on physical medium and maybe even 3D print)
    imageio.imwrite("images/michael_scanimation.png", res)
    imageio.imwrite("images/michael_scanimation_grid.png", (255*grid[0]).astype(np.uint8))

    # Save at various dpi levels
    # TODO ...

max_stars_repo_name: RyanHartzell/scanimate

max_issues_repo_name: RyanHartzell/scanimate

max_forks_repo_name: RyanHartzell/scanimate


text:
#// Author Jack Hales, 17-year-old Australian kas.nsw.edu.au Student
#// --
#// BitMap manipulation, BitMap refers to the process of creating a
#// black-background image that contains a prominent set of colours
#// that are significant in meaning.

import time, random, math
import numpy as np
from PIL import Image, ImageDraw
import pprint as pprint
import os, sys

pp = pprint.PrettyPrinter(indent=4)
pp = pp.pprint

# Pathfinding!
# from pathfinding.core.diagonal_movement import DiagonalMovement
# from pathfinding.core.grid import Grid
# from pathfinding.finder.a_star import AStarFinder

#// Manual Colour-coding array
#// To be in JSON file one day

#// Requirement splitting for Vision rules
#// % value, 0.25 = 25% means splitting into 4 parts x 4 parts

# Meta retention-data,
# not defined by the type of image,
# but the self-manipulation.
Complexity, by = 1, 16
split = 1 / by
squares  = int(1 / split)
squaresx = squares
totalSquares = int(squares*squares)
xs, ys = 0, 0

class PixelArray(object):
    # - dynamically set variables
    # Squares
    # OGPixelArray is a numpy array

    def SetBy(self, toSet):
        global by
        by = toSet

    def __init__(self, PA, bySet = False):
        # cp-string for data collection and retention
        global xs, ys, xSplit, ySplit, squarePixels, totalSquares, squaresx, squares, split, by
        by = bySet if bySet != False else by

        # global re-distributer
        split = 1 / by
        squares  = int(1 / split)
        squaresx = squares
        totalSquares = int(squares*squares)
        self.SquareMeta = np.zeros(totalSquares, dtype=dict)

        self.OGPixelArray = PA
        self.BitMap = Image.fromarray(PA)
        xs, ys = self.BitMap.size[0], self.BitMap.size[1]

        squarePixels = (xs*ys)/squares
        xSplit = int((xs) / (1 / split))
        ySplit = int((ys) / (1 / split))

        self.Squares = []
        # - convert the "pixel array" location
        #   into the square information and insert
        #   seq
        self.PixelArrCache = np.zeros(shape=(squares*xSplit*xSplit+1, squares*ySplit*ySplit+1), dtype=list)

    # | Quick anti-alias removal, black & white ONLY,
    # | lesser is reduced to black, and higher is
    # | ceiling'd to white
    def AARemove(self):
        global xs, ys, xSplit, ySplit, squarePixels, totalSquares, squaresx, squares, split, by
        PixelArr = self.OGPixelArray
        for i in range(0, xs):
            for b in range(0, ys):
                if PixelArr[i, b][0] < 120: PixelArr[i, b] = [0, 0, 0]
                else: PixelArr[i, b] = [255, 255, 255]
        self.OGPixelArray = PixelArr

    def SortSquares(self):
        PixelArr = self.OGPixelArray
        # print("from ss: ", str(squares), )
        Squares = [None for f in range(0, squares*squares)]
        startX, startY, iteration = 0, 0, 0
        # squares are stored from 0 -> squares*squares
        for i in range(int(squares*squares)):
            # print(startX, startY)
            iteration += 1
            Squares[i] = self.SquareBundlerStartFrom(i, startX, startY)
            if iteration % int(xs/xSplit) == 0:
                startY = startY + ySplit
                startX = 0
            else:
                startX = startX + xSplit
        self.Squares = Squares
        self.GenerateSquareReports()

    def GenerateSquareReports(self):
        for squarex, square in enumerate(self.Squares):
            self.SquareMeta[squarex] = {"empty": self.SquareEmpty(squarex)}

    def SquareBundlerStartFrom(self, squareNo, xStart, yStart):
        # Presetting the values for the bundled variables in a
        # 2D Numpy Array
        bundled = np.zeros(shape=(xSplit, ySplit), dtype=object)

        # Where we will start from for finding the X/Y values
        PixelArr = self.OGPixelArray
        storeX, storeY = xStart, yStart

        # For Inserting
        insertX, insertY, iteration = 0, 0, 0
        for i in range(0, (xSplit * ySplit)):
            iteration = i+1
            #// Storing into the "bundled" array to return later
            #// storing in the bundle return [Y, X]
            pixel = PixelArr[storeY, storeX]
            bundled[insertY, insertX] = pixel

            #// Storing data for later converts, such as generalized
            #// information store such as empty, etc...
            self.PixelArrCache[storeY, storeX] = {"square": squareNo, "insertY": storeY, "insertX": storeX, "squareLocX": insertY, "squareLocY": insertX}
            if (iteration % int(xSplit) == 0):
                storeX = xStart
                storeY +=  1
                insertX = 0
                insertY += 1
            else:
                storeX += 1
                insertX += 1
        # bundled[0, 0] = [255, 255, 255]
        return bundled

    def SquareToPixelArr(self, SquareNo, Coords):
        # print(self.OGPixelArray)
        # - coords = [X, Y]
        for eachInfoBundle in self.PixelArrCache:
            for cacheInformation in eachInfoBundle:
                if (cacheInformation != 0) and (SquareNo == cacheInformation["square"]) and (Coords[0] == cacheInformation["squareLocX"]) and (Coords[1] == cacheInformation["squareLocY"]):
                    return cacheInformation

    # | Return the colours in the square number
    def AllColourInSquare(self, squareNo, colour=[255, 255, 255]):
        # print(squareNo)
        area = []
        for eachInfoBundle in self.PixelArrCache:
            for cacheInformation in eachInfoBundle:
                if (cacheInformation != 0) and (squareNo == cacheInformation["square"]):
                    x, y = cacheInformation["insertX"], cacheInformation["insertY"]
                    # y, x
                    # print(self.OGPixelArray[y, x], colour)
                    if str(self.OGPixelArray[y, x]) == str(colour).replace(',', ''):
                        area.append(cacheInformation)
        return area

    def Square2PAEditor(self, SquareNo, Coords):
        # print(SquareNo, Coords, "\n")
        CInformation = self.SquareToPixelArr(SquareNo, Coords)
        return [CInformation["insertY"], CInformation["insertX"]]

    # - editing functions

    def Change(self, src="PixelArr", xy=[], squareno=False, to=[]):
        # - src = PixelArr or Squares,
        # - Each edit has to edit the OGPixelArray
        x, y = xy[0], xy[1]
        # - Default format: [Y, X] when accessing
        if src == "Squares":
            SquareNo = squareno
            EditInfo = self.Square2PAEditor(SquareNo, [x, y])
            self.OGPixelArray[EditInfo[0], EditInfo[1]] = to

    def SQCHange(self, s=0, x=0, y=0, to=[255, 255, 255]):
        self.Change(src="Squares", xy=[x, y], squareno=s, to=to)

    # - all displaying and saving functions

    def ShowOG(self):
        OGPA = self.OGPixelArray
        im = Image.fromarray(OGPA)
        im.show()
        return self

    def SaveOG(self):
        OGPA = self.OGPixelArray
        im = Image.fromarray(OGPA)
        im.save("og.png")
        return self

    #// - takes a square, and creates a new image and fills it
    #// - with the area data provided by the following functions.
    def ShowSquareAlone(self, squareNo):
        fa = Image.new(mode="RGB", size=[16, 16])
        faa = np.array(fa)
        for pixelParcel in self.AllColourInSquare(squareNo, [255, 255, 255]):
            for x, y in self.Unpack(pixelParcel, "squareLoc"):
                faa[x, y] = [255, 255, 255]
                # print(x, y)
        faai = Image.fromarray(faa)
        faai.show()

    #// real checker
    def SquareEmpty(self, squareNo):
        square = self.Squares[squareNo]
        flat = np.hstack(np.hstack(square))
        if 255 in flat: return False
        else: return True

    #// checking from the meta
    def Empty(self, squareNo):
        # print(squareNo)
        return self.SquareMeta[squareNo]["empty"]

    def Empty_OldAlgorithm(self, squareNo):
        square = self.Squares[squareNo]
        for a in range(len(square)):
            for b in range(len(square[a])):
                if square[a][b][0] == 255:
                    return False
                    # isEmpty = 0
        return True

    def Contains(self, squareNo, highlight="no"):
        if squareNo == highlight: return 2
        return not self.Empty(squareNo)

    #// inserting all the "contains" info into an array that can be edited
    def ContainsArray(self, cast=bool):
        access = 0
        mapArray = [[]]
        for i in range(len(self.Squares)):
            ic = i + 1
            if access > (len(mapArray)-1): mapArray.append([])
            mapArray[access].append(cast(self.Contains(i)))
            if ic % squares == 0: access += 1
        # print("finished at {0}".format(i))
        return mapArray

    #// inserting all the "contains" info into an array that can be edited
    def ContainsArrayPoints(self):
        access = 0
        mapArray = [[]]
        for i in range(len(self.Squares)):
            ic = i + 1
            if access > (len(mapArray)-1): mapArray.append([])
            if self.Contains(i): mapArray[access].append(i)
            if ic % squares == 0: access += 1
        return mapArray

    def PrintSquareMap(self, highlight="no"):
        print("HIGHLIGHTING: {0}".format(highlight))
        for i in range(len(self.Squares)):
            ic = i + 1
            print(int(self.Contains(i, highlight=highlight)), end=" ")
            if ic % squares == 0: print("")

    def SquareXYCache(self):
        xtemp = 0
        ytemp = 0
        squarescache = [i for i in range(squaresx*squaresx)]

        for i in range(squaresx*squaresx): # 0 -> 63
            # want to increate X each run
            # when going Y up, X = 0, Y += 1
            squarescache[i] = [xtemp, ytemp]
            #             |
            #             |
            # x = ->, y = v

            if (xtemp+1) % squaresx == 0:
                xtemp =  0
                ytemp += 1
            else:
                xtemp += 1
        # pp(squarescache)
        return squarescache

    # - taking the square, and applying the direction sequence
    # - to where to apply each period
    def SquareBounds(self, squareNo, direction):

        # Creating an array converter, since
        # the dest-mathematical generality is just
        # too simple for the creation.
        SCHE = self.SquareXYCache()

        # print(str(squareNo) + " stopped by " + str(direction), end=" ")
        # X = math.ceil((squareNo+1)/squaresx)
        # Y = math.ceil((squareNo)%squaresx)+1
        # X = math.floor((squareNo)/squaresx)
        # Y = math.floor((squareNo)%squaresx)
        X, Y = SCHE[squareNo]
        X += 1
        Y += 1

        # right = the square's right-side
        TrueRight = xs - ((xSplit-1) * X)
        TrueDown  = ys - ((ySplit-1) * Y)
        Right     = xSplit * X
        Down      = ySplit * Y
        Left = Right - xSplit
        Up   = Down - ySplit

        # print("bounds for {0}, X,Y={1},{2} - ".format(squareNo, X, Y), Up, TrueDown, Left, TrueRight)
        if Left < 0: Left = 0
        if Up < 0: Up = 0

        DirectionConvert = {
            "O":   [0, 0, 0, 0],
            "L":   [0, 0, Left, 0],
            "D":   [0, TrueDown, 0, 0],
            "UR":  [Up, 0, 0, TrueRight],
            "UL":  [Up, 0, Left, 0],
            "BL":  [0, TrueDown, Left, 0],
            "BR":  [0, TrueDown, 0, TrueRight],
            "R":   [0, 0, 0, TrueRight],
            "U":   [Up, 0, 0, 0],
            }

        # print(" - " + str(DirectionConvert[direction]))
        return DirectionConvert[direction]

    def TestBounds(self):
        for i in range(squaresx*squaresx):
            self.SquareBounds(i, "U")
            if i%squaresx == 0: print()

    # - taking the square, and applying the direction sequence
    # - to where to apply each period
    def SquareCentre(self, squareNo):

        # X = math.ceil((squareNo+1)/squaresx)
        # Y = math.ceil((squareNo)%squaresx)+1
        SCHE = self.SquareXYCache()
        X, Y = SCHE[squareNo]
        X += 1
        Y += 1 # higher-ref, since it index starts at 0

        Right = xSplit * X
        Down = ySplit * Y
        Left = Right - xSplit
        Up = Down - ySplit
        if Left < 0: Left = 0
        if Up < 0: Up = 0

        Right -= round(xSplit/2)
        Down  -= round(ySplit/2)
        # print("centre position for {0}:".format(squareNo), Right, Down)
        # print()
        return [Right, Down]

    # - | finding the "best path" to
    # - | to follow in the writing algorithm,
    # - | writes all the path code and
    # - | does all the handling for passing
    # - | into the further creation realm
    def Path(self):
        # self.PrintSquareMap()
        # Compression = 5 # overall tests, for the "sloppiness"
        # Leveler     = 3 # the expected overhead of moves required
                        # to finally get to the ending area of
                        # total control.
        # Leveler     = (by/8)*(1.8+Complexity) if by >= 8 else 3
        # Leveler     = 10000 # fin
        # leveler-specific itemic values
        # if by == 8: Leveler = 1.5

        ConsiderableSquares = self.ContainsArrayPoints()
        ConsiderableSquares1d = [] # - a flat array
        for cs in ConsiderableSquares:
            for each in cs:
                ConsiderableSquares1d.append(each)

        # | - the parent, starting on a random
        # | - line then sub-bruting till it
        # | - finds a dead end and goes till
        # | - it covers every square and resolves
        # | - the path
        # Brute, OverallBrute = round(len(ConsiderableSquares1d)*Leveler), 10000*Compression # Brute=overall attempts, Overall=individual attempts at pathfinding
        # Brute, OverallBrute = 2, 1
        CoveredSquares = []
        AllCovered = False

        # FinishPath = 100000
        # Brute = FinishPath

        Movement, CurrentPlace = "O", random.choice(ConsiderableSquares1d)
        DirectionSequence = [["O", CurrentPlace]] # the operations of movement dir's
        CoveredArea = [CurrentPlace] # the covered area through pathfinding
        # recording the direction and meta info
        StepInformation, Original, LinearDirections = [], CurrentPlace, [CurrentPlace]

        while True:
            AllCovered = True
            for ConsLoc in ConsiderableSquares1d:
                if ConsLoc not in CoveredArea:
                    AllCovered = False
                    break

            if AllCovered == True:
                self.DirectionSequenceDone = True
                self.DirectionSequence = DirectionSequence
                # self.PrintSquareMap(highlight=Original)
                # print(self.DirectionSequence)
                # print("All covered!")
                return DirectionSequence

            Movements = {
                # "UU":  CurrentPlace-squaresx-squaresx,
                "U":   CurrentPlace-squaresx,
                "R":   (CurrentPlace+1)           if ((CurrentPlace+1)%squaresx!=0) else -1, #rc
                "UR":  CurrentPlace-squaresx+1    if ((CurrentPlace+1)%squaresx!=0) else -1, #rc
                "UL":  CurrentPlace-squaresx-1    if ((CurrentPlace)%squaresx!=0) else -1, #lc
                "BL":  CurrentPlace+squaresx-1    if ((CurrentPlace)%squaresx!=0) else -1, #lc
                "BR":  CurrentPlace+squaresx+1    if ((CurrentPlace+1)%squaresx!=0) else -1, #rc
                "L":   CurrentPlace-1             if ((CurrentPlace)%squaresx!=0) else -1, #lc
                "D":   CurrentPlace+squaresx}

            # Random sort.
            # for Movement, Place in Movements.items():
            # Movements = random.shuffle(Movements)
            # for Movement, Place in Movements.items():
            for Movement, Place in sorted(Movements.items(), key=lambda x: random.random()):
                if Place in ConsiderableSquares1d:
                    DirectionSequence.append([Movement, Place])
                    CoveredArea.append(Place)

                    # print(Movement, Place)

                    # StepInformation.append({
                    #     "before": CurrentPlace,
                    #     "after": Place,
                    #     "movements": Movements,
                    #     "newPlace": Place,
                    #     "beforePlace": CurrentPlace,
                    #     "direction": Movement
                    # })

                    CurrentPlace = Place

                    # placing into direction
                    LinearDirections.append("{0}:{1}".format(Movement, Place))
                    break
            else:
                print("No movement possible...")

    # - takes squareNo and applies where it will be after
    # - direction is applied
    def DirectionApplySquareNo(self, SquareNo, Direction):
        CurrentPlace = SquareNo
        Movements = {
            "O":   CurrentPlace,
            "U":   CurrentPlace-squaresx,
            "R":   (CurrentPlace+1)           if ((CurrentPlace+1)%squaresx!=0) else -1, #rc
            "UR":  CurrentPlace-squaresx+1    if ((CurrentPlace+1)%squaresx!=0) else -1, #rc
            "UL":  CurrentPlace-squaresx-1    if ((CurrentPlace)%squaresx!=0) else -1, #lc
            "BL":  CurrentPlace+squaresx-1    if ((CurrentPlace)%squaresx!=0) else -1, #lc
            "BR":  CurrentPlace+squaresx+1    if ((CurrentPlace+1)%squaresx!=0) else -1, #rc
            "L":   CurrentPlace-1             if ((CurrentPlace)%squaresx!=0) else -1, #lc
            "D":   CurrentPlace+squaresx
        }
        return Movements[Direction]

    # - | formating the self.DirectionSequence into
    # - | a fair format.
    def PathFormat(self):
        if self.DirectionSequenceDone == False:
            print("Path has not been established...")
            return False

        AngleConvert = {
            "O": 0,
            "U": 0,
            "D": 180,
            "L": 270,
            "R": 90,
            "UR": 45,
            "UL": 315,
            "BL": 225,
            "BR": 135}
        i, Series = 0, []

        for DirectionSpace in self.DirectionSequence:
            if i >= 1:
                Direction, SquareNo, PastDirection = DirectionSpace[0], DirectionSpace[1], self.DirectionSequence[i-1]
                # THIS IS THE MOST IMPORTANT KEYWORD THAT IS IN THIS ENTIRE CODEBASE VVV
                # we're going from PastDirection, to the SquareNo in Direction
                OldSquareNo = PastDirection[1]

                Angle = AngleConvert[Direction]
                Bounds = self.SquareBounds(SquareNo, Direction)
                X, Y   = self.SquareCentre(OldSquareNo)

                PackedFormatOBJ = {
                    "oldSquareX": X,
                    "oldSquareY": Y,
                    "stoppingBounds": Bounds,
                    "newSquareNo": SquareNo,
                    "oldSquareNo": OldSquareNo,
                    "direction": Direction,
                    "pastDirection": PastDirection,
                    "angle": Angle,
                }
                Series.append([X, Y, Angle, Bounds, OldSquareNo, Direction])
            i += 1

        return Series

    # - iter functions to make it easier to iterate
    # | - pixelParcel (parcel) is a pack provided by the cache
    # | - information series.

    def Unpack(self, parcel, pref="insert"):
        return zip([parcel["{0}X".format(pref)]], [parcel["{0}Y".format(pref)]])

    # - bitmap opening functions, to convert the array into an
    # - original piece

    def Open(self):
        PixelArr = self.OGPixelArray
        self.BitMap = Image.fromarray(PixelArr)
        # self.BitMapDraw = ImageDraw.Draw(self.BitMap)
        return ImageDraw.Draw(self.BitMap)

    def Draw(self):
        return ImageDraw.Draw(self.BitMap)

    def Close(self, removeAA=False):
        BitMap = self.BitMap
        self.OGPixelArray = np.array(BitMap)
        if removeAA: self.AARemove()

ext: py

reason: import numpy


text:
import numpy as np
import torch
import warnings

class batchless_VanillaLSTM_pytorch(object):
    """ Vanilla LSTM implementation using pytorch """

    def __init__(self, num_layers=2, cells_per_layer=50, dropout=0.5, seed=None, stateful=True, lag=5):
        """
        Initialise and build the model
        """
        self.num_layers = num_layers
        self.cells_per_layer = cells_per_layer
        self.dropout = dropout
        self.seed = seed

        self.stateful = stateful
        self.lag = lag

        if seed != None:
            np.random.seed(seed)
            torch.manual_seed(seed)

    def build(self, sequence, debug=False):
        """
        Build model
        """
        # Sequence either list of lists or a list.
        if sequence.ndim != 1:
            self.features = len(sequence[0])
        else:
            self.features = 1

        # Reshape and convert to torch tensor
        self.sequence = torch.FloatTensor(sequence).view(-1, 1, self.features)

        self.model = pytorch_LSTM(input_dim=self.features, hidden_dim=self.cells_per_layer, batch_size=1, output_dim=self.features, num_layers=self.num_layers, dropout=self.dropout)
        if self.features != 1:
            self.loss_fn = torch.nn.CrossEntropyLoss(reduction='sum')
            self.optimizer = torch.optim.Adam(self.model.parameters())
        else:
            self.loss_fn = torch.nn.MSELoss(size_average=False)
            self.optimizer = torch.optim.Adam(self.model.parameters())
        self.model.init_weights(self.model)

    def construct_training_index(self, debug=False):
        """
        Construct training index (compatible with model) from sequence of vectors of dimension d,
        """
        n = len(self.sequence)
        self.index = []
        if self.stateful:
            # Create groups
            self.num_augs = min(self.lag, n - self.lag)
            for el in range(self.num_augs):
                self.index.append(np.arange(el, n - self.lag, self.lag))
        else:
            self.num_augs = 1
            self.index = np.arange(0, n - self.lag, 1)

    def train(self, patience=10, max_epoch=1000, acceptable_loss=np.inf, batch_size=1, weight_restarts=False, debug=False):
        """
        Train the model on the constructed training data
        """
        if batch_size != 1:
            warnings.warn('batch_size must equal 1, setting to 1.')
        ########################################################################
        # Weight restarts
        ########################################################################
        states = self.model.initialise_states()
        if weight_restarts:
            weight_restarts = 10
            store_weights = [0]*weight_restarts
            initial_loss = [0]*weight_restarts
            for i in range(weight_restarts):
                # reset cell state
                states = self.model.initialise_states()

                y_pred, states = self.model(self.sequence[0:self.lag, :, :], (states[0].detach(), states[1].detach()))

                # calculate loss
                if self.features == 1:
                    self.loss = self.loss_fn(y_pred.view(-1, 1), self.sequence[self.lag, :, :])
                else:
                    target = self.sequence[self.lag, :, :].max(-1)[1]
                    self.loss = self.loss_fn(y_pred.reshape(1,-1), target)

                initial_loss[i] = self.loss.data
                store_weights[i] = self.model.state_dict()

                # Re initialise weights
                self.model.init_weights(self.model)
            m = np.argmin(initial_loss)
            self.model.load_state_dict(store_weights[int(m)])
            del store_weights

        ########################################################################
        # Train
        ########################################################################
        vec_loss = np.zeros(max_epoch)
        min_loss = np.inf
        min_loss_ind = np.inf
        losses = [0]*self.num_augs
        if self.stateful: # no shuffle and reset state manually
            for iter in range(max_epoch):
                rint = np.random.permutation(self.num_augs) # shuffle groups
                for r in rint: # run through groups
                    # reset cell state
                    states = self.model.initialise_states()

                    loss_sum = 0
                    for i in self.index[r]: # run through group
                        # Forward pass
                        y_pred, states = self.model(self.sequence[i:i+self.lag, :, :], (states[0].detach(), states[1].detach()))

                        # calculate loss
                        if self.features == 1:
                            self.loss = self.loss_fn(y_pred.view(-1, 1), self.sequence[i+self.lag, :, :])
                        else:
                            target = self.sequence[i+self.lag, :, :].max(-1)[1]
                            self.loss = self.loss_fn(y_pred.reshape(1,-1), target)

                        loss_sum += (float(self.loss.data))**2

                        # Backward pass
                        self.loss.backward(retain_graph=True)

                        # Update parameters
                        self.optimizer.step()
                        # clear gradients
                        self.model.zero_grad()

                    losses[r] = loss_sum/len(self.index[r])
                vec_loss[iter] = np.mean(losses)

                if vec_loss[iter] >= min_loss:
                    if iter - min_loss_ind >= patience and min_loss<acceptable_loss:
                        break
                else:
                    min_loss = vec_loss[iter]
                    old_weights = self.model.state_dict()
                    min_loss_ind = iter

        else: # shuffle in fit
            for iter in range(max_epoch):
                loss_sum = 0
                for i in np.random.permutation(len(self.index)):
                    states = self.model.initialise_states()

                    # Forward pass
                    y_pred, states = self.model.forward(self.sequence[i:i+self.lag, :, :], (states[0].detach(), states[1].detach()))

                    # calculate loss
                    if self.features == 1:
                        self.loss = self.loss_fn(y_pred.view(-1, 1), self.sequence[i+self.lag, :, :])
                    else:
                        target = self.sequence[i+self.lag, :, :].max(-1)[1]
                        self.loss = self.loss_fn(y_pred.reshape(1,-1), target)

                    loss_sum += (float(self.loss.data))**2

                    # Backward pass
                    self.loss.backward()

                    # Update parameters
                    self.optimizer.step()
                    # clear gradients
                    self.model.zero_grad()

                vec_loss[iter] = loss_sum/len(self.index)

                if vec_loss[iter] >= min_loss:
                    if iter - min_loss_ind >= patience and min_loss < acceptable_loss:
                        break
                else:
                    min_loss = vec_loss[iter]
                    old_weights = self.model.state_dict()
                    min_loss_ind = iter

        self.model.load_state_dict(old_weights)
        self.epoch = iter+1
        self.loss = vec_loss[0:iter+1]

    def forecast(self, k, randomize=False, debug=False):
        """
        Make k step forecast into the future.
        """
        self.model.eval()
        prediction = self.sequence.clone()

        # Recursively make k one-step forecasts
        for ind in range(len(self.sequence), len(self.sequence) + k):
            # Build data to feed into model
            if self.stateful:
                index = np.arange(ind%self.lag, ind, self.lag)
            else:
                index = [ind - self.lag]

            # Feed through model
            states = self.model.initialise_states()
            for i in index:
                p, states = self.model.forward(prediction[i:i+self.lag, :, :], (states[0].detach(), states[1].detach()))

            # Convert output
            if self.features != 1:
                softmax = torch.nn.Softmax(dim=-1)
                p = softmax(p).tolist()
                p = np.array(p)
                p /= p.sum()
                if randomize:
                    idx = np.random.choice(range(self.features), p=(p.ravel()))
                else:
                    idx = np.argmax(list(p), axis = 0)

                # Add forecast result to appropriate vectors.
                pred = torch.zeros([1, 1, self.features])
                pred[0, 0, idx] = 1
            else:
                pred = torch.zeros([1, 1, 1])
                pred[0, 0, 0] = p

            prediction = torch.cat([prediction, pred], dim=0)

        if self.features != 1:
            return prediction.view(-1, self.features).tolist()
        else:
            return prediction.view(-1).detach()

################################################################################
################################################################################
################################################################################

class pytorch_LSTM(torch.nn.Module):
    """
    Class to define LSTM model using pytorch.
    """
    def __init__(self, input_dim, hidden_dim, batch_size, output_dim, num_layers, dropout):
        super(pytorch_LSTM, self).__init__()
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        self.batch_size = batch_size
        self.num_layers = num_layers
        self.output_dim = output_dim

        self.states = (0, 0)

        self.lstm = torch.nn.LSTM(self.input_dim, self.hidden_dim, self.num_layers, dropout=dropout)

        self.dropout = torch.nn.Dropout(p=dropout)
        self.final = torch.nn.Linear(self.hidden_dim, self.output_dim)

    def init_weights(self, m):
        for n, p in m.named_parameters():
            if 'weight_ih' in n:
                for ih in p.chunk(4,0):
                    torch.nn.init.xavier_uniform_(ih)
            elif 'weight_hh' in n:
                for hh in p.chunk(4,0):
                    torch.nn.init.orthogonal_(hh)
            elif 'bias_ih' in n:
                torch.nn.init.zeros_(p)
            elif 'bias_hh' in n:
                torch.nn.init.zeros_(p)
            elif 'final.weight' in n:
                torch.nn.init.xavier_uniform_(p)
            elif 'final.bias' in n:
                torch.nn.init.zeros_(p)

    def initialise_states(self):
        """
        Reset both cell state and hidden state
        """
        return (torch.zeros(self.num_layers, self.batch_size, self.hidden_dim), torch.zeros(self.num_layers, self.batch_size, self.hidden_dim))

    def reset_hidden(self, states):
        """
        Reset hidden state
        """
        return (torch.zeros(self.num_layers, self.batch_size, self.hidden_dim), states[1])

    def reset_cell(self):
        """
        Reset cell state
        """
        return (states[0], torch.zeros(self.num_layers, self.batch_size, self.hidden_dim))

    def forward(self, input, states):
        """
        Define forward pass through LSTM
        """
        # pass through LSTM layers
        lstm_out, states = self.lstm(input.view(len(input), self.batch_size, -1), states)
        # pass through linear layer
        y_pred = self.final(self.dropout(lstm_out[-1].view(self.batch_size, -1)))
        return y_pred.view(-1), states

max_stars_repo_name: nla-group/LSTMs

max_issues_repo_name: nla-group/ABBA-LSTM

max_forks_repo_name: nla-group/ABBA-LSTMs


text:
import rsa.utils.pytorch_utils as ptu

import torch
import torch.nn as nn
import torch.nn.functional as F

import copy
import os
import numpy as np

class MCTD3:
    def __init__(self, params):

        self.max_action = params['max_action']
        self.discount = params['discount']
        self.tau = params['tau']
        self.policy_noise = params['policy_noise']
        self.noise_clip = params['noise_clip']
        self.policy_freq = params['policy_freq']
        self.batch_size = params['batch_size']
        self.batch_size_demonstrator = params['batch_size_demonstrator']
        self.bc_weight = params['bc_weight']
        self.bc_decay = params['bc_decay']
        self.do_drtg_bonus = params['do_drtg_bonus']
        self.plot_drtg_maxes = params['plot_drtg_maxes']
        self.do_bc_loss = params['do_bc_loss']
        self.do_q_filter = params['do_q_filter']

        self.total_it = 0
        self.running_risk = 1

        self.actor = Actor(params['d_obs'], params['d_act'], self.max_action).to(ptu.TORCH_DEVICE)
        self.actor_target = copy.deepcopy(self.actor)
        self.actor_optimizer = torch.optim.Adam(self.actor.parameters(), lr=params['lr_actor'])

        self.critic = Critic(params['d_obs'], params['d_act']).to(ptu.TORCH_DEVICE)
        self.critic_target = copy.deepcopy(self.critic)
        self.critic_optimizer = torch.optim.Adam(self.critic.parameters(), lr=params['lr_critic'])

        self.drtg_buffer = set()
        self.bellman_buffer = set()

    def select_action(self, state):
        state = torch.FloatTensor(state.reshape(1, -1)).to(ptu.TORCH_DEVICE)
        return self.actor(state).cpu().data.numpy().flatten()

    def update(self, replay_buffer):
        # Sample from replay buffer
        out_dict = replay_buffer.sample(self.batch_size)
        obs, action, next_obs, reward, mask, drtg, expert = out_dict['obs'], out_dict['act'], \
                                                            out_dict['next_obs'], out_dict['rew'], \
                                                            out_dict['mask'], out_dict['drtg'], \
                                                            out_dict['expert']

        obs, action, next_obs, reward, mask, drtg, expert = \
            ptu.torchify(obs, action, next_obs, reward, mask, drtg, expert)

        info = {}

        ###############################################################################################

        with torch.no_grad():
            # Select action according to policy and add clipped noise
            noise = (torch.randn_like(action) * self.policy_noise)\
                .clamp(-self.noise_clip, self.noise_clip)

            next_action = (self.actor_target(next_obs) + noise)\
                .clamp(-self.max_action, self.max_action)

            # Compute the target Q value
            target_Q1, target_Q2 = self.critic_target(next_obs, next_action)
            target_Q = torch.min(target_Q1, target_Q2).squeeze()
            target_Q = reward + mask * self.discount * target_Q
            target_Q = target_Q.squeeze()

            if self.do_drtg_bonus:
                if self.plot_drtg_maxes:
                    # print(len(self.drtg_buffer), len(self.bellman_buffer))
                    obs_np = ptu.numpify(obs)
                    nqv_np = ptu.numpify(target_Q)
                    drtg_np = ptu.numpify(drtg)
                    drtg_bigger = drtg_np > nqv_np
                    # print(nqv_np[np.logical_not(drtg_bigger)].shape, drtg_np[drtg_bigger])
                    self.drtg_buffer.update([tuple(x) for x in obs_np[drtg_bigger]])
                    self.bellman_buffer.update([tuple(x) for x in obs_np[np.logical_not(drtg_bigger)]])
                    print(len(self.drtg_buffer), len(self.bellman_buffer))

                target_Q = torch.max(target_Q, drtg)

        # print('--------------------------------------------------')
        # print(drtg)
        # print(reward)
        # print(target_Q)

        # Get current Q estimates
        current_Q1, current_Q2 = self.critic(obs, action)
        current_Q1, current_Q2 = current_Q1.squeeze(), current_Q2.squeeze()

        # Compute critic loss
        critic_loss = F.mse_loss(current_Q1, target_Q) + F.mse_loss(current_Q2, target_Q)

        # Optimize the critic
        self.critic_optimizer.zero_grad()
        critic_loss.backward()
        self.critic_optimizer.step()
        info['critic_loss'] = critic_loss.item()
        info['Q1'] = current_Q1.mean().item()
        info['Q2'] = current_Q2.mean().item()

        # Delayed policy updates
        if self.total_it % self.policy_freq == 0:

            # Compute actor losse
            actor_q_loss = -self.critic.Q1(obs, self.actor(obs)).mean()

            # Behavior cloning auxiliary loss, inspired by DDPGfD paper
            if self.do_bc_loss:
                # Sample expert actions from the replay buffer
                out_dict = replay_buffer.sample_positive(self.batch_size_demonstrator, 'expert')
                obs, action = out_dict['obs'], out_dict['act']
                obs, action = ptu.torchify(obs, action)

                # Calculate loss as negative log prob of actions
                act_hat = self.actor(obs)
                losses = F.mse_loss(act_hat, action, reduction='none')

                # Optional Q filter
                if self.do_q_filter:
                    with torch.no_grad():
                        q_agent = self.critic.Q1(obs, act_hat)
                        q_expert = self.critic.Q1(obs, action)
                        q_filter = torch.gt(q_expert, q_agent).float()

                    if torch.sum(q_filter) > 0:
                        bc_loss = torch.sum(losses * q_filter) / torch.sum(q_filter)
                    else:
                        bc_loss = actor_q_loss * 0
                else:
                    bc_loss = torch.mean(losses)

            else:
                bc_loss = actor_q_loss * 0

            lambda_bc = self.bc_decay ** self.total_it * self.bc_weight
            actor_loss = actor_q_loss + lambda_bc * bc_loss

            # Optimize the actor
            self.actor_optimizer.zero_grad()
            actor_loss.backward()
            self.actor_optimizer.step()
            info['actor_loss'] = actor_loss.item()
            info['actor_q_loss'] = actor_q_loss.item()
            info['actor_bc_loss'] = bc_loss.item()

            # Update the frozen target models
            ptu.soft_update(self.critic, self.critic_target, 1 - self.tau)
            ptu.soft_update(self.actor, self.actor_target, 1 - self.tau)

        ############################################################################################################

        # Compute targets using bellman backup and target function
        # with torch.no_grad():
        #     next_state_action, next_state_log_pi, _ = self.policy.sample(next_obs)
        #     qf_list_next_target = self.critic_target(next_obs, next_state_action)
        #     min_qf_next_target = torch.min(torch.cat(qf_list_next_target, dim=1), dim=1)[0] \
        #                          - self.alpha * next_state_log_pi.squeeze()
        #     min_qf_next_target = min_qf_next_target.squeeze()
        #     next_q_value = reward + mask * self.discount * min_qf_next_target
        #
        #     if self.do_drtg_bonus:
        #         if self.plot_drtg_maxes:
        #             # print(len(self.drtg_buffer), len(self.bellman_buffer))
        #             obs_np = ptu.numpify(obs)
        #             nqv_np = ptu.numpify(next_q_value)
        #             drtg_np = ptu.numpify(drtg)
        #             drtg_bigger = drtg_np > nqv_np
        #             # print(nqv_np[np.logical_not(drtg_bigger)].shape, drtg_np[drtg_bigger])
        #             self.drtg_buffer.update([tuple(x) for x in obs_np[drtg_bigger]])
        #             self.bellman_buffer.update([tuple(x) for x in obs_np[np.logical_not(drtg_bigger)]])
        #
        #         next_q_value = torch.max(next_q_value, drtg)
        #
        #
        #
        # # Compute Q losses
        # # Two Q-functions to mitigate positive bias in the policy improvement step
        # qf_list = self.critic(obs, action)
        # # JQ = 𝔼(st,at)~D[0.5(Q1(st,at) - r(st,at) - γ(𝔼st+1~p[V(st+1)]))^2]
        # qf_losses = [
        #     F.mse_loss(qf.squeeze(), next_q_value)
        #     for qf in qf_list
        # ]
        # # JQ = 𝔼(st,at)~D[0.5(Q1(st,at) - r(st,at) - γ(𝔼st+1~p[V(st+1)]))^2]
        # qf_loss = sum(qf_losses)
        #
        # # Q function backward pass
        # self.critic_optim.zero_grad()
        # qf_loss.backward()
        # self.critic_optim.step()
        #
        # # Sample from policy, compute minimum Q value of sampled action
        # pi, log_pi, _ = self.policy.sample(obs)
        # qf_list_pi = self.critic(obs, pi)
        # min_qf_pi = torch.min(torch.cat(qf_list_pi, dim=1), dim=1)[0]
        #
        # # Calculate policy loss
        # # Jπ = 𝔼st∼D,εt∼N[α * logπ(f(εt;st)|st) − Q(st,f(εt;st))]
        # policy_q_loss = ((self.alpha * log_pi) - min_qf_pi).mean()
        #
        # # Behavior cloning auxiliary loss, inspired by DDPGfD paper
        # if self.do_bc_loss:
        #     # Sample expert actions from the replay buffer
        #     out_dict = replay_buffer.sample_positive(self.batch_size_demonstrator, 'expert')
        #     obs, action = out_dict['obs'], out_dict['act']
        #     obs, action = ptu.torchify(obs, action)
        #
        #     # Calculate loss as negative log prob of actions
        #     mean, log_std = self.policy(obs)
        #     dist = torch.distributions.Normal(mean, torch.exp(log_std))
        #     losses = -torch.mean(dist.log_prob(action), dim=1)
        #
        #     # Optional Q filter
        #     if self.do_q_filter:
        #         with torch.no_grad():
        #             q_agent = self.critic.Q1(obs, dist.sample())
        #             q_expert = self.critic.Q1(obs, action)
        #             q_filter = torch.gt(q_expert, q_agent).float()
        #
        #         if torch.sum(q_filter) > 0:
        #             bc_loss = torch.sum(losses * q_filter) / torch.sum(q_filter)
        #         else:
        #             bc_loss = policy_q_loss * 0
        #     else:
        #         bc_loss = torch.mean(losses)
        #
        # else:
        #     bc_loss = policy_q_loss * 0
        #
        # # Lambda from DDPGfD paper I believe
        # lambda_bc = self.bc_decay ** self.total_it * self.bc_weight
        # policy_loss = policy_q_loss + lambda_bc * bc_loss
        #
        # # Policy backward pass
        # self.policy_optim.zero_grad()
        # policy_loss.backward()
        # self.policy_optim.step()
        #
        # # Automatic entropy tuning
        # if self.automatic_entropy_tuning:
        #     alpha_loss = -(self.log_alpha * (log_pi + self.target_entropy).detach()).mean()
        #
        #     self.alpha_optim.zero_grad()
        #     alpha_loss.backward()
        #     self.alpha_optim.step()
        #
        #     self.alpha = self.log_alpha.exp()
        #     alpha_tlogs = self.alpha.clone()  # For TensorboardX logs
        # else:
        #     alpha_loss = torch.tensor(0.).to(ptu.TORCH_DEVICE)
        #     alpha_tlogs = torch.tensor(self.alpha)  # For TensorboardX logs
        #
        # # if init:
        # #     ptu.hard_update(self.critic, self.critic_target)
        # if self.total_it % self.target_update_interval == 0:
        #     ptu.soft_update(self.critic, self.critic_target, 1 - self.tau)

        self.total_it += 1
        return info

    def save(self, folder):
        os.makedirs(folder, exist_ok=True)

        torch.save(self.critic.state_dict(), os.path.join(folder, "critic.pth"))
        torch.save(self.critic_optimizer.state_dict(), os.path.join(folder, "critic_optimizer.pth"))

        torch.save(self.actor.state_dict(), os.path.join(folder, "actor.pth"))
        torch.save(self.actor_optimizer.state_dict(), os.path.join(folder, "actor_optimizer.pth"))

    def load(self, folder):
        self.critic.load_state_dict(torch.load(os.path.join(folder, "critic.pth")))
        self.critic_optimizer.load_state_dict(
            torch.load(os.path.join(folder, "critic_optimizer.pth")))
        self.critic_target = copy.deepcopy(self.critic)

        self.actor.load_state_dict(torch.load(os.path.join(folder, "actor.pth")))
        self.actor_optimizer.load_state_dict(
            torch.load(os.path.join(folder, "actor_optimizer.pth")))

class Actor(nn.Module):
    def __init__(self, state_dim, action_dim, max_action):
        super(Actor, self).__init__()

        self.l1 = nn.Linear(state_dim[0], 256)
        self.l2 = nn.Linear(256, 256)
        self.l3 = nn.Linear(256, action_dim[0])

        self.max_action = max_action

    def forward(self, state):
        a = F.relu(self.l1(state))
        a = F.relu(self.l2(a))
        return self.max_action * torch.tanh(self.l3(a))

class Critic(nn.Module):
    def __init__(self, state_dim, action_dim):
        super(Critic, self).__init__()

        # Q1 architecture
        self.l1 = nn.Linear(state_dim[0] + action_dim[0], 256)
        self.l2 = nn.Linear(256, 256)
        self.l3 = nn.Linear(256, 1)

        # Q2 architecture
        self.l4 = nn.Linear(state_dim[0] + action_dim[0], 256)
        self.l5 = nn.Linear(256, 256)
        self.l6 = nn.Linear(256, 1)

    def forward(self, state, action):
        sa = torch.cat([state, action], 1)

        q1 = F.relu(self.l1(sa))
        q1 = F.relu(self.l2(q1))
        q1 = self.l3(q1)

        q2 = F.relu(self.l4(sa))
        q2 = F.relu(self.l5(q2))
        q2 = self.l6(q2)
        return q1, q2

    def Q1(self, state, action):
        sa = torch.cat([state, action], 1)

        q1 = F.relu(self.l1(sa))
        q1 = F.relu(self.l2(q1))
        q1 = self.l3(q1)
        return q1

max_stars_repo_name: zaynahjaved/AWAC

max_issues_repo_name: zaynahjaved/AWAC

max_forks_repo_name: zaynahjaved/AWAC


text:
import numpy as np
import matplotlib.pyplot as plt
from sklearn.discriminant_analysis import LinearDiscriminantAnalysis
import pandas as pd
from sklearn.decomposition import PCA
from sklearn import svm

sum = np.zeros((1,9))
for i in range(5):
    data = pd.read_csv("datasets/lifestyle.csv")
    data = data.sample(frac=1).reset_index(drop=True)
    x_columns = ["age", "bp", "grav", "htn", "dm", "cad", "appet", "pe", "ane"]
    x = data[x_columns]
    y = data['class']

    X = np.array(x)
    y = np.array(y)
    clf = svm.SVC(kernel='linear', probability=True)
    clf.fit(x, y)

    print(clf.coef_.shape)
    sum += clf.coef_

    plt.show()
    #print(lda.predict([[4, 1, 6]]))
avg = sum/5
print(avg)


text:
import os
import re
import json
import numpy as np

SKETCH_RAW_FOLDER = os.path.join(os.path.realpath('..'), 'data/raw_data')
SKETCH_HISTORY_FOLDER = os.path.join(os.path.realpath('..'), 'data/sketch_history')
SKETCH_GRAPH_FOLDER = os.path.join(os.path.realpath('..'), 'data/sketch_graph')
CURVES_FOLDER = os.path.join(os.path.realpath('..'), 'data/curves')

def get_file_path(root_folder, user_id, system_id, model_id, extension='json'):
    correct_file = None

    regex_string = r"{:02}-{}-{}".format(user_id, system_id, model_id)

    extension = '.' + extension

    # r=root, d=directories, f = files
    for r, d, f in os.walk(root_folder):
        for file in f:
            if extension in file:
                if re.search(regex_string, file) is not None:
                    correct_file = file

    if correct_file is None:
        print("file not found")
        return

    return os.path.join(root_folder, correct_file)

def try_load_data(full_path):
    if full_path is None:
        return None

    # Reading the json as a dict
    with open(full_path) as json_data:
        data = json.load(json_data)

    return data

def load_curves_data(path):
    if path is None:
        return None
    full_path = os.path.join(CURVES_FOLDER, path)

    polylines = []
    current_polyline = np.empty((0,3))
    target_vertex_count = -1

    with open(full_path) as f:
        for line in f:
            els = line.split()
            if els[0] == 'v':
                # Start a new polyline
                current_polyline = np.empty((0,3))
                target_vertex_count = int(els[1])

            else:
                if len(current_polyline) >= target_vertex_count:
                    print("error reading file, vertex count and actual vertex data is inconsistent.")
                    break
                # Parse vertex 3D position and add it to current polyline
                pos = np.array([float(els[i]) for i in range(3)])
                current_polyline = np.row_stack([current_polyline, pos])

                if len(current_polyline) == target_vertex_count:
                    polylines.append(current_polyline)

    return polylines

# Support different encoding of vec3, either a list or an object with the coordinates

def parse_vec3(vec3_object):
    if isinstance(vec3_object, list):
        return vec3_object
    else:
        return [vec3_object['x'], vec3_object['y'], vec3_object['z']]

def parse_ctrl_pts(ctrl_pts):
    return [[parse_vec3(p) for p in pts] for pts in ctrl_pts]

max_stars_repo_name: V-Sekai/cassie-data

max_issues_repo_name: V-Sekai/cassie-data

max_forks_repo_name: V-Sekai/cassie-data


text:
#!/usr/bin/env python

import sys

import numpy as np
np.set_printoptions(threshold=5)

#from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor
from frovedis.mllib.tree import DecisionTreeClassifier, DecisionTreeRegressor

# initializing the Frovedis server
argvs = sys.argv
argc = len(argvs)
if (argc < 2):
    print ('Please give frovedis_server calling command as the first argument \n(e.g. "mpirun -np 2 /opt/nec/frovedis/ve/bin/frovedis_server")')
    quit()

from frovedis.exrpc.server import FrovedisServer
FrovedisServer.initialize(argvs[1])

# classification data
from sklearn.datasets import load_breast_cancer
mat, lbl = load_breast_cancer(return_X_y=True)

# fitting input matrix and label on DecisionTree Classifier object
dtc = DecisionTreeClassifier(criterion='gini', max_depth=5)
dtc.fit(mat,lbl)
#dtc.debug_print()

# predicting on train model
print("predicting on DecisionTree classifier model: ")
print(dtc.predict(mat))
print("predicting probability on DecisionTree classifier model: ")
print (dtc.predict_proba(mat))
print("prediction accuracy: %.4f" % (dtc.score(mat, lbl)))

# regression data
from sklearn.datasets import load_boston
mat, lbl = load_boston(return_X_y=True)

# fitting input matrix and label on DecisionTree Regressor object
dtr = DecisionTreeRegressor(criterion='mse', max_depth=5)
dtr.fit(mat, lbl)
#dtr.debug_print()

# predicting on train model
print("predicting on  DecisionTree Regressor model: ")
print(dtr.predict(mat))
print("prediction score: %.4f" % (dtr.score(mat, lbl)))

#clean-up
#dtc.release()
#dtr.release()
FrovedisServer.shut_down()

max_stars_repo_name: XpressAI/frovedis

max_issues_repo_name: XpressAI/frovedis

max_forks_repo_name: XpressAI/frovedis


text:
import copy
import glob
import os
import time
from collections import deque

# TODO remove warning for tensorflow and for pytorch nonzero

import gym
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

from a2c_ppo_acktr import algo, utils
from a2c_ppo_acktr.envs import make_vec_envs
from a2c_ppo_acktr.model import Policy
from a2c_ppo_acktr.storage import RolloutStorage
from evaluation import evaluate

import argparse

def get_args():
    parser = argparse.ArgumentParser(description='RL')
    parser.add_argument(
        '--algo', default='a2c', help='algorithm to use: a2c | ppo ')
    parser.add_argument(
        '--lr', type=float, default=1e-4, help='learning rate (default: 1e-4)')
    parser.add_argument(
        '--eps',
        type=float,
        default=1e-5,
        help='RMSprop optimizer epsilon (default: 1e-5)')
    parser.add_argument(
        '--alpha',
        type=float,
        default=0.99,
        help='RMSprop optimizer alpha (default: 0.99)')
    parser.add_argument(
        '--gamma',
        type=float,
        default=0.99,
        help='discount factor for rewards (default: 0.99)')
    parser.add_argument(
        '--use-gae',
        action='store_true',
        default=True,
        help='use generalized advantage estimation')
    parser.add_argument(
        '--gae-lambda',
        type=float,
        default=0.95,
        help='gae lambda parameter (default: 0.95)')
    parser.add_argument(
        '--entropy-coef',
        type=float,
        default=0.01,
        help='entropy term coefficient (default: 0.01)')
    parser.add_argument(
        '--value-loss-coef',
        type=float,
        default=0.5,
        help='value loss coefficient (default: 0.5)')
    parser.add_argument(
        '--max-grad-norm',
        type=float,
        default=0.5,
        help='max norm of gradients (default: 0.5)')
    parser.add_argument(
        '--seed', type=int, default=1, help='random seed (default: 1)')
    parser.add_argument(
        '--cuda-deterministic',
        action='store_true',
        default=False,
        help="sets flags for determinism when using CUDA (potentially slow!)")
    parser.add_argument(
        '--training',
        action='store_true',
        default=False,
        help="Whether or not to train the model")
    parser.add_argument(
        '--save-experimental-data',
        action='store_true',
        default=False,
        help="Whether or not to save experimental data for analysis")
    parser.add_argument(
        '--num-processes',
        type=int,
        default=32,
        help='how many training CPU processes to use (default: 16)')
    parser.add_argument(
        '--num-steps',
        type=int,
        default=40,
        help='number of forward steps in A2C (default: 5)')
    parser.add_argument(
        '--shift-len',
        type=int,
        default=10,
        help='number of steps to shift the rollout forward each update' +
             ' (default: 10)')
    parser.add_argument(
        '--ppo-epoch',
        type=int,
        default=4,
        help='number of ppo epochs (default: 4)')
    parser.add_argument(
        '--num-mini-batch',
        type=int,
        default=32,
        help='number of batches for ppo (default: 32)')
    parser.add_argument(
        '--clip-param',
        type=float,
        default=0.2,
        help='ppo clip parameter (default: 0.2)')
    parser.add_argument(
        '--log-interval',
        type=int,
        default=10,
        help='log interval, one log per n updates (default: 10)')
    parser.add_argument(
        '--save-interval',
        type=int,
        default=100,
        help='save interval, one save per n updates (default: 100)')
    parser.add_argument(
        '--load-id',
        type=str,
        help='The ID of the model to load (only if a model is to be loaded)')
    parser.add_argument(
        '--eval-interval',
        type=int,
        default=None,
        help='eval interval, one eval per n updates (default: None)')
    parser.add_argument(
        '--num-env-steps',
        type=int,
        default=10e8,
        help='number of environment steps to train (default: 10e8)')
    parser.add_argument(
        '--env-name',
        default="Bandit-v0", # Or try "Pendulum-v0" | 'CartPole-v0' etc.
        help='environment to train on (default: Bandit-v0)')
    parser.add_argument(
        '--no-cuda',
        action='store_true',
        default=False,
        help='disables CUDA training')
    parser.add_argument(
        '--use-proper-time-limits',
        action='store_true',
        default=True,
        help='compute returns taking into account time limits')
    parser.add_argument(
        '--recurrent-policy',
        action='store_true',
        default=True,
        help='use a recurrent policy')
    parser.add_argument(
        '--use-linear-lr-decay',
        action='store_true',
        default=False,
        help='use a linear schedule on the learning rate')
    args = parser.parse_args()

    args.cuda = not args.no_cuda and torch.cuda.is_available()

    assert args.algo in ['a2c', 'ppo']

    return args

def main():
    args = get_args()

    if not args.training:
        print("Warning! Training is turned off!")

    torch.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)

    if args.cuda and torch.cuda.is_available() and args.cuda_deterministic:
        torch.backends.cudnn.benchmark = False
        torch.backends.cudnn.deterministic = True

    # Set up the directories and names for saving stuff
    session_name = utils.datetimenow(subseconds=True)
    if args.load_id:
        loading_id = args.load_id
        cutoff_idx = loading_id.find('_s')
        model_name = loading_id[1:cutoff_idx]
    else:
        model_name = session_name  #i.e. if new model, model_name is same as session_name

    unique_id = 'm'+ model_name + '_s' + session_name

    print("The unique ID for this model and session combination " + \
          "is %s" % str(unique_id))

    # Make dirs to log experimental data, models
    exp_dir = '../exps'
    if not os.path.isdir(exp_dir):
        os.mkdir(exp_dir)

    data_logs_dir = os.path.join(exp_dir, 'data_logs')
    if not os.path.isdir(data_logs_dir):
        os.mkdir(data_logs_dir)

    data_logs_dir_uniq = os.path.join(data_logs_dir, unique_id)
    if not os.path.isdir(data_logs_dir_uniq):
        os.mkdir(data_logs_dir_uniq)
        print("Data will be logged to %s. " % data_logs_dir_uniq + \
              "Ignore the tmp/openai logging statement below.")

    models_dir = os.path.join(exp_dir, 'models')
    if not os.path.isdir(models_dir):
        os.mkdir(models_dir)
        print("New model will be saved at %s" % models_dir)

    torch.set_num_threads(1)
    device = torch.device("cuda:0" if args.cuda else "cpu")

    # Save args into csv for record keeping
    utils.save_configs_to_csv(args, session_name=session_name,
                              model_name=model_name, unique_id=unique_id)

    # Set up envs and model etc
    envs = make_vec_envs(args.env_name, args.seed, args.num_processes,
                         args.gamma, data_logs_dir_uniq, device, False)

    if args.load_id:
        loaded_id = str(args.load_id)
        path = '../exps/models/' + loaded_id + '.pt'
        actor_critic = torch.load(path)
    else:
        actor_critic = Policy(
            obs_shape=envs.observation_space.shape,
            action_space=envs.action_space,
            base_kwargs={'recurrent': args.recurrent_policy})
    actor_critic.to(device)

    if args.algo == 'a2c':
        agent = algo.A2C(
            args,
            actor_critic,
            args.value_loss_coef,
            args.entropy_coef,
            lr=args.lr,
            eps=args.eps,
            alpha=args.alpha,
            max_grad_norm=args.max_grad_norm)
    elif args.algo == 'ppo':
        agent = algo.PPO(
            args, # todo change ppo script to include args (which was added
            # to help make the training fully episodic and with overlapping
            # segments)
            actor_critic,
            args.clip_param,
            args.ppo_epoch,
            args.num_mini_batch,
            args.value_loss_coef,
            args.entropy_coef,
            lr=args.lr,
            eps=args.eps,
            max_grad_norm=args.max_grad_norm)

    rollouts = RolloutStorage(200, args.num_processes,
                              envs.observation_space.shape, envs.action_space,
                              actor_critic.recurrent_hidden_state_size)

    obs = envs.reset()
    rollouts.obs[0].copy_(obs)
    rollouts.to(device)

    episode_rewards = deque(maxlen=10)

    start = time.time()
    num_episodes = int(10e9)

    for j in range(num_episodes):

        if args.use_linear_lr_decay:
            # decrease learning rate linearly
            utils.update_linear_schedule(
                agent.optimizer, j, num_episodes, args.lr)

        for step in range(200):
            # Sample actions
            with torch.no_grad():
                value, action, action_log_prob, recurrent_hidden_states = \
                    actor_critic.act(rollouts.obs[step],
                                     rollouts.recurrent_hidden_states[step],
                                     rollouts.masks[step])

            # Obser reward and next obs
            # envs.render()
            obs, reward, done, infos = envs.step(action)

            for info in infos:
                if 'episode' in info.keys():  # Only added at the end of an epi
                    episode_rewards.append(info['episode']['r'])

            p_dists = torch.FloatTensor(
                [info['p_dist'] if 'p_dist' in info.keys() else np.zeros(2)
                 for info in infos])

            # If done then clean the history of observations.
            masks = torch.FloatTensor(
                [[0.0] if done_ else [1.0] for done_ in done])
            bad_masks = torch.FloatTensor(
                [[0.0] if 'bad_transition' in info.keys() else [1.0]
                 for info in infos])

            rollouts.insert(obs, recurrent_hidden_states, action,
                            action_log_prob, value, reward, masks,
                            bad_masks, p_dists)

        with torch.no_grad():
            next_value = actor_critic.get_value(
                rollouts.obs[-1], rollouts.recurrent_hidden_states[-1],
                rollouts.masks[-1]).detach()

        rollouts.compute_returns(next_value, args.use_gae, args.gamma,
                                 args.gae_lambda, args.use_proper_time_limits)

        if args.training:
            value_loss, action_loss, dist_entropy = agent.update(rollouts)

        if args.save_experimental_data:
            rollouts.save_experimental_data(save_dir=data_logs_dir_uniq)

        if "Bandit" in args.env_name:
            reset_hxs_every_episode = True
        else:
            reset_hxs_every_episode = False
        rollouts.after_update(reset_hxs_every_episode)

        # save for every interval-th episode or for the last epoch
        if (j % args.save_interval == 0 or j == num_episodes - 1):
            torch.save(actor_critic,
                       os.path.join(models_dir, unique_id + ".pt"))

        if j % args.log_interval == 0 and len(episode_rewards) > 1 and args.training:
            total_num_steps = (j + 1) * args.num_processes * 200
            end = time.time()
            print(
                "Episodes {}, num timesteps {}, FPS {}. Entropy: {:.4f} , Value loss: {:.4f}, Policy loss: {:.4f}, \n Last {} training episodes: mean/median reward {:.1f}/{:.1f}, min/max reward {:.1f}/{:.1f}\n"
                .format(j, total_num_steps,
                        int(total_num_steps / (end - start)), dist_entropy,
                        value_loss, action_loss,
                        len(episode_rewards), np.mean(episode_rewards),
                        np.median(episode_rewards), np.min(episode_rewards),
                        np.max(episode_rewards)))
        elif j % args.log_interval == 0 and len(episode_rewards) > 1:
            total_num_steps = (j + 1) * args.num_processes * 200
            end = time.time()
            print(
                "Episodes {}, num timesteps {}. \n"
                .format(j, total_num_steps))

        if (args.eval_interval is not None and len(episode_rewards) > 1
                and j % args.eval_interval == 0):
            ob_rms = utils.get_vec_normalize(envs).ob_rms
            evaluate(actor_critic, ob_rms, args.env_name, args.seed,
                     args.num_processes, data_logs_dir_uniq, device)

if __name__ == "__main__":
    main()


text:
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.

from __future__ import division
from __future__ import print_function

import numpy as np
import pandas as pd
from typing import Text, Union
import copy
from ...utils import get_or_create_path
from ...log import get_module_logger

import torch
import torch.nn as nn
import torch.optim as optim

from .pytorch_utils import count_parameters
from ...model.base import Model
from ...data.dataset import DatasetH
from ...data.dataset.handler import DataHandlerLP
from ...contrib.model.pytorch_lstm import LSTMModel
from ...contrib.model.pytorch_gru import GRUModel

class IGMTF(Model):
    """IGMTF Model

    Parameters
    ----------
    d_feat : int
        input dimension for each time step
    metric: str
        the evaluate metric used in early stop
    optimizer : str
        optimizer name
    GPU : str
        the GPU ID(s) used for training
    """

    def __init__(
        self,
        d_feat=6,
        hidden_size=64,
        num_layers=2,
        dropout=0.0,
        n_epochs=200,
        lr=0.001,
        metric="",
        early_stop=20,
        loss="mse",
        base_model="GRU",
        model_path=None,
        optimizer="adam",
        GPU=0,
        seed=None,
        **kwargs
    ):
        # Set logger.
        self.logger = get_module_logger("IGMTF")
        self.logger.info("IMGTF pytorch version...")

        # set hyper-parameters.
        self.d_feat = d_feat
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.dropout = dropout
        self.n_epochs = n_epochs
        self.lr = lr
        self.metric = metric
        self.early_stop = early_stop
        self.optimizer = optimizer.lower()
        self.loss = loss
        self.base_model = base_model
        self.model_path = model_path
        self.device = torch.device("cuda:%d" % (GPU) if torch.cuda.is_available() and GPU >= 0 else "cpu")
        self.seed = seed

        self.logger.info(
            "IGMTF parameters setting:"
            "\nd_feat : {}"
            "\nhidden_size : {}"
            "\nnum_layers : {}"
            "\ndropout : {}"
            "\nn_epochs : {}"
            "\nlr : {}"
            "\nmetric : {}"
            "\nearly_stop : {}"
            "\noptimizer : {}"
            "\nloss_type : {}"
            "\nbase_model : {}"
            "\nmodel_path : {}"
            "\nvisible_GPU : {}"
            "\nuse_GPU : {}"
            "\nseed : {}".format(
                d_feat,
                hidden_size,
                num_layers,
                dropout,
                n_epochs,
                lr,
                metric,
                early_stop,
                optimizer.lower(),
                loss,
                base_model,
                model_path,
                GPU,
                self.use_gpu,
                seed,
            )
        )

        if self.seed is not None:
            np.random.seed(self.seed)
            torch.manual_seed(self.seed)

        self.igmtf_model = IGMTFModel(
            d_feat=self.d_feat,
            hidden_size=self.hidden_size,
            num_layers=self.num_layers,
            dropout=self.dropout,
            base_model=self.base_model,
        )
        self.logger.info("model:\n{:}".format(self.igmtf_model))
        self.logger.info("model size: {:.4f} MB".format(count_parameters(self.igmtf_model)))

        if optimizer.lower() == "adam":
            self.train_optimizer = optim.Adam(self.igmtf_model.parameters(), lr=self.lr)
        elif optimizer.lower() == "gd":
            self.train_optimizer = optim.SGD(self.igmtf_model.parameters(), lr=self.lr)
        else:
            raise NotImplementedError("optimizer {} is not supported!".format(optimizer))

        self.fitted = False
        self.igmtf_model.to(self.device)

    @property
    def use_gpu(self):
        return self.device != torch.device("cpu")

    def mse(self, pred, label):
        loss = (pred - label) ** 2
        return torch.mean(loss)

    def loss_fn(self, pred, label):
        mask = ~torch.isnan(label)

        if self.loss == "mse":
            return self.mse(pred[mask], label[mask])

        raise ValueError("unknown loss `%s`" % self.loss)

    def metric_fn(self, pred, label):

        mask = torch.isfinite(label)

        if self.metric == "ic":
            x = pred[mask]
            y = label[mask]

            vx = x - torch.mean(x)
            vy = y - torch.mean(y)
            return torch.sum(vx * vy) / (torch.sqrt(torch.sum(vx**2)) * torch.sqrt(torch.sum(vy**2)))

        if self.metric == ("", "loss"):
            return -self.loss_fn(pred[mask], label[mask])

        raise ValueError("unknown metric `%s`" % self.metric)

    def get_daily_inter(self, df, shuffle=False):
        # organize the train data into daily batches
        daily_count = df.groupby(level=0).size().values
        daily_index = np.roll(np.cumsum(daily_count), 1)
        daily_index[0] = 0
        if shuffle:
            # shuffle data
            daily_shuffle = list(zip(daily_index, daily_count))
            np.random.shuffle(daily_shuffle)
            daily_index, daily_count = zip(*daily_shuffle)
        return daily_index, daily_count

    def get_train_hidden(self, x_train):
        x_train_values = x_train.values
        daily_index, daily_count = self.get_daily_inter(x_train, shuffle=True)
        self.igmtf_model.eval()
        train_hidden = []
        train_hidden_day = []

        for idx, count in zip(daily_index, daily_count):
            batch = slice(idx, idx + count)
            feature = torch.from_numpy(x_train_values[batch]).float().to(self.device)
            out = self.igmtf_model(feature, get_hidden=True)
            train_hidden.append(out.detach().cpu())
            train_hidden_day.append(out.detach().cpu().mean(dim=0).unsqueeze(dim=0))

        train_hidden = np.asarray(train_hidden, dtype=object)
        train_hidden_day = torch.cat(train_hidden_day)

        return train_hidden, train_hidden_day

    def train_epoch(self, x_train, y_train, train_hidden, train_hidden_day):

        x_train_values = x_train.values
        y_train_values = np.squeeze(y_train.values)

        self.igmtf_model.train()

        daily_index, daily_count = self.get_daily_inter(x_train, shuffle=True)

        for idx, count in zip(daily_index, daily_count):
            batch = slice(idx, idx + count)
            feature = torch.from_numpy(x_train_values[batch]).float().to(self.device)
            label = torch.from_numpy(y_train_values[batch]).float().to(self.device)
            pred = self.igmtf_model(feature, train_hidden=train_hidden, train_hidden_day=train_hidden_day)
            loss = self.loss_fn(pred, label)

            self.train_optimizer.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_value_(self.igmtf_model.parameters(), 3.0)
            self.train_optimizer.step()

    def test_epoch(self, data_x, data_y, train_hidden, train_hidden_day):

        # prepare training data
        x_values = data_x.values
        y_values = np.squeeze(data_y.values)

        self.igmtf_model.eval()

        scores = []
        losses = []

        daily_index, daily_count = self.get_daily_inter(data_x, shuffle=False)

        for idx, count in zip(daily_index, daily_count):
            batch = slice(idx, idx + count)
            feature = torch.from_numpy(x_values[batch]).float().to(self.device)
            label = torch.from_numpy(y_values[batch]).float().to(self.device)

            pred = self.igmtf_model(feature, train_hidden=train_hidden, train_hidden_day=train_hidden_day)
            loss = self.loss_fn(pred, label)
            losses.append(loss.item())

            score = self.metric_fn(pred, label)
            scores.append(score.item())

        return np.mean(losses), np.mean(scores)

    def fit(
        self,
        dataset: DatasetH,
        evals_result=dict(),
        save_path=None,
    ):

        df_train, df_valid = dataset.prepare(
            ["train", "valid"],
            col_set=["feature", "label"],
            data_key=DataHandlerLP.DK_L,
        )
        if df_train.empty or df_valid.empty:
            raise ValueError("Empty data from dataset, please check your dataset config.")

        x_train, y_train = df_train["feature"], df_train["label"]
        x_valid, y_valid = df_valid["feature"], df_valid["label"]

        save_path = get_or_create_path(save_path)
        stop_steps = 0
        train_loss = 0
        best_score = -np.inf
        best_epoch = 0
        evals_result["train"] = []
        evals_result["valid"] = []

        # load pretrained base_model
        if self.base_model == "LSTM":
            pretrained_model = LSTMModel()
        elif self.base_model == "GRU":
            pretrained_model = GRUModel()
        else:
            raise ValueError("unknown base model name `%s`" % self.base_model)

        if self.model_path is not None:
            self.logger.info("Loading pretrained model...")
            pretrained_model.load_state_dict(torch.load(self.model_path, map_location=self.device))

        model_dict = self.igmtf_model.state_dict()
        pretrained_dict = {
            k: v for k, v in pretrained_model.state_dict().items() if k in model_dict  # pylint: disable=E1135
        }
        model_dict.update(pretrained_dict)
        self.igmtf_model.load_state_dict(model_dict)
        self.logger.info("Loading pretrained model Done...")

        # train
        self.logger.info("training...")
        self.fitted = True

        for step in range(self.n_epochs):
            self.logger.info("Epoch%d:", step)
            self.logger.info("training...")
            train_hidden, train_hidden_day = self.get_train_hidden(x_train)
            self.train_epoch(x_train, y_train, train_hidden, train_hidden_day)
            self.logger.info("evaluating...")
            train_loss, train_score = self.test_epoch(x_train, y_train, train_hidden, train_hidden_day)
            val_loss, val_score = self.test_epoch(x_valid, y_valid, train_hidden, train_hidden_day)
            self.logger.info("train %.6f, valid %.6f" % (train_score, val_score))
            evals_result["train"].append(train_score)
            evals_result["valid"].append(val_score)

            if val_score > best_score:
                best_score = val_score
                stop_steps = 0
                best_epoch = step
                best_param = copy.deepcopy(self.igmtf_model.state_dict())
            else:
                stop_steps += 1
                if stop_steps >= self.early_stop:
                    self.logger.info("early stop")
                    break

        self.logger.info("best score: %.6lf @ %d" % (best_score, best_epoch))
        self.igmtf_model.load_state_dict(best_param)
        torch.save(best_param, save_path)

        if self.use_gpu:
            torch.cuda.empty_cache()

    def predict(self, dataset: DatasetH, segment: Union[Text, slice] = "test"):
        if not self.fitted:
            raise ValueError("model is not fitted yet!")
        x_train = dataset.prepare("train", col_set="feature", data_key=DataHandlerLP.DK_L)
        train_hidden, train_hidden_day = self.get_train_hidden(x_train)
        x_test = dataset.prepare(segment, col_set="feature", data_key=DataHandlerLP.DK_I)
        index = x_test.index
        self.igmtf_model.eval()
        x_values = x_test.values
        preds = []

        daily_index, daily_count = self.get_daily_inter(x_test, shuffle=False)

        for idx, count in zip(daily_index, daily_count):
            batch = slice(idx, idx + count)
            x_batch = torch.from_numpy(x_values[batch]).float().to(self.device)

            with torch.no_grad():
                pred = (
                    self.igmtf_model(x_batch, train_hidden=train_hidden, train_hidden_day=train_hidden_day)
                    .detach()
                    .cpu()
                    .numpy()
                )

            preds.append(pred)

        return pd.Series(np.concatenate(preds), index=index)

class IGMTFModel(nn.Module):
    def __init__(self, d_feat=6, hidden_size=64, num_layers=2, dropout=0.0, base_model="GRU"):
        super().__init__()

        if base_model == "GRU":
            self.rnn = nn.GRU(
                input_size=d_feat,
                hidden_size=hidden_size,
                num_layers=num_layers,
                batch_first=True,
                dropout=dropout,
            )
        elif base_model == "LSTM":
            self.rnn = nn.LSTM(
                input_size=d_feat,
                hidden_size=hidden_size,
                num_layers=num_layers,
                batch_first=True,
                dropout=dropout,
            )
        else:
            raise ValueError("unknown base model name `%s`" % base_model)
        self.lins = nn.Sequential()
        for i in range(2):
            self.lins.add_module("linear" + str(i), nn.Linear(hidden_size, hidden_size))
            self.lins.add_module("leakyrelu" + str(i), nn.LeakyReLU())
        self.fc_output = nn.Linear(hidden_size * 2, hidden_size * 2)
        self.project1 = nn.Linear(hidden_size, hidden_size, bias=False)
        self.project2 = nn.Linear(hidden_size, hidden_size, bias=False)
        self.fc_out_pred = nn.Linear(hidden_size * 2, 1)

        self.leaky_relu = nn.LeakyReLU()
        self.d_feat = d_feat

    def cal_cos_similarity(self, x, y):  # the 2nd dimension of x and y are the same
        xy = x.mm(torch.t(y))
        x_norm = torch.sqrt(torch.sum(x * x, dim=1)).reshape(-1, 1)
        y_norm = torch.sqrt(torch.sum(y * y, dim=1)).reshape(-1, 1)
        cos_similarity = xy / (x_norm.mm(torch.t(y_norm)) + 1e-6)
        return cos_similarity

    def sparse_dense_mul(self, s, d):
        i = s._indices()
        v = s._values()
        dv = d[i[0, :], i[1, :]]  # get values from relevant entries of dense matrix
        return torch.sparse.FloatTensor(i, v * dv, s.size())

    def forward(self, x, get_hidden=False, train_hidden=None, train_hidden_day=None, k_day=10, n_neighbor=10):
        # x: [N, F*T]
        device = x.device
        x = x.reshape(len(x), self.d_feat, -1)  # [N, F, T]
        x = x.permute(0, 2, 1)  # [N, T, F]
        out, _ = self.rnn(x)
        out = out[:, -1, :]
        out = self.lins(out)
        mini_batch_out = out
        if get_hidden is True:
            return mini_batch_out

        mini_batch_out_day = torch.mean(mini_batch_out, dim=0).unsqueeze(0)
        day_similarity = self.cal_cos_similarity(mini_batch_out_day, train_hidden_day.to(device))
        day_index = torch.topk(day_similarity, k_day, dim=1)[1]
        sample_train_hidden = train_hidden[day_index.long().cpu()].squeeze()
        sample_train_hidden = torch.cat(list(sample_train_hidden)).to(device)
        sample_train_hidden = self.lins(sample_train_hidden)
        cos_similarity = self.cal_cos_similarity(self.project1(mini_batch_out), self.project2(sample_train_hidden))

        row = (
            torch.linspace(0, x.shape[0] - 1, x.shape[0])
            .reshape([-1, 1])
            .repeat(1, n_neighbor)
            .reshape(1, -1)
            .to(device)
        )
        column = torch.topk(cos_similarity, n_neighbor, dim=1)[1].reshape(1, -1)
        mask = torch.sparse_coo_tensor(
            torch.cat([row, column]),
            torch.ones([row.shape[1]]).to(device) / n_neighbor,
            (x.shape[0], sample_train_hidden.shape[0]),
        )
        cos_similarity = self.sparse_dense_mul(mask, cos_similarity)

        agg_out = torch.sparse.mm(cos_similarity, self.project2(sample_train_hidden))
        # out = self.fc_out(out).squeeze()
        out = self.fc_out_pred(torch.cat([mini_batch_out, agg_out], axis=1)).squeeze()
        return out

max_stars_repo_name: SunsetWolf/qlib

max_issues_repo_name: SunsetWolf/qlib

max_forks_repo_name: SunsetWolf/qlib


text:
#
# OtterTune - test_cluster.py
#
# Copyright (c) 2017-18, Carnegie Mellon University Database Group
#
import unittest
import numpy as np
from sklearn import datasets

from analysis.cluster import KMeans, KMeansClusters, create_kselection_model

class TestKMeans(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        super(TestKMeans, cls).setUpClass()
        iris = datasets.load_iris()
        cls.model = KMeans()
        cls.model.fit(iris.data, 5, iris.target,
                      estimator_params={'n_init': 50, 'random_state': 42})

    def test_kmeans_n_clusters(self):
        self.assertEqual(self.model.n_clusters_, 5)

    def test_kmeans_cluster_inertia(self):
        self.assertAlmostEqual(self.model.cluster_inertia_, 46.535, 2)

    def test_kmeans_cluster_labels(self):
        expected_labels = [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
                           1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
                           1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 3, 3, 3, 2, 3, 3, 3,
                           2, 3, 2, 2, 3, 2, 3, 2, 3, 3, 2, 3, 2, 3, 2, 3, 3, 3, 3,
                           3, 3, 3, 2, 2, 2, 2, 3, 2, 3, 3, 3, 2, 2, 2, 3, 2, 2, 2,
                           2, 2, 3, 2, 2, 4, 3, 0, 4, 4, 0, 2, 0, 4, 0, 4, 4, 4, 3,
                           4, 4, 4, 0, 0, 3, 4, 3, 0, 3, 4, 0, 3, 3, 4, 0, 0, 0, 4,
                           3, 3, 0, 4, 4, 3, 4, 4, 4, 3, 4, 4, 4, 3, 4, 4, 3]
        for lab_actual, lab_expected in zip(self.model.cluster_labels_, expected_labels):
            self.assertEqual(lab_actual, lab_expected)

    def test_kmeans_sample_labels(self):
        for lab_actual, lab_expected in zip(self.model.sample_labels_, datasets.load_iris().target):
            self.assertEqual(lab_actual, lab_expected)

    def test_kmeans_cluster_centers(self):
        expected_centers = [[7.475, 3.125, 6.300, 2.050],
                            [5.006, 3.418, 1.464, 0.244],
                            [5.508, 2.600, 3.908, 1.204],
                            [6.207, 2.853, 4.746, 1.564],
                            [6.529, 3.058, 5.508, 2.162]]
        for row_actual, row_expected in zip(self.model.cluster_centers_, expected_centers):
            for val_actual, val_expected in zip(row_actual, row_expected):
                self.assertAlmostEqual(val_actual, val_expected, 2)

class TestKSelection(unittest.TestCase):

    def setUp(self):
        np.random.seed(seed=42)

    @classmethod
    def setUpClass(cls):
        super(TestKSelection, cls).setUpClass()

        # Load Iris data
        iris = datasets.load_iris()
        cls.matrix = iris.data
        cls.kmeans_models = KMeansClusters()
        cls.kmeans_models.fit(cls.matrix,
                              min_cluster=1,
                              max_cluster=10,
                              sample_labels=iris.target,
                              estimator_params={'n_init': 50, 'random_state': 42})

    def test_detk_optimal_num_clusters(self):
        # Compute optimal # cluster using det-k
        detk = create_kselection_model("det-k")
        detk.fit(self.matrix, self.kmeans_models.cluster_map_)
        self.assertEqual(detk.optimal_num_clusters_, 2)

    def test_gap_statistic_optimal_num_clusters(self):
        # Compute optimal # cluster using gap-statistics
        gap = create_kselection_model("gap-statistic")
        gap.fit(self.matrix, self.kmeans_models.cluster_map_)
        self.assertEqual(gap.optimal_num_clusters_, 8)

    def test_silhouette_optimal_num_clusters(self):
        # Compute optimal # cluster using Silhouette Analysis
        sil = create_kselection_model("s-score")
        sil.fit(self.matrix, self.kmeans_models.cluster_map_)
        self.assertEqual(sil.optimal_num_clusters_, 2)

max_stars_repo_name: newopscn/ottertune

max_issues_repo_name: yangdsh/ottertune

max_forks_repo_name: yangdsh/ottertune