<reponame>yifan-you-37/rl_swiss
import joblib

import numpy as np
from numpy.random import choice, randint

from rlkit.data_management.env_replay_buffer import get_dim as gym_get_dim
from rlkit.data_management.simple_replay_buffer import SimpleReplayBuffer
from rlkit.envs.maze_envs.trivial_grid import TrivialGrid
from rlkit.envs.maze_envs.pogrid import PartiallyObservedGrid
from rlkit.envs.maze_envs.mem_map_grid import MemoryGrid

class RandomGridPolicy():
    def __init__(self, max_num_consecutive):
        self.max_num_consecutive = max_num_consecutive
        self.cur_act = randint(4)
        self.num_left = randint(max_num_consecutive) + 1

    def get_action(self, obs, *args):
        if self.num_left == 0:
            self.num_left = randint(self.max_num_consecutive) + 1
            self.cur_act = randint(4)
        self.num_left -= 1
        return self.cur_act

    def reset(self):
        pass

class ListPolicy():
    def __init__(self, act_list):
        self.act_list = act_list
        self.ptr = 0

    def get_action(self, obs, *args):
        a = self.act_list[self.ptr]
        self.ptr = (self.ptr + 1) % len(self.act_list)
        return a

    def reset(self):
        self.ptr = 0

def generate_transitions(policy, env, num_timesteps_total, max_steps_per_episode, save_path):
    buff = SimpleReplayBuffer(
        num_timesteps_total, env.observation_space.shape,
        gym_get_dim(env.action_space), discrete_action_dim=True
    )

    cur_total = 0
    steps_left_in_episode = 0
    while cur_total != num_timesteps_total:
        if steps_left_in_episode == 0:
            steps_left_in_episode = max_steps_per_episode
            obs = env.reset()

        act = policy.get_action(obs)
        next_obs, rew, done, _ = env.step(act)
        buff.add_sample(obs, act, rew, done, next_obs)

        obs = next_obs
        cur_total += 1
        steps_left_in_episode -= 1

    save_dict = dict(
        observations=buff._observations,
        actions=buff._actions,
        rewards=buff._rewards,
        terminals=buff._terminals,
        next_observations=buff._next_obs,
    )
    joblib.dump(save_dict, save_path)

    # debug
    from scipy.misc import imsave
    actions = buff._actions
    observations = buff._observations
    for i in range(1000):
        a = actions[i]
        obs = observations[i]
        print(a)
        imsave('junk_vis/tiny/mem_grid_{}.png'.format(i), np.transpose(obs, (1,2,0)))

    # for i in range(90, 110):
    #     a = actions[i]
    #     obs = observations[i]
    #     print(a)
    #     imsave('junk_vis/maze_{}.png'.format(i), np.transpose(obs, (1,2,0)))

    # for i in range(70, 90):
    #     a = actions[i]
    #     obs = observations[i]
    #     print(a)
    #     imsave('junk_vis/maze_{}.png'.format(i), np.transpose(obs, (1,2,0)))

    # for i in range(110, 130):
    #     a = actions[i]
    #     obs = observations[i]
    #     print(a)
    #     imsave('junk_vis/maze_{}.png'.format(i), np.transpose(obs, (1,2,0)))

if __name__ == '__main__':
    # env_specs = {
    #     'flat_repr': False,
    #     'one_hot_repr': False,
    #     'maze_h': 4,
    #     'maze_w': 4,
    #     'scale': 1,
    # }
    # env = TrivialGrid(env_specs)
    # policy = RandomGridPolicy(1)

    # env_specs = {
    #     'flat_repr': False,
    #     'one_hot_repr': False,
    #     'maze_h': 9,
    #     'maze_w': 9,
    #     'obs_h': 5,
    #     'obs_w': 5,
    #     'scale': 4,
    #     'num_objs': 10
    # }
    # act_list = [1, 0, 3, 2]
    # env = MemoryGrid(env_specs)
    # policy = ListPolicy(act_list)

    env_specs = {
        'flat_repr': False,
        'one_hot_repr': False,
        'maze_h': 9,
        'maze_w': 9,
        'obs_h': 5,
        'obs_w': 5,
        'scale': 4,
        'num_objs': 10
    }
    act_list = [1, 0, 3, 2]
    env = PartiallyObservedGrid(env_specs)
    policy = RandomGridPolicy(1)

    generate_transitions(policy, env, 50000, 8, '/ais/gobi6/kamyar/oorl_rlkit/maze_trans_data/pogrid_len_8_scale_4')

# 32, 128, 512 for 3x3
# 32, 512, 2048 for 5x5


<reponame>RedHeadM/pytorch_segmentation<gh_stars>0
# Originally written by <NAME>
# https://github.com/kazuto1011/deeplab-pytorch

from base import BaseDataSet, BaseDataLoader
from utils import palette
import numpy as np
import os
import scipy
import torch
from PIL import Image
import cv2
from torch.utils.data import Dataset
from torchvision import transforms

from multiview.video.datasets import ViewPairDataset
class MuiltivwDataset(BaseDataSet):
    """
    Pascal Voc dataset
    http://host.robots.ox.ac.uk/pascal/VOC/voc2012/VOCtrainval_11-May-2012.tar
    """
    def __init__(self, number_views, view_idx, **kwargs):
        # self.num_classes = 3
        self.num_classes = 2+10
        self.palette = palette.get_voc_palette(self.num_classes)
        self.number_views = number_views
        self.view_idx = view_idx
        if not isinstance(view_idx,int):
            raise ValueError('view_idx: {}'.format(view_idx))
        self.view_key_img = "frames views " + str(self.view_idx)
        self.view_key_seg = "seg "+str(self.view_idx)
        assert isinstance(view_idx, int) and isinstance(number_views, int)
        super(MuiltivwDataset, self).__init__(**kwargs)
        print('data dir {}, view idx {}, num views'.format(self.root, view_idx, number_views))

    def _set_files(self):
        def data_len_filter(comm_name,frame_len_paris):
            if len(frame_len_paris)<2:
                return frame_len_paris[0]>10
            return min(*frame_len_paris)>10
        self.mvbdata = ViewPairDataset(self.root.strip(),
					    segmentation= True,
                                            transform_frames= None,
					    number_views=self.number_views,
					    filter_func=data_len_filter)

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

    def _load_data(self, index):
        s = self.mvbdata[index]
        label = s[self.view_key_seg]
        image = s[self.view_key_img]
        image = np.asarray(image, dtype=np.float32)
        label = np.asarray(label, dtype=np.int32)
        return image, label, s["id"]

class MVB(BaseDataLoader):
    def __init__(self, data_dir, batch_size, split, crop_size=None, base_size=None, scale=True, num_workers=1, val=False,
                    shuffle=False, flip=False, rotate=False, blur= False, augment=False, val_split= None, return_id=False,
                    number_views=1, view_idx=None):

        self.MEAN = [0.45734706, 0.43338275, 0.40058118]
        self.STD = [0.23965294, 0.23532275, 0.2398498]
        assert view_idx is not None, "set view idx in config"
        kwargs = {
            'root': data_dir,
            'split': split,
            'mean': self.MEAN,
            'std': self.STD,
            'augment': augment,
            'crop_size': crop_size,
            'base_size': base_size,
            'scale': scale,
            'flip': flip,
            'blur': blur,
            'rotate': rotate,
            'return_id': return_id,
            'val': val
        }

        self.dataset = MuiltivwDataset(number_views,view_idx,**kwargs)
        super(MVB, self).__init__(self.dataset, batch_size, shuffle, num_workers, val_split)


# ===============================================
#            Foreword
# ===============================================

""" This script works as a demonstration how a given audio file can be transformed into speaker embeddings. When working with multiple audio files,
consider first creating the MFCC files as described in "feature extraction" and then embed the files separately.
"""

# ===============================================
#            Headers
# ===============================================

import argparse
import os
import wavefile
import numpy as np
from keras.models import Model, load_model
from keras.layers import *
from keras.constraints import *
import tensorflow as tf
import keras
from tensorflow.python.keras import backend as K
from librosa.util import frame
from librosa.feature import mfcc
from sklearn import preprocessing
from scipy.optimize import brentq
from sklearn.metrics import roc_curve
from scipy.interpolate import interp1d

# ===============================================
#           Functions
# ===============================================

class VLAD(keras.engine.Layer):

    """
    NetVLAD implementation by <NAME> based on
    the "NetVLAD: CNN architecture for weakly supervised place recognition" paper
    """

    def __init__(self, k_centers=8, kernel_initializer='glorot_uniform', **kwargs):
        self.k_centers = k_centers
        self.kernel_initializer = initializers.get(kernel_initializer)
        super(VLAD, self).__init__(**kwargs)

    def build(self, input_shape):
        self.w = self.add_weight(name='kernel',
                                      shape=(input_shape[2], self.k_centers),
                                      initializer=self.kernel_initializer,
                                      trainable=True)
        self.b = self.add_weight(name='kernel',
                                      shape=(self.k_centers, ),
                                      initializer=self.kernel_initializer,
                                      trainable=True)
        self.c = self.add_weight(name='kernel',
                                      shape=(input_shape[2], self.k_centers),
                                      initializer=self.kernel_initializer,
                                      trainable=True)
        super(VLAD, self).build(input_shape)

    def call(self, x):

        Wx_b = K.dot(x, self.w)+self.b
        a = tf.nn.softmax(Wx_b)

        rows = []

        for k in range(self.k_centers):
            error = x-self.c[:, k]

            row = K.batch_dot(a[:, :, k],error)
            row = tf.nn.l2_normalize(row,dim=1)
            rows.append(row)

        output = tf.stack(rows)
        output = tf.transpose(output, perm = [1, 0, 2])
        output = tf.reshape(output, [tf.shape(output)[0], tf.shape(output)[1]*tf.shape(output)[2]])

        return output

    def compute_output_shape(self, input_shape):
        return (input_shape[0], self.k_centers*input_shape[2])

    def get_config(self):
        config = super(VLAD, self).get_config()
        config['k_centers'] = self.k_centers
        config['kernel_initializer'] = initializers.serialize(self.kernel_initializer)
        return config

def amsoftmax_loss(y_true, y_pred, scale = 30, margin = 0.35):
    y_pred = y_pred - y_true*margin
    y_pred = y_pred*scale
    return K.categorical_crossentropy(y_true, y_pred, from_logits = True)

def feature_extractor(s, fs = 16000):
   mfcc_feat = mfcc(s, n_mfcc = 30, sr = fs, n_fft=512, hop_length=160)
   mfcc_feat = preprocessing.scale(mfcc_feat, axis = 1)
   return mfcc_feat

def EER_calc(cos_dists, labels):
    fpr, tpr, thresholds = roc_curve(labels, cos_dists, pos_label=1)
    EER = brentq(lambda x : 1. - x - interp1d(fpr, tpr)(x), 0., 1.)
    thresh = interp1d(fpr, thresholds)(EER)
    return EER, threshI

# ===============================================
#                   MAIN
# ===============================================

