<filename>growspace/envs/spotlight.py<gh_stars>10-100
import os
import random
import sys
from enum import IntEnum
from random import sample

import cv2
import gym
import numpy as np
import tqdm
from numpy.linalg import norm
from scipy.spatial import distance
import random
import torchvision
from torchvision import datasets
import growspace.plants.tree

np.set_printoptions(threshold=sys.maxsize)
# customizable variables by user

BRANCH_THICCNESS = .036 # before was 0.036 in 28 x 28 and .015 in 28 x 28
MAX_BRANCHING = 10
DEFAULT_RES = 28
LIGHT_WIDTH = 0.25
LIGHT_DIF = 200
LIGHT_DISPLACEMENT = 0.1
LIGHT_W_INCREMENT = 0.1
MIN_LIGHT_WIDTH = 0.1
MAX_LIGHT_WIDTH = 0.5

BRANCH_LENGTH = (1 / 9) * DEFAULT_RES

def to_int(v):
    return int(round(v))

ir = to_int  # shortcut for function calld

FIRST_BRANCH_HEIGHT = ir(0.1 * DEFAULT_RES)

def unpack(w):
    return map(list, zip(*enumerate(w)))

def load_mnist(digit):
    mnist_dataset = datasets.MNIST('./data', train=True, download=True)
    if digit == 'partymix':
        return mnist_dataset
    if digit == 'curriculum':
        return mnist_dataset
    else:
        return idx_digit(mnist_dataset,digit)

def idx_digit(mnist_dataset, digit):
    idx = mnist_dataset.train_labels == int(digit) #and mnist_dataset.train_labels  == 6  # digit comes as string
    mnist_digit = mnist_dataset.train_data[idx]

    return mnist_digit

def sample_digit(digit_set):
    digit_set = digit_set.data.numpy()
    digit = np.random.randint(len(digit_set)-1)
    im1 = np.zeros((28, 28), dtype=np.uint8)
    im2 = np.zeros((28, 28), dtype=np.uint8)
    im3 = np.zeros((28, 28, 3), dtype=np.uint8)
    number = np.where(digit_set[digit] == 0, im2, 1 * 255)
    img = np.dstack((im1,im2, number))
    A = np.array((im3, img))
    final_img = np.float32(np.sum(A, axis=0))

    return final_img

def stage_curriculum(episode, mnist_dataset):
    _0 = idx_digit(mnist_dataset, 0)
    _1 = idx_digit(mnist_dataset, 1)
    _2 = idx_digit(mnist_dataset, 2)
    _3 = idx_digit(mnist_dataset, 3)
    _4 = idx_digit(mnist_dataset, 4)
    _5 = idx_digit(mnist_dataset, 5)
    _6 = idx_digit(mnist_dataset, 6)
    _7 = idx_digit(mnist_dataset, 7)
    _8 = idx_digit(mnist_dataset, 8)
    _9 = idx_digit(mnist_dataset, 9)
    curriculum = []
    if episode < 500:
        curriculum.append(_3)
        curriculum.append(_6)

        if episode == 1:
            print('check1')

    if 500 <= episode < 1000:
        curriculum.append(_3)
        curriculum.append(_6)
        curriculum.append(_2)
        if episode == 500:
             print('check2')

    if 1000 <= episode < 1500:
        curriculum.append(_3)
        curriculum.append(_6)
        curriculum.append(_2)
        curriculum.append(_1)

        if episode == 1000:
            print('check3')

    if 1500 <= episode < 2500:
        curriculum.append(_3)
        curriculum.append(_6)
        curriculum.append(_2)
        curriculum.append(_1)
        curriculum.append(_4)

        if self.episode == 1500:
            print('check4')

    if 2500 <= episode < 3500:
        curriculum.append(_3)
        curriculum.append(_6)
        curriculum.append(_2)
        curriculum.append(_1)
        curriculum.append(_4)
        curriculum.append(_5)

        if self.episode == 2500:
            print('check5')

    if 3500 <= episode < 5000:
        curriculum.append(_3)
        curriculum.append(_6)
        curriculum.append(_2)
        curriculum.append(_1)
        curriculum.append(_4)
        curriculum.append(_5)
        curriculum.append(_7)

        if self.episode == 3500:
            print('check6')

    if 5000 <= episode < 6500:
        curriculum.append(_3)
        curriculum.append(_6)
        curriculum.append(_2)
        curriculum.append(_1)
        curriculum.append(_4)
        curriculum.append(_5)
        curriculum.append(_7)
        curriculum.append(_8)

        if episode == 5000:
            print('check7')

    if 6500 <= episode < 10000:
        curriculum.append(_3)
        curriculum.append(_6)
        curriculum.append(_2)
        curriculum.append(_1)
        curriculum.append(_4)
        curriculum.append(_5)
        curriculum.append(_7)
        curriculum.append(_8)
        curriculum.append(_0)
        curriculum.append(_9)

    idx = len(curriculum)
    id = np.random.randint(idx)
    x = sample_digit(curriculum[id])
    return x

class Features(IntEnum):
    light = 0
    scatter = 1

class GrowSpaceEnvSpotlightMnist(gym.Env):
    metadata = {'render.modes': ['human', 'rgb_array']}  # Required, otherwise gym.Monitor disables itself.

    def __init__(self, width=DEFAULT_RES, height=DEFAULT_RES, digit='curriculum'):
        self.width = width
        self.height = height
        self.seed()
        self.action_space = gym.spaces.Discrete(5)  # L, R, keep of light paddle, or increase, decrease
        self.feature_maps = np.zeros((len(Features), self.height, self.width), dtype=np.uint8)

        self.observation_space = gym.spaces.Box(0, 255, shape=(self.height, self.width, 3), dtype=np.uint8)
        self.digit = digit
        self.mnist_digit = load_mnist(self.digit)

        self.focus_point = None
        self.focus_radius = None

        self.branches = None
        self.target = None
        self.steps = None
        self.new_branches = None
        self.tips_per_step = None
        self.tips = None
        self.episode = -1

    def render(self, mode='human', debug_show_scatter=False):
        """
        @mode: ['mode', 'rgb_array']

        """
        img = self.get_observation(debug_show_scatter)
        #dsize = (84,84)
        #img = cv2.resize(img, dsize)
        # if self.obs_type == 'Binary':
        #     image = img.astype(np.uint8)
        #     img = image * 255

        if mode == "human":
            cv2.imshow('plant', img)  # create opencv window to show plant
            cv2.waitKey(1)  # this is necessary or the window closes immediately
        else:
            return img

    def seed(self, seed=None):
        return [np.random.seed(seed)]

    def light_scatter(self):
        filter_ = np.logical_and(self.feature_maps[Features.light], self.feature_maps[Features.scatter])
        return np.argwhere(filter_)

    def tree_grow(self, activated_photons, mindist, maxdist):
        branches_trimmed = self.branches
        # number of possible scatters, check if they allow for branching with min_dist
        for i in range(len(activated_photons) - 1, 0, -1):
            closest_branch = 0
            dist = 1 * self.width

            if len(self.branches) > MAX_BRANCHING:
                branches_trimmed = sample(self.branches, MAX_BRANCHING)
            else:
                branches_trimmed = self.branches

            for branch in branches_trimmed:
                photon_ptx = np.flip(activated_photons[i])  # flip was necessary bc coordinate systems are inverted -
                tip_to_scatter = norm(photon_ptx - np.array(branch.tip_point))  # one is np.array, one is tuple
                if tip_to_scatter < dist:
                    dist = tip_to_scatter
                    closest_branch = branch

            # removes scatter points if reached

            if dist < mindist:
                activated_photons = np.delete(activated_photons, i)

            # when distance is greater than max distance, branching occurs to find other points.
            elif dist < maxdist:
                closest_branch.grow_count += 1
                branch_length = BRANCH_LENGTH / dist
                photon = np.flip(activated_photons[i])
                g = (photon - closest_branch.tip_point) * branch_length
                closest_branch.grow_direction += np.round(g).astype(np.int)

        for branch in branches_trimmed:
            if branch.grow_count > 0:
                (x2, y2) = branch.tip_point + branch.grow_direction / branch.grow_count
                x2 = np.clip(x2, 0, self.width - 1)
                y2 = np.clip(y2, 0, self.height - 1)

                newBranch = growspace.plants.tree.PixelBranch(
                    branch.x2, ir(x2), branch.y2, ir(y2), self.width, self.height
                )
                self.branches.append(newBranch)
                branch.child.append(newBranch)
                branch.grow_count = 0
                branch.grow_direction.fill(0)

        # increase thickness of first elements added to tree as they grow

        self.branches[0].update_width()

        branch_coords = []

        # sending coordinates out
        for branch in self.branches:
            branch_coords.append(branch.tip_point)

        self.tips = branch_coords
        return branch_coords

    def get_observation(self, debug_show_scatter=False):
        img = np.zeros((self.height, self.width, 3), dtype=np.uint8)

        yellow = (0, 128, 128)  # RGB color (dark yellow)

        img[self.feature_maps[Features.light].nonzero()] = yellow
        cv2.circle(img, tuple(self.to_image(self.focus_point)), int(self.focus_radius * self.height), (0, 128, 128), thickness=2) # only contour

        if debug_show_scatter:
            pts = self.light_scatter()
            for p in pts:
                x, y = p
                cv2.circle(img, center=(y, x), radius=2, color=(255, 0, 0), thickness=-1)

        # Draw plant as series of lines (1 branch = 1 line)
        for branch in self.branches:
            thiccness = ir(branch.width * BRANCH_THICCNESS * self.width)
            cv2.line(img, pt1=branch.p, pt2=branch.tip_point, color=(0, 255, 0), thickness=thiccness)

        img = cv2.flip(img,0)
        z = np.where(self.mnist_shape < 255, img, 150)

        return z

    def reset(self):
        random_start = random.randint(self.width - (self.width*3/4), self.width - 1 - (self.width*1/4))
        self.branches = [
            growspace.plants.tree.PixelBranch(
                x=random_start,
                x2=random_start,
                y=0,
                y2=FIRST_BRANCH_HEIGHT,
                img_width=self.width,
                img_height=self.height,
            )
        ]

        self.episode += 1
        if self.digit == 'curriculum':
            self.mnist_shape = stage_curriculum(self.episode,self.mnist_digit)

        else:
            self.mnist_shape = sample_digit(self.mnist_digit)

        self.focus_point = np.array([random_start / self.width, FIRST_BRANCH_HEIGHT / self.height])
        self.focus_radius = 0.1

        x_scatter = np.random.randint(0, self.width, LIGHT_DIF)
        y_scatter = np.random.randint(0, self.height, LIGHT_DIF)
        self.feature_maps[Features.scatter].fill(False)

        self.feature_maps[Features.scatter][y_scatter, x_scatter] = True

        self.steps = 0
        self.new_branches = 0
        self.tips_per_step = 0
        self.tips = [
            self.branches[0].tip_point,
        ]

        self.draw_spotlight()

        self.mnist_pixels = (self.get_observation()[:, :, 2] / 150)  # binary map of mnist shape

        plant_stem = (self.get_observation()[:, :, 1] / 255)
        plant_stem[plant_stem>0.6] =1              # filter for green
        self.plant_original = plant_stem.astype(int)

        return self.get_observation()

    def step(self, action):
        if action == 0:
            self.focus_radius = min(0.2, self.focus_radius + 0.05)

        if action == 1:
            self.focus_radius = max(0.03, self.focus_radius - 0.05)

        if action == 2:
            self.focus_point[0] = max(0, self.focus_point[0] - 0.1)

        if action == 3:
            self.focus_point[0] = min(1, self.focus_point[0] + 0.1)

        if action == 4:
            self.focus_point[1] = min(1, self.focus_point[1] + 0.1)

        if action == 5:
            self.focus_point[1] = max(0, self.focus_point[1] - 0.1)

        if action == 6:
            pass
        #self.draw_spotlight()

        pts = self.light_scatter()
        tips = self.tree_grow(pts, 0.01 * self.width, 0.15 * self.width)
        self.draw_spotlight()
        observation = self.get_observation()  #image

        plant = (observation[:,:,1]/255) # binary map of plant
        pixel_plant = np.sum(plant)

        plant[plant>0.6] =1              # filter for green
        plant = plant.astype(int)
        true_plant = np.subtract(plant,self.plant_original)

        mnist = (observation[:,:,2]/150)     # binary map of mnist

        mnist[mnist>0.5] =1
        mnist = mnist.astype(int)

        check = np.sum((true_plant, mnist), axis=0)
        intersection = np.sum(np.where(check < 2, 0, 1))

        union = np.sum(np.where(check<2,check,1))

        reward = intersection / union

        print('reqward',reward)

        done = False  # because we don't have a terminal condition
        misc = {"tips": tips, "target": self.target, "light": None}

        if self.steps == 0:
            self.new_branches = len(tips)
            misc["new_branches"] = self.new_branches

        else:
            new_branches = len(tips) - self.new_branches
            misc["new_branches"] = new_branches
            self.new_branches = len(tips)  # reset for future step

        misc["img"] = observation
        misc["plant_pixel"] = pixel_plant
        self.steps += 1
        return observation, float(reward), done, misc

    def draw_spotlight(self):
        self.feature_maps[Features.light].fill(False)
        cv2.circle(
            self.feature_maps[Features.light],
            tuple(self.to_image(self.focus_point)),
            int(self.focus_radius * self.height),
            True,
            thickness=-1,

        )

    def to_image(self, p):
        if hasattr(p, "normalized_array"):
            return np.around((self.height, self.width) * p.normalized_array[:-1]).astype(np.int32)
        else:
            y, x = p
            return np.around((self.height * y, self.width * x)).astype(np.int32)

def enjoy():
    #gse = gym.make("GrowSpaceSpotlight-Mnist1-v0")
    gse = GrowSpaceEnvSpotlightMnist()

    def key2action(key):
        if key == ord("+"):
            return 0  # move left
        elif key == ord("-"):
            return 1  # move right
        elif key == ord("a"):
            return 2
        elif key == ord("d"):
            return 3
        elif key == ord("w"):
            return 4
        elif key == ord("s"):
            return 5
        elif key == ord("x"):
            return 6
        else:
            return None

    while True:
        gse.reset()
        img = gse.get_observation(debug_show_scatter=False)
        cv2.imshow("plant", img)
        cv2.waitKey(-1)
        rewards = []
        for _ in range(50):
            action = key2action(cv2.waitKey(-1))
            if action is None:
                quit()

            b, t, c, f = gse.step(action)
            print(f["new_branches"])
            rewards.append(t)
            cv2.imshow("plant", gse.get_observation(debug_show_scatter=False))
        total = sum(rewards)

        print("amount of rewards:", total)  # cv2.waitKey(1)  # this is necessary or the window closes immediately
        # else:
        # dreturn img

def profile():
    #gse = gym.make("GrowSpaceSpotlight-Mnist1-v0")
    gse = GrowSpaceEnvSpotlightMnist()
    gse.reset()

    def do_step():
        a = gse.action_space.sample()
        s, r, d, i = gse.step(a)
        if d:
            gse.reset()

    for _ in tqdm.trange(100000):
        do_step()
    print("hi")

if __name__ == "__main__":
    # import matplotlib.pyplot as plt
    #
    # env = GrowSpaceEnvSpotlightMnist()
    # env.reset()
    # fig, axs = plt.subplots(1, 2)
    # axs[0].imshow(env.feature_maps[Features.light])
    # axs[1].imshow(env.feature_maps[Features.scatter])
    # plt.show()
    #
    # env.draw_spotlight()
    # pts = env.light_scatter()
    # print(pts.shape)
    # fig, axs = plt.subplots(1, 2)
    # axs[0].imshow(env.feature_maps[Features.light])
    # img = np.zeros((71, 71), dtype=np.float)
    # img[pts[:, 0], pts[:, 1]] = True
    # axs[1].imshow(img)
    # plt.show()

    enjoy()


import warnings
import numpy as np
import pandas as pd
import scipy.stats as st
import statsmodels as sm
import matplotlib
import matplotlib.pyplot as plt

import numpy as np
from statsmodels.base.model import GenericLikelihoodModel

from scipy.stats import gamma
shape = 12; loc = 0.71; scale = 0.0166
data = record["peaks"]
params = gamma.fit(data) # params close to but not the same as (shape, loc, scale)
# HOW TO ESTIMATE/GET ERRORS FOR EACH PARAM?

print(params)
print('\n')

class Gamma(GenericLikelihoodModel):

    nparams = 3

    def loglike(self, params):
        return gamma.logpdf(self.endog, *params).sum()

res = Gamma(data).fit(start_params=params)
res.df_model = len(params)
res.df_resid = len(data) - len(params)
print(res.summary())


"""
    Common processing functions for grayscale images: (x, y), where x,y in range (0,255)
"""

def _memoryview_safe(x):
    """Make array safe to run in a Cython memoryview-based kernel. These
    kernels typically break down with the error ``ValueError: buffer source
    array is read-only`` when running in dask distributed.
    SEE: https://github.com/dask/distributed/issues/1978
    """
    if not x.flags.writeable:
        if not x.flags.owndata:
            x = x.copy(order='C')
        x.setflags(write=True)
    return x

def local_contrast(im_array, box_size, DEBUG = False):
    """ REF: https://scikit-image.org/docs/dev/auto_examples/color_exposure/plot_local_equalize.html
    """
    try:
        from skimage.filters import rank
        from skimage.morphology import disk
    except:
        print(" ERROR :: scikit-image not installed, try:")
        print("     pip install scikit-image")
        return

    ## ensure the input array is writable
    im_array = _memoryview_safe(im_array)

    footprint = disk(box_size * 2)
    if DEBUG:
        print("=======================================")
        print(" image_handler :: local_contrast")
        print("---------------------------------------")
        print("  input img dim = ", im_array.shape)
        print("  box_size = %s px" % box_size)
        print("  local contrast footprint = ", footprint.shape)
        print("=======================================")
    im = rank.equalize(im_array, footprint)
    return im

def auto_contrast(im_array, DEBUG = True):
    """ Rescale the image intensity levels to a reasonable range using the top/bottom 2 percent
        of the data to define the intensity levels
    """
    import numpy as np
    ## avoid hotspot pixels by looking at a group of pixels at the extreme ends of the image
    minval = np.percentile(im_array, 2)
    maxval = np.percentile(im_array, 98)

    if DEBUG:
        print("=======================================")
        print(" image_handler :: auto_contrast")
        print("---------------------------------------")
        print("  input img dim = ", im_array.shape)
        print("  original img min, max = (%s, %s)" % (np.min(im_array), np.max(im_array)))
        print("  stretch to new min, max = (%s %s)" % (minval, maxval))
        print("=======================================")

    ## remove pixles above/below the defined limits
    im_array = np.clip(im_array, minval, maxval)
    ## rescale the image into the range 0 - 255
    im_array = ((im_array - minval) / (maxval - minval)) * 255

    return im_array

def sigma_contrast(im_array, sigma, DEBUG = True):
    """ Rescale the image intensity levels to a range defined by a sigma value (the # of
        standard deviations to keep). Can perform better than auto_contrast when there is
        a lot of dark pixels throwing off the level balancing.
    """
    import numpy as np
    stdev = np.std(im_array)
    mean = np.mean(im_array)
    minval = mean - (stdev * sigma)
    maxval = mean + (stdev * sigma)

    if DEBUG:
        print("=======================================")
        print(" image_handler :: sigma_contrast")
        print("---------------------------------------")
        print("  input img dim = ", im_array.shape)
        print("  img mean, stdev = (%s, %s)" % (mean, stdev))
        print("  stretch to new min, max = (%s %s)" % (minval, maxval))
        print("=======================================")

    ## remove pixles above/below the defined limits
    im_array = np.clip(im_array, minval, maxval)
    ## rescale the image into the range 0 - 255
    im_array = ((im_array - minval) / (maxval - minval)) * 255

    return im_array

def whiten_outliers(im_array, min, max):
    """ Set any pixels outside of a defined intensity range to 255 (white)
    """
    import numpy as np
    ## use the input mean/max to clip the image before increasing contrast
    im_array = np.where(im_array < min, 255, im_array)
    im_array = np.where(im_array > max, 255, im_array)
    return im_array

def extract_boxes(im_array, box_size, coords, DEBUG = True):
    """
    PARAMETERS
        im_array = np array (0 - 255)
        box_size = int(); pixel size of the box to extract
        coords = list( tuple(x, y), ... ); centered coordinates in pixels (top left == 0,0 by convention)
    RETURNS
        extracted_imgs = list( np arrays of dimension box_size , ... )
    """
    extracted_imgs = []
    ## sanity check that not too many coordinates are being asked to be extracted
    if len(coords) > 500:
        print(" ERROR :: extracted_boxes is capped at 500 coordinates to avoid memory issues, remove some and re-run")
        return extracted_imgs

    box_size_halfwidth = int( box_size / 2)

    for coordinate in coords:
        x0 = coordinate[0] - box_size_halfwidth
        y0 = coordinate[1] - box_size_halfwidth
        x1 = coordinate[0] + box_size_halfwidth
        y1 = coordinate[1] + box_size_halfwidth

        extracted_img = im_array[y0:y1,x0:x1]
        extracted_imgs.append(extracted_img)

    if DEBUG:
        print(" Extracted %s boxes from image" % len(extracted_imgs))

    return extracted_imgs

def find_intensity_range(im_arrays, DEBUG = True):
    """
        im_arrays = list( 2d numpy arrays, ...)

    RETURNS:
        min = int()
        max = int()
    """
    import numpy as np
    import statistics
    mins = []
    maxs = []
    ## get the min/max value for each image we are loading
    for im_array in im_arrays:
        mins.append(np.min(im_array))
        maxs.append(np.max(im_array))
    ## get the mean value for the min/maxes
    min = int(statistics.mean(mins))
    max = int(statistics.mean(maxs))
    if DEBUG:
        print(" find_intensity_range :: %s imgs" % len(im_arrays))
        print("    ... (min, max) = (%s, %s)" % (min, max))
    return (min, max)

def gaussian_blur(im_array, sigma, DEBUG = True):
    import scipy.ndimage as ndimage
    if DEBUG:
        print("=======================================")
        print(" image_handler :: gaussian_blur")
        print("---------------------------------------")
        print("  input img dim = ", im_array.shape)
        print("  sigma = ", sigma)
        print("=======================================")

    blurred_img = ndimage.gaussian_filter(im_array, sigma)
    return blurred_img

def image2array(file, DEBUG = True):
    """
        Import an image into a grayscal 2d numpy array with values from (0 - 255), where
            0 == black
            255 == white
    """
    from PIL import Image as PIL_Image
    import numpy as np

    im = PIL_Image.open(file).convert('L') # 'L' == convert to grayscale data
    # convert image to numpy array
    im_data = np.asarray(im)

    if DEBUG:
        print("===================================================")
        print(" IMPORT IMAGE :: %s" % file)
        print("===================================================")
        print("   >> %s px, min = %s, max = %s" % (im_data.shape, np.min(im_data), np.max(im_data)))

    return im_data

def display_img(im_array, coords = None, box_size = 1):
    """
        box_size = int(); pixel size of the particle
    """
    from PIL import Image as PIL_Image
    from PIL import ImageTk

    box_size_halfwidth = int( box_size / 2 )

    root = Tk()
    canvas = Canvas(root, width = im_array.shape[0], height = im_array.shape[1])
    canvas.pack()
    img = PIL_Image.fromarray(im_array).convert('L')
    img = ImageTk.PhotoImage(img)
    canvas.create_image(0, 0, anchor=NW, image=img)

    if coords is None:
        pass
    else:
        for coordinate in coords:
            # break
            ## each coordinate is the center of a box, thus we need to offset by half the gif_box_width pixel length to get the bottom left and top right of the rectangle
            x0 = coordinate[0] - box_size_halfwidth
            y0 = coordinate[1] - box_size_halfwidth
            x1 = coordinate[0] + box_size_halfwidth
            y1 = coordinate[1] + box_size_halfwidth #y0 - img_box_size # invert direction of box to take into account x0,y0 are at bottom left, not top left
            canvas.create_rectangle(x0, y0, x1, y1, outline='red', width=1, tags='particle_positions')

    root.mainloop()

def gaussian_disk(size):
    """ Creates a soft gaussian grayscale image of given pixel size with values in range 0 -- 255
    """
    import numpy as np

    size = int(size)
    x, y = np.meshgrid(np.linspace(-1,1, size), np.linspace(-1,1, size))
    d = np.sqrt(x*x+y*y)
    sigma, mu = 0.4, 0.0
    g = np.exp(-( (d-mu)**2 / ( 2.0 * sigma**2 ) ) )
    ## invert color to we match to dark pixels
    g = 1 - g

    g = g * 255
    return g