if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    parser.add_argument("--signal", help="Input signal file", default = "/teamwork/t40511_asr/p/spherediar/data/augmented_wavs/signal.npy")
    parser.add_argument("--frame_len", type = int, default = 2000, help="Frame length in milliseconds") # Ensure that the frame length is less or equal to the signal duration
    parser.add_argument("--hop_len",  type = int, default = 500, help="Hop length in milliseconds")
    parser.add_argument("--mode",  type = int, default = 1, help="1: Save all embeddings, 0: Save only the average embedding.")
    parser.add_argument("--dest",  default = "/teamwork/t40511_asr/p/spherediar/data/augmented_wavs", help="Directory to save the embeddings")
    parser.add_argument("--model",  default = "/teamwork/t40511_asr/p/spherediar/models/current_best.h5", help="Model file.")
    args = parser.parse_args()

    # ===============================================
    #           Feature extraction
    # ===============================================

    # Get the type of the signal file
    file_name = args.signal.split("/")[-1]
    file_format = file_name.split(".")[1]

    # Load signal - for now, only works with wav or numpy files
    if file_format == "npy":
        signal = np.load(args.signal)
    else:
        (rate,sig) = wavefile.load(args.signal)
        signal = sig[0]

    # Frame and compute MFCCs
    S = np.transpose(frame(signal, int(args.frame_len*16), int(args.hop_len*16)))  # For now, only 16kHz sampling rate can be used
    X = list(map(lambda s: feature_extractor(s, 16000), S))
    X = np.array(np.swapaxes(X,1, 2))
    X = X.astype(np.float16) # Compression to save memory, 16-bit MFCCs have also been used in the training of the current_best.h5
    num_timesteps = X.shape[1]

    # ===============================================
    #           Embedding extraction
    # ===============================================

    emb_model = load_model(args.model, custom_objects={'VLAD': VLAD, 'amsoftmax_loss': amsoftmax_loss})

    # Modify input shape if necessary
    if num_timesteps != 201:
        emb_model.layers.pop(0)
        new_input = Input(batch_shape=(None,num_timesteps,30))
        new_output = emb_model(new_input)
        emb_model = Model(new_input, new_output)

    # Create embeddings
    embs = emb_model.predict(X)

    # ===============================================
    #           Save
    # ===============================================

    # Get the name of the signal file
    sig_name = file_name.split(".")[0]

    if args.mode == 0:
        avg_emb = np.mean(embs, axis = 0)
        avg_emb = avg_emb/np.sqrt(sum(avg_emb**2))
        dest_file = os.path.join(args.dest, "_".join([sig_name, "single_emb", str(args.frame_len), str(args.hop_len)]))
        np.save(dest_file, avg_emb)
    else:
        dest_file = os.path.join(args.dest, "_".join([sig_name, "embs", str(args.frame_len), str(args.hop_len)]))
        np.save(dest_file, embs)


# MIT License
#
# Copyright (c) 2022 Quandela
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import concurrent.futures
import logging
import random
import sys
import threading
import time

import numpy as np
from scipy.optimize import minimize

import perceval as pcvl
import perceval.lib.phys as phys

mapping = {
    "HH": pcvl.FockState([1, 1, 0, 1, 1, 1, 0]),
    "HL": pcvl.FockState([1, 1, 0, 1, 1, 0, 1]),
    "LH": pcvl.FockState([1, 0, 1, 1, 1, 1, 0]),
    "LL": pcvl.FockState([1, 0, 1, 1, 1, 0, 1])
}

istates = {v: k for k, v in mapping.items()}
simulator_backend = pcvl.BackendFactory().get_backend("Naive")

def f_to_minimize(circuit, params, params_value, prob=1):
    for idx, p in enumerate(params_value):
        params[idx].set_value(p)

    sim = simulator_backend(circuit.compute_unitary(use_symbolic=False), n=5, mask=["1  11  "])
    sim.compile(istates.keys())
    ca = pcvl.CircuitAnalyser(sim, istates, "*")
    ca.compute(expected={"LH": "LL", "LL": "LH", "HH": "HH", "HL": "HL"})
    loss = np.sqrt((0.1*(prob-ca.performance))**2+(0.9*ca.error_rate)**2)
    return loss

def discover(circuit, p, params=None, method=None, init_params=None, bounds=None):
    if init_params is None:
        init_params = np.random.randn(len(params))

    minimize(lambda x: f_to_minimize(circuit, params, x, p), init_params, method=method, bounds=bounds)

    sim = simulator_backend(circuit.compute_unitary(use_symbolic=False), n=5, mask=["1  11  "])
    ca = pcvl.CircuitAnalyser(sim, istates, "*")
    ca.compute(expected={"LH": "LL", "LL": "LH", "HH": "HH", "HL": "HL"})
    performance = ca.performance
    ber = ca.error_rate
    return ber, performance, ca.pdisplay()

global_lock = threading.Lock()
result_file = sys.stdout  # open("result-discovery.txt", "w")

def run_a_discovery(name):
    logging.info("Discovery thread %s: starting", name)
    start = time.time()

    # generic interferometer
    n = 7
    gen_rect = pcvl.Circuit.generic_interferometer(
        n,
        lambda i: random.randint(0, 1) and phys.BS(R=pcvl.P("R%d" % i), phi_b=np.pi, phi_d=0)
            or phys.BS(R=pcvl.P("R%d" % i)),
        shape="rectangular",
        depth=4
    )

    params = gen_rect.get_parameters()
    bounds = [p.bounds for p in params]

    (end_ber, end_performance, cadisplay) = discover(gen_rect,
                                                     0.5,
                                                     params=params,
                                                     init_params=[random.random() for _ in params],
                                                     bounds=bounds)
    with global_lock:
        result_file.write("time=%f, performance=%f, ber=%f\n" % (time.time()-start, end_performance, end_ber))
        result_file.write(cadisplay+"\n")
        result_file.write(str(gen_rect.get_parameters())+"\n\n")

    logging.info("Discovery thread %s: ending in %f seconds", name, time.time()-start)

log_format = "%(asctime)s: %(message)s"
logging.basicConfig(format=log_format, level=logging.INFO,
                    datefmt="%H:%M:%S")

# max_works is the number of parallel workers that will be launched
with concurrent.futures.ThreadPoolExecutor(max_workers=8) as executor:
    # launch N iterations
    executor.map(run_a_discovery, range(10))


import skimage.io
import skvideo.io
import os
import h5py
from sklearn.externals import joblib
from sklearn.linear_model import LinearRegression
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import f1_score
import scipy.misc
import scipy.signal
import numpy as np
from sporco import util
import matplotlib.pyplot as plt
import pylab as py
import glob
from PIL import Image
import cv2
import sys

def get_postrainpatches(hdf5_im, hdf5_lab, hdf5_trainset, idx=0, traintest=0):
  return genericpospatcher(hdf5_im, hdf5_lab, hdf5_trainset, idx=idx, traintest=traintest)

def genericpospatcher(hdf5_im, hdf5_lab, hdf5_trainset, idx=0, traintest=0):
  width = 256+40
  height = 256+40

  lst480p = np.array(glob.glob("/mnt/hd3/scenes/480p/*avi"))
  lst1080p = np.array(glob.glob("/mnt/hd3/scenes/1080p/*avi"))
  lst480p = np.sort(lst480p)
  lst1080p = np.sort(lst1080p)

  if traintest == 0:
    lst = np.hstack((lst480p[:575], lst1080p[:215]))
  else:
    lst = np.hstack((lst480p[575:], lst1080p[215:]))

  n_samples = len(lst)
  for jjj, fname in enumerate(lst):
    print jjj, n_samples
    vid_pris = skvideo.io.vread(fname, as_grey=True).astype(np.float32)
    T, H, W, C = vid_pris.shape

    adj_h = H - height
    adj_w = W - width
    iv, jv = np.meshgrid(np.arange(adj_h), np.arange(adj_w), sparse=False, indexing='ij')
    iv = iv.reshape(-1)
    jv = jv.reshape(-1)

    jdx = np.random.permutation(adj_h*adj_w)

    iv = iv[jdx]
    jv = jv[jdx]
    tv = np.arange(0, T-1)

    limit = 0
    for (y, x) in zip(iv, jv):
      np.random.shuffle(tv)
      t = tv[0]
      goodpatch = vid_pris[t:t+2, y:y+height, x:x+width, 0]
      badpatch = goodpatch.copy()

      # random amount of change
      amt = np.random.randint(3, 6)
      badpatch /= 2**amt
      badpatch = np.floor(badpatch)
      badpatch *= 2**amt

      diff = np.mean((badpatch[0] - goodpatch[0])**2)
      if diff < 1.5:
        continue
      if np.std(badpatch[0]) < 20:
        continue
      if np.std(badpatch[1]) < 20:
        continue
      #print diff
      #skimage.io.imsave("extract/test_%d.png" % (idx,), np.hstack((goodpatch[0].astype(np.uint8), badpatch[0].astype(np.uint8))))

      #preprocess = preprocess_image(badpatch)

      #preprocess = preprocess[:, 5:-5, 5:-5]
      badpatch = badpatch[:, 20:-20, 20:-20]

      hdf5_im[idx] = badpatch
      hdf5_lab[idx] = 1
      hdf5_trainset[idx] = traintest

      #skimage.io.imsave("extract/%d.png" % (idx,), patch)
      limit += 1
      idx += 1
      if limit >= 25:
        break

  return idx

def get_negtrainpatches(image_patches, labels, trainset, idx=0, traintest=0):
  return genericnegpatcher(image_patches, labels, trainset, idx=idx, traintest=traintest)

def genericnegpatcher(hdf5_im, hdf5_lab, hdf5_trainset, idx=0, traintest=0):
  width = 256+40
  height = 256+40

  lst480p = np.array(glob.glob("/mnt/hd3/scenes/480p/*avi"))
  lst1080p = np.array(glob.glob("/mnt/hd3/scenes/1080p/*avi"))
  lst480p = np.sort(lst480p)
  lst1080p = np.sort(lst1080p)

  if traintest == 0:
    lst = np.hstack((lst480p[:575], lst1080p[:215]))
  else:
    lst = np.hstack((lst480p[575:], lst1080p[215:]))

  n_samples = len(lst)
  for jjj, fname in enumerate(lst):
    print jjj, n_samples
    vid_pris = skvideo.io.vread(fname, as_grey=True).astype(np.float32)
    T, H, W, C = vid_pris.shape

    adj_h = H - height
    adj_w = W - width
    iv, jv = np.meshgrid(np.arange(adj_h), np.arange(adj_w), sparse=False, indexing='ij')
    iv = iv.reshape(-1)
    jv = jv.reshape(-1)

    jdx = np.random.permutation(adj_h*adj_w)

    iv = iv[jdx]
    jv = jv[jdx]
    tv = np.arange(0, T-1)

    limit = 0
    for (y, x) in zip(iv, jv):
      np.random.shuffle(tv)
      t = tv[0]
      goodpatch = vid_pris[t:t+2, y:y+height, x:x+width, 0]

      #preprocess = preprocess[:, 5:-5, 5:-5]
      goodpatch = goodpatch[:, 20:-20, 20:-20]
      #if np.std(goodpatch[0]) < 20:
      #  continue
      #if np.std(goodpatch[1]) < 20:
      #  continue

      hdf5_im[idx] = goodpatch
      hdf5_lab[idx] = 0
      hdf5_trainset[idx] = traintest

      #skimage.io.imsave("extract/%d.png" % (idx,), patch)
      limit += 1
      idx += 1
      if limit >= 25:
        break

  return idx

# get the number of patches
np.random.seed(12345)

n_total_images = 64925#12000
patch_height = 256
patch_width = 256
n_channels = 2

# sf = single frame
# fd = frame diff
f = h5py.File('bandingdataset_sf.hdf5', mode='w')

image_patches = f.create_dataset('image_patches', (n_total_images, n_channels, patch_height, patch_width), dtype='float')
image_patches.dims[0].label = 'batch'
image_patches.dims[1].label = 'channel'
image_patches.dims[2].label = 'height'
image_patches.dims[3].label = 'width'

labels = f.create_dataset('labels', (n_total_images,), dtype='uint8')
trainset = f.create_dataset('set', (n_total_images,), dtype='uint8')

n_idx = 0
n_idx = get_postrainpatches(image_patches, labels, trainset, n_idx, 0)
n_idx = get_negtrainpatches(image_patches, labels, trainset, n_idx, 0)
n_idx = get_postrainpatches(image_patches, labels, trainset, n_idx, 1)
n_idx = get_negtrainpatches(image_patches, labels, trainset, n_idx, 1)

print n_idx, n_total_images

f.flush()
f.close()


import MWPotential2014Likelihood
import astropy.units as u
import gd1_util
import hypothesis
import numpy
import numpy as np
import numpy as np
import pickle
import torch

from galpy.orbit import Orbit
from galpy.potential import MWPotential2014, turn_physical_off, vcirc
from galpy.util import bovy_conversion, bovy_coords, save_pickles, bovy_plot
from gd1_util_MWhaloshape import lb_to_phi12
from scipy import integrate, interpolate
from scipy.integrate import quad
from torch.distributions.uniform import Uniform

def allocate_prior_stream_age():
    lower = torch.tensor(3).float().to(hypothesis.accelerator)
    upper = torch.tensor(7).float().to(hypothesis.accelerator)

    return Uniform(lower, upper)

def allocate_prior_wdm_mass():
    lower = torch.tensor(1).float().to(hypothesis.accelerator)
    upper = torch.tensor(50).float().to(hypothesis.accelerator)

    return Uniform(lower, upper)

def load_observed_gd1(path, phi, degree=1):
    data = np.genfromtxt(path, names=True)
    phi_max = max(phi) + 5 # For stability in fitting the splines
    phi_min = min(phi) - 5 # For stability in fitting the splines
    phi_data = data["phi1mid"]
    if phi_min < min(phi_data) or phi_max > max(phi_data):
        raise ValueError("Angles not supported by observation.")
    indices = (phi_data <= phi_max) & (phi_data >= phi_min)
    phi_data = phi_data[indices]
    linear_density = data["lindens"][indices]
    error = data["e_lindens"][indices]
    trend = np.polyfit(phi_data, linear_density, deg=degree)
    fitted = np.poly1d(trend)(phi_data)
    error /= fitted
    linear_density /= fitted
    # Fit a spline and extract the requested values
    l = np.array(linear_density)
    fit_density = interpolate.InterpolatedUnivariateSpline(phi_data, linear_density)
    fit_error = interpolate.InterpolatedUnivariateSpline(phi_data, error)
    linear_density = fit_density(phi)
    trend = np.polyfit(phi, linear_density, deg=degree)
    fitted = np.poly1d(trend)(phi)
    linear_density /= fitted
    error = fit_error(phi)
    error /= fitted

    return linear_density, error

h=0.6774
ro=8.
vo=220.

def parse_times(times,age):
    if 'sampling' in times:
        nsam= int(times.split('sampling')[0])
        return [float(ti)/bovy_conversion.time_in_Gyr(vo,ro)
                for ti in np.arange(1,nsam+1)/(nsam+1.)*age]
    return [float(ti)/bovy_conversion.time_in_Gyr(vo,ro)
            for ti in times.split(',')]

def parse_mass(mass):
    return [float(m) for m in mass.split(',')]

def nsubhalo(m):
    return 0.3*(10.**6.5/m)

def rs(m,plummer=False,rsfac=1.):
    if plummer:
        #print ('Plummer')
        return 1.62*rsfac/ro*(m/10.**8.)**0.5
    else:
        return 1.05*rsfac/ro*(m/10.**8.)**0.5

def alpha(m_wdm):
    return (0.048/h)*(m_wdm)**(-1.11) #in Mpc , m_wdm in keV

def lambda_hm(m_wdm):
    nu=1.12
    return 2*numpy.pi*alpha(m_wdm)/(2**(nu/5.) - 1.)**(1/(2*nu))

def M_hm(m_wdm):
    Om_m=0.3089
    rho_c=1.27*10**11 #Msun/Mpc^3
    rho_bar=Om_m*rho_c
    return (4*numpy.pi/3)*rho_bar*(lambda_hm(m_wdm)/2.)**3

def Einasto(r):
    al=0.678 #alpha_shape
    rm2=199 #kpc, see Erkal et al 1606.04946 for scaling to M^1/3
    return numpy.exp((-2./al)*((r/rm2)**al -1.))*4*numpy.pi*(r**2)

def dndM_cdm(M,c0kpc=2.02*10**(-13),mf_slope=-1.9):
    #c0kpc=2.02*10**(-13) #Msun^-1 kpc^-3 from Denis' paper
    m0=2.52*10**7 #Msun from Denis' paper
    return c0kpc*((M/m0)**mf_slope)

def fac(M,m_wdm):
    beta=-0.99
    gamma=2.7
    return (1.+gamma*(M_hm(m_wdm)/M))**beta

def dndM_wdm(M,m_wdm,c0kpc=2.02*10**(-13),mf_slope=-1.9):
    return fac(M,m_wdm)*dndM_cdm(M,c0kpc=c0kpc,mf_slope=mf_slope)

def nsub_cdm(M1,M2,r=20.,c0kpc=2.02*10**(-13),mf_slope=-1.9):
    #number density of subhalos in kpc^-3
    m1=10**(M1)
    m2=10**(M2)
    return integrate.quad(dndM_cdm,m1,m2,args=(c0kpc,mf_slope))[0]*integrate.quad(Einasto,0.,r)[0]*(8.**3.)/(4*numpy.pi*(r**3)/3) #in Galpy units

def nsub_wdm(M1,M2,m_wdm,r=20.,c0kpc=2.02*10**(-13),mf_slope=-1.9):
    m1=10**(M1)
    m2=10**(M2)
    return integrate.quad(dndM_wdm,m1,m2,args=(m_wdm,c0kpc,mf_slope))[0]*integrate.quad(Einasto,0.,r)[0]*(8.**3)/(4*numpy.pi*(r**3)/3) #in Galpy units

def simulate_subhalos_mwdm(sdf_pepper,m_wdm,mf_slope=-1.9,c0kpc=2.02*10**(-13),r=20.,Xrs=5.,sigma=120./220.):

    Mbin_edge=[5.,6.,7.,8.,9.]
    Nbins=len(Mbin_edge)-1
    #compute number of subhalos in each mass bin
    nden_bin=np.empty(Nbins)
    rate_bin=np.empty(Nbins)
    for ll in range(Nbins):
        nden_bin[ll]=nsub_wdm(Mbin_edge[ll],Mbin_edge[ll+1],m_wdm=m_wdm,r=r,c0kpc=c0kpc,mf_slope=mf_slope)
        Mmid=10**(0.5*(Mbin_edge[ll]+Mbin_edge[ll+1]))
        rate_bin[ll]=sdf_pepper.subhalo_encounters(sigma=sigma,nsubhalo=nden_bin[ll],bmax=Xrs*rs(Mmid,plummer=True))

    rate = np.sum(rate_bin)

    Nimpact= numpy.random.poisson(rate)

    norm= 1./quad(lambda M : fac(M,m_wdm)*((M)**(mf_slope +0.5)),10**(Mbin_edge[0]),10**(Mbin_edge[Nbins]))[0]

    def cdf(M):
        return quad(lambda M : norm*fac(M,m_wdm)*(M)**(mf_slope +0.5),10**Mbin_edge[0],M)[0]

    MM=numpy.linspace(Mbin_edge[0],Mbin_edge[Nbins],10000)

    cdfl=[cdf(i) for i in 10**MM]
    icdf= interpolate.InterpolatedUnivariateSpline(cdfl,10**MM,k=1)
    timpact_sub= numpy.array(sdf_pepper._uniq_timpact)[numpy.random.choice(len(sdf_pepper._uniq_timpact),size=Nimpact,
                                p=sdf_pepper._ptimpact)]
    # Sample angles from the part of the stream that existed then
    impact_angle_sub= numpy.array([sdf_pepper._icdf_stream_len[ti](numpy.random.uniform())
            for ti in timpact_sub])

    sample_GM=lambda: icdf(numpy.random.uniform())/bovy_conversion.mass_in_msol(vo,ro)
    GM_sub= numpy.array([sample_GM() for a in impact_angle_sub])
    rs_sub= numpy.array([rs(gm*bovy_conversion.mass_in_msol(vo,ro)) for gm in GM_sub])
    # impact b
    impactb_sub= (2.*numpy.random.uniform(size=len(impact_angle_sub))-1.)*Xrs*rs_sub
    # velocity
    subhalovel_sub= numpy.empty((len(impact_angle_sub),3))
    for ii in range(len(timpact_sub)):
        subhalovel_sub[ii]=sdf_pepper._draw_impact_velocities(timpact_sub[ii],sigma,impact_angle_sub[ii],n=1)[0]
    # Flip angle sign if necessary
    if not sdf_pepper._gap_leading: impact_angle_sub*= -1.

    return impact_angle_sub,impactb_sub,subhalovel_sub,timpact_sub,GM_sub,rs_sub

def compute_obs_density_no_interpolation(phi1, apars, dens_apar):
    apar_edge=[]
    phi1_edge=[]

    abw0=apars[1]-apars[0]
    apar_edge.append(apars[0]-(abw0/2.))

    phi1bw0=phi1[1]-phi1[0]
    phi1_edge.append(phi1[0]-(phi1bw0/2.))

    for ii in range(len(apars)-1):
        abw=apars[ii+1]-apars[ii]
        phi1bw=phi1[ii+1]-phi1[ii]
        apar_edge.append(apars[ii]+abw/2.)
        phi1_edge.append(phi1[ii]+phi1bw/2.)

    abw_last=apars[len(apars)-1]-apars[len(apars)-2]
    apar_edge.append(apars[len(apars)-1]+(abw_last/2.))

    phi1bw_last=phi1[len(phi1)-1]-phi1[len(phi1)-2]
    phi1_edge.append(phi1[len(phi1)-1]+(phi1bw_last/2.))

    #compute the Jacobian d(apar)/d(phi1) using finite difference method
    dapar_dphi1=np.fabs(numpy.diff(apar_edge)/numpy.diff(phi1_edge))
    density = dens_apar * dapar_dphi1

    return density