def template_cross_correlate(im_array, template, threshold, DEBUG = False):
    """
    PARAMETERS
        im_array = np array of grayscale img in range 0 - 255
        template = np array of grayscale template img in range 0 - 255
        threshold = int(); peak cutoff in range 0 - 1
    RETURNS
        cc = cross correlation image as a grayscale (0 - 255), note peaks represent positions aligned with top-right of template!
    """
    import numpy as np
    from scipy import signal

    if DEBUG:
        print("Template info: %s, min = %s, max %s" % (template.shape, np.min(template), np.max(template)))
    cc = signal.correlate2d(im_array, template, boundary='symm', mode='same')
    ## determine the threshold at which to keep peaks
    cc_min, cc_max = (np.min(cc), np.max(cc))
    cc_range = cc_max - cc_min
    cc_threshold_cutoff = cc_min + (threshold * cc_range)
    print("cc min, max, range, threshold value = (%s, %s, %s, %s)" % ( np.min(cc), np.max(cc), cc_range, cc_threshold_cutoff))
    ## remove signal below peaks
    cc = np.where(cc < cc_threshold_cutoff, 0, 255)
    return cc

def bool_img(im_array, threshold, DEBUG = False):
    """
        For a given threshold value (intensity, i.e. between 0 - 255), make any pixels below the
        threshold equal to 255 (white) and any above 0 (black)
    PARAMETERS
        im_array = np array of a grayscale image (0 - 255)
    RETURNS
        im_array = np array as grayscale image (0 - 255)
    """
    import numpy as np
    if DEBUG:
        print("=======================================")
        print(" image_handler :: bool_img")
        print("=======================================")
        print("  intensity cutoff = ", threshold)

    im_array = np.where(im_array >= threshold, 255, 0)
    return im_array

def find_local_peaks(im_array, min_area, max_area, INVERT = False, DEBUG = False):
    """
    """
    try:
        from skimage.measure import label, regionprops
    except:
        print(" ERROR :: scikit-image not installed, try:")
        print("     pip install scikit-image")
        return
    if DEBUG:
        print("=======================================")
        print(" image_handler :: find_local_peaks")
        print("=======================================")
        print("  input img dim = ", im_array.shape)
        print("  min_area = %s px" % min_area)
        print("  max_area = %s px" % max_area)

    if INVERT:
        if DEBUG:
            print("  invert color scale for autopicking")
        im_array = 255 - im_array

    labeled_img = label(im_array) #, connectivity = 2)
    regions = regionprops(labeled_img)
    coordinates = []

    # min_area = int(box_size / 2) ## minimum # of pixels for a labeled feature to be added as a coordinate
    # max_area = int(box_size * box_size * 1.5)
    for props in regions:
        area = getattr(props, 'area')
        if area >= min_area:
            if area <= max_area:
                y0, x0 = props.centroid
                coordinates.append((x0, y0))

    if DEBUG:
        print("  >> %s coordinates found!" % len(coordinates))
        print("=======================================")

    return coordinates

#############################################
##  RUN BLOCK
#############################################
if __name__ == "__main__":

    import sys, os
    from tkinter import *

    ## allow functions of this script to be tested by supplying in an image on the commandline via:
    ##      $ image_handler.py  <img name>
    fname = sys.argv[1]

    ## edit the functions that run in this block to test proper execution
    im = image2array(fname)

    ## box_size and some coordinates for img B1g1...Exp_10.jpg
    coords = [(444,138), (452, 124), (466, 117)]
    box_size = 17

    im = local_contrast(im, box_size)
    # im = sigma_contrast(im, 1.5)
    # im = auto_contrast(im)

    im = gaussian_blur(im, 1)

    # extracted_imgs = extract_boxes(im, box_size, coords, DEBUG = True)
    # min, max = find_intensity_range(extracted_imgs)
    # im = whiten_outliers(im, min, max)

    ## boolean image
    # cutoff = 100 #int(max - min / 1000000) + min
    # im = bool_img(im, cutoff)

    # im_inverted = 255 - im ## local maxima uses 0 as background so we need particle peaks to be white
    # coords = find_local_maxima(im_inverted, box_size, DEBUG = True)

    display_img(im, coords, box_size)


<filename>hidrocomp/comparasion/genpareto.py
from scipy.stats import genpareto
import pandas as pd

from hidrocomp.comparasion.bootstrap_build import BootstrapBuild

class BootsGenPareto(BootstrapBuild):

    def __init__(self, shape, location, scale, size):
        super().__init__(shape, location, scale, size)

    def fit_resample(self):
        resample = genpareto.rvs(self.shape, self.location, self.scale,
                                 self.size)
        return genpareto.fit(resample)

    def fits_resamples(self, quantity):
        list_fits = list()
        for i in range(quantity):
            list_fits.append(self.fit_resample())

        return list_fits

    def magnitudes_resamples(self, quantity):
        dic_magns = {0.001: list(), 0.01: list(), 0.1: list(),
                     0.5: list(), 0.9: list(), 0.99: list(), 0.999: list()
                     }
        for i in range(quantity):
            fit = self.fit_resample()
            for j in dic_magns:
                mag = genpareto.ppf(j, fit[0], fit[1], fit[2])
                dic_magns[j].append(mag)

        return pd.DataFrame(dic_magns)


import numpy as np
from sigpipes.sigcontainer import SigContainer
from sigpipes.sigoperator import SigOperator
from warnings import warn
from typing import Sequence, List
from scipy.signal import correlate, convolve
from deprecated import deprecated

class Joiner(SigOperator):
    def __init__(self, *branches):
        self.subop = branches

    def fromSources(self) -> SigContainer:
        container = self.prepare_container(self.subop[0])
        return self.join(container, self.subop[1:])

    def apply(self, container: SigContainer) -> SigContainer:
        container = self.prepare_container(container)
        containers = [inp if isinstance(inp, SigContainer) else container | inp for inp in self.subop]
        # containers = container | Alternatives(*self.subop)
        assert all(isinstance(c, SigContainer) for c in containers), "Join over non containers"
        self.assertion(container, containers)
        return self.join(container, containers)

    def prepare_container(self, container: SigContainer) -> SigContainer:
        return SigContainer(container.d.deepcopy(shared_folders=["annotations"],
                                                 empty_folders=["meta"]))

    def join(self, output: SigContainer, inputs: Sequence[SigContainer]) -> SigContainer:
        raise NotImplementedError("abstract method")

    def assertion(self, output: SigContainer, inputs: Sequence[SigContainer]) -> None:
        assert all(c.signals.shape == output.signals.shape for c in inputs)
        if any(c.d["signals/fs"] != output.d["signals/fs"] for c in inputs):
            warn("Join operation on signals with incompatible frequencies")

class Merge(Joiner):
    def __init__(self, ufunc: np.ufunc, *branches) -> None:
        super().__init__(*branches)
        self.ufunc = ufunc

    def join(self, output: SigContainer, inputs: Sequence[SigContainer]) -> SigContainer:
        result = np.copy(output.signals)
        for inc in inputs:
            self.ufunc(result, inc.signals, out=result)
        output.d["signals/data"] = result
        print(output.d["signals/channels"])
        output.d["signals/channels"] = [
            f"{self.ufunc.__name__}({', '.join(input.d['signals/channels'][i] for input in [output] + inputs)})"
            for i in range(output.channel_count)]
        return output

    def log(self):
        return f"M@{self.ufunc.__name__}"

@deprecated(reason='more generalized and efficient version in Merge joiner')
class Sum(Joiner):
    def __init__(self, *branches) -> None:
        super().__init__(*branches)

    def join(self, outcontainer: SigContainer, incontainers: Sequence[SigContainer]) -> SigContainer:
        result = np.copy(outcontainer.signals)
        for inc in incontainers:
            result += inc.signals
        outcontainer.d["signals/data"] = result
        return outcontainer

class JoinChannels(Joiner):
    def __init__(self, *branches):
        super().__init__(*branches)

    def join(self, output: SigContainer, inputs: Sequence[SigContainer]) -> SigContainer:
        for input in inputs:
            output.d["signals/data"] = np.vstack((output.signals, input.signals))
            output.d["signals/channels"].extend(input.d["signals/channels"])
            output.d["signals/units"].extend(input.d["signals/units"])
        return output

    def assertion(self, output: SigContainer, inputs: Sequence[SigContainer]) -> None:
        assert all(c.signals.shape[1] == output.signals.shape[1] for c in inputs)
        if any(c.d["signals/fs"] != output.d["signals/fs"] for c in inputs):
            warn("Join operation on signals with incompatible frequencies")

class Concatenate(Joiner):
    def __init__(self, *branches, channel_names = None):
        super().__init__(*branches)
        self.names = channel_names

    def join(self, output: SigContainer, inputs: Sequence[SigContainer]) -> SigContainer:
        if self.names is not None:
            output.d["signals/channels"] = self.names
        for input in inputs:
            output.d["signals/data"] = np.hstack((output.signals, input.signals))
            if self.names is None:
                output.d["signals/channels"] = [
                    output.d["signals/channels"][i] + " + " + input.d["signals/channels"][i]
                        for i in range(len(output.d["signals/channels"]))]
        return output

    def assertion(self, output: SigContainer, inputs: Sequence[SigContainer]) -> None:
        assert all(c.signals.shape[0] == output.signals.shape[0] for c in inputs)
        if any(c.d["signals/fs"] != output.d["signals/fs"] for c in inputs):
            warn("Join operation on signals with incompatible frequencies")

class AssymetricJoiner(Joiner):
    def assertion(self, output: SigContainer, inputs: Sequence[SigContainer]) -> None:
        assert all(c.signals.shape[0] == output.signals.shape[0] for c in inputs)
        if any(c.d["signals/fs"] != output.d["signals/fs"] for c in inputs):
            warn("Join operation on signals with incompatible frequencies")

    def crossChannelNames(self, c1: SigContainer, c2: SigContainer) -> List[str]:
        return [
             f"{name1} x {name2}" if name1 !=name2 else f"{name1} (cross)"
             for name1, name2 in zip(c1.d["/signals/channels"], c2.d["/signals/channels"])
        ]

    def crossUnit(self, c1: SigContainer, c2: SigContainer) -> List[str]:
        return [
             f"{unit1} x {unit2}" if unit1 !=unit2 else f"{unit1}$^2$"
             for unit1, unit2 in zip(c1.d["/signals/units"], c2.d["/signals/units"])
        ]

class CrossCorrelate(AssymetricJoiner):
    def __init__(self, *branches, mode:str = 'full', method: str = 'auto'):
        super().__init__(*branches)
        self.mode = mode
        self.method = method

    def join(self, output: SigContainer, inputs: Sequence[SigContainer]) -> SigContainer:
        assert len(inputs) <= 1, "Cross corelation with more than two signal is not supported"
        in1 = output
        in2 = inputs[0] if inputs else in1
        result = np.vstack([
            correlate(in1.signals[i, :], in2.signals[i, :], self.mode, self.method)
            for i in range(in1.signals.shape[0])
        ])
        output.d["signals/data"] = result
        output.d["signals/channels"] = self.crossChannelNames(in1, in2)
        output.d["signals/units"] = self.crossUnit(in1, in2)
        output.d["signals/lag"] = output.sample_count // 2
        return output

class Convolve(AssymetricJoiner):
        def __init__(self, *branches, mode: str = 'full', method: str = 'auto'):
            super().__init__(*branches)
            self.mode = mode
            self.method = method

        def join(self, output: SigContainer, inputs: Sequence[SigContainer]) -> SigContainer:
            assert len(inputs) <= 1, "Convolution with more than two signal is not supported"
            in1 = output
            in2 = inputs[0] if inputs else in1
            result = np.vstack([
                convolve(in1.signals[i, :], in2.signals[i, :], self.mode, self.method)
                for i in range(in1.signals.shape[0])
            ])
            output.d["signals/data"] = result
            output.d["signals/lag"] = output.sample_count // 2
            return output


<gh_stars>0
import numpy as np
import netCDF4 as nc
import xarray as xr
from scipy.interpolate import griddata, interp1d
from salishsea_tools import (nc_tools, gsw_calls,viz_tools)
# import matplotlib.pyplot as plt
# import matplotlib.cm as cm

fname = '/home/ssahu/saurav/JP_BC/cat_42_days_T.nc'

gridT = xr.open_dataset(fname)