def compute_obs_density(phi1,apars,dens_apar,Omega):

    apar_edge=[]
    phi1_edge=[]

    abw0=apars[1]-apars[0]
    apar_edge.append(apars[0]-(abw0/2.))

    phi1bw0=phi1[1]-phi1[0]
    phi1_edge.append(phi1[0]-(phi1bw0/2.))

    for ii in range(len(apars)-1):
        abw=apars[ii+1]-apars[ii]
        phi1bw=phi1[ii+1]-phi1[ii]
        apar_edge.append(apars[ii]+abw/2.)
        phi1_edge.append(phi1[ii]+phi1bw/2.)

    abw_last=apars[len(apars)-1]-apars[len(apars)-2]
    apar_edge.append(apars[len(apars)-1]+(abw_last/2.))

    phi1bw_last=phi1[len(phi1)-1]-phi1[len(phi1)-2]
    phi1_edge.append(phi1[len(phi1)-1]+(phi1bw_last/2.))

    #compute the Jacobian d(apar)/d(phi1) using finite difference method
    dapar_dphi1=np.fabs(numpy.diff(apar_edge)/numpy.diff(phi1_edge))
    #print (dapar_dphi1)

    #Interpolate dens(apar)
    ipdens_apar= interpolate.InterpolatedUnivariateSpline(apars,dens_apar)

    #Interpolate apar(phi1)
    if phi1[1] < phi1[0] : # ad-hoc way of checking whether increasing or decreasing
        ipphi1= interpolate.InterpolatedUnivariateSpline(phi1[::-1],apars[::-1])
        #Interpolate Jacobian
        ipdapar_dphi1=interpolate.InterpolatedUnivariateSpline(phi1[::-1],dapar_dphi1[::-1])
        #Interpolate density(phi1) by multiplying by jacobian
        dens_phi1=interpolate.InterpolatedUnivariateSpline(phi1[::-1],ipdens_apar(ipphi1(phi1[::-1]))*ipdapar_dphi1(phi1[::-1]))

    else :
        ipphi1= interpolate.InterpolatedUnivariateSpline(phi1,apars)
        #Interpolate Jacobian
        ipdapar_dphi1=interpolate.InterpolatedUnivariateSpline(phi1,dapar_dphi1)
        #Interpolate density(phi1) by multiplying by jacobian
        dens_phi1=interpolate.InterpolatedUnivariateSpline(phi1,ipdens_apar(ipphi1(phi1))*ipdapar_dphi1(phi1))

    return (dens_phi1(phi1))


<reponame>ClusterDuck123/ringity
import scipy
import numpy as np

def slope(rho, rate, a):
    mu_S = mean_similarity(rate,a)
    if rho <= mu_S:
        return rho/mu_S
    else:
        const = 1/np.sinh(np.pi*rate)
        def integral(k):
            term1 = np.sinh((1 + 2*a*(1/k-1))*np.pi*rate)
            term2 = (k*np.sinh((a*np.pi*rate)/k)*np.sinh(((a+k-2*a*k)*np.pi*rate)/k))/(a*np.pi*rate)
            return term1-term2
        return scipy.optimize.newton(
            func = lambda k: const*integral(k) + (1-cdf_similarity(1/k, rate, a)) - rho,
            x0 = rho/mu_S)


import sys
sys.path.append('C:/python scripts/ciecam02 plot')
import Read_Meredith as rm
# from scipy.ndimage import binary_dilation
# from scipy.stats import circstd
# import scipy.fftpack as fftpack
# from scipy.linalg import solve_banded
import vispol
import numpy as np
from scipy.sparse.linalg import spsolve
# from scipy.sparse.linalg import spilu
import matplotlib.pyplot as plt
from scipy.signal import convolve2d
# from scipy.ndimage import gaussian_filter
from scipy.ndimage import grey_erosion
from scipy.signal import medfilt2d
from scipy.signal import wiener

def ang_diff(D):
    P = np.pi
    return 1 - 2/P * np.abs((D + P/2) % (2 * P) - P)

vispol.register_cmaps()
data_folder = 'c:/users/z5052714/documents/unsw/unsw/data_sets/'
# filename = data_folder + 'forscottfrommeredith/162041.L1B2.v006.hdf5'
# filename = data_folder + 'forscottfrommeredith/162651.L1B2.v006.hdf5'
filename = data_folder + 'forscottfrommeredith/094821.L1B2.v006.hdf5'
# I, P, A = rm.getIPA(filename, start=[2, 32], end=[3900, 1403])
_, P, A = rm.getIPA(filename, start=[2, 32], end=[3000, 1403])
# I, P, A = rm.getIPA(filename, start=[1000, 800], end=[1600, 1300])
# I, _, A = rm.getIPA(filename, start=[2000, 150], end=[2200, 450]) # looks great
# _, _, A = rm.getIPA(filename, start=[1500, 150], end=[2200, 450]) # looks good with scaling
# _, _, A = rm.getIPA(filename, start=[2, 150], end=[2600, 1000])
# I = np.clip(I/np.percentile(I,99), 0, 1)
A *= np.pi/180.0
A[A > np.pi] -= np.pi

delta = vispol.delta_aop(A)
A45 = A + np.pi/8
A45[A45 > np.pi] -= np.pi
Aneg45 = A - np.pi/8
Aneg45[Aneg45 < 0 ] += np.pi
# plt.plot(np.linspace(-np.pi, np.pi, 256), ang_diff(np.linspace(-np.pi, np.pi, 256)))

# delta_patch = delta[2520:2620, 1200:1300].reshape((-1, 1))
# A_patch = A[2520:2620, 1200:1300].reshape((-1, 1))
# P_patch = P[2520:2620, 1200:1300].reshape((-1, 1))
# hist1, hist2, edges = np.histogram2d(A_patch, delta_patch, bins='fd')

# f, ax = plt.subplots(3)
# ax[0].scatter(A_patch, delta_patch)
# ax[1].scatter(A_patch, P_patch)
# ax[2].scatter(P_patch, delta_patch)
# print(np.mean(P_patch))
# print(np.std(P_patch))
# print(vispol.circular_mean(A_patch))
# print(np.sqrt(-2 * np.log(np.hypot(np.mean(np.sin(2 * A_patch)), np.mean(np.cos(2 * A_patch))))))
# print(np.mean(delta_patch))
# plt.show()

cap = 95
sigma = 2
# delta = grey_erosion(delta, size=(5, 5))
# delta = medfilt2d(delta, 7)
# delta = wiener(delta, 5)
# A, _ = vispol.histogram_eq(A,
#                                weighted=True,
#                                min_change=0.25,
#                                element=5,
#                                deltas = delta,
#                                suppress_noise=True,
#                                interval=[0.0,np.pi])#,
#                                # box=[[1100, A.shape[0]], [0, A.shape[1]]])
# plt.imsave("C:/users/z5052714/documents/weekly_meetings/28-06-2019/AoP_rot.png", A, cmap="AoP", vmin=0, vmax=np.pi)

f, ax = plt.subplots(1, 4)
# ax[0].imshow(delta, vmin=0, vmax=1)
# ax[1].imshow(A, vmin=0, vmax=np.pi, cmap="AoP")
# ax[2].imshow(P, vmin=0, vmax=1)
ax[0].imshow(np.cos(2 * A), cmap="gray")
ax[1].imshow(np.cos(2 * A45), cmap="gray")
ax[2].imshow(np.sin(2 * A), cmap="gray")
ax[3].imshow(np.cos(2 * Aneg45), cmap="gray")
plt.show()
# kernel = np.array([[0, 1, 0], [1, -4, 1], [0, 1, 0]])

# Ld = convolve2d(delta, kernel, mode='same', boundary='symm')
# Ld /= np.percentile(np.abs(Ld), cap)
# Ld = np.clip(Ld, -1, 1)

x_neg = np.array([[0, 0, 0],
                  [0, -1, 1],
                  [0, 0, 0]])
x_pos = np.array([[0, 0, 0],
                  [1, -1, 0],
                  [0, 0, 0]])
y_neg = np.array([[0, 0, 0],
                  [0, -1, 0],
                  [0, 1, 0]])
y_pos = np.array([[0, 1, 0],
                  [0, -1, 0],
                  [0, 0, 0]])

L = np.zeros_like(A)

# plt.imshow(medfilt2d(delta, 7))
# plt.show()

# f, ax = plt.subplots(1,4)
# cosA = np.cos(2 * A)
# sinA = np.sin(2 * A)
# ax[0].imshow(cosA, vmin=-1, vmax=1, cmap="BlueWhiteRed")
# ax[2].imshow(sinA, vmin=-1, vmax=1, cmap="BlueWhiteRed")
# filt_size = 5
# cosA = wiener(cosA, filt_size)
# sinA = wiener(sinA, filt_size)
# cosA = medfilt2d(cosA, filt_size)
# sinA = medfilt2d(sinA, filt_size)
# ax[1].imshow(cosA, vmin=-1, vmax=1, cmap="BlueWhiteRed")
# ax[3].imshow(sinA, vmin=-1, vmax=1, cmap="BlueWhiteRed")
# plt.show()
close_to_zero = np.abs(np.cos(2 * A) - 1) < 0.000005

for kernel in [x_neg, x_pos, y_neg, y_pos]:
# for kernel in [x_neg, y_neg]:
    # Lsub = np.sin(convolve2d(A, kernel, mode='same', boundary='symm'))
    Lsub0 = ang_diff(convolve2d(A, kernel, mode='same', boundary='symm'))
    Lsub45 = ang_diff(convolve2d(A45, kernel, mode='same', boundary='symm'))
    f, ax = plt.subplots(1, 5)
    ax[0].imshow(Lsub0, vmin=-1, vmax=1, cmap="BlueWhiteRed")
    ax[1].imshow(Lsub45, vmin=-1, vmax=1, cmap="BlueWhiteRed")
    ax[2].imshow(close_to_zero)
    ax[3].imshow(Lsub0 - Lsub45, cmap="BlueWhiteRed", vmin=-0.1, vmax=0.1)
    Lsub = Lsub0
    Lsub[close_to_zero] = Lsub45[close_to_zero]

    ax[4].imshow(Lsub - Lsub0, vmin=-.1, vmax=.1, cmap="BlueWhiteRed")
    # plt.show()
    # cos_arr = convolve2d(cosA, kernel, mode='same', boundary='symm')
    # sin_arr = convolve2d(sinA, kernel, mode='same', boundary='symm')
    # Lsub = cos_arr
    # Lsub[np.abs(cos_arr) < np.abs(sin_arr)] = sin_arr[np.abs(cos_arr) < np.abs(sin_arr)]
    L += Lsub

# L[500,500] = 0
# L = np.sin(np.pi/2 * L)

# from scipy.special import erf, erfinv
# endpoint = 0.99999999999
# factor = erfinv(endpoint)
# L = erf(factor * L) / endpoint
plt.figure()
plt.imshow(L, cmap="BlueWhiteRed", vmin=-1, vmax=1)

# plt.show()

n, m = A.shape

# kernel = np.array([[0, 1, 0], [1, -4, 1], [0, 1, 0]])
# L_F = fftpack.fft2(L)
# K_F = fftpack.fft2(kernel, shape=L.shape)
# f, ax = plt.subplots(1, 2)
# ax[0].imshow(np.real(L_F))
# ax[1].imshow(np.real(K_F))
# plt.show()
# U = fftpack.ifft2(L_F / K_F)
M = vispol.construct_matrix(A.shape, type='laplacian')
# M = vispol.construct_matrix(A.shape, type='lap_xy')

U = spsolve(M, L.reshape(n * m, 1)).reshape((n, m))
# U -= np.median(U)
#
# U /= 2 * np.max(np.abs(U))
# U += 0.5
#

U = (U - np.min(U)) / (np.max(U) - np.min(U))
U_enhanced = (U - np.percentile(U, 1)) / (np.percentile(U, 99) - np.percentile(U, 1))
U_enhanced = np.clip(wiener(U_enhanced, 3), 0, 1)
U = np.clip(U, 0, 1)
Ushrink = U_enhanced * (0.8 - 0.2) + 0.2

plt.figure()
A_slice = A[1000,:600]/np.pi
# A_slice[A_slice > 0.93] -= 1
# A_slice[564:] -= 1
U_slice = U[1000,:600]
slopes_A = ang_diff(np.convolve(A_slice, [-1, 1], mode='same'))
slopes_U = np.convolve(U_slice, [-1, 1], mode='same')
# A_slice = (A_slice - np.min(A_slice)) / (np.max(A_slice) - np.min(A_slice))
# U_slice = (U_slice - np.min(U_slice)) / (np.max(U_slice) - np.min(U_slice))
# A_delta = np.convolve(A_slice, [1, -1], mode='same')
# A_slice[A_delta > 0.5] -= 0.5
# A_slice[A_delta < -0.5] += 0.5
plt.plot(range(599), slopes_A[1:])
plt.plot(range(599), slopes_U[1:])
# plt.plot(range(600), slopes_A / slopes_U)
# plt.plot(range(600), np.abs(slopes_A / slopes_U))

P = 1 - np.clip(delta / np.max(delta), 0, 1)
# mask_params ={'thresh':0.4,
#               'smooth':True,
#               'morph':True}
# delta_params = {'mask_on':True, 'mask_params':mask_params}
# dmask = vispol.dmask(delta, thresh=0.3, smooth=True, morph=True)

# U *= dmask
# A += np.pi/2
RGB = vispol.IPAtoRGB(I = Ushrink, P=P, A=A, dependent="P")

f, ax = plt.subplots(1, 2)
# plt.imsave("C:/users/z5052714/documents/weekly_meetings/28-06-2019/U.png", U, cmap="gray")
# plt.imsave("C:/users/z5052714/documents/weekly_meetings/28-06-2019/U_enhanced.png", U_enhanced, cmap="gray")
# plt.imsave("C:/users/z5052714/documents/weekly_meetings/28-06-2019/RGB_rot.png", RGB)
ax[0].imshow(U_enhanced, cmap="gray")
ax[1].imshow(RGB)
# for idx, sig in enumerate([30, 20, 10, 5]):
#     ax[idx + 1].imshow(U - gaussian_filter(U, sigma=sig), cmap="gray")
    # ax[idx].imshow(Usig, cmap="gray")
plt.show()


# -*- coding: utf-8 -*-
"""
Simple Detection Code
"""

import matplotlib
matplotlib.use('agg')
import glob
import sys
import numpy as np
import os.path as op
import fitsio
import argparse as ap

from astropy.io import fits
from astropy.convolution import convolve, Gaussian1DKernel, interpolate_replace_nans
from input_utils import setup_logging
from scipy.interpolate import interp1d
from astropy.stats import mad_std, biweight_location
from utils import biweight_midvariance
from astropy.table import Table
from scipy.signal import savgol_filter, medfilt

def get_script_path():
    return op.dirname(op.realpath(sys.argv[0]))

DIRNAME = get_script_path()

def grab_attribute(filename, args, attributes=[], amps=['LL', 'LU', 'RU',
                   'RL']):
    ''' grab specified attributes from multi* file '''
    basename = filename[:-8]
    s = [[] for a in attributes]
    norm = []
    for amp in amps:
        name = basename + '_%s.fits' % amp
        data = fitsio.read(name, 'twi_spectrum')
        xl = data.shape[1] / 3
        xh = 2 * data.shape[1] / 3
        norm.append(np.median(data[:, xl:xh], axis=1))
    for amp in amps:
        name = basename + '_%s.fits' % amp
        for i, attribute in enumerate(attributes):
            s[i].append(fitsio.read(name, attribute))
    X = [np.vstack(si) for si in s]
    X.append(np.hstack(norm))
    return X

def put_attribute(filename, args, data, attributes=[]):
    ''' put specified attributes into multi* file '''
    try:
        for i, attribute in enumerate(attributes):
            F = fitsio.FITS(filename, 'rw')
            F.write(data[i], extname=attribute+'_1')
    except IOError:
        for i, attribute in enumerate(attributes):
            args.log.warning('%s not found to add %s' % attribute)

def rectify(wave, spec, lims, mask=None, fac=1.0):
    N, D = wave.shape
    rect_wave = np.linspace(lims[0], lims[1], int(D*fac))
    rect_spec = np.zeros((N, len(rect_wave)))
    G = Gaussian1DKernel(1.5 * fac)
    for i in np.arange(N):
        dw = np.diff(wave[i])
        dw = np.hstack([dw[0], dw])
        if mask is None:
            x = wave[i]
            y = spec[i] / dw
        else:
            x = wave[i]
            y = (spec[i] / dw)
            y[mask[i]] = np.nan
            y = interpolate_replace_nans(y, G)
        I = interp1d(x, y, kind='quadratic',
                     bounds_error=False, fill_value=-999.)
        rect_spec[i, :] = I(rect_wave)
    return rect_wave, rect_spec

def build_weight_matrix(x, y, sig=1.5):
    d = np.sqrt((x - x[:, np.newaxis])**2 + (y - y[:, np.newaxis])**2)
    G = np.exp(-0.5 * (d / sig)**2)
    G = G / G.sum(axis=0)[:, np.newaxis]
    return G.swapaxes(0, 1)

def convolve_spatially(x, y, spec, wave, mask, sig_spatial=0.7, sig_wave=1.5):
    W = build_weight_matrix(x, y, sig=sig_spatial)
    Z = spec * 1.
    G = Gaussian1DKernel(sig_wave)
    for i in np.arange(spec.shape[0]):
        Z[i, :] = convolve(Z[i, :], G, nan_treatment='fill', fill_value=0.0)
    for i in np.arange(spec.shape[1]):
        Z[:, i] = np.dot(Z[:, i], W)
    return Z

def safe_division(num, denom, eps=1e-8, fillval=0.0):
    good = np.isfinite(denom) * (np.abs(denom) > eps)
    div = num * 0.
    if num.ndim == denom.ndim:
        div[good] = num[good] / denom[good]
        div[~good] = fillval
    else:
        div[:, good] = num[:, good] / denom[good]
        div[:, ~good] = fillval
    return div

def mask_cosmics(error, trace, cosmic_avoidance=4):
    cyind, cxind = np.where(error == -1)
    mask = np.array(trace * 0., dtype=bool)
    for xind, yind in zip(cxind, cyind):
        trace_a = trace[:, xind]
        fibs = np.where(np.abs(trace_a - yind) < cosmic_avoidance)[0]
        for fib in fibs:
            lx = (xind-cosmic_avoidance)
            hx = (xind+cosmic_avoidance) + 1
            mask[fib, lx:hx] = True
    return mask

def dummy_test(image):
    y = savgol_filter(image, 315, 1, axis=1)
    s = np.zeros((y.shape[1], 4))
    for i in np.arange(y.shape[1]):
        chunks = np.array_split(y[:, i], 4)
        avg = chunks[-1]
        n = [biweight_location(avg / chunk) for chunk in chunks]
        s[i, :] = np.array(n)

    chunks = np.array_split(y, 4, axis=0)
    norm = np.zeros((3*448, chunks[0].shape[1]))
    for k in np.arange(3):
        for i in np.arange(y.shape[1]):
            x = np.arange(chunks[0].shape[0])
            z = (chunks[k] * s[:, k] / chunks[-1])[:, i]
            test = z - medfilt(z, 51)
            threshold = 3. * np.nanmedian(np.abs(test))
            mask = np.abs(test) < threshold
            xchunk, zchunk, mchunk = [np.array_split(j, 4, axis=0) for j in [x, z, mask]]
            for xc, zc, mc in zip(xchunk, zchunk, mchunk):
                p = np.polyfit((xc/448.)[mc], zc[mc], 2)
                norm[xc+448*k, i] = np.polyval(p, xc/448.)
    return norm

def main():
    parser = ap.ArgumentParser(add_help=True)

    parser.add_argument("-f", "--filename",
                        help='''Filename that contains list of files''',
                        type=str, default=None)
    parser.add_argument("-ac", "--spatial_conv_size",
                        help='''Spatial Convolution Kernel Sigma (")''',
                        type=float, default=0.6)
    parser.add_argument("-ec", "--spectral_conv_size",
                        help='''Spectral Convolution Kernel Sigma (A)''',
                        type=float, default=2.5)
    parser.add_argument("-cc", "--spectral_cont_conv_size",
                        help='''Spectral Continuum Convolution Kernel Sigma (A)''',
                        type=float, default=25.)
    parser.add_argument("-t", "--threshold",
                        help='''Detection Threshold''',
                        type=float, default=5.0)
    parser.add_argument("-o", "--outdir",
                        help='''Out directory for detections''',
                        type=str, default='temp')
    args = parser.parse_args(args=None)
    args.log = setup_logging(logname='detection')

    filenames = [line.rstrip('\n').split()
                 for line in open(args.filename, 'r')]
    allwave, allspec, allifupos, allmask, alltwi, allmodel, allftf = ([], [], [], [], [], [], [])
    for filename in filenames:
        args.log.info('Reading in %s' % filename[0][:-8])
        dither = np.array([float(filename[2]), float(filename[3])])
        amps = ['LL', 'LU', 'RU', 'RL']
        attributes = ['wavelength', 'spectrum', 'twi_spectrum',
                      'ifupos', 'error', 'trace', 0, 'flat_image']
        w, s, f, i, e, t, T, m, n  = grab_attribute(filename[0], args,
                                             attributes=attributes, amps=amps)
        mask = mask_cosmics(e, t)
        norm = (n / np.median(n))[:, np.newaxis]
        allwave.append(w)
        allspec.append(s)#safe_division(s, f * norm))
        allifupos.append(i + dither)
        allmask.append(mask)
        alltwi.append(T)
        allmodel.append(m)
        allftf.append(f)
    allwave, allspec, allifupos, allmask, alltwi, allmodel, allftf = [np.array(np.vstack(x), dtype='float64')
                                            for x in [allwave, allspec,
                                                      allifupos, allmask, alltwi, allmodel, allftf]]
    args.log.info('Rectifying sky subtracted spectra')
    allmask = np.array(allmask, dtype=bool)
    rw, rs = rectify(allwave, allspec, [3500., 5500.], mask=allmask, fac=1.5)
    args.log.info('Convolving sky subtracted spectra for continuum')
    Zc = convolve_spatially(allifupos[:, 0], allifupos[:, 1], rs, rw, allmask,
                           sig_spatial=args.spatial_conv_size,
                           sig_wave=(args.spectral_cont_conv_size / (rw[1]-rw[0])))
    noise = biweight_midvariance(Zc, axis=(0, ))
    SNc = Zc / noise
    F1 = fits.PrimaryHDU(alltwi)
    args.log.info('Convolving sky subtracted spectra for emission')
    Ze = convolve_spatially(allifupos[:, 0], allifupos[:, 1], rs, rw, allmask,
                           sig_spatial=args.spatial_conv_size,
                           sig_wave=(args.spectral_conv_size / (rw[1]-rw[0])))
    noise = biweight_midvariance(Ze-Zc, axis=(0, ))
    SNe = (Ze-Zc) / noise
    F2 = fits.ImageHDU(allmodel)
    #norm = dummy_test(np.vstack([allspec, alltwi[:448,:]]))
    F3 = fits.ImageHDU(allftf)
    fits.HDUList([F1, F2, F3]).writeto('test.fits', overwrite=True)
    # peaks_fib, peaks_wave = np.where(SN > args.threshold)

if __name__ == '__main__':
    main()


<filename>python_scripts/opencv_utils.py<gh_stars>0
# -*- coding: utf-8 -*-
"""
Created on Fri Jun 23 10:45:49 2017

@author: Aditi
"""
from collections import deque
import numpy as np
import cv2
import math
import os
from scipy.spatial import distance as dist
import copy