mask = nc.Dataset('/ocean/ssahu/CANYONS/wcvi/grid/meshmask_files/JP_mesh_mask.nc')

tmask_JP = mask.variables['tmask'][0,:,55:420,446:730]

votemper_JP = np.array(gridT['votemper'][1,:,55:420,446:730]) #Our 1st day of NEMO run (set in namelist and atmos files from that day)  is second day of data which starts from beginning of month
votemper_JP[...] = np.ma.masked_array(votemper_JP[...], mask = tmask_JP[...])
votemper_JP[votemper_JP == 0] =['Nan']

vosaline_JP = np.array(gridT['vosaline'][1,:,55:420,446:730])
vosaline_JP[...] = np.ma.masked_array(vosaline_JP[...], mask = tmask_JP[...])
vosaline_JP[vosaline_JP == 0] = ['Nan']

glamt_bc_JP = np.array(gridT['nav_lon'][55:420,446:730])
gphit_bc_JP = np.array(gridT['nav_lat'][55:420,446:730])
deptht_JP = np.array(gridT['deptht'][:])

print("Removing Nan values from JP's data, since we replaced the zero masks with Nans")

for i in np.arange(1,votemper_JP.shape[0]):
    for p in np.arange(votemper_JP.shape[1]):
        for l in np.arange(votemper_JP.shape[2]):
            if np.isnan(votemper_JP[i,p,l]):
                votemper_JP[i,p,l] = votemper_JP[i-1,p,l]
            else:
                continue

for i in np.arange(1,vosaline_JP.shape[0]):
    for p in np.arange(vosaline_JP.shape[1]):
        for l in np.arange(vosaline_JP.shape[2]):
            if np.isnan(vosaline_JP[i,p,l]):
                vosaline_JP[i,p,l] = vosaline_JP[i-1,p,l]
            else:
                continue

for i in np.arange(votemper_JP.shape[0]):
    for p in np.arange(votemper_JP.shape[1]):
        for l in np.arange(votemper_JP.shape[2]):
            if np.isnan(votemper_JP[i,p,l]):
                votemper_JP[i,p,l] = votemper_JP[i,p,l-1]
            else:
                continue

for i in np.arange(vosaline_JP.shape[0]):
    for p in np.arange(vosaline_JP.shape[1]):
        for l in np.arange(vosaline_JP.shape[2]):
            if np.isnan(vosaline_JP[i,p,l]):
                vosaline_JP[i,p,l] = vosaline_JP[i,p,l-1]
            else:
                continue

for i in np.arange(vosaline_JP.shape[1]):
    for j in np.arange(vosaline_JP.shape[2]):
        if np.isnan(vosaline_JP[0,i,j]):
            vosaline_JP[0,i,j] = vosaline_JP[1,i,j]
        else:
            continue

for i in np.arange(votemper_JP.shape[1]):
    for j in np.arange(votemper_JP.shape[2]):
        if np.isnan(votemper_JP[0,i,j]):
            votemper_JP[0,i,j] = votemper_JP[1,i,j]
        else:
            continue

#fname_wcvi = '/ocean/ssahu/CANYONS/wcvi/grid/coordinates.nc'

fname_wcvi = '/ocean/ssahu/CANYONS/wcvi/grid/coordinates_westcoast_seagrid_high_resolution_truncated_wider_west_bdy.nc'

with nc.Dataset(fname_wcvi, 'r') as coord:
    gphit_wcvi = coord.variables['gphit'][0,...];
    glamt_wcvi =  coord.variables['glamt'][0,...];

X = glamt_bc_JP.flatten();

Y = gphit_bc_JP.flatten();

points = (X[:],Y[:]);

xi = (glamt_wcvi.flatten(), gphit_wcvi.flatten());

votemper_ic = np.empty((votemper_JP.shape[0], glamt_wcvi.shape[0], glamt_wcvi.shape[1]));
vosaline_ic = np.empty((vosaline_JP.shape[0], glamt_wcvi.shape[0], glamt_wcvi.shape[1]));

for i in np.arange(votemper_JP.shape[0]):
    votemper_ic[i,...] = np.reshape(griddata(points, votemper_JP[i,...].flatten(), xi, method= 'linear'), glamt_wcvi.shape)
    vosaline_ic[i,...] = np.reshape(griddata(points, vosaline_JP[i,...].flatten(), xi, method= 'linear'), glamt_wcvi.shape)

print("Interpolation to WCVI horizontal points successful")

print("Calling GSW tools to convert to Conservative Temperature and Reference Salinity")

lat = np.empty_like(gphit_wcvi)
lon = np.empty_like(gphit_wcvi)
depth = np.empty_like(deptht_JP)

lat[:] = gphit_wcvi[:]
lon[:] = glamt_wcvi[:]

depth[:] = deptht_JP[:]

z = np.multiply(depth[:],-1)

votemper_PT = np.empty_like(votemper_ic)
vosaline_PSU =np.empty_like(vosaline_ic)

votemper_PT[:] = votemper_ic[:]
vosaline_PSU[:] = vosaline_ic[:]

pressure = np.empty(z.shape)
lats = np.empty(pressure.shape)
lons = np.empty(pressure.shape)
lats[:] = np.mean(lat)
lons[:] = np.mean(lon)

vosaline_SA = np.empty(vosaline_PSU.shape)
vosaline_RS = np.empty(vosaline_PSU.shape)
votemper_CT = np.empty(votemper_PT.shape)

pressure = gsw_calls.generic_gsw_caller('gsw_p_from_z.m', [z, np.mean(lat)])

print("Converted z to p: first GSW call successful")

for i in np.arange(vosaline_SA.shape[0]):
    vosaline_SA[i,...] = gsw_calls.generic_gsw_caller('gsw_SA_from_SP', [vosaline_PSU[i,...],pressure[i],lons[i],lats[i]])

print("Got SA from SP: GSW ran successfully inside loop")

vosaline_RS[:] = gsw_calls.generic_gsw_caller('gsw_SR_from_SP', [vosaline_PSU[:]])

print("Reference Salinity obtained from PS: one more GSW call left")

votemper_CT[:] = gsw_calls.generic_gsw_caller('gsw_CT_from_pt', [vosaline_SA[:], votemper_PT[:]])

print("GSW Calls successfull")

mask_wcvi = nc.Dataset('/ocean/ssahu/CANYONS/wcvi/grid/meshmask_files/mesh_mask_high_res.nc')

NEMO_depth = mask_wcvi.variables['gdept_0'][0,:,0,0]
tmask_WCVI = mask_wcvi.variables['tmask'][:]

salinity_function = interp1d(depth, vosaline_RS, axis = 0, bounds_error=False, fill_value='extrapolate')
temperature_function = interp1d(depth, votemper_CT, axis = 0, bounds_error=False, fill_value='extrapolate')

vosaline_NEMO = np.empty((NEMO_depth.shape[0], vosaline_RS.shape[1], vosaline_RS.shape[2]));
votemper_NEMO = np.empty((NEMO_depth.shape[0], vosaline_RS.shape[1], vosaline_RS.shape[2]));

for indx in np.arange(NEMO_depth.shape[0]):
    vosaline_NEMO[indx,...] = salinity_function(NEMO_depth[indx]);
    votemper_NEMO[indx,...] = temperature_function(NEMO_depth[indx]);

for i in np.arange(votemper_NEMO.shape[0]):
    for p in np.arange(votemper_NEMO.shape[1]):
        for l in np.arange(votemper_NEMO.shape[2]):
            if np.isnan(votemper_NEMO[i,p,l]):
                votemper_NEMO[i,p,l] = votemper_NEMO[i-1,p,l]
            else:
                continue

for i in np.arange(vosaline_NEMO.shape[0]):
    for p in np.arange(vosaline_NEMO.shape[1]):
        for l in np.arange(vosaline_NEMO.shape[2]):
            if np.isnan(vosaline_NEMO[i,p,l]):
                vosaline_NEMO[i,p,l] = vosaline_NEMO[i-1,p,l]
            else:
                continue

print("Vertical Interpolation to WCVI depth levels successful")

print("Now writing into a binary file to be used as IC for NEMO")

file_temp = nc.Dataset('/ocean/ssahu/CANYONS/wcvi/initial_conditions/West_coast_NEMO_IC_high_resolution_wide_bdy.nc', 'w', zlib=True)
# dataset attributes
nc_tools.init_dataset_attrs(
    file_temp,
    title='Temperature and salinity Initial Condition',
    notebook_name='Making_IC_from_JP',
    nc_filepath='/ocean/ssahu/CANYONS/wcvi/initial_conditions/West_coast_NEMO_IC_high_resolution.nc',
    comment='Temperature and salinity from JP Model, high_resolution__grid; used at all grid points and interpolated vertically')

file_temp.createDimension('xb', votemper_NEMO.shape[2]);
file_temp.createDimension('yb', votemper_NEMO.shape[1]);
file_temp.createDimension('deptht', votemper_NEMO.shape[0]);
file_temp.createDimension('time_counter', None);

nav_lat = file_temp.createVariable('nav_lat', 'float32', ('yb','xb'));
nav_lat.long_name = 'Latitude';
nav_lat.units = 'degrees_north';

nav_lon = file_temp.createVariable('nav_lon', 'float32', ('yb','xb'));
nav_lon.long_name = 'Longitude';
nav_lon.units = 'degrees_east';

deptht = file_temp.createVariable('deptht', 'float32', ('deptht'));
deptht.long_name = 'Vertical T Levels';
deptht.units = 'm';
deptht.positive = 'down';

time_counter = file_temp.createVariable('time_counter', 'float32', ('time_counter'));
time_counter.units = 's';
time_counter.long_name = 'time';
time_counter.calendar = 'noleap';

votemper = file_temp.createVariable('votemper', 'float32', ('time_counter','deptht','yb','xb'));
votemper.units = 'degC'
votemper.long_name = 'Conservative Temperature (CT)';
votemper.grid = 'WCVI';

vosaline = file_temp.createVariable('vosaline', 'float32', ('time_counter','deptht','yb','xb'));
vosaline.units = 'g/Kg';
vosaline.long_name = 'Reference Salinity (SR)';
vosaline.grid = 'WCVI';

nav_lat[:] = lat[:];
nav_lon[:] = lon[:];
deptht[:] = NEMO_depth[:];
time_counter[0] = 1;
votemper[0,...] = votemper_NEMO[:]
vosaline[0,...]= vosaline_NEMO[:]

file_temp.close()

print("File written: Thank you")


<reponame>johertrich/Proximal_Neural_Networks<filename>conv_denoiser/test_conv_denoiser_BSD68.py
# This code belongs to the paper
#
# <NAME>, <NAME> and <NAME>.
# Convolutional Proximal Neural Networks and Plug-and-Play Algorithms.
# Linear Algebra and its Applications, vol 631 pp. 203-234, 2021.
#
# Please cite the paper if you use this code.
#
from core.stiefel_network import *
from core.layers import *
import numpy as np
import numpy.random
from PIL import Image
from conv_denoiser.readBSD import *
import bm3d
import pickle
from scipy.io import savemat