COLORS = {  'plains_color': np.array([42,161,152], dtype=np.uint8),
            'wetlands_color': np.array([108,113,196], dtype=np.uint8),
            'desert_color': np.array([181,137,0], dtype=np.uint8),
            'jungle_color': np.array([211,54,130], dtype=np.uint8),
            'reservoir_color': np.array([232, 236, 237], dtype=np.uint8),
            'mountain_valley_color': np.array([232, 236, 237], dtype=np.uint8),
            'water_color': np.array([0,255,255], dtype=np.uint8),
            'logs_color': np.array([255,0,0], dtype=np.uint8) }

#to capture the position / cordinates of the mini video
def know_coordinates(event,x,y,flags,param):
    if event == cv2.EVENT_LBUTTONDBLCLK:
        print (x ,",", y)

# finds the angle of the contour of water path and writes it into a file
def detect_water_encoding(min_video_frame,frame_num,directory):
    angle_matrix=[]
    #to decode the color information theo embeded
    hsv = cv2.cvtColor(min_video_frame, cv2.COLOR_BGR2HSV)
    # define range of green?? color in HSV
    lower_green = np.array([29,86,6])
    upper_green = np.array([64,255,255])

    # Threshold the HSV image to get only green colors
    mask = cv2.inRange(hsv, lower_green, upper_green)
    # Bitwise-AND mask and original image
    res = cv2.bitwise_and(min_video_frame,min_video_frame, mask= mask)
   # cv2.imshow('res',res)
    #cv2.imwrite('frame.png',res)
    name = str(directory)+"\\water_path_angle"+str(frame_num)+".csv"
    rows, cols, channel = res.shape
    handle = open(name, 'wb')
    blank = np.zeros(res.shape,res.dtype)
    #get green and red channel and convert into useful information
    for row in range(0,rows):
        for col in range(0,cols):
            k = res[row,col]    #B=0 G=1 R=2
            #green channel
            green =  res.item(row,col,1)-120
            #red channel
            red = res.item(row,col,2)-120
            handle.write("\n")
            handle.write(str(row))
            handle.write(",")
            handle.write(str(col))
            handle.write(",")
            #Blue channel which was made 0 above
            handle.write(str(res.item(row,col,0)))
            handle.write(",")
            handle.write(str(green))
            handle.write(",")
            handle.write(str(red))
            handle.write(",")
            #all channels
            handle.write(str(k))
            """ red+120 <0 and green+120 < 0 first quadrant
                if green negative and red==0 y axis decide
                direction based on green negative or poistive
                if red < 0 ( third quadrant) add 180
                else check green negative
            """
            if(red + 120 < 20 and green + 120 < 20):
                angle = -1000
            elif(red==0):
                if((-1)*green < 0):#downwards y axis negative
                    angle =270
                    handle.write(",")
                    handle.write(str(angle))
                elif((-1)*green > 0): #upwards y axis positive
                    angle =90
                    handle.write(",")
                    handle.write(str(angle))
                else:
                    angle= -1000
            else:
                if(red < 0): #second quadrant left negative x axis
                    angle= 180 + math.degrees(math.atan2((-1)*green,red))
                    handle.write(",")
                    handle.write(str(angle))
                else: #positive x axis right
                    if((-1)*green < 0): #first quadrant upwards negative y axis
                        angle= 360 + math.degrees(math.atan2((-1)*green,red))
                        handle.write(",")
                        handle.write(str(angle))
                    else:#    fourth quadrant
                        angle=  math.degrees(math.atan2((-1)*green,red))
                        handle.write(",")
                        handle.write(str(angle))
            if(angle >= 0 and angle < 90):
                cv2.circle(blank,(col,row),1,(0,255,0),-1)
            elif(angle >=90 and angle <180):
                cv2.circle(blank,(col,row),1,(0,0,255),-1)
            elif(angle>= 180 and angle<270):
                cv2.circle(blank,(col,row),1,(255,0,0),-1)
            elif(angle>=270 and angle<=360):
                cv2.circle(blank,(col,row),1,(0,255,255),-1)
            angle_matrix.append([row,col,angle])
            #cv2.imshow("blank",blank)
    handle.close()
    return angle_matrix

#just identify water flow path for drawing graphs
def detect_water(min_video_frame):
    hsv = cv2.cvtColor(min_video_frame, cv2.COLOR_BGR2HSV)
    # define range of green/yellow color in HSV
    lower_green = np.array([29,86,6])
    upper_green = np.array([64,255,255])
    th3 = extract_boundary(min_video_frame,hsv,lower_green, upper_green,0)
    store = th3
    # morphing to get the skeletal structure/ medial line of the water flow
    size = np.size(th3)
    skel = np.zeros(th3.shape,np.uint8)
    element = cv2.getStructuringElement(cv2.MORPH_CROSS,(3,3))
    done = False

    while(not done):
        eroded = cv2.erode(th3,element)
        temp = cv2.dilate(eroded,element)
        temp = cv2.subtract(th3,temp)
        skel = cv2.bitwise_or(skel,temp)
        th3 = eroded.copy()

        zeros = size - cv2.countNonZero(th3)
        if zeros==size:
            done = True
    return store,skel

"""checks if a particular log is active or inactive, i.e. if
log touches water active else inactive"""
def check_logs(log,water):
    res = np.logical_and(log, water)
    ans = np.where(res == True)
    if len(ans[0])==0 and len(ans[1])==0:
        return "inactive"
    return "active"

def detect_logs(min_video_frame, filename, thresh_water):
    hsv = cv2.cvtColor(min_video_frame, cv2.COLOR_BGR2HSV)
    # define range of blue color in HSV
    lower_blue = np.array([110,50,50])
    upper_blue = np.array([130,255,255])

    th3 = extract_boundary(min_video_frame,hsv,lower_blue, upper_blue,0)

    #smooth the logs (current version very fat lines)
    image ,contours, heirarchy = cv2.findContours(th3,1,2)#cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE)

    #Draw log contour + bonding rects
    colored = cv2.cvtColor(image,cv2.COLOR_GRAY2BGR)
    count =0
    black = np.zeros(colored.shape)
    centers=[]
    for contour in contours:
       # returns top left (x,y), width, height, and angle of rotation
       rect = cv2.minAreaRect(contour)
       box = cv2.boxPoints(rect)
       box = np.int0(box)
       """ rotated rectangle angle in opencv is calculated fro -0 to
       -90 in degrees as the angle the lowest vertex of the bounding
       rectangle makes with the horizontal counter clockwise direction, (the vertices are
       considered clockwise from the lowest vertex), width is the
       distance between second and third vertex, and height is the
       distance between first and second vertex.
       The following "if" converts the angle to o to -180 counterclockwise."""
       if(rect[1][0] < rect[1][1]):
           angle = rect[2] - 90
       else:
           angle = rect[2]
       filename.write(str(count))
       filename.write(",")
       filename.write(str(abs(angle)))
       filename.write(",")
       filename.write(str(rect[0][0]))
       filename.write(",")
       filename.write(str(rect[0][1]))
       filename.write(",")

       #to check if log is active/ inactive
       blank = np.zeros(thresh_water.shape,thresh_water.dtype)
       log = cv2.drawContours(blank,[contour],-1,(255,0,255),-1)
       res = check_logs(log,thresh_water)
       filename.write(res)
       filename.write(",")
       count= count+1
       image=cv2.drawContours(colored,[box],-1,(0,0,255),-1)
       #returns centroid of the bounding rectangle same  as contour centroid

       cx= int(rect[0][0])
       cy= int(rect[0][1])
       image = cv2.circle(black,(cx,cy),1,(0,255,0),-1)
       if(res is "active"):
           centers.append([cx,cy])
    return image,filename,centers

#def check_validity(frame, frame_number, handle_valid):
#    handle_valid.write("\n")
#    rows,cols, channels = frame.shape
#    print frame.shape
#    for row in range(0,rows):
#        for col in range(0,cols):
#            if(row>= 5 and row <=9 and col >=215 and col<=219):
#                k = frame[row,col]
#                handle_valid.write(str(row))
#                handle_valid.write(",")
#                handle_valid.write(str(col))
#                handle_valid.write(",")
#                handle_valid.write(str(frame_number))
#                handle_valid.write(",")
#                handle_valid.write(str(k))
#                handle_valid.write(",")
#    return handle_valid

def extract_boundary(original,hsv_image, lower, upper, flag):
    # need end points of the boundary too
    mask = cv2.inRange(hsv_image, lower, upper)
    # Bitwise-AND mask and original image
    res = cv2.bitwise_and(original,original,mask=mask)
    #boundaries in gray scale
    gray = cv2.cvtColor(res,cv2.COLOR_BGR2GRAY)
    # Otsu's thresholding and gaussian filtering  to make the logs white and the background black for better detection
    ret2,th2 = cv2.threshold(gray,0,255,cv2.THRESH_BINARY+cv2.THRESH_OTSU)
    # Otsu's thresholding after Gaussian filtering
    blur = cv2.GaussianBlur(gray,(5,5),0)
    #logs will be white in th3
    ret3,th3 = cv2.threshold(blur,0,255,cv2.THRESH_BINARY+cv2.THRESH_OTSU)
    if(flag==1):
        black, extLeft, extRight, cx,cy = find_contour(th3,original)
        return black,extLeft,extRight,cx,cy
    return th3

def find_contour(th3,original):
    black = np.zeros(original.shape,original.dtype)
    image ,contours, heirarchy = cv2.findContours(th3,1,2)
    for contour in contours:
        M = cv2.moments(contour)
        cx = int(M['m10']/M['m00'])
        cy = int(M['m01']/M['m00'])
        black = cv2.circle(black,(cx,cy),3,(255,255,255),-1)
        extLeft = tuple(contour[contour[:, :, 0].argmin()][0])
        extRight = tuple(contour[contour[:, :, 0].argmax()][0])
        black = cv2.circle(black,extLeft,3,(255,255,255),-1)
        black = cv2.circle(black,extRight,3,(255,255,255),-1)
    return black,extLeft,extRight,cx,cy

def fix_color(matrix, lower, upper, new_color, roll=2):
    mask = np.zeros_like(matrix[:,:,0])
    mask += matrix[:,:,0]<=upper[0]
    mask += matrix[:,:,0]>=lower[0]
    mask += matrix[:,:,1]<=upper[1]
    mask += matrix[:,:,1]>=lower[1]
    mask += matrix[:,:,2]<=upper[2]
    mask += matrix[:,:,2]>=lower[2]

    matrix[mask==6] = new_color
    matrix[np.roll(mask, roll, axis=0)==6] = new_color
    matrix[np.roll(mask, -roll, axis=0)==6] = new_color
    matrix[np.roll(mask, roll, axis=1)==6] = new_color
    matrix[np.roll(mask, -roll, axis=1)==6] = new_color
    return matrix