def run(scale=1.99,noise_level=25./255.,num=None,pretrained_weights=True):
    # declare network
    act=tf.keras.activations.relu
    num_filters=64
    max_dim=128
    num_layers=8
    sizes=[None]*(num_layers)
    conv_shapes=[(num_filters,max_dim)]*num_layers
    filter_length=5
    if scale is None:
        model=StiefelModel(sizes,None,convolutional=True,filter_length=filter_length,dim=2,conv_shapes=conv_shapes,activation=act,scale_layer=False)
    else:
        model=StiefelModel(sizes,None,convolutional=True,filter_length=filter_length,dim=2,conv_shapes=conv_shapes,activation=act,scale_layer=scale)

    pred=model(tf.random.normal((10,40,40)))
    model.fast_execution=True

    if scale is None:
        # load weights
        if pretrained_weights:
            file_name='data/pretrained_weights/free_noise_level'+str(noise_level)+'.pickle'
        else:
            if num is None:
                file_name='results_conv/free_noise_level'+str(noise_level)+'/adam.pickle'
            else:
                file_name='results_conv/free_noise_level'+str(noise_level)+'/adam'+str(num)+'.pickle'
        with open(file_name,'rb') as f:
            trainable_vars=pickle.load(f)
        for i in range(len(model.trainable_variables)):
            model.trainable_variables[i].assign(trainable_vars[i])
    else:
        # load weights
        if pretrained_weights:
            file_name='data/pretrained_weights/scale'+str(scale)+'_noise_level'+str(noise_level)+'.pickle'
        else:
            if num is None:
                file_name='results_conv/scale'+str(scale)+'_noise_level'+str(noise_level)+'/adam.pickle'
            else:
                file_name='results_conv/scale'+str(scale)+'_noise_level'+str(noise_level)+'/adam'+str(num)+'.pickle'
        with open(file_name,'rb') as f:
            trainable_vars=pickle.load(f)
        for i in range(len(model.trainable_variables)):
            model.trainable_variables[i].assign(trainable_vars[i])
        beta=1e8
        project=True
        if project:
            # projection of the convolution matrices onto the Stiefel manifold
            for i in range(len(model.stiefel)):
                convs=model.stiefel[i].convs
                smaller=convs.shape[0]<convs.shape[1]
                if smaller:
                    convs=transpose_convs(convs)
                iden=np.zeros((convs.shape[1],convs.shape[1],4*filter_length+1,4*filter_length+1),dtype=np.float32)
                for j in range(convs.shape[1]):
                    iden[j,j,2*filter_length,2*filter_length]=1
                iden=tf.constant(iden)
                C=tf.identity(convs)
                def projection_objective(C):
                    return 0.5*beta*tf.reduce_sum((conv_mult(transpose_convs(C),C)-iden)**2)+.5*tf.reduce_sum((C-convs)**2)
                for iteration in range(100):
                    with tf.GradientTape(persistent=True) as tape:
                        tape.watch(C)
                        val=projection_objective(C)
                        grad=tape.gradient(val,C)
                        grad_sum=tf.reduce_sum(grad*grad)
                    hess=tape.gradient(grad_sum,C)
                    hess*=0.5/tf.sqrt(grad_sum)
                    C-=grad/tf.sqrt(tf.reduce_sum(hess*hess))
                if smaller:
                    C=transpose_convs(C)
                model.stiefel[i].convs.assign(C)

    np.set_printoptions(threshold=sys.maxsize)

    # set parameters
    test_directory='data/BSD68'
    fileList=os.listdir(test_directory+'/')
    fileList.sort()
    img_names=fileList
    sig=25.
    sig/=255.
    if scale is None:
        save_path='results_conv/denoise_results_free_noise_level'+str(noise_level)
    else:
        save_path='results_conv/denoise_results_scale'+str(scale)+'_noise_level'+str(noise_level)
    if not os.path.isdir('results_conv'):
        os.mkdir('results_conv')
    if not os.path.isdir(save_path):
        os.mkdir(save_path)
    if not os.path.isdir(save_path+'/noisy_data'):
        os.mkdir(save_path+'/noisy_data')
    if not os.path.isdir(save_path+'/other'):
        os.mkdir(save_path+'/other')
    psnr_sum=0.
    psnr_noisy_sum=0.
    psnr_bm3d_sum=0.
    error_sum=0.
    error_bm3d_sum=0.
    counter=0
    my_shift=10
    myfile=open(save_path+"/psnrs.txt","w")
    myfile.write("PSNRs:\n")
    myfile.close()
    np.random.seed(30)
    # Denoising
    for name in img_names:
        counter+=1
        img=Image.open(test_directory+'/'+name)
        img=img.convert('L')
        img_gray=1.0*np.array(img)
        img_gray/=255.0
        img_gray-=.5
        img_gray_noisy=img_gray+sig*np.random.normal(size=img_gray.shape)
        savemat(save_path+'/noisy_data/'+name[:-4]+'_noisy.mat',{'noisy': (img_gray_noisy+0.5)*255})
        img_gray_pil=Image.fromarray((img_gray+.5)*255.0)
        img_gray_pil=img_gray_pil.convert('RGB')
        img_gray_pil.save(save_path+'/original'+name)
        pred=model(tf.expand_dims(tf.constant(img_gray_noisy,dtype=tf.float32),0))
        noisy=(img_gray_noisy+.5)*255
        reconstructed=((img_gray_noisy-np.reshape(pred.numpy(),img_gray.shape))+.5)*255.
        img_gray=(img_gray+.5)*255.
        res_bm3d=(bm3d.bm3d(img_gray_noisy,sig)+0.5)*255.
        error_sum+=tf.reduce_sum(((reconstructed-img_gray)/255.)**2).numpy()
        error_bm3d_sum+=tf.reduce_sum(((res_bm3d-img_gray)/255.)**2).numpy()
        psnr=meanPSNR(tf.keras.backend.flatten(reconstructed).numpy()/255.0,tf.keras.backend.flatten(img_gray).numpy()/255.0,one_dist=True)
        psnr_noisy=meanPSNR(tf.keras.backend.flatten(noisy).numpy()/255.0,tf.keras.backend.flatten(img_gray).numpy()/255.0,one_dist=True)
        psnr_bm3d=meanPSNR(tf.keras.backend.flatten(res_bm3d).numpy()/255.0,tf.keras.backend.flatten(img_gray).numpy()/255.0,one_dist=True)
        print('PSNR of '+name+':                    '+str(psnr))
        print('PSNR of BM3D '+name+':               '+str(psnr_bm3d))
        print('PSNR of noisy '+name+':              '+str(psnr_noisy))
        psnr_sum+=psnr
        psnr_noisy_sum+=psnr_noisy
        psnr_bm3d_sum+=psnr_bm3d
        print('Mean PSNR PPNN:                      '+str(psnr_sum/counter))
        print('Mean PSNR BM3D:                      '+str(psnr_bm3d_sum/counter))
        print('Mean Error PPNN:                     '+str(error_sum/counter))
        print('Mean Error BM3D:                     '+str(error_bm3d_sum/counter))
        print('Mean PSNR noisy:                     '+str(psnr_noisy_sum/counter))
        myfile=open(save_path+"/psnrs.txt","a")
        myfile.write('PSNR of '+name+':                    '+str(psnr)+'\n')
        myfile.write('PSNR of BM3D '+name+':               '+str(psnr_bm3d)+'\n')
        myfile.write('PSNR of noisy '+name+':              '+str(psnr_noisy)+'\n')
        myfile.close()
        img=Image.fromarray(noisy)
        img=img.convert('RGB')
        img.save(save_path+'/noisy'+name)
        img=Image.fromarray(reconstructed)
        img=img.convert('RGB')
        img.save(save_path+'/reconstructed'+name)
        img=Image.fromarray(res_bm3d)
        img=img.convert('RGB')
        img.save(save_path+'/other/BM3D'+name)
    print('Mean PSNR on images: '+str(psnr_sum/len(img_names)))
    print('Mean PSNR on noisy images: '+str(psnr_noisy_sum/len(img_names)))


# -*- coding: utf-8 -*-
"""
Created on Mon Dec 04 10:13:10 2017

@author: PJ
"""

import numpy as np
import matplotlib.pyplot as plt
import scipy.io as sio

def generate(popSize,m):
    '''
    generate a random Individual for HLO
    '''
    Bpop = np.random.randint(2,size=(popSize,m))
    return Bpop

def B2R(popSize,Bpop,dim,bit,bound):
    '''
    convert binary indivdual to real code format
    '''
    Rpop = np.zeros((popSize,dim))
    for i in range(0,popSize):
        for d in range(0,dim):
            for j in range(d*bit,(d+1)*bit):
                Rpop[i,d] = Rpop[i,d] + (Bpop[i,j]*2**((d+1)*bit-j-1))
            Rpop[i,d] = bound[0]+(bound[1]-bound[0])*Rpop[i,d]/(2**bit-1)
    return Rpop

def Evfit(Rpop,dim):
    """
    sphere function
    """
    o=sio.loadmat('sphere_func_data')
    o = o['o'].reshape(-1)
    f = np.zeros((len(Rpop)))
    for i in range(0,len(Rpop)):
        for j in range(0,dim):
            f[i] = f[i] + (Rpop[i,j]-o[j])**2
    return f

def printResult(Gmax,trace):
    '''
    plot the result of the HLO
    '''
    x = np.arange(0, Gmax)
    y1 = trace[:, 0]
    plt.plot(x, y1, 'r', label='optimal value')
    plt.xlabel("Iteration")
    plt.ylabel("function value")
    plt.title("HLO for function optimization")
    plt.legend()
    plt.show()


<filename>intrinsic/outlier_detection.py
from numpy import mean, std, sqrt, log, exp
from scipy.stats import norm, shapiro, anderson
import math

# (nj) TODO: Come up with different (better?) ways of
# (1) Combining "probabilities" for a single passage. Right now we just mulitply
# the probabilities we get from each of the features. Done in:
# _combine_feature_probs(prob_vector)
# (2) Building a "confidence" score based on the relative "probability" of
# a passage being plagiarized vs. not-plagiarized. Done in:
# _get_confidence(plag_prob, non_plag_prob)
# (3) Scaling the "confidence" scores across an entire document (set of
# stylometric features). Done in:
# _scale_confidences(confs)

MIN_PROB = 10**(-30)
# The confidence we use when we are unable to calculate anything about the
# likelihood of plag
IMPURITY_ASSUMPTION = .2

def density_based(stylo_vectors, center_at_mean=True, num_to_ignore=1, impurity=.2, feature_confidence_weights=None):
    '''
    Implements the algorithm described in <NAME>'s
    "Intrinsic Plagiarism Analysis", Section 2.4 "Outlier Detection"

    Estimates the distribution of each feature's non-plag. portion as a normal distribution
    simply using the MLE (i.e. a normal distribution with sample mean
    and sample std as parameters). Note that the min and max of the observed
    features are removed before computing sample mean/std (and this could perhaps
    be extended to remove the <k> largest/smallest observations before computing
    sample mean/std)

    if <center_at_mean>, then normal dist is centered at feature's mean
    if not, then normal dist at feature's median

    When calculating the mean, the minimum and maximum <num_to_ignore> elements
    are ignored.

    NOTE <impurity> argument is ignored right now.
    '''
    transpose = _rotate_vectors(stylo_vectors)
    confidences = []

    means, stds, mins, medians, maxs = [], [], [], [], []
    normality_pvals = []

    for row in transpose:
        cur_mean, cur_std, cur_min, cur_median, cur_max = _get_distribution_features(row, num_to_ignore)
        means.append(cur_mean)
        stds.append(cur_std)
        mins.append(cur_min)
        medians.append(cur_median)
        maxs.append(cur_max)

        # Print some data about the normality of our features
        norm_p = _test_normality(row)
        normality_pvals.append(norm_p)

    #print 'Normality pvals have min, max, mean:', min(normality_pvals), max(normality_pvals), mean(normality_pvals)

    for i in xrange(len(stylo_vectors)):
        vec = stylo_vectors[i]
        # For current <vec>,
        # featurewise_plag_prob[i] == prob. that feature <i> was plagiarized in <vec>
        # NOTE that these are taken from PDFs, so they don't actually correspond
        # to real probabilities
        featurewise_nonplag_prob = []
        featurewise_plag_prob = []
        featurewise_confs = []

        for feat_num in xrange(len(vec)):
            # TODO plag_prob is just constant -- precompute this
            cur_val = vec[feat_num]
            cur_center = means[feat_num] if center_at_mean else medians[feat_num]
            cur_std = stds[feat_num]
            cur_min, cur_max = mins[feat_num], maxs[feat_num]

            # A std of 0 => the feature is constant, and therefore won't
            # help us distinguish anything!
            if cur_std != 0.0 and not _in_uncertainty_interval(cur_val, cur_center, cur_std):

                cur_norm_prob = _get_norm_prob(cur_val, cur_center, cur_std)
                if math.isnan(cur_norm_prob) or cur_norm_prob == 0.0:
                    cur_norm_prob = MIN_PROB
                    #print 'Norm prob was nan or 0: %f. Using MIN_PROB' % cur_norm_prob

                cur_unif_prob = _get_unif_prob(cur_val, cur_min, cur_max)
                if math.isnan(cur_unif_prob) or cur_unif_prob == 0.0:
                    cur_unif_prob = MIN_PROB
                    print 'Unif prob was nan or 0: %f. Using MIN_PROB' % cur_unif_prob

                featurewise_confs.append(_get_confidence(cur_unif_prob, cur_norm_prob))
            # TODO what happens if all points are in uncertainty interval??
        confidences.append(_combine_feature_probs(featurewise_confs))

    # No more scaling -- we're already returning probs between 0 and 1, so should be all good
    # scaled = _scale_confidences(confidences)
    # return scaled

    return confidences

def _combine_feature_probs(prob_vector):
    '''
    Returns the Naive Bayes version of combining probabilites: just multiply them.
    Note that we take the sum of the log of each probability and exponentiate
    in hopes of avoiding underflow when multiplying many small numbers

    If the probability vector is empty, then all features were in the uncertainty
    interval. Return the IMPURITY_ASSUMPTION (unless something else logical comes
    about)
    '''
    if len(prob_vector) == 0:
        return IMPURITY_ASSUMPTION
    else:
        return exp(sum(log(prob_vector)))

def _get_confidence(plag_prob, non_plag_prob):
    '''
    Returns some notion of confidence:
    If we think there's plag., return the Naive Bayes estimated prob of plag
    If not, return the negative of the Naive Bayes estimate prob of NOT plag

    Other options:
    if plag_prob > non_plag_prob:
        return plag_prob
    else:
        return -non_plag_prob

    The above is an old notion of confidence, which eventually gets
    scaled. It worked well at the start, but also doesn't make a ton
    of sense...

    Note that these values are scaled later on to be between 0 and 1
    '''
    return plag_prob / (plag_prob + non_plag_prob)

def _scale_confidences(confs):
    '''
    NOT BEING USED ANYMORE. The way we combine combine plag/non-plag probs
    should already give confs between 0 and 1.

    Scales all "confidences" to (0, 1) interval simply by dividing by
    the maximum "confidence"

    If <confs> is constant (i.e. contains all the same values), return
    a vector of all IMPURITY_ASSUMPTIONs.
    '''
    # offset will be either 0 or some negative number, in which case
    # we subtract the negative offset (i.e. add)
    offset = min(min(confs), 0.0)
    max_conf_with_offset = max(confs) - offset

    if max_conf_with_offset == 0.0:
        return [IMPURITY_ASSUMPTION for x in confs]
    else:
        return [(x - offset) / max_conf_with_offset for x in confs]

def _get_distribution_features(row, extremes_to_ignore):
    '''
    <row> corresponds to all the observed values of feature

    Per Section 2.4 in Stein, Lipka, Prettenhofer's "Intrinsic Plagiarism Analysis",
    removes min and max when calculating parameters of Gaussian distribution
    '''
    sorted_row = sorted(row)

    min_val = sorted_row[0]
    max_val = sorted_row[-1]
    median_val = sorted_row[len(sorted_row) / 2]

    row_without_extremes = sorted_row[extremes_to_ignore : -extremes_to_ignore]
    mean_of_row = mean(row_without_extremes)
    std_of_row = std(row_without_extremes)

    return mean_of_row, std_of_row, min_val, median_val, max_val

def _in_uncertainty_interval(x, center, sd_dist):
    '''
    Returns whether <x> is between 1 and 2 standard deviations away from
    <center> in either direction. Such points are dubbed "uncertain"
    in Stein et. al.'s paper
    '''
    right_interval = (center + sd_dist, center + 2 * sd_dist)
    left_interval = (center - 2 * sd_dist, center - sd_dist)

    if right_interval[0] <= x <= right_interval[1] or \
        left_interval[0] <= x <= left_interval[1]:
        return True

    return False

def _get_norm_prob(x, loc, scale):
    '''
    Returns normal PDF evaluated at <x> for a normal dist centered at <loc>
    with SD == <scale>
    '''
    return norm.pdf(x, loc, scale)

def _get_unif_prob(v, min_val, max_val):
    '''
    Returns uniform PDF (probability density function)
    '''
    diff = max_val - min_val
    if diff == 0:
        return 1.0
    else:
        return 1.0 / (max_val - min_val)

def _rotate_vectors(mat):
    '''
    This is more verbose than it needs to be, but also perhaps
    more readable. Transposes <mat> and returns it
    '''
    rotated = []

    for col_num in range(len(mat[0])):
        # Append one column at a time
        rotated.append([mat[row_num][col_num] for row_num in range(len(mat))])

    return rotated

def _test_normality(data, threshold=.05):
    '''
    '''
    if len(data) < 4:
        return 0
    else:
        test_stat, pval = shapiro(data)

        return pval

def _test():
    '''
    '''
    dim1_mean = 0
    dim1_std = 3

    dim2_mean = 20
    dim2_std = 3
    # 20 normal RVs with mean=0, std=
    dim1 = list(norm.rvs(dim1_mean, dim1_std, size=20))
    # Add a couple of obvious outliers
    dim1.append(-10)
    dim1.append(10)

    dim2 = list(norm.rvs(dim2_mean, dim2_std, size=20))
    dim2.append(10)
    dim2.append(30)

    data = zip(dim1, dim2)

    confs = density_based(data)

    print 'Dim1 params:', dim1_mean, dim1_std
    print 'Dim2 params:', dim2_mean, dim2_std
    for d, conf in zip(data, confs):
        print d, conf

if __name__ == '__main__':
    _test()


#!/usr/bin/python
# -*- coding: latin-1 -*-

import numpy as np

# Definimos una función senoidal simple.
def ondasimple(t):
  A = 1.0    # amplitud
  f = 500.0  # frequencia
  Phi = 0.0  # fase
  return A * np.sin(2 * np.pi * f * t + Phi)

# Generamos 16000 puntos a 16kHz.
ts = np.arange(16000.0) / 16000.0

# Armamos una onda senoidal discretizada.
mionda = []
for t in ts:
  mionda.append(ondasimple(t))
mionda = np.array(mionda)

# Graficamos la onda.
import matplotlib.pyplot as pyplot
pyplot.clf()
pyplot.plot(ts[0:100], mionda[0:100])
pyplot.savefig('mionda.png')

# La guardamos como wav.
import scipy.io.wavfile
#MAL: wavdata = np.array(mionda, dtype=np.int16) * 10000
wavdata = np.array(mionda * 10000.0, dtype=np.int16)
scipy.io.wavfile.write('mionda.wav', 16000, wavdata)

# Ejercicios:
#
# 1. Generar un archivo wav para cada nota musical Do, Re, Mi,
#    Fa, Sol, La, Si. Consultar las frecuencias en
#    http://www.phy.mtu.edu/~suits/notefreqs.html
#    Tomar como referencia La = 440Hz.
#
# 2. Buscar la frecuencia más aguda y más grave que pueden percibir.
#
# 3. Percepcion relativa. Escuchar la diferencia entre dos tonos graves
#    separados por 100Hz (ej: 200 y 300Hz) y dos tonos agudos separados
#    también por 100Hz (ej: 1200 y 1300Hz).
#
# 4. Crear una onda cuadrada a 500 Hz, modificando ondasimple(t) de modo
#    que devuelva solamente 1 o -1. Generar un wav y comparar con una
#    senoidal de la misma frecuencia.
#
# 5. Repetir el punto anterior para 100Hz y para 1000Hz. ¿En algún caso
#    suenan parecidas las ondas senoidales y cuadradas? (Más allá de las
#    diferencias de volumen).


<gh_stars>0
import numpy as np
from scipy.stats import zscore
import warnings

def cofluct(X, quantiles: tuple = (0, 1), return_mat=True):
    """
    Computes cofluctuation time-series (per edge) for a nodes x timepoints matrix X.
    Based on https://www.pnas.org/content/early/2020/10/21/2005531117

    Parameters
    ----------
    quantiles: tuple,default=(0,1)
        list of lowest/highest quantile of edge events to use
        [0, 1]: all events = pearson corr; [0, .05]: bottom 5% of events; [.95, 1]: top 5% of events
    return_mat: bool,default=True
        Whether to return a connectivity matrix (True) or dictionary (False). The dict edge contains cofluctuation
        time-series (pairs_of_nodes x timepoints) and event timeseries.

    Returns
    -------
    float
        edge cofluctuation time-series dict (pairs_of_nodes x timepoints) and event timeseries as dict

    """
    X = np.float32(X)
    # get z-value time series (mean 0, std 1)
    z_ts = zscore(X, axis=1)

    # multiply z-scored time series to get edge time series (co-fluctuation)
    edge_ids = np.triu_indices(X.shape[0], k=1)
    edge_ts = np.einsum('ik,jk->ijk', z_ts, z_ts)[edge_ids]     # get upper triangular values

    # get event_ts as Root Mean Square over nodes
    event_ts = np.sqrt(np.mean(np.square(edge_ts), axis=0))

    # get correlations for quantile intervals of events
    q_low = np.quantile(np.abs(event_ts), q=quantiles[0])
    q_high = np.quantile(np.abs(event_ts), q=quantiles[1])
    edge_ts = edge_ts[:, (event_ts >= q_low) & (event_ts <= q_high)]

    # compute correlation/connectivity vector or matrix
    cors = np.mean(edge_ts, axis=1)
    if return_mat:
        tmp = np.eye(X.shape[0])
        tmp[np.triu_indices(tmp.shape[0], k=1)] = cors
        cors = tmp + tmp.T
        np.fill_diagonal(cors, 1)

    # return as dict
    else:
        warnings.warn("dicts can not be handled by photonai."
                      "Dicts are implemented for use outside of Photon.")
        edge_ids_str = [str(m) + '_' + str(n) for m, n in zip(edge_ids[0], edge_ids[1])]
        edge_ts_dict = {key: val for val, key in zip(edge_ts, edge_ids_str)}
        edge_ts_dict['event_ts'] = event_ts
        return edge_ts_dict

    return cors


# -*- coding: utf-8 -*-
""" Code for training and evaluating Self-Explaining Neural Networks.
Copyright (C) 2018 <NAME> <<EMAIL>>

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License,
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program. If not, see <https://www.gnu.org/licenses/>.
"""

import os
import pdb
import numpy as np
from scipy import optimize

import matplotlib.pyplot as plt
import pprint # For feature explainer

import torch
from torch.autograd import Variable
from torchvision.utils import make_grid
import matplotlib.gridspec as gridspec

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

def generate_dir_names(dataset, args, make = True):
    if args.h_type == 'input':
        suffix = '{}_H{}_Th{}_Reg{:0.0e}_LR{}'.format(
                    args.theta_reg_type,
                    args.h_type,
                    args.theta_arch,
                    args.theta_reg_lambda,
                    args.lr,
                    )
    else:
        suffix = '{}_H{}_Th{}_Cpts{}_Reg{:0.0e}_Sp{}_LR{}'.format(
                    args.theta_reg_type,
                    args.h_type,
                    args.theta_arch,
                    args.nconcepts,
                    args.theta_reg_lambda,
                    args.h_sparsity,
                    args.lr,
                    )

    model_path     = os.path.join(args.model_path, dataset, suffix)
    log_path       = os.path.join(args.log_path, dataset, suffix)
    results_path   = os.path.join(args.results_path, dataset, suffix)

    if make:
        for p in [model_path, results_path]: #, log_path,
            if not os.path.exists(p):
                os.makedirs(p)

    return model_path, log_path, results_path

class AverageMeter(object):
    """Computes and stores the average and current value"""
    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count

### Animation Utils

# animation function
def animate_fn(i, xx,yy,Ts, Cs):
    t = Ts[i]
    C = Cs[i]
    cont = plt.contourf(xx, yy, C, 25, cmap = plt.cm.RdBu)
    plt.title(r'Iter = %i' % t)
    return cont

def animate_training(Steps, Cs, X_train, y_train):
    cm = plt.cm.RdBu
    cm_bright = ListedColormap(['#FF0000', '#0000FF'])
    h = .02  # step size in the mesh
    x_min, x_max = X_train[:, 0].min() - .5, X_train[:, 0].max() + .5
    y_min, y_max = X_train[:, 1].min() - .5, X_train[:, 1].max() + .5
    xx, yy = np.meshgrid(np.arange(x_min, x_max, h),
                         np.arange(y_min, y_max, h))

    fig = plt.figure()
    ax = fig.add_subplot(111)
    line, = ax.plot([],[], '-')
    line2, = ax.plot([],[],'--')
    ax.set_xlim(np.min(xx), np.max(xx))
    ax.set_xlim(np.min(yy), np.max(yy))
    ax.scatter(X_train[:, 0], X_train[:, 1], c=y_train, cmap=cm_bright,
       edgecolors='k')
    #ax.contourf(xx, yy, Cs[0].reshape(xx.shape))

    anim = animation.FuncAnimation(fig, animate_fn, frames = len(Steps), fargs = (xx,yy,Steps,Cs,), interval = 200, blit = False)

    return anim