def mark_boundaries():
     boundary_ends=[]
     center_biomes =[]

     convert_boundaries = cv2.imread("CW-BiomeBoundary.png")
     hsv = cv2.cvtColor(convert_boundaries, cv2.COLOR_BGR2HSV)

     # mountain valley
     lower = np.array([0,100,0])
     upper = np.array([5,255,5])
     convert_boundaries = fix_color(convert_boundaries, lower, upper, COLORS['mountain_valley_color'])

     #Plains
     lower = np.array([170,0,0])
     upper = np.array([255,5,5])
     convert_boundaries = fix_color(convert_boundaries, lower, upper, COLORS['plains_color'])

     #jungle - megenta
     lower=np.array([100, 0, 100])
     upper=np.array([255,5,255])
     convert_boundaries = fix_color(convert_boundaries, lower, upper, COLORS['jungle_color'])

     #wetlands
     lower=np.array([0, 150, 150])
     upper=np.array([5,255,255])
     convert_boundaries = fix_color(convert_boundaries, lower, upper, COLORS['wetlands_color'])

     #reservoir
     lower=np.array([135, 135, 0])
     upper=np.array([255,255,5])
     convert_boundaries = fix_color(convert_boundaries, lower, upper, COLORS['reservoir_color'])

     #desert
     lower = np.array([0,0,170])
     upper = np.array([5,5,255])
     convert_boundaries = fix_color(convert_boundaries, lower, upper, COLORS['desert_color'])

     temp = np.float64(convert_boundaries)
     return temp, boundary_ends, center_biomes

def find_biome_sector(shape,dtype,start,end,x,y):
    blank = np.zeros(shape,dtype)
    pts=np.array([start,end,[x,y]])
    rect = cv2.minAreaRect(pts)
    box = cv2.boxPoints(rect)
    upper = dist.euclidean(box[0],box[1])
    lower = dist.euclidean(box[1], box[2])
    #if longest side is not necessarily the 0 and 1 coordinates
    if lower > upper:
        temp = lower
        lower= upper
        upper= temp
    box = np.int0(box)
    cv2.drawContours(blank,[box],0,(255,255,255),2)
    filled = cv2.fillPoly(blank, [pts], color=(255,255,255))
    return filled, lower, upper

#def find_long_side(x,y,start,end):
#    dist1 = dist.euclidean([x,y],start)
#    dist2 = dist.euclidean([x,y],end)
#    limit = abs(dist1-dist2)
#    if(dist1 >= dist2):
#        return dist1,limit
#    else:
#        return dist2,limit

def log2log(biome,log_centers,flag_logs):
    i =0
    for x,y in log_centers:
        if(biome[y,x] == 255 and flag_logs[i]==0):
            join =[x,y]
            break
        i=i+1
    if(i >= len(log_centers)):
         join = [-1000,-1000]
    return join

#def customsort(log_centers,logs):
#    log_x=[]
#    log_y=[]
#    for log in logs:
#        x,y = log_centers[log]
#        log_x.append(x)
#        log_y.append(y)
#    sorted_logs = np.lexsort((log_y,log_x))
#    return_logs = [(log_x[i],log_y[i]) for i in sorted_logs]
#    print return_logs
#    return return_logs

def draw_graph(boundaries,skeleton,log_centers,boundary_ends,center_biome):
    log_list = []
    flag=0
    for log in log_centers:
        #a flag along with the coordinates to indcate log=0 or wall=1
        log_list.append([])
        x,y = log
        log_list[flag].append([(x,y),0])
        flag=flag+1

    flag_logs =np.zeros(len(log_centers))
    blank = np.zeros(boundaries.shape, boundaries.dtype)
    flag =0
    for x,y in log_centers:
            flag_logs[flag]=1
            cv2.circle(blank,(x,y),3,[0,0,255],-1)
            count=0
            for i in range(0,len(boundary_ends),2):
                filled_biome, lower, upper = find_biome_sector(skeleton.shape,skeleton.dtype,boundary_ends[i],boundary_ends[i+1],x,y)
                res = cv2.bitwise_and(np.float32(filled_biome),np.float32(skeleton))
                #cv2.imshow("result",res)
                if(cv2.countNonZero(res) > lower):#lower and cv2.countNonZero(res) <= upper):  #threshhold of  pixels
                    #indices of logs in the same sector=1 else 0
                    join = log2log(filled_biome,log_centers,flag_logs)
                    x_join , y_join = join
                    #separate join and if it is not -1000 then draw line from log to end
                    if(x_join!=-1000 and y_join!=-1000):
                        log_list[flag].append([(x_join,y_join),0])
                        cv2.line(blank.astype(np.int32),(x,y),(x_join,y_join),[0,255,255],1)
                       # cv2.imshow("middle",blank)#check for upper asending may not work
                    else: #(other check can be checking if point of intersection is between which boundaries ends and join for adjacent boundary ends)
                        #check point of intersection in which sector this or neighbouring
                        if(count==1):
                            center_count = check_bound_intersection(boundaries,res,boundary_ends,1)
                        elif(count==5):
                            center_count = check_bound_intersection(boundaries,res,boundary_ends,2)
                        else:
                            center_count = count
                        center_count = int(center_count)
                        if(center_count!=13):     # if 13 no line to be drawn
                            cb = (int(center_biomes[center_count][0]), int(center_biomes[center_count][1]))
                            log_list[flag].append([cb,1])
                            cv2.line(blank,(x,y),cb,[0,255,255],1)
                count=count+1
            flag=flag+1
    #checking cycles without directions is difficult
    #so for each two logs that are neighbours of each other measure distance from the common wall
    #check for cycles remove redundant edges  (can i remove longest edge??)

    #log_list = remove_redundant_edges(log_list)

    #convert into alphabets and write into a file framewise
    #blanknew = cv2.add(boundaries,blank)
   # cv2.imshow('blanknew',blanknew)
    return log_list

def remove_redundant_edges(log_list):
    count =0
    for log in log_list:
        neighbour_logs = find_neighbour_log(log)
        head = neighbour_logs.popleft()
        while (len(neighbour_logs)!=0):
            tail = neighbour_logs.popleft()
            index = locate_tail(tail,log_list)

            common = findcommon(log_list[count],log_list[index])
            for each_common in common:
                discard = find_discard(head,tail,each_common)
                if(discard =="head"):
                    log_list[count].remove(each_common)
                else:
                    log_list[index].remove(each_common)
        count =count+1
    return log_list

def locate_tail(tail, log_list):
    all_heads = [item[0] for item in log_list]
    for i in range(0,len(all_heads)):
        if(all_heads[i]==tail):
            return i

def find_discard(log1,log2,common):
    [log1_x,log1_y], log1_flag = log1
    [log2_x,log2_y], log2_flag = log2
    [common_x,common_y], common1_flag = common
    log1_common = dist.euclidean([log1_x,log1_y],[common_x,common_y])
    log2_common = dist.euclidean([log2_x,log2_y],[common_x,common_y])
    if(log1_common > log2_common):
        return "head"
    else:
        return "tail"

def findcommon(log1, log2):
    temp1 = copy.deepcopy(log1)
    temp2 = copy.deepcopy(log2)
    del temp1[0]
    del temp2[0]
    return [x for x in temp1 if x in temp2]

def find_neighbour_log(log):
    d = deque()
    for element in log:
        (x,y),flag = element
        if (flag==0):
           d.append([(x,y),flag])
    return d

def check_bound_intersection(boundaries,res,boundary_ends,flag):
    #may not work when  some times the water doesnt intersect the walls like in desert goes little out side
     boundaries_thresh = cv2.cvtColor(np.float32(boundaries),cv2.COLOR_BGR2GRAY)
     res = cv2.bitwise_and(boundaries_thresh,np.float32(res))
     index = np.where(res)
     index_x = index[1]
     index_y = index[0]
     if(flag==1):
         start = 0
         finish = 5
     elif(flag==2):
         start=8
         finish=13
     for ends in range(start, finish,2):
         if(ends>=0 and ends<=5):
             lower= boundary_ends[ends]
             upper = boundary_ends[ends+1]
         elif(ends >=8  and ends <=13):
             lower= boundary_ends[ends +1]
             upper = boundary_ends[ends]
        # if(x >= lower[0] and y>=lower[1] and x<= upper[0] and y<=upper[1]):
         check = np.where(np.greater_equal(index_x,lower[0]) & np.less_equal(index_x,upper[0]) & np.greater_equal(index_y,lower[1]) & np.less_equal( index_y,upper[1]))
         if(len(check[0])!=0):
             return ends/2   # only first point may not be a good idea

             #may be calculate maximum coverage

     return 13   # some dummy value if intersection doesnt fall in the range

#def check_angle(angle_matrix, row,col):
#    for row_elem, col_elem,angle in angle_matrix:
#         if row == row_elem and col == col_elem:
#            return angle
#    return 0

def add_waterfall(boundaries, boundary_ends, center_biomes):
    water_fall_start = boundary_ends[6]
    water_fall_end = boundary_ends[7]
    x_start= water_fall_start[0]
    y_start = water_fall_start[1]
    x_end = water_fall_end[0]
    y_end = water_fall_end[1]
    x_mid = (x_start + x_end) / 2
    y_mid = (y_start + y_end) / 2
    center_biomes.insert(3,((x_mid,y_mid)))
    boundaries = cv2.line(boundaries,boundary_ends[6],boundary_ends[7],[255,255,255],3)
    return boundaries,center_biomes

def biome_alpha(index):
    if(index==0):
        return "D"
    elif(index==1):
        return "MV"
    elif(index==2):
        return "P"
    elif(index==3):
        return "WF"
    elif(index==4):
        return "J"
    elif(index==5):
        return "R"
    elif(index==6):
        return "W"

def write_log_list(image,file_write,log_list,frame, center_biomes):
    file_write.write("\n"+str(frame))
    file_write.write(",")
    unique_logs=[item[0] for item in log_list]
    for lists in log_list:
        head = lists[0]
        head_pts, flag = head
        if(flag==0):
            file_write.write(str(head_pts)+": log"+str(unique_logs.index(head))+"->")
        for tail in range(1,len(lists)):
            tail_pts, flag =  lists[tail]
            if(flag==0):
                file_write.write(str(tail_pts)+": log"+str(unique_logs.index(lists[tail]))+"->")
            else:
                index = center_biomes.index(tail_pts)
                alpha= biome_alpha(index)
                file_write.write(str(tail_pts)+":"+ alpha+ "->" )
            cv2.line(image,head_pts,tail_pts,[75,0,130],3)
    cv2.imshow("finalize",image)
    return file_write,image