# Got these two from scikit learn embedding example
def make_meshgrid(x, y, h=.02):
    """Create a mesh of points to plot in

    Parameters
    ----------
    x: data to base x-axis meshgrid on
    y: data to base y-axis meshgrid on
    h: stepsize for meshgrid, optional

    Returns
    -------
    xx, yy : ndarray
    """
    x_min, x_max = x.min() - 1, x.max() + 1
    y_min, y_max = y.min() - 1, y.max() + 1
    xx, yy = np.meshgrid(np.arange(x_min, x_max, h),
                         np.arange(y_min, y_max, h))
    return xx, yy

def plot_contours(model, X, **params):
    """Plot the decision boundaries for a classifier.

    Parameters
    ----------
    ax: matplotlib axes object
    clf: a classifier
    xx: meshgrid ndarray
    yy: meshgrid ndarray
    params: dictionary of params to pass to contourf, optional
    """
    fig, ax = plt.subplots()
    Z = model(X)
    Z = Z.data.numpy()
    C = np.argmax(Z,axis=1)
    C = C.reshape(xx.shape)
    out = ax.contourf(xx, yy, C, **params)
    return out

def plot_embedding(X,y,Xp, title=None):
    """ Scale and visualize the embedding vectors """
    x_min, x_max = np.min(Xp, 0), np.max(Xp, 0)
    Xp = (Xp - x_min) / (x_max - x_min)

    plt.figure(figsize=(20,10))
    ax = plt.subplot(111)
    for i in range(Xp.shape[0]):
        plt.text(Xp[i, 0], Xp[i, 1], str(y[i]),
                 color=plt.cm.Set1(y[i] / 10.),
                 fontdict={'weight': 'bold', 'size': 9})

    if hasattr(offsetbox, 'AnnotationBbox'):
        # only print thumbnails with matplotlib > 1.0
        shown_images = np.array([[1., 1.]])  # just something big
        for i in range(X.shape[0]):
            dist = np.sum((Xp[i] - shown_images) ** 2, 1)
            if np.min(dist) < 4e-3:
                # don't show points that are too close
                continue
            shown_images = np.r_[shown_images, [Xp[i]]]
            imagebox = offsetbox.AnnotationBbox(
                offsetbox.OffsetImage(X[i,:].reshape(28,28), cmap=plt.cm.gray_r),
                Xp[i])
            ax.add_artist(imagebox)
    plt.xticks([]), plt.yticks([])
    if title is not None:
        plt.title(title)

def _explain_class(model, x_raw, x,k,typ='pos',thresh = 0.5,recompute=True):
    """
        Given an input x and class index k, explain f(x) by returning indices of
        features in x that have highest positive impact on predicting class k.
    """
    if recompute:
        y = model(x) # y = self.model(x)
    B_k = model.params[0,k,:].data.numpy()
    if typ == 'pos':
        Mask = (B_k > thresh).astype(np.int).reshape(x.size()).squeeze()
    elif typ == 'neg':
        Mask = (B_k < -thresh).astype(np.int).reshape(x.size()).squeeze()
    else:
        # Return weights instead of mask
        return B_k.reshape(x.size()).squeeze()
    Masked_x = Mask*x_raw.numpy().squeeze()
    return Masked_x

def explain_digit(model, x_raw, x, thresh = 0.5, save_path = None):
    """
        Given an input x, explain f(x) by returning indices of
        features in x that have highest positive impact on predicting each class.

        x_raw is passed for plotting purposes
    """
    plt.imshow(x_raw.squeeze().numpy())
    plt.title('Input:')
    plt.xticks([])
    plt.yticks([])
    if save_path:
        plt.savefig(save_path+'_input.pdf',  bbox_inches = 'tight', format='pdf', dpi=300)
    plt.show()
    y_pred = model(x)

    pred_class = np.argmax(y_pred.data.numpy())
    print('Predicted: ',pred_class)

    fig, ax = plt.subplots(3,model.dout,figsize=(1.5*model.dout,1.5*3))
    for i in range(model.dout):
        #print('Class {}:'.format(i))

        # Positive
        x_imask = _explain_class(model, x_raw, x,i,typ ='pos', recompute=False, thresh = thresh)
        ax[0,i].imshow(x_imask)
        ax[0,i].set_xticks([])
        ax[0,i].set_yticks([])
        ax[0,i].set_title('Class: {}'.format(i))

        # Negative
        x_imask = _explain_class(model, x_raw, x,i,  typ ='neg', recompute=False, thresh = thresh)
        ax[1,i].imshow(x_imask)
        ax[1,i].set_xticks([])
        ax[1,i].set_yticks([])

        # Combined
        x_imask = _explain_class(model, x_raw, x,i,  typ ='both', recompute=False)
        ax[2,i].imshow(x_imask, cmap = plt.cm.RdBu)
        ax[2,i].set_xticks([])
        ax[2,i].set_yticks([])
        #print(np.linalg.norm(x_imask))
        #print(x_imask[:5,:5])

        if i == 0:
            ax[0,0].set_ylabel('Pos. Feats.')
            ax[1,0].set_ylabel('Neg. Feats.')
            ax[2,0].set_ylabel('Combined')

    if save_path:
        plt.savefig(save_path + '_expl.pdf', bbox_inches = 'tight', format='pdf', dpi=300)
    plt.show()

def plot_text_explanation(words, values, n_cols = 6, save_path = None):
    import seaborn as sns
    # Get some pastel shades for the colors
    #colors = plt.cm.BuPu(np.linspace(0, 0.5, len(rows)))
    cmap = sns.cubehelix_palette(light=1, as_cmap=True)

    n_rows = int(min(len(values), len(words)) / n_cols) + 1

    # Plot bars and create text labels for the table
    if type(words) is str:
        words = words.split(' ')

    cellcolours = np.empty((n_rows, n_cols), dtype='object')
    celltext    = np.empty((n_rows, n_cols), dtype='object')

    for r in range(n_rows):
        for c in range(n_cols):
            idx = (r * n_cols + c)
            val =  values[idx] if (idx < len(values)) else 0
            cellcolours[r,c] = cmap(val)
            celltext[r,c] = words[idx] if (idx < len(words)) else ''

    fig, ax = plt.subplots()#figsize=(n_cols, n_rows))

    # Hide axes
    ax.xaxis.set_visible(False)
    ax.yaxis.set_visible(False)

    # Add a table at the bottom of the axes
    tab = plt.table(cellText=celltext,
                          cellColours = cellcolours,
                          rowLabels=None,
                          rowColours=None,
                          colLabels=None,
                          cellLoc='center',
                          loc='center')

    for key, cell in tab.get_celld().items():
        cell.set_linewidth(0)

    tab.set_fontsize(14)
    tab.scale(1.5, 1.5)  # may help

    # Adjust layout to make room for the table:
    #plt.subplots_adjust(left=0.2, bottom=0.2)

    #plt.ylabel("Loss in ${0}'s".format(value_increment))
    plt.yticks([])
    plt.xticks([])
    plt.title('')
    plt.axis('off')
    plt.grid('off')
    if save_path:
        plt.savefig(save_path + '_expl.pdf', bbox_inches = 'tight', format='pdf', dpi=300)
    plt.show()

class FeatureInput_Explainer():
    """
        Explainer for classification task models that take vector of features
        as input and and return class probs.

        Arguments:

    """
    def __init__(self, feature_names, binary = False, sort_rows = True, scale_values = True):
        super(FeatureInput_Explainer, self).__init__()
        self.features = feature_names
        self.binary   = binary # Whether it is a binary classif task
        self.sort_rows = sort_rows
        self.scale_values = scale_values

    def explain(self, model, x, thresh = 0.5, save_path = None):
        np.set_printoptions(threshold=15, precision = 2)
        #print('input: {}'.format(x.data.numpy()))
        print('Input:')
        pp.pprint(dict(zip(self.features, *x.data.numpy())))
        print('')

        np.set_printoptions()
        y_pred = model(x)
        pred_class = np.argmax(y_pred.data.numpy())
        print('Predicted: ',pred_class)

        # Get data-dependent params
        B = model.thetas[0,:,:].data.numpy() # class x feats

        Pos_Mask = (B > thresh).astype(np.int)#.reshape(x.size()).squeeze()
        Neg_Mask = (B < thresh).astype(np.int)#.reshape(x.size()).squeeze()

        title = r'Relevance Score $\theta(x)$' + (' (Scaled)' if self.scale_values else '')
        if self.binary:
            d = dict(zip(self.features, B[:,0])) # Change to B[0,:] when B model is truly binary
            A = plot_dependencies(d, title= title,
                                     scale_values = self.scale_values,
                                     sort_rows = self.sort_rows)
        else:
            Pos_Feats = {}
            for k in range(B.shape[0]):
                d = dict(zip(self.features, B[k,:])) # Change to B[0,:] when B model is truly binary
                A = plot_dependencies(d, title= title,
                                         scale_values = self.scale_values,
                                         sort_rows = self.sort_rows)
                Neg_Feats = list(compress(self.features, B[k,:] < -thresh))
                Pos_Feats = list(compress(self.features, B[k,:] > thresh))
                print('Class:{:5} Neg: {}, Pos: {}'.format(k, ','.join(Neg_Feats), ','.join(Pos_Feats)))
        if save_path:
            plt.savefig(save_path, bbox_inches = 'tight', format='pdf', dpi=300)
        plt.show()
        print('-'*60)

    def _explain_class(self, x_raw, x,k,typ='pos',feat_names = None, thresh = 0.5,recompute=True):
        """
            Given an input x and class index k, explain f(x) by returning indices of
            features in x that have highest positive impact on predicting class k.
        """
        if recompute:
            y = model(x) # y = self.model(x)
        B_k = model.params[0,k,:].data.numpy()
        #print((B_k > thresh).astype(np.int))

        if feat_names and typ == 'pos':
            # Return masked features instead of values
            return list(compress(feat_names, B_k > thresh ))
        elif feat_names and typ == 'neg':
            return list(compress(feat_names, B_k < thresh ))
        if typ == 'pos':
            Mask = (B_k > thresh).astype(np.int).reshape(x.size()).squeeze()
        elif typ == 'neg':
            Mask = (B_k < -thresh).astype(np.int).reshape(x.size()).squeeze()
        else:
            # Return weights instead of mask
            return B_k.reshape(x.size()).squeeze()

        Masked_x = Mask*x_raw.numpy().squeeze()
        return Masked_x

def plot_dependencies(dictionary_values,
                      pos_color="#ff4d4d",
                      negative_color="#3DE8F7",
                      reverse_values=False,
                      sort_rows =True,
                      scale_values = True,
                      title="",
                      fig_size=(4, 4),
                      ax = None,
                      x = None,
                      digits = 1, prediction_text = None,
                      show_table = False, ax_table = None):
    """ This function was adapted form the fairml python package

        x needed only if show_table = True

        digits: (int) significant digits in table
    """
    # add check to make sure that dependence features are not zeros
    if np.sum(np.array(dictionary_values.values())) == 0.0:
        print("Feature dependence for all attributes equal zero."
              " There is nothing to plot here. ")
        return None

    column_names = list(dictionary_values.keys())
    coefficient_values = list(dictionary_values.values())

    # get maximum
    maximum_value = np.absolute(np.array(coefficient_values)).max()
    if scale_values:
        coefficient_values = ((np.array(coefficient_values) / maximum_value) * 100)

    if sort_rows:
        index_sorted = np.argsort(np.array(coefficient_values))
    else:
        index_sorted = range(len(coefficient_values))[::-1]

    sorted_column_names = list(np.array(column_names)[index_sorted])
    sorted_column_values = list(np.array(coefficient_values)[index_sorted])
    pos = np.arange(len(sorted_column_values)) + 0.7

    # rearrange this at some other point.
    def assign_colors_to_bars(array_values,
                              pos_influence=pos_color,
                              negative_influence=negative_color,
                              reverse=reverse_values):

        # if you want the colors to be reversed for positive
        # and negative influences.
        if reverse:
            pos_influence, negative_influence = (negative_influence,
                                                 pos_influence)

        # could rewrite this as a lambda function
        # but I understand this better
        def map_x(x):
            if x > 0:
                return pos_influence
            else:
                return negative_influence
        bar_colors = list(map(map_x, array_values))
        return bar_colors

    bar_colors = assign_colors_to_bars(coefficient_values, reverse=True)
    bar_colors = list(np.array(bar_colors)[index_sorted])

    #pdb.set_trace()
    if ax is None and not show_table:
        #pdb.set_trace()
        fig, ax = plt.subplots(figsize=fig_size)
    elif ax is None and show_table:
        fig, axes = plt.subplots(1, 2, figsize=fig_size)
        ax_table, ax = axes

    ax.barh(pos, sorted_column_values, align='center', color=bar_colors)
    ax.set_yticks(pos)
    ax.set_yticklabels(sorted_column_names)
    if scale_values:
        ax.set_xlim(-105, 105)
    else:
        pass
        #ax.set_xlim(-1.05, 1.05)
    if title:
        ax.set_title("{}".format(title))

    if show_table and ax_table:
        cell_text = [[('%1.' + str(digits) + 'f') % v] for v in x]
        if prediction_text is None:
            ax_table.axis('off')
        else:
            print('here')
            ax_table.set_xticklabels([])
            ax_table.set_yticklabels([])
            ax_table.set_yticks([])
            ax_table.set_xticks([])
            for side in ['top', 'right', 'bottom', 'left']:
                ax_table.spines[side].set_visible(False)
            ax_table.set_xlabel(prediction_text)

        ax_table.table(cellText=cell_text,
                                  rowLabels=sorted_column_names[::-1],
                                  rowColours=bar_colors[::-1],
                                  colLabels=None,#['Value'],
                                  colWidths=[1],
                                  loc='left', cellLoc = 'right',
                                  bbox=[0.2, 0.025, 0.95, 0.95])
        ax_table.set_title('Input Value')
        return ax, ax_table

    return ax

def plot_theta_stability(model, input, pert_type = 'gauss', noise_level = 0.5,
                         samples = 5, save_path = None):
    """ Test stability of relevance scores theta for perturbations of an input.

        If model is of type 1 (i.e. theta is of dimension nconcepts x nclass), visualizes
        the perturbations of dependencies with respect to predicted class.

        If model is of type 1/3 (theta is a vector of size nconcepts), then there's only
        one dimension of theta to visualize.

        Args:
            model (GSENN): GSENN model.
            inputs (list of tensors): Inputs over which stability will be tested. First one is "base" input.

        Returns:
            stability: scalar.

        Displays plot also.

    """
    def gauss_perturbation(x, scale = 1):
        noise = Variable(scale*torch.randn(x.size()), volatile = True)
        if x.is_cuda:
            noise = noise.cuda()
        return x + noise

    model.eval()

    # Generate perturbations
    inputs = [input]
    for i in range(samples):
        inputs.append(gauss_perturbation(input, scale=noise_level))

    fig, ax = plt.subplots(2,len(inputs),figsize=(2*len(inputs),1.5*3))

    # Map Them
    thetas = []
    dists  = []
    for i,x in enumerate(inputs):
        pred = model(x)
        ax[0,i].imshow(x.data.cpu().numpy().squeeze())#, cmap = 'Greys', interpolation = 'nearest')
        ax[0,i].set_xticks([])
        ax[0,i].set_yticks([])
        if i == 0:
            ax[0,i].set_title('Original'.format(i))
        else:
            ax[0,i].set_title('Perturbation {}'.format(i))

        theta = model.thetas.data.cpu().numpy().squeeze()
        if theta.shape[1] > 1:
            # Means this is model 1, scalar h and theta_i vector-sized. Choose which one ti visualize
            klass = pred.data.max(1)[1] # Predicted class
            deps = theta[:,klass].squeeze()
            thetas.append(deps)
        else:
            deps = theta
            thetas.append(deps)
        classes = ['C' + str(i) for i in range(theta.shape[0])]
        d = dict(zip(classes, deps))
        A = plot_dependencies(d, title= 'Dependencies', sort_rows = False, ax = ax[1,i])
        #ax[1,i].locator_params(axis = 'y', nbins=10)

        # max_yticks = 10
        # yloc = plt.MaxNLocator(max_yticks)
        # ax[1,i].yaxis.set_major_locator(yloc)
        #print(thetas[-1])
        if i > 0:
            dists.append(np.linalg.norm(thetas[0] - deps))

    dists = np.array(dists)
    plt.tight_layout()
    #print(dists.max())
    if save_path:
        plt.savefig(save_path, bbox_inches = 'tight', format='pdf', dpi=300)
    #plt.show(block=False)

def concept_grid(model, data_loader, cuda=False, top_k = 6, layout = 'vertical', return_fig=False, save_path = None):
    """
        Finds examples in data_loader that are most representatives of concepts.

        For scalar concepts, activation is simply the value of the concept.
        For vector concepts, activation is the norm of the concept.

    """
    print('Warning: make sure data_loader passed to this function doesnt shuffle data!!')
    all_norms = []
    num_concepts = model.parametrizer.nconcept
    concept_dim  = model.parametrizer.dout

    top_activations = {k: np.array(top_k*[-1000.00]) for k in range(num_concepts)}
    top_examples = {k: top_k*[None] for k in range(num_concepts)}
    all_activs = []
    for idx, (data, target) in enumerate(data_loader):
        # get the inputs
        if cuda:
            data, target = data.cuda(), target.cuda()
        data, target = Variable(data, volatile=True), Variable(target)
        output = model(data)
        concepts = model.concepts.data
        #pdb.set_trace()
        #concepts[concepts < 0] = 0.0 # This is unncessary if output of H is designed to be > 0.
        if concepts.shape[-1] > 1:
            print('ERROR')
            print(asd.asd)
            activations = np.linalg.norm(concepts, axis = 2)
        else:
            activations = concepts

        all_activs.append(activations)
        # if idx == 10:
        #     break

    all_activs = torch.cat(all_activs)
    top_activations, top_idxs = torch.topk(all_activs, top_k, 0)
    top_activations = top_activations.squeeze().t()
    top_idxs = top_idxs.squeeze().t()
    top_examples = {}
    for i in range(num_concepts):
        top_examples[i] = data_loader.dataset.test_data[top_idxs[i]]
    #top_examples =

    # Before, i was doing this manually :
        # for i in range(activations.shape[0]):
        #     #pdb.set_trace()
        #     for j in range(num_concepts):
        #         min_val  = top_activations[j].min()
        #         min_idx  = top_activations[j].argmin()
        #         if activations[i,j] >  min_val:
        #             # Put new one in place of min
        #             top_activations[j][min_idx]  = activations[i,j]
        #             top_examples[j][min_idx] = data[i, :, :, :].data.numpy().squeeze()
        #     #pdb.set_trace()
    # for k in range(num_concepts):
    #     #print(k)
    #     Z = [(v,e) for v,e in sorted(zip(top_activations[k],top_examples[k]),  key=lambda x: x[0], reverse = True)]
    #     top_activations[k], top_examples[k] = zip(*Z)

    if layout == 'horizontal':
        num_cols = top_k
        num_rows = num_concepts
        figsize=(num_cols, 1.2*num_rows)
    else:
        num_cols = num_concepts
        num_rows = top_k
        figsize=(1.4*num_cols, num_rows)

    fig, axes  = plt.subplots(figsize=figsize, nrows=num_rows, ncols=num_cols )

    for i in range(num_concepts):
        for j in range(top_k):
            pos = (i,j) if layout == 'horizontal' else (j,i)

            l = i*top_k + j
            #print(i,j)
            #print(top_examples[i][j].shape)
            axes[pos].imshow(top_examples[i][j], cmap='Greys',  interpolation='nearest')
            if layout == 'vertical':
                axes[pos].axis('off')
                if j == 0:
                    axes[pos].set_title('Cpt {}'.format(i+1), fontsize = 24)
            else:
                axes[pos].set_xticklabels([])
                axes[pos].set_yticklabels([])
                axes[pos].set_yticks([])
                axes[pos].set_xticks([])
                for side in ['top', 'right', 'bottom', 'left']:
                    axes[i,j].spines[side].set_visible(False)
                if i == 0:
                    axes[pos].set_title('Proto {}'.format(j+1))
                if j == 0:
                    axes[pos].set_ylabel('Concept {}'.format(i+1), rotation = 90)

    print('Done')

    # cols = ['Prot.{}'.format(col) for col in range(1, num_cols + 1)]
    # rows = ['Concept # {}'.format(row) for row in range(1, num_rows + 1)]
    #
    # for ax, col in zip(axes[0], cols):
    #     ax.set_title(col)
    #
    # for ax, row in zip(axes[:,0], rows):
    #     ax.set_ylabel(row, rotation=0, size='large')
    #plt.tight_layout()

    if layout == 'vertical':
        fig.subplots_adjust(wspace=0.01, hspace=0.1)
    else:
        fig.subplots_adjust(wspace=0.1, hspace=0.01)

    if save_path is not None:
        plt.savefig(save_path, bbox_inches = 'tight', format='pdf', dpi=300)
    plt.show()
    if return_fig:
        return fig, axes

def plot_prob_drop(attribs, prob_drop, save_path = None):

    ind = np.arange(len(attribs))
    column_names = [str(j) for j in range(1,22)]

    width = 0.65

    fig, ax1 = plt.subplots(figsize=(8,4))

    color1 = '#377eb8'
    ax1.bar(ind+width+0.35, attribs, 0.45, color=color1)
    ax1.set_ylabel(r'Feature Relevance $\theta(x)_i$',color=color1, fontsize = 14)
    #ax1.set_ylim(-1,1)
    ax1.set_xlabel('Feature')
    ax1.tick_params(axis='y', colors=color1)

    color2 = '#ff7f00'
    ax2 = ax1.twinx()
    ax2.ticklabel_format(style='sci',scilimits=(-2,2),axis = 'y')
    ax2.plot(ind+width+0.35, prob_drop, 'bo', linestyle='dashed', color=color2)
    ax2.set_ylabel('Probability Drop', color = color2, fontsize = 14)
    ax2.tick_params(axis='y', colors=color2)

    ax1.set_xticks(ind+width+(width/2))
    ax1.set_xticklabels(column_names)

    fig.tight_layout()
    if save_path:
        plt.savefig(save_path, bbox_inches = 'tight', format='pdf', dpi=300)

    plt.show()

def noise_stability_plots(model, dataset, cuda, save_path):
    # find one example of each digit:
    examples = {}
    i = 0
    while (not len(examples.keys()) == 10) and (i < len(dataset)):
        if dataset[i][1] not in examples:
            examples[dataset[i][1]] = dataset[i][0].view(1,1,28,28)
        i += 1

    for i in range(10):
        x = Variable(examples[i], volatile = True)
        if cuda:
            x = x.cuda()
        plot_theta_stability(model, x, noise_level = 0.5,
                save_path=save_path + '/noise_stability_{}.pdf'.format(i))

def plot_digit(x, ax = None, greys = True):
    if ax is None:
        fig, ax = plt.subplots()
    if type(x) is torch.Tensor:
        x = x.numpy()
    if x.ndim > 2:
        x = x.squeeze()
    if greys:
        ax.imshow(x, cmap='Greys')
    else:
        ax.imshow(x)
    ax.axis('off')
    ax.set_xticklabels([])
    ax.set_yticklabels([])

def opt_concept_func(x, model, i, accumulator):
    x = torch.from_numpy(x).float()
    x = x.reshape(28,28)[None, None, :, :]
    accumulator.append(x)
    # printing the process
    if len(accumulator) % 5000 == 0:
        print(len(accumulator), " rounds so far.")
    return torch.sum(model.conceptizer.encode(x).squeeze()[:, None]).numpy() - model.conceptizer.encode(x).squeeze()[:, None][i].numpy()

def get_concept_vis(model, nconcepts):
    convis = list()
    for i in range(nconcepts):
        print("Concept: (" , str(i+1), "/", str(nconcepts), ")")
        # normalized random image as specified in the load minst function in main_mnist.py
        randimg = torch.normal(0.1307, 0.3081, size=(28,28))[None, :, :]
        # Start with zeros
        # randimg = torch.zeros(28,28)[None, :, :]
        accumulator = list()
        optimize.minimize(opt_concept_func, randimg, args=(model, i, accumulator))
        convis.append(accumulator[-1])
        print("Concept: " , str(i), " done! ", str(len(accumulator)), " rounds needed.")

    return convis

def plot_concept_vis(model, nconcepts, cols, rows):
    convis = get_concept_vis(model, nconcepts)
    fig=plt.figure(figsize=(8, 8))

    for i in range(nconcepts):
        ax = fig.add_subplot(cols, rows, i+1)
        ax.imshow(convis[i].squeeze())
    plt.show()