# filename = 'C:\\Aditi\\phd\\Log_Files\\2017-05-05\\theo emails\\17-43-36-May_17_Highschool-miniVid-PlantCreatureInfo\\17-43-36-May_17_Highschool-miniVid-PlantCreatureInfo\\0-Section_One-17-43-36-502.mov'
# #filename = 'C:\\Aditi\\phd\\Log_Files\\2017-05-05\\theo emails\\june 6 email\\16-21-27-May_17_Highschool-StoicBabyMigrate\\16-21-27-May_17_Highschool-StoicBabyMigrate\\0-Section_One-16-21-27-528.mov'
# cap = cv2.VideoCapture(filename)
# ex_filename = filename.strip(".mov")
# directory=ex_filename
# ex_filename = str(ex_filename)+".csv"
#
# handle = open(ex_filename, 'w')
# handle_valid= open('validity_file.csv','wb')
# """ per frame angle should be stored, so for each video file
# for each frame create a water_path angles file giving water contour angles
# """
# name = str(directory)+"\\adjacency_list.txt"
# file_writer= open(name,'w')
# handle.write("frame position")
#
# for i in range(0,8):
#     handle.write(",")
#     handle.write("index of log")
#     handle.write(",")
#     handle.write("angle with horizontal")
#     handle.write(",")
#     handle.write("centroid x")
#     handle.write(",")
#     handle.write("centroid y")
#     handle.write(",")
#     handle.write("status: active/inactive")
#
# cv2.namedWindow('video',0)
# boundaries, boundary_ends, center_biomes = mark_boundaries()
# complete_boundary,center_biomes = add_waterfall(boundaries,boundary_ends, center_biomes)
# while(cap.isOpened()):
#     flag, frame = cap.read()
#     if flag:
#         #display only mini video
#         min_video = frame[864:1080,0:384]
#
#         cv2.imshow('video',min_video)
#         handle.write("\n")
#         temp = frame[864:1080,0:384]
#
#         #break
#         # 1= will be frame number of next frame
#         pos_frame = cap.get(1)
#         #handle_valid=check_validity(min_video,pos_frame-1,handle_valid)
#
#         if not os.path.exists(directory):
#             os.makedirs(directory)
#         angle_matrix = detect_water_encoding(temp,(pos_frame-1),directory)
#         handle.write(str(pos_frame-1))
#         handle.write(",")
#         thresh_water, skeleton = detect_water(min_video)
#         logs,handle,centers = detect_logs(min_video, handle,thresh_water)
#         #cv2.imwrite("logs.png",logs)
#         color_water = cv2.cvtColor(skeleton,cv2.COLOR_GRAY2BGR)
#         water = np.float64(color_water)
#         #cv2.imwrite('snap.png',thresh_water)
#         dst = cv2.add(water,logs)
#         dst = cv2.add(complete_boundary, dst)
#
#         log_list = draw_graph(dst,skeleton,centers,boundary_ends,center_biomes)
#         file_writer,image = write_log_list(dst,file_writer,log_list,str(pos_frame-1),center_biomes)
#         #The frame is ready and already captured
#         cv2.imshow('video', image)
#
#
# #        break #added to process just one frame
#         cap.set(3,240)
#         cap.set(4,640)
#         cv2.setMouseCallback('video',know_coordinates)
#         cv2.waitKey(1000)
#     else:
#         # The next frame is not ready, so we try to read it again
#         cap.set(1, pos_frame-1)
#         print ("frame is not ready")
#         # It is better to wait for a while for the next frame to be ready
#         cv2.waitKey(1000)
#
#     if cv2.waitKey(10) == 27:
#         break
#     if cap.get(1) == cap.get(7):
#         # If the number of captured frames is equal to the total number of frames,
#         # we stop
#         break
# #handle_valid.close()
# file_writer.close()
# handle.close()
# cap.release()
# cv2.destroyAllWindows()


import ctypes
import numpy
import os
import time
from scipy.linalg.blas import daxpy, saxpy
import functools
from numba import jit, prange, njit
import sys

n_axpby = numpy.frompyfunc(lambda x,y,a,b: a*x + b*y, 4,1)

print ("current dir ", os.path.abspath(__file__))
N = 1048
if len(sys.argv) > 1:
    N = int(sys.argv[1])
shape = (N, N)

A = numpy.float32(2.)
B = numpy.float32(1.)

numpy.random.seed(1)
a = numpy.random.random(shape)
b = numpy.random.random(shape)
a = 2.5 * numpy.ones(shape)
b = 1. * numpy.ones(shape)

dll = os.path.abspath(os.path.join(
         os.path.abspath(os.path.dirname(__file__)),
         'libfdiff.dll')
)

dll = 'libfdiff.so'

print ("dll location", dll)
fdiff = ctypes.cdll.LoadLibrary(dll)

def axpby(a,b,out,A,B,dtype=numpy.float64):

    c_float_p = ctypes.POINTER(ctypes.c_float)
    c_double_p = ctypes.POINTER(ctypes.c_double)

    if a.dtype != dtype:
        a = a.astype(dtype)
    if b.dtype != dtype:
        b = b.astype(dtype)

    if dtype == numpy.float32:
        a_p = a.ctypes.data_as(c_float_p)
        b_p = b.ctypes.data_as(c_float_p)
        out_p = out.ctypes.data_as(c_float_p)

    elif dtype == numpy.float64:
        a = a.astype(numpy.float64)
        b = b.astype(numpy.float64)
        a_p = a.ctypes.data_as(c_float_p)
        b_p = b.ctypes.data_as(c_float_p)
        out_p = out.ctypes.data_as(c_float_p)
    else:
        raise TypeError('Unsupported type {}. Expecting numpy.float32 or numpy.float64'.format(dtype))

    out = numpy.empty_like(a)

    # int psaxpby(float * x, float * y, float * out, float a, float b, long size)
    fdiff.saxpby.argtypes = [ctypes.POINTER(ctypes.c_float), # pointer to the first array
                              ctypes.POINTER(ctypes.c_float), # pointer to the second array
                              ctypes.POINTER(ctypes.c_float), # pointer to the third array
                              ctypes.c_float,                 # type of A (float)
                              ctypes.c_float,                 # type of B (float)
                              ctypes.c_long]                  # type of size of first array
    fdiff.daxpby.argtypes = [ctypes.POINTER(ctypes.c_double), # pointer to the first array
                              ctypes.POINTER(ctypes.c_double), # pointer to the second array
                              ctypes.POINTER(ctypes.c_double), # pointer to the third array
                              ctypes.c_double,                 # type of A (c_double)
                              ctypes.c_double,                 # type of B (c_double)
                              ctypes.c_long]                  # type of size of first array

    if dtype == numpy.float32:
        return fdiff.saxpby(a_p, b_p, out_p, A, B, a.size)
    elif dtype == numpy.float64:
        return fdiff.daxpby(a_p, b_p, out_p, A, B, a.size)

c_float_p = ctypes.POINTER(ctypes.c_float)
c_double_p = ctypes.POINTER(ctypes.c_double)
dtype = numpy.float32
a = a.astype(dtype)
b = b.astype(dtype)

if dtype == numpy.float32:
    a_p = a.ctypes.data_as(c_float_p)
    b_p = b.ctypes.data_as(c_float_p)
    out = numpy.empty_like(a)
    out_p = out.ctypes.data_as(c_float_p)
    # int psaxpby(float * x, float * y, float * out, float a, float b, long size)
    fdiff.saxpby.argtypes = [ctypes.POINTER(ctypes.c_float), # pointer to the first array
                              ctypes.POINTER(ctypes.c_float), # pointer to the second array
                              ctypes.POINTER(ctypes.c_float), # pointer to the third array
                              ctypes.c_float,                 # type of A (float)
                              ctypes.c_float,                 # type of B (float)
                              ctypes.c_long]                  # type of size of first array

#out_nfrom = n_axpby(a,b,A,B)

@jit(nopython=True, parallel=True)
def numba_axpby(x,y,out,ca,cb):
    for i in prange(x.size):
        out.flat[i] = ca*x.flat[i] + cb*y.flat[i]

N = 100

fdiff.saxpby(a_p, b_p, out_p, A, B, a.size)
t0 = time.time()
for i in range(N):
    fdiff.saxpby(a_p, b_p, out_p, A, B, a.size)
t1 = time.time()
print ("lib saxpby", t1-t0)

axpby(a, b, out, A, B, numpy.float32)
t0 = time.time()
for i in range(N):
    axpby(a, b, out, A, B, numpy.float32)
t1 = time.time()
print ("lib axpby", t1-t0)

fdiff.saxpby(a_p, b_p, out_p, A, B, a.size)
t0 = time.time()
for i in range(N):
    fdiff.saxpby(a_p, b_p, out_p, A, B, a.size)
t1 = time.time()
print ("lib saxpby", t1-t0)

out_numpy = numpy.empty_like(a)
out_numpy2 = numpy.empty_like(a)
numpy.subtract(a,b, out=out_numpy)
t2 = time.time()
for i in range(N):
    numpy.multiply (a, A, out=out_numpy )
    numpy.multiply (b, B, out=out_numpy2)
    numpy.add(out_numpy,out_numpy2, out=out_numpy)
t3 = time.time()
print ("numpy memopt", t3-t2)

t2 = time.time()
for i in range(N):
    out_numpy2 = A*a + B*b
t3 = time.time()
print ("numpy no memopt", t3-t2)

t2 = time.time()
for i in range(N):
    out_scipy = saxpy(a,b,a=A)
t3 = time.time()
print ("scipy saxpy", t3-t2)

af = numpy.asfortranarray(a)
bf = numpy.asfortranarray(b)
print ("af.shape", af.shape, "a.shape", a.shape, "af.dtype", af.dtype)
t2 = time.time()
for i in range(N):
    out_scipy_f = daxpy(af,bf,a=A)
t3 = time.time()
print ("scipy saxpy fortran", t3-t2)

out_nfrom = n_axpby(a,b,A,B)
t2 = time.time()
for i in range(1):
    out_nfrom = n_axpby(a,b,A,B)
t3 = time.time()
print ("numpy frompyfunc", t3-t2)

from itertools import chain

out_map =  numpy.fromiter(chain.from_iterable(map(lambda y: A * y[0] + B * y[1], zip(a,b))), dtype=a.dtype).reshape(a.shape)
t2 = time.time()
for i in range(N):
    out_map = numpy.fromiter(chain.from_iterable(map(lambda y: A * y[0] + B * y[1], zip(a,b))), dtype=a.dtype).reshape(a.shape)
t3 = time.time()
print ("map", t3-t2)

out_numba = numpy.empty_like(a)
numba_axpby(a,b,out_numba,A,B)
t4 = time.time()
for i in range(N):
    numba_axpby(a,b,out_numba,A,B)
t5 = time.time()
print ("numba", t5-t4)

out_scipy_f_c = numpy.ascontiguousarray(out_scipy_f)
print ("out_scipy_f.shape", out_scipy_f.shape, "out_scipy_f_c.shape", out_scipy_f_c.shape)
numpy.testing.assert_array_equal(out, out_numpy)
numpy.testing.assert_array_equal(out_numpy, out_scipy)
numpy.testing.assert_array_almost_equal(out_numpy, numpy.ascontiguousarray(out_scipy_f))
numpy.testing.assert_array_almost_equal(out_numpy, out_nfrom)
numpy.testing.assert_array_almost_equal(out_numpy, numpy.asarray(list(out_map)))

numpy.testing.assert_array_almost_equal(out_numpy, out_numba)


<reponame>ifsm/apollon
"""hmm_test.py
Unit test for HMM implementation."""

import numpy as np
from scipy.stats import poisson
import unittest

from apollon.hmm.poisson import PoissonHmm

class TestHMM_utilities(unittest.TestCase):
    def setUp(self):
        # Arbitrary transition probability matrix
        self.A = np.array([[1., 0, 0], [.2, .3, .5], [.1, .3, .6]])

        # Wrong number of dimensions
        self.B1 = np.array([1., 0, 0, 0])
        self.B2 = np.array([[[1., 0, 0], [.2, .3, .5], [.1, .3, .6]]])

        # Not quadratic
        self.C1 = np.array([[1., 0, 0], [.2, .3, .5]])
        self.C2 = np.array([[1.0], [.5, .5], [.2, .8]])

        # Rows do not sum up to one
        self.D = np.array([[.2, .3, .5], [.5, .4, .2], [1., 0, 0]])

    def test_success(self):
        mus = [20, 40, 80, 120, 40]
        m = len(mus)
        data = np.concatenate([poisson(mu).rvs(30) for mu in mus])
        hmm = PoissonHmm(data, m)
        hmm.fit(data)
        self.assertTrue(hmm.success)

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