text:
#! /usr/bin/python3 -u

# Copyright © 2017 Sami Mezhoud <sami.mezhoud@ulb.ac.be>
# This work is free. You can redistribute it and/or modify it under the
# terms of the Do What The Fuck You Want To Public License, Version 2,
# as published by Sam Hocevar. See the COPYING file for more details.

""" ...

Usage:
    slam.py [-v|-vv] [-f|] <PICKEL_FILE>
    slam.py [-v|-vv] [-f|] <PICKEL_FILE> -p <FIGURE_DIR>

Arguments:
    <PICKEL_FILE>         Pickle file containing processed data.
    <FIGURE_DIR>          Directory for saved figures.

Options:
    -p <FIGURE_DIR>       Directory where to save figures.
    -f                    Force overwrite.
    -v                    Verbose mode.
    -vv                   Very Verbose mode.
    -h, --help
"""

from docopt import docopt
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import math
import copy
import os
import shutil
import warnings
import pickle

# RanSaC parameters
D_RANSAC = 20  # Degree range
S_RANSAC = 8  # Nb of samples
X_RANSAC = 0.4   # Distance to wall in m
N_RANSAC = 40  # Max trials
C_RANSAC = 0.6  # Consensus = C *len(data)
C_RANSAC_MIN = 50  # 50 # Minimum consensus

# Landmark validation parameters
D_VALIDATION = 4  # Distance max between 2 observations of a landmarks in m
N_VALIDATION = 2  # Nb of observation of a landmark need before data adjustment

# EKF parameters
TIME_WINDOW = 3  # Time window size in s
TIME_STEP = 3  # Sliding window step in s
ODOMETRY_NOISE = 5  # noise on the odometry process
RANGE_LIDAR_NOISE = 0.1  # distances noise on the LiDAR
BEARING_LIDAR_NOISE = 1 * math.pi/180  # angle noise on the LiDAR

# Figure parameters
PLOT_RED = 'tomato'
PLOT_BLUE = 'cornflowerblue'
PLOT_GREEN = 'forestgreen'
PLOT_ORANGE = 'darkorange'
PLOT_YELLOW = 'gold'
PLOT_LIGHT_GREY = 'gainsboro'
PLOT_COLORS = [PLOT_RED, PLOT_ORANGE, PLOT_YELLOW, 'navy'] * 10
PLOT_AXIS = [-12, 12, -12, 12]
PLOT_DRONE_HISTORY = 10  # seconds

IS_VERBOSE_MODE = docopt(__doc__)['-v'] >= 1
IS_HARD_VERBOSE_MODE = docopt(__doc__)['-v'] >= 2
IS_PLOTTING_MODE = docopt(__doc__)['-p'] is not None
FILE = docopt(__doc__)['<PICKEL_FILE>']
IS_OVERWRITE_MODE = docopt(__doc__)['-f'] >= 1
if IS_PLOTTING_MODE:
    FIGURE_DIR = docopt(__doc__)['-p']

# Ignore the warnings related to polyfit
warnings.filterwarnings('ignore')

def main():

    # If in plotting mode, assure that the directory to save figure exist
    if IS_PLOTTING_MODE:
        if os.path.isdir(FIGURE_DIR):
            if not IS_OVERWRITE_MODE:
                i = input("Figure directory already exist. Erase ?[y/n] ")
            if IS_OVERWRITE_MODE or i == "y":
                shutil.rmtree(FIGURE_DIR)
            else:
                print("Exit")
                exit()

        os.mkdir(FIGURE_DIR)
        os.mkdir(os.path.join(FIGURE_DIR, "pdf"))
        os.mkdir(os.path.join(FIGURE_DIR, "png"))
        os.mkdir(os.path.join(FIGURE_DIR, "data"))

    data = pd.read_pickle(FILE)

    correctedData, X = SLAM(data)

    if IS_PLOTTING_MODE:
        s = FIGURE_DIR + "/data/{}"

        dataFile = open(s.format("origData.pickle"), "wb")
        correctedDataFile = open(s.format("correctedData.pickle"), "wb")

        pickle.dump(data, dataFile)
        pickle.dump(correctedData, correctedDataFile)

def SLAM(origData):
    """ Apply a SLAM algorithm on LiDAR data.

    It recognize walls as landmarks using the RANSAC algorithm. Then landmarks
    are associated with a nearest-neighbor approach. The SLAM problem is solved
    by an EKF procedure.

    origData is a DataFrame object from the pandas module corresponding to
    post-precessed data acquired by the drone.
    """

    correctedData = pd.DataFrame(columns=origData.columns)

    timeIdx = 0
    offset = np.array([0.0, 0.0, 0.0])

    # Initialization of SLAM relevant matrix
    X, P, R, Jxr, Jz = initSlamMatrix()
    while timeIdx < origData.iloc[-1].time - TIME_WINDOW:

        # Reducing the data to the current time window
        curData = origData.copy()
        curData = curData[(curData.time > timeIdx) &
                          (curData.time < timeIdx + TIME_WINDOW)]
        timeIdx += TIME_STEP

        # Correction
        curData = dataCorrection(curData, offset)

        # Saving data for figures
        if IS_PLOTTING_MODE:
            previousCurData = curData.copy()  # Saving for plot
            previousX = copy.deepcopy(X)

        # Consensus law
        if C_RANSAC*len(curData) >= C_RANSAC_MIN:
            Consensus = len(curData) * C_RANSAC
        else:
            Consensus = C_RANSAC_MIN

        # Detecting landmarks using the RanSaC algorithm
        localLandmarks, trial = RanSaC(curData.copy(), D_RANSAC, S_RANSAC,
                                       X_RANSAC, Consensus, N_RANSAC)

        # Distinguishing which landmarks are observed for the first time
        localLandmarks = [makeLandmark(m) for m in localLandmarks]
        newLandmarkds, reobservedLandmark = \
            lmAssociation(X, localLandmarks)

        # Updating drone position using EKF methodology
        # STEP 1: update state using odometry data
        if len(curData) > 0:
            P, X, Jxr, Jz = updateFromOdometry(curData, X, P, Jxr, Jz)

        # STEP 2: update state using re-observed landmarks
        if len(reobservedLandmark) > 0:
            X, P, offset = updateFromReobservedLandmarks(reobservedLandmark,
                                                         P, X, R)

            curData = dataCorrection(curData, offset)

        # STEP 3: add new landmarks to the state
        if len(newLandmarkds) > 0:
            P, X = addNewLandmarks(newLandmarkds, X, P, Jxr, Jz, R)

        for i in curData.index:
            correctedData.loc[i] = curData.loc[i]

        if IS_VERBOSE_MODE:
            s = "\nTime window {:2.1f}s-{:2.1f}s: {} Data."
            section = s.format(timeIdx, timeIdx+TIME_WINDOW, len(curData))

            if len(newLandmarkds) + len(reobservedLandmark) > 0:
                print(section)
                for m in newLandmarkds:
                    s = "\tNew wall index {:d}."
                    print(s.format(m['index']))
                for m in reobservedLandmark:
                    s = "\tWall {:d} re-observed {:d} times."
                    print(s.format(m['index'], X.lmOccurrence[m['index']]))
                if len(reobservedLandmark) > 0:
                    s = "\tSLAM offset:\n"
                    s += "\t\tx: {:.3f} m\n"
                    s += "\t\ty: {:.3f} m\n"
                    s += "\t\tyaw: {:.3f} rad\n"
                    print(s.format(*offset))

            elif IS_HARD_VERBOSE_MODE and trial == 0:
                print(section, "\nNot enough data.")

            elif IS_HARD_VERBOSE_MODE and trial == N_RANSAC:
                print(section, "\nNo consensus.")

        if IS_PLOTTING_MODE:

            if len(curData) > 0:
                plotStep1(correctedData, previousCurData, previousX,
                          timeIdx, localLandmarks)

            # Re-observation
            if len(reobservedLandmark) > 0:
                plotStep2(correctedData, curData, previousCurData, X,
                          previousX, timeIdx)
            # New landmarks
            #if len(newLandmarkds) > 0:
            if len(curData) > 0:
                plotStep3(correctedData, previousCurData, curData, X, timeIdx)

    if IS_VERBOSE_MODE:
        print("-"*40, "\nSummary:\nFound {} walls".format(X.lmLength))
        for m in range(X.lmLength):
            s = "\tWall {:d} observed {:d} times"
            print(s.format(m, X.lmOccurrence[m]))

        print("Distances:")
        for m in range(X.lmLength):
            for i in range(m+1, X.lmLength):
                s = "\tFrom wall {:d} to {:d}: {:5.2f} m."
                print(s.format(m, i, distPointToPoint(X.getLm(m), X.getLm(i))))

    if IS_PLOTTING_MODE:
        plt.clf()
        plt.scatter(origData.lidarX, origData.lidarY, color=PLOT_LIGHT_GREY,
                    marker='.', s=3, label='Raw data')
        plt.scatter(correctedData.lidarX, correctedData.lidarY, marker='.',
                    color=PLOT_BLUE, label='Corrected data', s=3)

        for idx in range(X.lmLength):
            plotLandmark(X, idx)

        if X.lmLength == 2:
            plotDistance(X.getLm(0), X.getLm(1))
        elif X.lmLength == 4:
            plotDistance(X.getLm(0), X.getLm(2), textRatio=0.15)
            plotDistance(X.getLm(1), X.getLm(3), textRatio=0.8)

        plt.xlabel('x [m]')
        plt.ylabel('y [m]')
        plt.legend()
        plt.axis(PLOT_AXIS)
        plt.savefig("{}/png/Final.png".format(FIGURE_DIR), bbox_inches='tight')
        plt.savefig("{}/pdf/Final.pdf".format(FIGURE_DIR), bbox_inches='tight')

    return correctedData, X

def plotDroneMarker(row, color='black'):
    angleOffset = row.yaw * 180/math.pi - 45

    plt.scatter(row.droneX, row.droneY, color=color, facecolor=color, s=100,
                marker=(3, 0, angleOffset), zorder=3)

def plotStep1(correctedData, previousCurData, previousX, time, localLandmarks):
    plt.clf()

    # Drone trajectory
    correctedData = \
        correctedData[time - correctedData.time < PLOT_DRONE_HISTORY]
    correctedData = correctedData.loc[:previousCurData.iloc[0].name][:-1]
    correctedData = pd.concat([correctedData, previousCurData])

    plt.plot(correctedData.droneX, correctedData.droneY,
             color=PLOT_LIGHT_GREY, zorder=1)
    plt.plot(previousCurData.droneX, previousCurData.droneY,
             color='gray', zorder=2)
    plotDroneMarker(previousCurData.iloc[-1])

    # Corrected points
    plt.scatter(previousCurData.lidarX, previousCurData.lidarY,
                color=PLOT_BLUE, s=1)

    # Landmarks
    for idx in range(previousX.lmLength):
        plotLandmark(previousX, idx)

    # Ransac landmarks
    for m in localLandmarks:
        plotLandmark(*m['param'], mkr='--', color='darkmagenta')

    # Figure properties
    plt.xlabel('x [m]')
    plt.ylabel('y [m]')
    plt.axis(PLOT_AXIS)
    s = "{}/png/{}-01.png"
    plt.savefig(s.format(FIGURE_DIR, time), bbox_inches='tight')
    s = "{}/pdf/{}-01.pdf"
    plt.savefig(s.format(FIGURE_DIR, time), bbox_inches='tight')

def plotStep2(correctedData, curData, previousCurData, X, previousX, time):
    plt.clf()

    # Drone trajectory
    history = \
        correctedData[time - correctedData.time < PLOT_DRONE_HISTORY]
    plt.plot(history.droneX, history.droneY, color=PLOT_LIGHT_GREY, zorder=1)
    plt.plot(previousCurData.droneX, previousCurData.droneY,
             color=PLOT_RED, zorder=2)
    plotDroneMarker(previousCurData.iloc[-1], color=PLOT_RED)

    plt.plot(curData.droneX, curData.droneY,
             color=PLOT_GREEN, zorder=2)
    plotDroneMarker(curData.iloc[-1], color=PLOT_GREEN)

    # Before and after point correction
    plt.scatter(previousCurData.lidarX, previousCurData.lidarY,
                color=PLOT_RED, s=2)
    plt.scatter(curData.lidarX, curData.lidarY, color=PLOT_GREEN, s=2)

    # Landmarks
    # Old
    for idx in range(previousX.lmLength):
        plotLandmark(previousX, idx, color=PLOT_RED)
    # Updated
    for idx in range(X.lmLength):
        plotLandmark(X, idx, color=PLOT_GREEN)

    # Figure properties
    plt.xlabel('x [m]')
    plt.ylabel('y [m]')
    plt.axis(PLOT_AXIS)
    s = "{}/png/{}-02.png"
    plt.savefig(s.format(FIGURE_DIR, time), bbox_inches='tight')
    s = "{}/pdf/{}-02.pdf"
    plt.savefig(s.format(FIGURE_DIR, time), bbox_inches='tight')

def plotStep3(correctedData, previousCurData, curData, X, time):
    plt.clf()

    # Drone trajectory
    correctedData = \
        correctedData[time - correctedData.time < PLOT_DRONE_HISTORY]
    correctedData = correctedData.loc[:previousCurData.iloc[0].name][:-1]
    correctedData = pd.concat([correctedData, previousCurData])

    plt.plot(correctedData.droneX, correctedData.droneY,
             color=PLOT_LIGHT_GREY, zorder=1)
    plt.plot(previousCurData.droneX, previousCurData.droneY,
             color='gray', zorder=2)
    plotDroneMarker(previousCurData.iloc[-1])

    # Corrected points
    plt.scatter(curData.lidarX, curData.lidarY, color=PLOT_BLUE, s=1)

    # Landmarks
    for idx in range(X.lmLength):
        plotLandmark(X, idx)

    # Figure properties
    plt.xlabel('x [m]')
    plt.ylabel('y [m]')
    plt.axis(PLOT_AXIS)
    s = "{}/png/{}-03.png"
    plt.savefig(s.format(FIGURE_DIR, time), bbox_inches='tight')
    s = "{}/pdf/{}-03.pdf"
    plt.savefig(s.format(FIGURE_DIR, time), bbox_inches='tight')

def plotLandmark(p1, p2, color=False, mkr='-', label=''):
    """ param = [p1, p2] or [X, idx]. """

    if type(p1).__name__ == StateMatrix.__name__:
        label = "Wall {:d}: {:d} times".format(p2, p1.lmOccurrence[p2])
        if not color:
            color = PLOT_COLORS[p2]
        coordLm = p1.getLm(p2)
        p1, p2 = lmPointToParam(*coordLm)
    else:
        if not color:
            color = 'purple'
        coordLm = lmParamToPoint(p1, p2)

    l1, l2 = getLineCircleIntersection([p1, p2], coordLm, 2)
    plt.plot(l1, l2, mkr, color=color, label=label)
    plt.scatter(l1, l2, color=color, marker='x', s=50)
    plt.scatter(*coordLm, color=color, marker='o', s=100)

def plotDistance(p1, p2, textRatio=0.5):

    d = distPointToPoint(p1, p2)
    text = "{:.2f} m".format(d)

    textPos = [(p1[0] - p2[0])*textRatio, (p1[1] - p2[1])*textRatio]
    textPos = [p1[0] - (p1[0] - p2[0])*textRatio - (p1[1] - p2[1])*0.08,
               p1[1] - (p1[1] - p2[1])*textRatio + (p1[0] - p2[0])*0.08]

    plt.annotate('', xy=p1, xytext=p2, arrowprops={'arrowstyle': '<->'})
    plt.annotate(text, xy=p1, xytext=textPos)

def getLineCircleIntersection(line, center, R):
    """ Return the intersection of a line and circle.

    line is a vector (a, c) such as the line equation is y = ax + c.
    center is the coordinate of the center of the circle.
    R is the radius of the circle.
    """

    a, c = line
    xl, yl = center

    A = a**2 + 1
    B = 2*a*(c - yl) - 2*xl
    C = xl**2 + (c - yl)**2 - R**2

    Delta = B**2 - 4*A*C

    x1 = (-B + math.sqrt(Delta)) / (2*A)
    x2 = (-B - math.sqrt(Delta)) / (2*A)
    y1 = a*x1 + c
    y2 = a*x2 + c

    return [x1, x2], [y1, y2]

def dataCorrection(data, offset):
    """ Add a position offset to relevant column of data.

    data is a DataFrame object from the pandas module.
    offset is list of offset such as [x_offset, y_offset, yaw_offset].
    """

    data['droneX'] = data['droneX'].map(lambda x: x-offset[0])
    data['droneY'] = data['droneY'].map(lambda x: x-offset[1])
    data['yaw'] = data['yaw'].map(lambda x: x-offset[2])
    data['lidarX'] = data['lidarX'].map(lambda x: x-offset[0])
    data['lidarY'] = data['lidarY'].map(lambda x: x-offset[1])

    return data

def initSlamMatrix():
    """ Return intial value of all SLAM relevant matrices.

    Return:
    X as a StateMatrix object.
    P as a CovarianceMatrix object.
    Jxr as the prediction model jacobian with respect to position.
    Jz as the prediction model jacobian with respect to landmarks observations.
    """

    # State Matrix
    X = StateMatrix()

    # Covariance matrix
    P = CovarianceMatrix()

    # Modified prediction model jacobian
    Jxr = np.array([[1.0, 0.0, 0.0],
                    [0.0, 1.0, 0.0]])

    # Modified prediction model jacobian
    Jz = np.zeros([2, 2])

    # LiDAR Noise
    V = np.identity(2)
    R = V @ np.identity(2) * [RANGE_LIDAR_NOISE, BEARING_LIDAR_NOISE] @ V.T

    return X, P, R, Jxr, Jz

def updateFromOdometry(data, X, P, Jxr, Jz):
    """ 1st step: update state from odometry data.

    Return the modified X, P and prediction model jacobian  matrices.
    data is a DataFrame object from the pandas module.
    X is a StateMatrix object.
    P is a CovarianceMatrix object.
    Jxr is the prediction model jacobian with respect to position.
    Jz is the prediction model jacobian with respect to landmarks observations.
    """

    dx = data.iloc[0].droneX - data.iloc[-1].droneX
    dy = data.iloc[0].droneY - data.iloc[-1].droneY
    dt = data.iloc[0].yaw - data.iloc[-1].yaw

    # State matrix
    X.X[0] += dx
    X.X[1] += dy
    X.X[2] += dt

    # Prediction model jacobian
    A = np.identity(3)
    A[0, 2] = -dy
    A[1, 2] = dx

    # Process noise matrix
    W = np.identity(3) * [dx, dy, dt]
    Q = np.ones([3, 3]) * ODOMETRY_NOISE**2

    # Covariance on the robot position
    P.Prr = A @ P.Prr @ A.T + W @ Q @ W.T

    # Cross-correlations covariance
    for i in range(int((len(P.P)-3)/2)):
        Pri = P.getE(i)
        E = A @ Pri
        D = E.T

        P.setD(i, D)
        P.setE(i, E)

    # Modified prediction model jacobian with respect to position
    Jxr[0, 2] = -dy
    Jxr[1, 2] = dx

    # Modified prediction model jacobian with respect to landmarks observations
    Jz[0, 0] = math.cos(X.X[2] + dt)
    Jz[1, 0] = math.sin(X.X[2] + dt)
    Jz[0, 1] = -math.sin(X.X[2] + dt) * dt
    Jz[1, 1] = math.cos(X.X[2] + dt) * dt

    return P, X, Jxr, Jz

def updateFromReobservedLandmarks(reobservedLandmark, P, X, R):
    """ 2nd step: update state from re-observed landmarks.

    Return the modified X matrix and the offset to apply on measurments.
    reobservedLandmark is a list of re-observe landmarks define by the landmark
    association process.
    X is a StateMatrix object.
    P is a CovarianceMatrix object.
    R is the LiDAR noise.
    """

    for landmark in reobservedLandmark:

        # Landmark index
        idx = 3+landmark['index']*2

        # Compute the Jacobian of the measurement model
        r = distPointToLine(X.X[:2], landmark['point'])
        H11 = (X.X[0] - landmark['point'][0]) / r
        H12 = (X.X[1] - landmark['point'][1]) / r
        H21 = (landmark['point'][1] - X.X[1]) / (r**2)
        H22 = (landmark['point'][0] - X.X[0]) / (r**2)
        H = np.zeros([2, len(X.X)])
        H[:2, :3] = [[H11, H12, 0],
                     [H21, H22, -1]]
        H[:, idx:idx+2] = [[-H11, -H12],
                           [-H21, -H22]]

        # Compute innovation covariance
        d = distPointToPoint(X.X[0:2], landmark['point'])
        V = np.identity(2)
        R = np.identity(2) * [RANGE_LIDAR_NOISE*d, BEARING_LIDAR_NOISE]
        S = H @ P.P @ H.T + V @ R @ V.T

        # Compute Kalman gain
        K = P.P @ H.T @ np.linalg.inv(S)

        # Compute range and bearing of the observed landmark and its previous
        z = np.array(getRelativePos(X.X[:3], X.getLm(landmark['index'])))
        h = np.array(getRelativePos(X.X[:3], landmark['point']))

        # Update X state
        corr = K @ (z-h)
        X.X += corr

        # TODO make sure we don't need this
        # This instruction is not in 'SLAM for dummies' but according to EKF
        # algorithm it should be. Anyway it doesn't work.
        # I = np.identity(3)
        # P.Prr = (I - K[:3, :3] @ H[:3, :3]) @ P.Prr

        # TODO the offset is overwritten by last landmark
        offset = corr[:3]

    return X, P, offset

def addNewLandmarks(newLandmarkds, X, P, Jxr, Jz, R):
    """ 3rd step: Adding new landmarks to state and covariance matrices.

    Return the modified X and P matrices.
    newLandmarks is a list of new landmarks define by the landmark association
    process.
    X is a StateMatrix object.
    P is a CovarianceMatrix object.
    Jxr is the prediction model jacobian with respect to position.
    Jz is the prediction model jacobian with respect to landmarks observations.
    R is the LiDAR noise.
    """

    for m in newLandmarkds:
        # State matrix
        X.addLandmark(m['point'])

        P.addLandmark()

        # Landmark covariance
        C = Jxr @ P.Prr @ Jxr.T
        d = distPointToPoint(X.X[0:2], m['point'])
        R = np.identity(2) * [RANGE_LIDAR_NOISE*d, BEARING_LIDAR_NOISE]
        C += Jz @ R @ Jz.T

        P.setC(m['index'], C)

        # Robot-Landmark covariance
        E = P.Prr @ Jxr.T
        D = E.T

        P.setD(m['index'], D)
        P.setE(m['index'], E)

        # Landmark-Landmark covariance
        for i in range(m['index']):
            Pri = P.getD(i)
            F = Jxr @ Pri.T
            G = F.T

            P.setF(m['index'], i, F)
            P.setG(m['index'], i, G)

    return P, X

def getRelativePos(dr, lm):
    """ Return polar coordinate of a landmark relative to the drone.

        dr is a the pose of the drone such as dr = [x, y, yaw].
        lm is the cartesian coordinate of the landmark.
    """

    r = np.linalg.norm([lm[0]-dr[0], lm[1]-dr[1]])
    theta = math.atan2((lm[1]-dr[1]), (lm[0]-dr[0]))-dr[2]

    return [r, theta]

class StateMatrix(object):
    """ Represent a matrix containing the drone pose and landmarks position.

    Getter and setter functions are set up for easier access.
    """

    def __init__(self):
        self.X = np.array([0.0, 0.0, 0.0])
        self.lmOccurrence = []

    def addLandmark(self, point):
        """ Add new landmark."""

        self.X = np.concatenate((self.X, point), axis=0)
        self.lmOccurrence += [1]

    def getLm(self, i):
        """" Getter: landmark point (2x1). """

        index = 3+i*2

        return self.X[index:index+2]

    def setLm(self, i, value):
        """" Setter: landmark point (2x1). """

        index = 3+i*2

        self.X[index:index+2] = value

    @property
    def lmLength(self):
        """ Number of landmarks. """
        return len(self.lmOccurrence)

class CovarianceMatrix(object):
    """ Represent a matrix containing all relevant covariance matrix:
        - Covariance on the robot position
        - Landmark covariance
        - Robot - landmark covariance
        - Landmark - robot covariance
        - Landmark - landmark  covariance

    Getter and setter functions are set up for easier access.
    """

    def __init__(self):
        self.P = np.identity(3)

    @property
    def Prr(self):
        """ Covariance on the robot position. """
        return self.P[:3, :3]

    @Prr.setter
    def Prr(self, value):
        """Setter: Covariance on the robot position. """
        self.P[:3, :3] = value

    def addLandmark(self):
        """ Add rows and columns for new landmarks."""
        z = np.zeros((2, len(self.P)+2), dtype=self.P.dtype)

        # add two rows
        self.P = np.concatenate((self.P, z[:, :-2]), axis=0)

        # add two columns
        self.P = np.concatenate((self.P, z.T), axis=1)

    def getC(self, i):
        """" Getter: landmark covariance (2x2). """
        index = 3 + 2*i

        return self.P[index:index+2, index:index+2]

    def setC(self, i, value):
        """" Setter: landmark covariance (2x2). """
        index = 3 + 2*i

        self.P[index:index+2, index:index+2] = value

    def getD(self, i):
        """" Getter: robot - landmark covariance (2x3). """
        index = 3 + 2*i

        return self.P[index:index+2, :3]

    def setD(self, i, value):
        """" Setter: robot - landmark covariance (2x3). """
        index = 3 + 2*i

        self.P[index:index+2, :3] = value

    def getE(self, i):
        """" Getter: landmark - robot covariance (3x2). """
        index = 3 + 2*i

        return self.P[:3, index:index+2]

    def setE(self, i, value):
        """" Setter: landmark - robot covariance (3x2). """
        index = 3 + 2*i

        self.P[:3, index:index+2] = value

    def getF(self, i, j):
        """" Getter: landmark - landmark covariance on the row (2x2). """
        rowIdx = 3 + 2*i
        colIdx = 3 + 2*j

        return self.P[rowIdx:rowIdx+2, colIdx:colIdx+2]

    def setF(self, i, j, value):
        """" Setter: landmark - landmark covariance on the row (2x2). """
        rowIdx = 3 + 2*i
        colIdx = 3 + 2*j

        self.P[rowIdx:rowIdx+2, colIdx:colIdx+2] = value

    def getG(self, i, j):
        """" Getter: landmark - landmark covariance on the columns (2x2). """

        return self.getF(j, i)

    def setG(self, i, j, value):
        """" Setter: landmark - landmark covariance on the columns (2x2). """

        self.setF(j, i, value)

def RanSaC(data, D, S, X, C, N, plt=False):
    """ Detect wall inside a set of data using the RanSaC algorithm.

    data is a DataFrame object from the pandas module.
    D is distance max in degree between min and max angle in the random set.
    S is the number of samples to compute initial line.
    X is max distance a point may be from wall to get associated to wall.
    C is the number of points that must lie on a line for it to be validated.
    N is the max number of trials to attempt to find a lines.
    plt is an optional matplotlib object for plotting.
    """

    landmarks = []
    trial = count = 0
    while len(data) > C and trial < N:
        trial += 1

        # Select a random set of S samples in a range of D degrees
        sampleSet = getSampleSet(data, S, D)

        # Find the nearest line and count all points within a distance X
        lineParam = fitLeastSquareLine(sampleSet.lidarX, sampleSet.lidarY)
        count, data = countPointsNearLine(data, lineParam, X)

        # If there is enough points (more than consensus C)
        if count > C:

            # Recompute the nearest line with all points near previous line
            x = data["lidarX"][(data["distToWall"] == True)]
            y = data["lidarY"][(data["distToWall"] == True)]
            wallParam = fitLeastSquareLine(x, y)

            # Remove points near the new line from the data set
            count, data = countPointsNearLine(data, wallParam, X)
            data = data[data.distToWall != True]

            # Add the line to the list of wall detected
            landmarks.append(wallParam)

            if plt:
                p = np.poly1d(wallParam)
                xp = np.linspace(-4, 5)
                xp = np.linspace(-8, -4)
                plt.scatter(x, y, color=PLOT_BLUE, s=2)
                plt.plot(xp, p(xp), color=PLOT_BLUE)
                plt.scatter(sampleSet.lidarX, sampleSet.lidarY, s=3,
                            color='tomato')

    return landmarks, trial

def lmAssociation(X, localLandmarks):
    """ Return two lists of new and re-observed landmarks.

    Landmarks are associated using a nearest-neighbor approach.
    localLandmarks is the list of landmarks observed.
    X is a StateMatrix object.
    """

    newLandmarkds, reobservedLandmark = [], []
    count = 0
    for candidate in localLandmarks:

        # Selecting the nearest landmark already observed
        minDistance = [10**9, -1]
        for j in range(X.lmLength):
            d = distPointToPoint(candidate['point'], X.getLm(j))
            if d < minDistance[0]:
                minDistance = [d, j]

        # If it is near enough it is a re-observation
        if minDistance[0] < D_VALIDATION:
            candidate['index'] = minDistance[1]
            reobservedLandmark.append(candidate)
            X.lmOccurrence[candidate['index']] += 1

        # If not it is a new landmark
        else:
            candidate['index'] = X.lmLength + count
            count += 1
            newLandmarkds.append(candidate)

    # Only keep landmarks that have been observe enough time
    def condition(m): return X.lmOccurrence[m['index']] >= N_VALIDATION
    reobservedLandmark = [m for m in reobservedLandmark if condition(m)]

    return newLandmarkds, reobservedLandmark

def makeLandmark(param):
    """ Return a dictionary characterizing a landmark from its parameters.

    Wall landmarks are characterized by the point corresponding to the
    intersection of the wall line and its perpendicular passing through the
    origin (0, 0). The wall line is characterized by a vector (a, c) such as
    its equation is given by y = ax + c.

    Return a dictionary with landmark coordinate, parameters and occurrence.
    """

    return {"point": lmParamToPoint(*param),
            "param": param,
            "occurrence": 1}

def lmParamToPoint(a, c):
    """ Return the coordinates of a landmark from its line parameters.

    Wall landmarks are characterized by the point corresponding to the
    intersection of the wall line and its perpendicular passing through the
    origin (0, 0). The wall line is characterized by a vector (a, c) such as
    its equation is given by y = ax + c.
    """

    xp = float(-c*a / (1+a**2))
    yp = float(c / (1+a**2))

    return [xp, yp]

def lmPointToParam(xp, yp):
    """ Return the line parameters of a landmark from its coordinates.

    Wall landmarks are characterized by the point corresponding to the
    intersection of the wall line and its perpendicular passing through the
    origin (0, 0). The wall line is characterized by a vector (a, c) such as
    its equation is given by y = ax + c.

    xp and yp are the landmark coordinate.
    """

    a = -xp/yp
    b = yp*(1+a**2)

    return [a, b]

def countPointsNearLine(data, line, X):
    """ Count the number of point near a given line.

    Return the number of point within a distance X from a given line and add
    the column 'distToWall' to the data set corresponding to if the point is
    inside the range or not.
    data is a DataFrame object from the pandas module.
    line is a vector (a, c) such as the line equation is y = ax + c.
    """

    data['distToWall'] = False

    count = 0
    for idx, row in data.iterrows():
        d = distPointToLine([row.lidarX, row.lidarY], line)

        if d < X:
            data.loc[idx, 'distToWall'] = True
            count += 1

    return count, data

def distPointToPoint(p1, p2):
    """ Return the euclidean distance between two points.

    p1 and p2 are coordinate vectors (x, y).
    """

    return np.linalg.norm([p1[0]-p2[0], p1[1]-p2[1]])

def distPointToLine(point, line):
    """ Return the euclidean distance between a point and a line.

    point is a coordinate vector (x, y).
    line is a vector (a, c) such as the line equation is y = ax + c.
    """

    [xp, yp] = point
    [a, c] = line
    b = -1

    return abs((a*xp + b*yp + c) / np.linalg.norm([a, b]))

def fitLeastSquareLine(x, y):
    """ Least squares single degree polynomial fit.

    Fit a polynomial y = ax + c to points (x, y) and returns the vector of
    coefficients [a, c] that minimizes the squared error.
    x and y or vectors of points.
    """

    return np.polyfit(x, y, 1)

def getSampleSet(data, S, D):
    """ Return a random set of S elements within D degrees inside a data set.

    data is a DataFrame object from the pandas module.
    S is the size of the random set.
    D is distance max in degree between min and max angle in the random set.
    """

    sampleSet = pd.DataFrame(columns=data.columns)

    # Getting a first random sample
    sample = getRandomSample(data)
    sampleSet.loc[0] = sample

    for i in range(S-1):

        # Appending a new sample within the D range
        s = getRandomSample(data, sample.angle-D/2, sample.angle+D/2)
        sampleSet.loc[len(sampleSet)] = s

    return sampleSet

def getRandomSample(data, minimum=-1, maximum=361):
    """ Return a random sample within a data set.

    data is a DataFrame object from the pandas module.
    minimum and maximum arguments are optional angles (degree) limits in which
    the random sample needs to be, by default there is no limits.
    """

    # Get a random sample
    sample = data.sample(n=1).iloc[0]  # take a random sample

    # Take another one if it is not in the limits
    while sample.angle < minimum or sample.angle > maximum:
        sample = data.sample(n=1).iloc[0]

    return sample

def printMatrix(*args):
    """ Print matrices in a readable way. """

    for M in args:
        if type(M).__module__ == np.__name__ or type(M) == "list":
            for row in M:
                s = ["{:8.3}"] * len(row)
                s = ", ".join(s)
                print(s.format(*[i for i in row]))
        else:
            print(M)

if __name__ == '__main__':
    main()


text:
#!/usr/bin/env python

# inst: university of bristol
# auth: jeison sosa
# mail: j.sosa@bristol.ac.uk / sosa.jeison@gmail.com

import sys
import getopt
import subprocess
import configparser
import numpy as np
from lfptools import shapefile
import gdalutils
import gdalutils.extras.haversine as haversine
from osgeo import osr
from scipy.ndimage import distance_transform_edt
from scipy.spatial.distance import cdist

def getbankelevs_shell(argv):

    myhelp = '''
LFPtools v0.1

Name
----
getbankelevs

Description
-----------
Get river banks elevations from a high resolution DEM by reductions method like
nearest neighbour, mean, min or meanmin. Additionally, an outlier detection
can be applied to before running the reduction method.

Usage
-----
>> lfp-getbankelevs -i config.txt

Content in config.txt
---------------------
[getbankelevs]
output   = Shapefile output file path
netf     = Target mask file path
proj     = Output projection in Proj4 format
outlier  = Outlier detection yes/no
method   = Reductions available near, mean, min, meanmin
hrnodata = NODATA value for high resolution DEM
thresh   = Saerching threshold in degrees
hrdemf   = High resolution DEM
'''

    try:
        opts, args = getopt.getopt(argv, "i:")
        for o, a in opts:
            if o == "-i":
                inifile = a
    except:
        print(myhelp)
        sys.exit(0)

    config = configparser.SafeConfigParser()
    config.read(inifile)

    output = str(config.get('getbankelevs', 'output'))
    netf = str(config.get('getbankelevs', 'netf'))
    hrdemf = str(config.get('getbankelevs', 'hrdemf'))

    try:
        outlier = str(config.get('getbankelevs', 'outlier'))
    except:
        pass

    proj = str(config.get('getbankelevs', 'proj'))
    method = str(config.get('getbankelevs', 'method'))
    hrnodata = np.float64(config.get('getbankelevs', 'hrnodata'))
    thresh = np.float64(config.get('getbankelevs', 'thresh'))

    getbankelevs(output,netf,hrdemf,proj,method,hrnodata,thresh,outlier)

def getbankelevs(output,netf,hrdemf,proj,method,hrnodata,thresh,outlier):

    print("    running getbankelevs.py...")

    fname = output

    w = shapefile.Writer(shapefile.POINT)
    w.field('x')
    w.field('y')
    w.field('elev')

    # coordinates for bank elevations are based in river network mask
    net = gdalutils.get_data(netf)
    geo = gdalutils.get_geo(netf)
    iy, ix = np.where(net > 0)
    x = geo[8][ix]
    y = geo[9][iy]

    for i in range(len(x)):

        xmin = x[i] - thresh
        ymin = y[i] - thresh
        xmax = x[i] + thresh
        ymax = y[i] + thresh

        dem, dem_geo = gdalutils.clip_raster(hrdemf, xmin, ymin, xmax, ymax)
        ddem = np.ma.masked_where(dem == hrnodata, dem)

        if method == 'near':
            nodata = dem_geo[11]
            dfdem = gdalutils.array_to_pandas(dem, dem_geo, nodata, 'gt')
            arr = haversine.haversine_array(np.array(dfdem['y'].values, dtype='float32'), np.float32(
                dfdem['x'].values), np.float32(y[i]), np.float32(x[i]))
            dfdem['dis'] = np.array(arr)
            dfdem.sort_values(by='dis', inplace=True)
            elev = dfdem.iloc[0, 2]

        elif method == 'meanmin':
            if outlier == "yes":
                ddem = check_outlier(dem, ddem, hrnodata, 3.5)
            elev = np.mean([ddem.mean(), ddem.min()])

        elif method == 'mean':
            if outlier == "yes":
                ddem = check_outlier(dem, ddem, hrnodata, 3.5)
            elev = ddem.mean()

        elif method == 'min':
            if outlier == "yes":
                ddem = check_outlier(dem, ddem, hrnodata, 3.5)
            elev = ddem.min()

        # Write final file in a shapefile

        if np.isfinite(elev):
            w.point(x[i], y[i])
            w.record(x[i], y[i], elev)

    w.save("%s.shp" % fname)

    # Write .prj file
    prj = open("%s.prj" % fname, "w")
    srs = osr.SpatialReference()
    srs.ImportFromProj4(proj)
    prj.write(srs.ExportToWkt())
    prj.close()

    fmt = "GTiff"
    nodata = -9999
    bnkname1 = output+".shp"
    bnkname2 = output+".tif"
    subprocess.call(["gdal_rasterize", "-a_nodata", str(nodata), "-of", fmt, "-tr", str(geo[6]), str(geo[7]), "-a",
                     "elev", "-a_srs", proj, "-te", str(geo[0]), str(geo[1]), str(geo[2]), str(geo[3]), bnkname1, bnkname2])

def nearivpixel(ddem, rriv, ddsx, ddsy, XA):
    """
    Nearest river pixel when is possible if not
    take value from land
    """
    nodata = -9999
    _ds = np.where(rriv > 0)

    # if there are river pixels in the window
    if _ds[0].size > 0:
        XB = np.vstack((ddsy[_ds[0]], ddsx[_ds[1]])).T
        ind = np.int(cdist(XA, XB, metric='euclidean').argmin())
        elev = ddem[_ds[0][ind], _ds[1][ind]]

    # otherwise take nearest value from land
    elif np.where(rriv == 0)[0].size > 0:
        _ds = np.where(rriv == 0)
        XB = np.vstack((ddsy[_ds[0]], ddsx[_ds[1]])).T
        ind = np.int(cdist(XA, XB, metric='euclidean').argmin())
        elev = ddem[_ds[0][ind], _ds[1][ind]]

    # should be checked
    else:
        elev = nodata

    return elev

def avgrivpixel(ddem, rriv):
    """
    Average the mean and min of river pixels
    """
    nodata = -9999
    _ds = np.where(rriv > 0)

    # if there are river pixels in the window
    if _ds[0].size > 0:
        elev = np.mean([np.ma.masked_where(rriv == 0, ddem).mean(),
                        np.ma.masked_where(rriv == 0, ddem).min()])
    # otherwise
    else:
        elev = nodata

    return elev

def avgedgpixel(ddem, rriv):
    """
    Average the mean and min of edge pixels
    """
    nodata = -9999
    _ds = np.where(rriv > 0)

    # if there are river pixels in the window
    if _ds[0].size > 0:
        euclidis = distance_transform_edt(1-rriv)
        elev = np.mean([np.ma.masked_where(euclidis == 1, ddem).mean(
        ), np.ma.masked_where(euclidis == 1, ddem).min()])
    # otherwise
    else:
        elev = nodata

    return elev

def check_outlier(dem, ddem, hrnodata, thresh):

    shape = dem.shape
    chk = is_outlier(ddem.reshape(-1, 1), thresh)
    arr = np.where(chk == True)
    if arr[0].size > 0:
        dem1 = dem.reshape(-1, 1)
        # dem1_tmp = np.copy(dem1)
        dem1[arr[0]] = hrnodata
        dem2 = dem1.reshape(shape)
        ddem = np.ma.masked_where(dem2 == hrnodata, dem2)

        # # DEBUG DEBUG DEBUG
        # xaxis = np.arange(0,len(dem1)).reshape(-1,1)
        # plt.scatter(xaxis[dem1>hrnodata], dem1[dem1>hrnodata],  color='black')
        # plt.scatter(xaxis[arr[0]], dem1_tmp[arr[0]], color='red')
        # plt.show()

    return ddem

def is_outlier(points, thresh=3.5):
    """
    Returns a boolean array with True if points are outliers and False
    otherwise.

    Parameters:
    -----------
        points : An numobservations by numdimensions array of observations
        thresh : The modified z-score to use as a threshold. Observations with
            a modified z-score (based on the median absolute deviation) greater
            than this value will be classified as outliers.

    Returns:
    --------
        mask : A numobservations-length boolean array.

    References:
    ----------
        Boris Iglewicz and David Hoaglin (1993), "Volume 16: How to Detect and
        Handle Outliers", The ASQC Basic References in Quality Control:
        Statistical Techniques, Edward F. Mykytka, Ph.D., Editor.
    """

    if len(points.shape) == 1:
        points = points[:, None]
    median = np.median(points, axis=0)
    diff = np.sum((points - median)**2, axis=-1)
    diff = np.sqrt(diff)
    med_abs_deviation = np.median(diff)

    modified_z_score = 0.6745 * diff / med_abs_deviation

    return modified_z_score > thresh

if __name__ == '__main__':
    getbankelevs_shell(sys.argv[1:])

ext: py

reason: import numpy,from scipy


text:
# ====================================================================================== #
# Calculation derived for MFT of 1D firm growth.
# Author: Eddie Lee, edlee@csh.ac.at
# ====================================================================================== #
from scipy.optimize import minimize

from .utils import *

def density_bounds(density, wi,
                   vo=.49,
                   ve=.5,
                   dt=.1,
                   exact=False):
    """THIS IS A BOUND, NOT THE ACTUAL VELOCITY.

    Min density bound for nnovation front as derived from MFT and compared with
    simulation results.

    Depends on obsolescence and growth rates. Make sure that density has already been
    cached into parquet for faster runtime.

    Parameters
    ----------
    density : list of ndarray
    wi : float
        Innovation probability at each attempt.
    vo : float, .49
        Obsolescence rate per unit time.
    ve : float, .5
        Firm growth rate per unit time (attempt at moving to the right).
    dt : float, .1
        Size of simulation time step.
    exact : bool, False
        If True, use exact formulation. Else use small dt approximation. Make sure that dt
        corresponds to the actual time step in the simulation and not the rate at which
        samples were recorded.

    Returns
    -------
    float
        MFT min density bound for innovation front to keep moving.
    list
        Avg density for each simulation.
    """

    # number of firms on right boundary for each random trajectory
    right_density = [[i[-1] for i in d] for d in density]

    # histogram this
    y = [np.bincount(d, minlength=1) for d in right_density]

    sim_density = np.array([i.dot(range(i.size))/i.sum() for i in y])
    # min density required to progress
    if exact:
        assert not dt is None
        mft_bound = np.log(1 - vo * dt) / np.log(1 - wi * ve * dt)
    else:
        mft_bound = vo / wi / ve

    return mft_bound, sim_density

def exp_density_f(x, a, b, wi):
    """
    a * (1 - np.exp(-b*(1-x)))

    Parameters
    ----------
    x : float
    a : float
    b : float
    wi : float
        Probability of successful innovation.
    """

    return np.exp(-b*(1-x)) / wi + a * (1 - np.exp(-b*(1-x)))

def fit_density(x, ydata, wi, full_output=False):
    """Fit MFT prediction of firm density.

    Parameters
    ----------
    x : ndarray
        Normalized lattice coordinate.
    ydata : ndarray
    wi : ndarray
        Probability of successful innovation.

    Returns
    -------
    ndarray
        (a,b) from a * (1-exp(-b*(1-x)))
        b is normalized by the typical width
    """

    assert x.min()>=0 and x.max()<=1

    def cost(args):
        a, b = np.exp(args)
        return np.linalg.norm(exp_density_f(x, a, b, wi) - ydata)

    soln = minimize(cost, (np.log(ydata[0]), np.log(10)))
    if full_output:
        return np.exp(soln['x']), soln
    return np.exp(soln['x'])

max_stars_repo_name: eltrompetero/innovation

max_issues_repo_name: eltrompetero/innovation

max_forks_repo_name: eltrompetero/innovation

reason: from scipy


text:
import unittest

import numpy as np

from constants import TEST_DIR, del_test_dir, gen_test_dir, solution_domain_setup
from perform.constants import REAL_TYPE
from perform.input_funcs import catch_list
from perform.system_solver import SystemSolver
from perform.solution.solution_domain import SolutionDomain
from perform.visualization.probe_plot import ProbePlot

class ProbePlotInitTestCase(unittest.TestCase):
    def setUp(self):

        # generate working directory
        gen_test_dir()

        # generate input text files
        solution_domain_setup(TEST_DIR)

        # set SystemSolver and SolutionDomain
        self.solver = SystemSolver(TEST_DIR)
        self.sol_domain = SolutionDomain(self.solver)

        # get some plotting variables
        param_dict = self.solver.param_dict
        self.vis_vars = catch_list(param_dict, "vis_var_0", [None])
        self.vis_x_bounds = catch_list(param_dict, "vis_x_bounds_0", [[None, None]], len_highest=len(self.vis_vars))
        self.vis_y_bounds = catch_list(param_dict, "vis_y_bounds_0", [[None, None]], len_highest=len(self.vis_vars))

    def tearDown(self):

        del_test_dir()

    def test_probe_plot_init(self):

        plot = ProbePlot(
            self.solver.image_output_dir,
            0,
            self.solver.sim_type,
            self.solver.probe_vars,
            self.vis_vars,
            1,
            self.solver.num_probes,
            self.vis_x_bounds,
            self.vis_y_bounds,
            self.sol_domain.gas_model.species_names,
        )

        self.assertEqual(plot.vis_type, "probe")
        self.assertEqual(plot.num_subplots, 4)
        self.assertEqual(plot.num_rows, 2)
        self.assertEqual(plot.num_cols, 2)

        for bound_idx, bound in enumerate(self.vis_x_bounds):
            self.assertTrue(np.array_equal(bound, plot.vis_x_bounds[bound_idx]))
            self.assertTrue(np.array_equal(self.vis_y_bounds[bound_idx], plot.vis_y_bounds[bound_idx]))

class ProbePlotMethodsTestCase(unittest.TestCase):
    def setUp(self):

        # generate working directory
        gen_test_dir()

        # generate input text files
        solution_domain_setup()

        # set SystemSolver and SolutionDomain
        self.solver = SystemSolver(TEST_DIR)
        self.sol_domain = SolutionDomain(self.solver)

    def tearDown(self):

        del_test_dir()

    def test_probe_plot_plot(self):

        sol_domain = self.sol_domain

        # set plotting variables
        vis_vars = [
            "pressure",
            "velocity",
            "temperature",
            "species-0",
            "density",
            "momentum",
            "energy",
            "density-species-0",
        ]
        vis_y_bounds = [
            [0.8e6, 1.1e6],
            [0.5, 2.5],
            [900, 1300],
            [0.3, 0.7],
            [1.8, 2.6],
            [1, 6],
            [-2e7, -1.5e7],
            [0.5, 1.75],
        ]
        vis_x_bounds = [[None, None]] * len(vis_vars)

        # overwrite SolutionDomain probe info
        sol_domain.probe_vars = vis_vars
        sol_domain.num_probe_vars = len(vis_vars)
        self.solver.probe_vars = vis_vars

        # loop probes
        for probe_idx in range(sol_domain.num_probes):

            start_idx = 0
            # loop probe vars
            for var_idx in range(len(vis_vars)):

                # reset probe values
                sol_domain.probe_vals = np.zeros(
                    (sol_domain.num_probes, sol_domain.num_probe_vars, self.solver.num_steps), dtype=REAL_TYPE
                )

                if var_idx >= 9:
                    start_idx += 1

                plot = ProbePlot(
                    self.solver.image_output_dir,
                    0,
                    self.solver.sim_type,
                    sol_domain.probe_vars,
                    vis_vars[start_idx : var_idx + 1],
                    probe_idx,
                    self.solver.num_probes,
                    vis_x_bounds[start_idx : var_idx + 1],
                    vis_y_bounds[start_idx : var_idx + 1],
                    sol_domain.gas_model.species_names,
                )

                # update probes
                for self.solver.iter in range(self.solver.num_steps):
                    sol_domain.update_probes(self.solver)

                # plot
                plot.plot(sol_domain.probe_vals, sol_domain.time_vals, self.solver.num_steps, "b-", True)

                # check that bounds and values were set correctly
                if var_idx == 0:
                    self.assertTrue(np.array_equal(vis_y_bounds[0], plot.ax.get_ylim()))
                    self.assertTrue(
                        np.array_equal(sol_domain.probe_vals[probe_idx, 0, :], plot.ax.lines[0].get_xydata()[:, 1])
                    )
                else:
                    for ax_idx in range(var_idx - start_idx):
                        ax = plot.ax.ravel()[ax_idx]
                        self.assertTrue(np.array_equal(vis_y_bounds[ax_idx + start_idx], ax.get_ylim()))
                        self.assertTrue(
                            np.array_equal(
                                sol_domain.probe_vals[probe_idx, ax_idx + start_idx, :], ax.lines[0].get_xydata()[:, 1]
                            )
                        )

                plot.fig.clf()

reason: import numpy


text:
from ProcessorSICrossVali import ProcessorSICrossVali
from SharedProcessors.const import SUB_AND_SI_TRIALS, SI_TRIALS, SUB_NAMES, TRIAL_NAMES
import matplotlib.pyplot as plt
import numpy as np
import copy

train = copy.deepcopy(SUB_AND_SI_TRIALS)
# train = {'190521GongChangyang': SI_TRIALS}

if __name__ == "__main__":
    date = '211206'
    my_SI_processor = ProcessorSICrossVali(train, ['l_foot'], date, strike_off_from_IMU=1,
                                           do_input_norm=True, tune_hp=True)
    step_num = len(my_SI_processor.train_all_data_list)
    print('Number of steps before removing: {}'.format(step_num))

    i_step = 0
    counter_bad = 0
    bad_si, bad_si_sub, bad_si_trial = [], {sub: 0 for sub in list(train.keys())}, {trial: 0 for trial in SI_TRIALS}
    input_list, _, output_list = my_SI_processor.train_all_data_list.get_input_output_list()
    sub_list = my_SI_processor.train_all_data_list.get_sub_id_list()
    trial_list = my_SI_processor.train_all_data_list.get_trial_id_list()
    while i_step < len(my_SI_processor.train_all_data_list):
        # delete steps without a valid strike index
        strikes = np.where(input_list[i_step][:, -1] == 1)[0]
        if np.max(output_list[i_step]) <= 0 or np.max(output_list[i_step]) >= 1:
            if np.max(output_list[i_step]) == 0.0:
                prntval = np.min(output_list[i_step])
            else:
                prntval = np.max(output_list[i_step])
            bad_si.append(prntval)
            bad_si_sub[SUB_NAMES[sub_list[i_step]]] += 1
            bad_si_trial[TRIAL_NAMES[trial_list[i_step]]] += 1
            my_SI_processor.train_all_data_list.pop(i_step)
            counter_bad += 1

        else:
            i_step += 1
    print(f"There were {counter_bad} bad steps here due to SI problems")
    # train_all_data_list = my_SI_processor.clean_all_data(my_SI_processor.train_all_data_list,
    #                                                      my_SI_processor.sensor_sampling_fre)

    for sub in list(train.keys()):
        if bad_si_sub[sub] > 0:
            print(sub, end=': ')
            print(bad_si_sub[sub])

    for trial in SI_TRIALS:
        if bad_si_trial[trial] > 0:
            print(trial, end=': ')
            print(bad_si_trial[trial])

    step_num = len(my_SI_processor.train_all_data_list)
    print('Number of steps after removing: {}'.format(step_num))

    bad_si = [round(t * 100, 1) for t in bad_si if t > -100]
    # print(max(bad_si))
    # print(min(bad_si))
    # print(bad_si)


text:
#   Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import absolute_import

import os
import copy
import collections
import numpy as np
from paddle2onnx.graph import Node, Graph
from paddle2onnx.constant import NodeDomain, PRODUCER, dtypes
from paddle2onnx.op_mapper import OpMapper
from paddle2onnx.onnx_helper import helper
from paddle2onnx.utils import check_model, logging

class ONNXNode(Node):
    def __init__(self, op_type, inputs, outputs, attrs, layer_name):
        super(ONNXNode, self).__init__(op_type, inputs, outputs, attrs,
                                       layer_name, NodeDomain.ONNX)
        self.onnx_node = self.make_onnx_node()

    def make_onnx_constant_node(self):
        dtype = self.attr('dtype')
        value = self.attr('value')
        if isinstance(value, list):
            dims = (len(value), )
        elif value is None:
            dims = ()
            value = []
        else:
            dims = ()
            value = [value]

        if 'dims' in self.attrs:
            dims = self.attrs['dims']

        tensor = helper.make_tensor(
            name=self.layer_name, data_type=dtype, dims=dims, vals=value)

        onnx_node = helper.make_node(
            self.type, inputs=self.inputs, outputs=self.outputs, value=tensor)

        return onnx_node

    def make_onnx_node(self):
        if self.type in ['Constant', 'ConstantOfShape']:
            onnx_node = self.make_onnx_constant_node()
        else:
            onnx_node = helper.make_node(
                self.type,
                inputs=self.inputs,
                outputs=self.outputs,
                name=self.layer_name,
                **self.attrs)
        return onnx_node

class ONNXGraph(Graph):
    def __init__(self, paddle_graph, opset_version, block=None):
        super(ONNXGraph, self).__init__()
        self.opset_version = opset_version
        self.ctx = paddle_graph

    def __str__(self):
        graph_str = 'graph { \n'
        for node in self.input_nodes:
            graph_str += " input: {} \n".format(node)
        for node in self.output_nodes:
            graph_str += " output: {} \n \n".format(node)
        for name, node in self.node_map.items():
            graph_str += node.__str__()
        graph_str += ' }'
        return graph_str

    def make_node(self,
                  op_type,
                  inputs=[],
                  outputs=[],
                  attrs=None,
                  layer_name=None,
                  **kw):
        if layer_name is None:
            layer_name = self.generate_node_name(op_type)

        if attrs is None:
            attrs = kw
        attrs.update(kw)

        if inputs is None:
            inputs = []

        real_outputs = None
        if outputs is None:
            real_outputs = [layer_name]
        elif isinstance(outputs, int):
            real_outputs = []
            for i in range(outputs):
                real_outputs.append(self.generate_node_name(op_type))
        elif isinstance(outputs, list):
            real_outputs = []
            if len(outputs) == 0:
                real_outputs = [layer_name]
            else:
                for opt in outputs:
                    if isinstance(opt, Node):
                        real_outputs.append(opt.layer_name)
                    elif isinstance(opt, int):
                        real_outputs.append(self.generate_node_name(op_type))
                    else:
                        real_outputs.append(opt)
        else:
            real_outputs = outputs

        node = ONNXNode(op_type, inputs, real_outputs, attrs, layer_name)

        self.insert_node(node)
        if len(node.outputs) == 1:
            return node.outputs[0]
        else:
            return node.outputs

    def update_node(self,
                    node,
                    op_type=None,
                    inputs=None,
                    outputs=None,
                    attrs=None,
                    **kw):
        if op_type is None:
            op_type = node.type
        if inputs is None:
            inputs = node.inputs
        if outputs is None:
            outputs = node.outputs
        if attrs is None:
            attrs = node.attrs
        attrs.update(kw)

        node = ONNXNode(op_type, inputs, outputs, attrs, node.layer_name)
        self.insert_node(node)
        return node

    def build_parameters(self, parameters):
        # build weight nodes
        for name, param in parameters.items():
            weight = param['data']
            if weight is not np.ndarray:
                weight = np.array(weight)
            tensor = helper.make_tensor(
                name=name,
                dims=param['shape'],
                data_type=dtypes.DTYPE_PADDLE_ONNX_MAP[param['dtype']],
                vals=weight.flatten().tolist())
            node = helper.make_node(
                'Constant', inputs=[], outputs=[name], value=tensor)
            self.parameters[name] = node

    def build_input_nodes(self, input_nodes):
        # build input nodes
        for ipt in input_nodes:
            self.add_input_node(ipt.layer_name,
                                ipt.attr('shape'), ipt.attr('dtype'))

    def build_output_nodes(self, output_nodes):
        # build output nodes
        for opt in output_nodes:
            self.add_output_node(opt.layer_name,
                                 opt.attr('shape'), opt.attr('dtype'))

    def build_op_nodes(self, node_map):
        OpMapper.check_support_status(node_map, self.opset_version)
        # build op nodes
        for name, node in list(node_map.items()):
            OpMapper.mapping(self, node)

    def make_value_info(self, name, shape, dtype):
        tensor_info = helper.make_tensor_value_info(
            name=name,
            shape=shape,
            elem_type=dtypes.DTYPE_PADDLE_ONNX_MAP[dtype])
        return tensor_info

    def add_input_node(self, name, shape, dtype):
        vi = self.make_value_info(name, shape, dtype)
        self.input_nodes.append(vi)

    def add_output_node(self, name, shape, dtype):
        vi = self.make_value_info(name, shape, dtype)
        self.output_nodes.append(vi)

    def export_proto(self, enable_onnx_checker=False):
        op_nodes = [node.onnx_node for node in self.node_map.values()]
        weight_nodes = [node for node in self.parameters.values()]

        onnx_graph = helper.make_graph(
            nodes=weight_nodes + op_nodes,
            name='paddle-onnx',
            initializer=[],
            inputs=self.input_nodes,
            outputs=self.output_nodes)

        opset_imports = [helper.make_opsetid("", self.opset_version)]
        onnx_proto = helper.make_model(
            onnx_graph, producer_name=PRODUCER, opset_imports=opset_imports)

        if enable_onnx_checker:
            check_model(onnx_proto)

        return onnx_proto

    @staticmethod
    def build(paddle_graph, opset_version, verbose=False):
        onnx_graph = ONNXGraph(paddle_graph, opset_version=opset_version)
        onnx_graph.build_parameters(paddle_graph.parameters)
        onnx_graph.build_input_nodes(paddle_graph.input_nodes)
        onnx_graph.build_output_nodes(paddle_graph.output_nodes)
        onnx_graph.build_op_nodes(paddle_graph.node_map)

        return onnx_graph


text:
#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import atexit
import math
import re
import signal
from typing import List, Tuple

import numpy as np
import psutil

MEMORY_SIZE_UNITS = {"K": 2**10, "M": 2**20, "G": 2**30, "T": 2**40}

# we use 4 bytes for block size, this means each block can contain
# 4294967296 records
BLOCK_SIZE_BIT = 32

def get_node_address() -> str:
    """
    Get the ip address used in ray.
    """
    pids = psutil.pids()
    for pid in pids:
        try:
            proc = psutil.Process(pid)
            # HACK: Workaround for UNIX idiosyncrasy
            # Normally, cmdline() is supposed to return the argument list.
            # But it in some cases (such as when setproctitle is called),
            # an arbitrary string resembling a command-line is stored in
            # the first argument.
            # Explanation: https://unix.stackexchange.com/a/432681
            # More info: https://github.com/giampaolo/psutil/issues/1179
            for arglist in proc.cmdline():
                for arg in arglist.split(" "):
                    if arg.startswith("--node-ip-address"):
                        addr = arg.split("=")[1]
                        return addr
        except psutil.AccessDenied:
            pass
        except psutil.NoSuchProcess:
            pass
    raise Exception("can't find any ray process")

def register_exit_handler(func):
    atexit.register(func)
    signal.signal(signal.SIGTERM, func)
    signal.signal(signal.SIGINT, func)

def random_split(df, weights, seed=None):
    """
    Random split the spark DataFrame or koalas DataFrame into given part
    :param df: the spark DataFrame or koalas DataFrame
    :param weights: list of doubles as weights with which to split the df.
                    Weights will be normalized if they don't sum up to 1.0.
    :param seed: The seed for sampling.
    """
    # convert to Spark DataFrame
    df, is_spark_df = convert_to_spark(df)
    splits = df.randomSplit(weights, seed)
    if is_spark_df:
        return splits
    else:
        # convert back to koalas DataFrame
        import databricks.koalas as ks
        return [ks.DataFrame(split) for split in splits]

def _df_helper(df, spark_callback, koalas_callback):
    try:
        import pyspark
    except Exception:
        pass
    else:
        if isinstance(df, pyspark.sql.DataFrame):
            return spark_callback(df)

    try:
        import databricks.koalas as ks
    except Exception:
        pass
    else:
        if isinstance(df, ks.DataFrame):
            return koalas_callback(df)

    raise Exception(f"The type: {type(df)} is not supported, only support "
                    "pyspark.sql.DataFrame and databricks.koalas.DataFrame")

def df_type_check(df):
    """
    Check whether the df is spark DataFrame or koalas DataFrame.
    :return True for spark DataFrame or Koalas DataFrame.
    :raise Exception when it is neither spark DataFrame nor Koalas DataFrame.
    """
    return _df_helper(df, lambda d: True, lambda d: True)

def convert_to_spark(df):
    """
    Do nothing if the df is spark DataFrame, convert to spark DataFrame if it is
    koalas DataFrame. Raise Exception otherwise.
    :return: a pair of (converted df, whether it is spark DataFrame)
    """
    return _df_helper(df, lambda d: (d, True), lambda d: (d.to_spark(), False))

def parse_memory_size(memory_size: str) -> int:
    """
    Parse the human readable memory size into bytes.
    Adapt from: https://stackoverflow.com/a/60708339
    :param memory_size: human readable memory size
    :return: convert to int size
    """
    memory_size = memory_size.strip().upper()
    if re.search(r"B", memory_size):
        # discard "B"
        memory_size = re.sub(r"B", "", memory_size)

    try:
        return int(memory_size)
    except ValueError:
        pass

    global MEMORY_SIZE_UNITS
    if not re.search(r" ", memory_size):
        memory_size = re.sub(r"([KMGT]+)", r" \1", memory_size)
    number, unit_index = [item.strip() for item in memory_size.split()]
    return int(float(number) * MEMORY_SIZE_UNITS[unit_index])

def divide_blocks(
        blocks: List[int],
        world_size: int,
        world_rank: int,
        shuffle: bool,
        pack_index: bool = True) -> Tuple[List[int], List]:
    """
    Divide the blocks into world_size partitions, and return the divided block indexes for the
    Given work_rank
    :param blocks: the blocks and each item is the given block size
    :param world_size: total world size
    :param world_rank: the current rank
    :param shuffle: whether we need to the shuffle the blocks index
    :param pack_index: whether pack the index
    :return: a pair, the first one is the selected block indexes. If the pack_index set to True,
             the second one return packed block index with block inner index, else return the
             selected block size.
    """
    # get the block size for each rank
    num_blocks = int(math.ceil(len(blocks) * 1.0 / world_size))
    # get the number of samples for each rank
    num_samples = int(math.ceil(sum(blocks) * 1.0 / world_size))
    total_block_size = num_blocks * world_size
    total_indexes = list(range(len(blocks)))

    # add extra samples to make it evenly divisible
    # we should use while here, because the len(total_indices) can be smaller than num_replicas
    while len(total_indexes) != total_block_size:
        total_indexes += total_indexes[: (total_block_size - len(total_indexes))]

    assert len(total_indexes) == total_block_size

    if shuffle:
        np.random.seed(0)
        np.random.shuffle(total_indexes)

    indexes = total_indexes[world_rank: total_block_size: world_size]
    assert len(indexes) == num_blocks

    def select(index: int, current_size: int, selected: List[Tuple[int, int]]) -> int:
        block_size = blocks[index]
        tmp = current_size + block_size
        if tmp < num_samples:
            selected.append((index, block_size))
            current_size = tmp
        elif tmp >= num_samples:
            selected.append((index, (num_samples - current_size)))
            current_size = num_samples
        return current_size

    total_size = 0
    selected_indices = []
    for i in indexes:
        total_size = select(i, total_size, selected_indices)
        if total_size == num_samples:
            break

    step = 1
    while total_size < num_samples:
        index = total_indexes[(world_rank + step) % len(total_indexes)]
        total_size = select(index, total_size, selected_indices)
        step += world_rank

    assert total_size == num_samples

    selected_indices = sorted(selected_indices)
    if pack_index:
        block_indexes = []
        packed_selected_indexes = []
        global BLOCK_SIZE_BIT
        for i, size in selected_indices:
            block_indexes.append(i)
            # we use 4 Bytes for the block inner index
            packed_selected_indexes.append([((i << BLOCK_SIZE_BIT) | j) for j in range(size)])
        return block_indexes, packed_selected_indexes
    else:
        block_indexes = []
        selected_sizes = []
        for i, size in selected_indices:
            block_indexes.append(i)
            selected_sizes.append(size)
        return block_indexes, selected_sizes

max_stars_repo_name: ConeyLiu/oap-raydp

max_issues_repo_name: ConeyLiu/oap-raydp

max_forks_repo_name: ConeyLiu/oap-raydp


text:
# -*- coding: utf-8 -*-
"""
This module provides classes to post process a model run output.
"""
__author__ = 'ajshajib'

import numpy as np
import matplotlib.pyplot as plt

from lenstronomy.Plots.model_plot import ModelPlot
from lenstronomy.Sampling.parameters import Param
from lenstronomy.Util.class_creator import create_im_sim

from dolphin.processor import Processor
from dolphin.processor.config import ModelConfig

class Output(Processor):
    """
    This class lets
    """
    def __init__(self, io_directory):
        """

        :param io_directory: path to the input/output directory. Should not
            end with slash.
        :type io_directory: `str`
        """
        super(Output, self).__init__(io_directory)

        self._fit_output = None
        self._kwargs_result = None
        self._model_settings = None
        self._samples_mcmc = None
        self._params_mcmc = None

    @property
    def fit_output(self):
        """
        The output from the `lenstronomy...Fitting_sequence.fit_sequence()`.

        :return:
        :rtype:
        """
        if self._fit_output is None:
            raise ValueError('Model output not specified!'
                             'Load an output using the `load_output()`'
                             'method!')
        else:
            return self._fit_output

    @property
    def kwargs_result(self):
        """
        The `kwargs_result` after running a model by calling
        `lenstronomy...Fitting_sequence.fit_sequence()`.

        :return:
        :rtype:
        """
        if self._kwargs_result is None:
            raise ValueError('Model output not specified!'
                             'Load an output using the `load_output()`'
                             'method!')
        else:
            return self._kwargs_result

    @property
    def model_settings(self):
        """
        The `kwargs_result` after running a model by calling
        `lenstronomy...Fitting_sequence.fit_sequence()`.

        :return:
        :rtype:
        """
        if self._model_settings is None:
            raise ValueError('Model output not specified!'
                             'Load an output using the `load_output()`'
                             'method!')
        else:
            return self._model_settings

    @property
    def samples_mcmc(self):
        """
        The array of MCMC samples from the model run.

        :return:
        :rtype:
        """
        if self._samples_mcmc is None:
            return []
        else:
            return self._samples_mcmc

    @property
    def params_mcmc(self):
        """
        The non-linear parameters sampled with MCMC.

        :return:
        :rtype:
        """
        if self._params_mcmc is None:
            return []
        else:
            return self._params_mcmc

    @property
    def num_params_mcmc(self):
        """
        Number of sampled non-linear parameters in MCMC.

        :return:
        :rtype:
        """
        if self._params_mcmc is None:
            return 0
        else:
            return len(self._params_mcmc)

    def swim(self, *args, **kwargs):
        """
        Override the `swim` method of the `Processor` class to make it
        not callable.
        """
        raise NotImplementedError

    def load_output(self, lens_name, model_id):
        """
        Load output from file and save in class variables.

        :param lens_name: lens name
        :type lens_name: `str`
        :param model_id: model identifier provided at run initiation
        :type model_id: `str`
        :return: output dictionary
        :rtype: `dict`
        """
        output = self.file_system.load_output(lens_name, model_id)

        self._model_settings = output['settings']
        self._kwargs_result = output['kwargs_result']
        self._fit_output = output['fit_output']

        if self.fit_output[-1][0] == 'EMCEE':
            self._samples_mcmc = self.fit_output[-1][1]
            self._params_mcmc = self.fit_output[-1][2]

        return output

    def get_model_plot(self, lens_name, model_id=None,
                       kwargs_result=None, band_index=0,
                       data_cmap='cubehelix'):
        """
        Get the `ModelPlot` instance from lenstronomy for the lens.

        :param lens_name: name of the lens
        :type lens_name: `str`
        :param model_id: model run identifier
        :type model_id: `str`
        :param kwargs_result: lenstronomy `kwargs_result` dictionary. If
            provided, it will be used to plot the model, otherwise the model
            will be plotted from the saved/loaded outputs for `lens_name` and
            `model_id`.
        :type kwargs_result: `dict`
        :param band_index: index of band to plot for multi-band case
        :type band_index: `int`
        :param data_cmap: colormap for image, reconstruction, and source plots
        :type data_cmap: `str` or `matplotlib.colors.Colormap`
        :return: `ModelPlot` instance, maximum pixel value of the image
        :rtype: `obj`, `float`
        """
        if model_id is None and kwargs_result is None:
            raise ValueError('Either the `model_id` or the `kwargs_result` '
                             'needs to be provided!')

        if kwargs_result is None:
            self.load_output(lens_name, model_id)
            kwargs_result = self.kwargs_result

        multi_band_list_out = self.get_kwargs_data_joint(
                                                lens_name)['multi_band_list']

        config = ModelConfig(settings=self.model_settings)

        mask = config.get_masks()
        kwargs_model = config.get_kwargs_model()

        v_max = np.log10(
            multi_band_list_out[band_index][0]['image_data'].max())

        model_plot = ModelPlot(multi_band_list_out, kwargs_model,
                               kwargs_result,
                               arrow_size=0.02, cmap_string=data_cmap,
                               likelihood_mask_list=mask,
                               multi_band_type='multi-linear')

        return model_plot, v_max

    def plot_model_overview(self, lens_name, model_id=None,
                            kwargs_result=None, band_index=0,
                            data_cmap='cubehelix', residual_cmap='RdBu',
                            convergence_cmap='afmhot',
                            magnification_cmap='viridis',
                            v_min=None, v_max=None):
        """
        Plot the model, residual, reconstructed source, convergence,
        and magnification profiles. Either `model_id` or `kwargs_result`
        needs to be provided. `kwargs_result` is prioritized for plotting if
        both are provided.

        :param lens_name: name of the lens
        :type lens_name: `str`
        :param model_id: model run identifier
        :type model_id: `str`
        :param kwargs_result: lenstronomy `kwargs_result` dictionary. If
            provided, it will be used to plot the model, otherwise the model
            will be plotted from the saved/loaded outputs for `lens_name` and
            `model_id`.
        :type kwargs_result: `dict`
        :param band_index: index of band to plot for multi-band case
        :type band_index: `int`
        :param data_cmap: colormap for image, reconstruction, and source plots
        :type data_cmap: `str` or `matplotlib.colors.Colormap`
        :param residual_cmap: colormap for noise residual plot
        :type residual_cmap: `str` or `matplotlib.colors.Colormap`
        :param convergence_cmap: colormap for convergence plot
        :type convergence_cmap: `str` or `matplotlib.colors.Colormap`
        :param magnification_cmap: colormap for magnification plot
        :type magnification_cmap: `str` or `matplotlib.colors.Colormap`
        :param v_min: minimum plotting scale for the model, data, & source plot
        :type v_min: `float` or `int`
        :param v_max: maximum plotting scale for the model, data, & source plot
        :type v_max: `float` or `int`
        :return: `matplotlib.pyplot.figure` instance with the plots
        :rtype: `matplotlib.pyplot.figure`
        """
        if v_max is None:
            model_plot, v_max = self.get_model_plot(
                                                 lens_name,
                                                 model_id=model_id,
                                                 kwargs_result=kwargs_result,
                                                 band_index=band_index,
                                                 data_cmap=data_cmap)
        else:
            model_plot = self.get_model_plot(lens_name,
                                             model_id=model_id,
                                             kwargs_result=kwargs_result,
                                             band_index=band_index,
                                             data_cmap=data_cmap)[0]

        fig, axes = plt.subplots(2, 3, figsize=(16, 8))

        model_plot.data_plot(ax=axes[0, 0], band_index=band_index,
                             v_max=v_max, v_min=v_min)
        model_plot.model_plot(ax=axes[0, 1], band_index=band_index,
                              v_max=v_max, v_min=v_min)
        model_plot.normalized_residual_plot(ax=axes[0, 2],
                                            band_index=band_index,
                                            cmap=residual_cmap, v_max=3,
                                            v_min=-3)
        model_plot.source_plot(ax=axes[1, 0], deltaPix_source=0.02, numPix=100,
                               band_index=band_index, v_max=v_max, v_min=v_min)
        model_plot.convergence_plot(ax=axes[1, 1], band_index=band_index,
                                    cmap=convergence_cmap)
        model_plot.magnification_plot(ax=axes[1, 2],
                                      band_index=band_index,
                                      cmap=magnification_cmap)
        fig.tight_layout()
        fig.subplots_adjust(left=None, bottom=None, right=None, top=None,
                            wspace=0., hspace=0.05)

        return fig

    def get_reshaped_emcee_chain(self, lens_name, model_id, walker_ratio,
                                 burn_in=-100, verbose=True):
        """

        :param lens_name:
        :type lens_name:
        :param model_id:
        :type model_id:
        :param walker_ratio:
        :type walker_ratio:
        :param burn_in:
        :type burn_in:
        :param verbose:
        :type verbose:
        :return:
        :rtype:
        """
        self.load_output(lens_name, model_id)

        num_params = self.num_params_mcmc  # self.samples_mcmc.shape[1]
        num_walkers = walker_ratio * num_params
        num_step = int(len(self.samples_mcmc) / num_walkers)

        chain = np.empty((num_walkers, num_step, num_params))

        for i in np.arange(num_params):
            samples = self.samples_mcmc[:, i].T
            chain[:, :, i] = samples.reshape((num_step, num_walkers)).T

        return chain

    def plot_mcmc_trace(self, lens_name, model_id, walker_ratio,
                        burn_in=-100, verbose=True, fig_width=16):
        """
        Plot the trace of MCMC walkers.

        :param lens_name: name of the lens
        :type lens_name: `str`
        :param model_id: model run identifier
        :type model_id: `str`
        :param walker_ratio: number of walkers per parameter in MCMC
        :type walker_ratio: `int`
        :param burn_in: number of burn-in steps to compute the medians after
            convergence of the MCMC chain
        :type: `int`
        :param verbose: if `True`, median values after burn-in will be printed
        :type verbose: `bool`
        :param fig_width: width of the figure
        :type fig_width: `float`
        :return: `matplotlib.pyplot.figure` instance with the plots
        :rtype: `matplotlib.pyplot.figure`
        """
        chain = self.get_reshaped_emcee_chain(lens_name, model_id,
                                              walker_ratio,
                                              burn_in=burn_in, verbose=verbose)

        num_params = self.num_params_mcmc
        num_step = chain.shape[1]

        mean_pos = np.zeros((num_params, num_step))
        median_pos = np.zeros((num_params, num_step))
        std_pos = np.zeros((num_params, num_step))
        q16_pos = np.zeros((num_params, num_step))
        q84_pos = np.zeros((num_params, num_step))

        # chain = np.empty((nwalker, nstep, ndim), dtype = np.double)
        for i in np.arange(num_params):
            for j in np.arange(num_step):
                mean_pos[i][j] = np.mean(chain[:, j, i])
                median_pos[i][j] = np.median(chain[:, j, i])
                std_pos[i][j] = np.std(chain[:, j, i])
                q16_pos[i][j] = np.percentile(chain[:, j, i], 16.)
                q84_pos[i][j] = np.percentile(chain[:, j, i], 84.)

        fig, ax = plt.subplots(num_params, sharex='all',
                               figsize=(fig_width, int(fig_width/8) *
                                        num_params))
        last = num_step
        medians = []

        for i in range(num_params):
            if verbose:
                print(self.params_mcmc[i],
                      '{:.4f} ± {:.4f}'.format(median_pos[i][last - 1],
                                               (q84_pos[i][last - 1] -
                                                q16_pos[i][last - 1]) / 2))
            # ax[i].plot(mean_pos[i][:3000], c='b')
            ax[i].plot(median_pos[i][:last], c='g')
            # ax[i].axhline(np.mean(mean_pos[i][burnin:2900]), c='b')
            ax[i].axhline(np.median(median_pos[i][burn_in:last]), c='r', lw=1)
            ax[i].fill_between(np.arange(last), q84_pos[i][:last],
                               q16_pos[i][:last], alpha=0.4)
            # ax[i].fill_between(np.arange(last), mean_pos[i][:last] \
            # +std_pos[i][:last], mean_pos[i][:last]-std_pos[i][:last],
            # alpha=0.4)
            ax[i].set_ylabel(self.params_mcmc[i], fontsize=10)
            ax[i].set_xlim(0, last)

            medians.append(np.median(median_pos[i][burn_in:last]))

        return fig

    def get_param_class(self, lens_name, model_id):
        """
        Get `Param` instance for the lens model.

        :param lens_name: name of the lens
        :type lens_name: `str`
        :param model_id: model run identifier
        :type model_id: `str`
        :return: `Param` instance
        :rtype: `obj`
        """
        self.load_output(lens_name, model_id=model_id)

        config = ModelConfig(settings=self._model_settings)
        kwargs_params = config.get_kwargs_params()
        kwargs_model = config.get_kwargs_model()
        kwargs_constraints = config.get_kwargs_constraints()

        param = Param(kwargs_model,
                      kwargs_params['lens_model'][2],
                      kwargs_params['source_model'][2],
                      kwargs_params['lens_light_model'][2],
                      kwargs_params['point_source_model'][2],
                      # kwargs_params['special'][2],
                      # kwargs_params['extinction_model'][2],
                      # kwargs_lens_init=kwargs_result['kwargs_lens'],
                      kwargs_lens_init=kwargs_params['lens_model'][0],
                      **kwargs_constraints
                      )
        return param

    def get_kwargs_from_args(self, lens_name, model_id, args, band_index=0,
                             linear_solve=False, param=None):
        """

        :param lens_name:
        :type lens_name:
        :param model_id:
        :type model_id:
        :return:
        :rtype:
        """
        if param is None:
            param = self.get_param_class(lens_name, model_id)

        kwargs = param.args2kwargs(args)

        if linear_solve:
            config = ModelConfig(settings=self._model_settings)

            # kwargs_numerics = config.get_kwargs_numerics()
            kwargs_model = config.get_kwargs_model()

            multi_band_list_out = self.get_kwargs_data_joint(
                lens_name)['multi_band_list']

            # kwargs_data = multi_band_list_out[band_index][0]
            # kwargs_psf = multi_band_list_out[band_index][1]

            im_sim = create_im_sim(multi_band_list_out,
                                   'multi-linear',
                                   kwargs_model,
                                   bands_compute=None,
                                   likelihood_mask_list=config.get_masks(),
                                   band_index=band_index)

            im_sim.image_linear_solve(**kwargs, inv_bool=True)

        return kwargs

max_stars_repo_name: chinyitan/dolphin

max_issues_repo_name: chinyitan/dolphin

max_forks_repo_name: chinyitan/dolphin


text:
#! python
# -*- coding: utf-8 -*-

import numpy as np
import cv2
import random
import librosa
from python_speech_features import logfbank, fbank, delta

NUM_PREVIOUS_FRAME = 0
NUM_NEXT_FRAME = 256
NUM_FRAMES = NUM_PREVIOUS_FRAME + NUM_NEXT_FRAME

USE_LOGSCALE = True
USE_DELTA = False
USE_SCALE = False
SAMPLE_RATE = 16000
TRUNCATE_SOUND_FIRST_SECONDS = 0.5
FILTER_BANK = 64

USING_LOG1P = False  # True log 离散化

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

def feature_normalize(v):
    return (v - v.mean(axis=0)) / (v.std(axis=0) + 2e-12)

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

def truncate_fea(frames_features, mode="train"):
    frames_features_t = frames_features.T
    freq, time = frames_features_t.shape

    if "train" == mode:
        try:
            j = random.randint(NUM_PREVIOUS_FRAME, time - NUM_NEXT_FRAME)
            spec_mag = frames_features_t[:, j -
                                         NUM_PREVIOUS_FRAME: j + NUM_NEXT_FRAME]
        except:
            spec_mag = cv2.resize(frames_features_t, (NUM_FRAMES, freq))
    else:
        spec_mag = frames_features_t

    final_data_t = feature_normalize(spec_mag)

    final_data = final_data_t.T

    final_feature = np.expand_dims(final_data, 0)
    return final_feature

def frames_fea(path, mode="train", feat_opt="melspec", double_length=False):
    try:
        (y, rate) = librosa.load(path, sr=SAMPLE_RATE)
    except Exception as e:
        print("Error audio file: ", path)
        return None

    # Trimming
    y, _ = librosa.effects.trim(y)

    # 预加重
    y = np.append(y[0], y[1:] - 0.97*y[:-1])  # TODO need check again!

    if double_length:
        if "train" == mode:
            sig = np.concatenate([y, y])
            # if np.random.random() < 0.3:
            # sig = sig[::-1]
        else:
            # sig = np.concatenate([y, y[::-1]])
            sig = np.concatenate([y, y])
    else:
        sig = y

    try:
        if "f26" == feat_opt:
            feat, energy = fbank(sig, rate, winfunc=np.hamming)
            filter_banks = np.log(feat)
            # print("f26: ", filter_banks.shape)
            return filter_banks
        elif "f40" == feat_opt:
            feat, energy = fbank(sig, rate, nfilt=40, winfunc=np.hamming)
            filter_banks = np.log(feat)
            # print("f40: ", filter_banks.shape)
            return filter_banks
        elif "melspec" == feat_opt:
            S = librosa.feature.melspectrogram(
                y=sig, sr=rate, window="hamm", n_fft=512, hop_length=160, win_length=400)
            feature = S.T
            # print("melspec: ", fea.shape)
            return feature
        elif "mel_pitch" == feat_opt:
            S = librosa.feature.melspectrogram(
                y=sig, sr=rate, window="hamm", n_fft=512, hop_length=160, win_length=400)
            # print("S: ", S.shape)
            pitches, magnitudes = librosa.piptrack(
                y=sig, sr=rate, window="hamm", n_fft=512, hop_length=160, win_length=400)
            # print("pitches: ", pitches.shape)
            # print("mags: ", magnitudes.shape)
            feature_t = np.concatenate([S, pitches])
            feature = feature_t.T
            return feature
        else:
            raise Exception("unsupported feat_opt: {}".format(feat_opt))

    except Exception:
        raise Exception("Error: %s" % path)

def get_truncated_features(wav_path, mode="train", feat_opt="stft"):
    frames_features = frames_fea(wav_path, mode=mode, feat_opt=feat_opt)
    truncated_features = truncate_fea(frames_features, mode=mode)
    return truncated_features


text:
import numpy as np
from gym import spaces
from pettingzoo import AECEnv
from pettingzoo.utils import agent_selector, wrappers

from .board import Board

def env(n_kewai=8):
    env = raw_env(n_kewai=n_kewai)
    env = wrappers.CaptureStdoutWrapper(env)
    env = wrappers.TerminateIllegalWrapper(env, illegal_reward=-1)
    env = wrappers.AssertOutOfBoundsWrapper(env)
    env = wrappers.OrderEnforcingWrapper(env)
    return env

class raw_env(AECEnv):
    metadata = {'render.modes': ['human'], "name": 'mutotere_v0'}

    def __init__(self, n_kewai=8):
        super().__init__()

        self.n_kewai = n_kewai
        self.board = Board(n_kewai=self.n_kewai)

        self.possible_agents = ['player_' + str(r) for r in range(2)]
        self.agents = self.possible_agents[:]

        self.action_spaces = {agent: spaces.Discrete(self.n_kewai + 1) for agent in self.possible_agents}
        self.observation_spaces = {agent: spaces.Dict({
            'observation': spaces.Box(low=0, high=1, shape=(self.n_kewai + 1, 2), dtype=np.int8),
            'action_mask': spaces.Box(low=0, high=1, shape=(self.n_kewai + 1,), dtype=np.int8)
        }) for agent in self.possible_agents}

        self._cumulative_rewards = {agent: 0 for agent in self.possible_agents}
        self.rewards = {agent: 0 for agent in self.possible_agents}
        self.dones = {agent: False for agent in self.possible_agents}
        self.infos = {agent: {} for agent in self.possible_agents}

        self._agent_selector = agent_selector(self.possible_agents)
        self.agent_selection = self._agent_selector.reset()

    def observation_space(self, agent):
        return self.observation_spaces[agent]

    def action_space(self, agent):
        return self.action_spaces[agent]

    def reset(self):
        self.board = Board(n_kewai=self.n_kewai)

        self.agents = self.possible_agents[:]
        self._cumulative_rewards = {agent: 0 for agent in self.agents}
        self.rewards = {agent: 0 for agent in self.agents}
        self.dones = {agent: False for agent in self.agents}
        self.infos = {agent: {} for agent in self.agents}

        self._agent_selector.reinit(self.agents)
        self.agent_selection = self._agent_selector.reset()

    def observe(self, agent):
        flat_board = self.board.kewai + [self.board.putahi]
        cur_player = self.agents.index(agent)
        opp_player = (cur_player + 1) % 2

        cur_p_stones = np.equal(flat_board, self.board.agent_stones[cur_player])
        opp_p_stones = np.equal(flat_board, self.board.agent_stones[opp_player])

        observation = np.stack([cur_p_stones, opp_p_stones], axis=-1).astype(np.int8)
        valid_moves = self.board.get_possible_actions(cur_player) if agent == self.agent_selection else []

        action_mask = np.zeros(self.n_kewai + 1, dtype=np.int8)
        for i in valid_moves:
            action_mask[i] = 1

        return {'observation': observation, 'action_mask': action_mask}

    def step(self, action):
        if self.dones[self.agent_selection]:
            return self._was_done_step(action)

        self.board.play_move(self.agents.index(self.agent_selection), action)

        next_agent = self._agent_selector.next()

        if len(self.board.get_possible_actions(self.agents.index(next_agent))) == 0:
            self.rewards[self.agent_selection] += 1
            self.rewards[next_agent] -= 1
            self.dones = {agent: True for agent in self.agents}

        self._cumulative_rewards[self.agent_selection] = 0
        self.agent_selection = next_agent
        self._accumulate_rewards()

    def render(self, mode='human'):
        self.board.render()

    def close(self):
        self.board.close()

max_stars_repo_name: Aroksak/MuTorere

max_issues_repo_name: Aroksak/MuTorere

max_forks_repo_name: Aroksak/MuTorere


text:
import src.Utils.raycast as rc
from scipy.stats import norm
import math

class RangeFinderBeamModel:
    def __init__(self, w_hit, w_short, w_max, w_rand, sigma_hit, lamda_short, z_max):
        self.w_hit = w_hit
        self.w_short = w_short
        self.w_max = w_max
        self.w_rand = w_rand
        self.sigma_hit = sigma_hit
        self.lamda_short = lamda_short
        self.z_max = z_max

        sum = w_hit + w_short + w_max + w_rand
        if math.fabs(sum - 1.0) > 0.0001:
            raise Exception("w_hit, w_short, w_max, w_rand should add up to 1.")

    # Calculate the probability of sensor measurements zt, given the map and robot pose.
    # ztArr: Array of sensor measurement [[theta, distance], [...], [...]]
    # grid: Occupancy grid.
    # pose: Robot pose [x, y, theta].
    def calculateProb(self, ztArr, pose, grid, resolution=1):
        prob = 1
        for i in range(len(ztArr)):
            # Info from sensor measurement i
            zt = ztArr[i]
            theta = zt[0]
            dist = zt[1]

            # Clamp the measurement value, only [0, z_max] range is allowed.
            if dist > self.z_max:
                dist = self.z_max

            if dist < 0:
                dist = 0

            # Calculate expected measurement.
            expec_dist = rc.distanceRaycast(grid, pose, theta, resolution=resolution, limit=self.z_max)

            # We cannot have expected distance exceeds z_max,
            # otherwise, normalizer in Phit can be very big.
            if expec_dist > self.z_max:
                expec_dist = self.z_max

            prob *= self._calculateTotalProb(dist, expec_dist)

        return prob

    def _calculateTotalProb(self, dist, expec_dist):
        # Calculate probabilities.
        Phit = self._calculatePhit(dist, expec_dist)
        Pshort = self._calculatePshort(dist, expec_dist)
        Pmax = self._calculatePmax(dist)
        Prand = self._calculatePrand(dist)

        # Sum up all the probabilities.
        return self.w_hit * Phit + self.w_short * Pshort + self.w_max * Pmax + self.w_rand * Prand

    # Calculate Phit probability of sensor given distance and expected distance.
    def _calculatePhit(self, dist, expec_dist):

        if dist > self.z_max or dist < 0:
            return 0

        prob = norm.pdf(dist, loc=expec_dist, scale=self.sigma_hit)

        # If the measured distance and expected distance are 3 standard deviation away,
        # then we treat the normalizer as 1.
        if math.fabs(expec_dist - dist) >= 3 * self.sigma_hit:
            normalizer = 1.0
        else:
            normalizer = 1.0 / \
                (norm.cdf(self.z_max, loc=expec_dist, scale=self.sigma_hit)
                    - norm.cdf(0, loc=expec_dist, scale=self.sigma_hit))

        return prob * normalizer

    def _calculatePshort(self, dist, expec_dist):

        if dist > expec_dist or dist < 0:
            return 0

        # Avoid division by zero.
        if expec_dist <= 0:
            expec_dist = 0.1

        prob = self.lamda_short * math.exp(-self.lamda_short * dist)
        normalizer = 1.0 / (1.0 - math.exp(-self.lamda_short * expec_dist))

        return prob * normalizer

    def _calculatePmax(self, dist):
        if dist >= self.z_max:
            return 1.0
        else:
            return 0

    def _calculatePrand(self, dist):

        if dist < 0 or dist > self.z_max:
            return 0

        return 1.0 / self.z_max


text:
# -*- coding: utf-8 -*-
"""
P-spline versions of Whittaker functions
----------------------------------------

pybaselines contains penalized spline (P-spline) versions of all of the
Whittaker-smoothing-based algorithms implemented in pybaselines. The reason
for doing so was that P-splines offer additional user flexibility when choosing
parameters for fitting and more easily work for unequally spaced data. This example
will examine the relationship of `lam` versus the number of data points when fitting
a baseline with the :func:`.arpls` function and its P-spline version,
:func:`.pspline_arpls`.

Note that the exact optimal `lam` values reported in this example are not of significant
use since they depend on many other factors such as the baseline curvature, noise, peaks,
etc.; however, the examined trends can be used to simplify the process of selecting `lam`
values for fitting new datasets.

"""
# sphinx_gallery_thumbnail_number = 2

from itertools import cycle

import matplotlib.pyplot as plt
import numpy as np

from pybaselines import spline, whittaker

# local import with setup code
from example_helpers import make_data, optimize_lam

# %%
# The baseline for this example is an exponentially decaying baseline, shown below.
# Other baseline types could be examined, similar to the
# :ref:`Whittaker lam vs data size example <sphx_glr_examples_whittaker_plot_lam_vs_data_size.py>`,
# which should give similar results.
plt.plot(make_data(1000, bkg_type='exponential')[0])

# %%
# For each function, the optimal `lam` value will be calculated for data sizes
# ranging from 500 to 20000 points. Further, the intercept and slope of the linear fit
# of the log of the data size, N, and the log of the `lam` value will be reported.
# The number of knots for the P-spline version is fixed at the default, 100 (the effect
# of the number of knots versus optimal `lam` is shown in another
# :ref:`example <sphx_glr_examples_spline_plot_lam_vs_num_knots.py>`).
print('Function, intercept & slope of log(N) vs log(lam) fit')
print('-' * 60)

show_plots = False  # for debugging
num_points = np.logspace(np.log10(500), np.log10(20000), 6, dtype=int)
symbols = cycle(['o', 's'])
_, ax = plt.subplots()
legend = [[], []]
for i, func in enumerate((whittaker.arpls, spline.pspline_arpls)):
    func_name = func.__name__
    best_lams = np.empty_like(num_points, float)
    min_lam = None
    for j, num_x in enumerate(num_points):
        y, baseline = make_data(num_x, bkg_type='exponential')
        # use a slightly lower tolerance to speed up the calculation
        min_lam = optimize_lam(y, baseline, func, min_lam, tol=1e-2, max_iter=50)
        best_lams[j] = min_lam

        if show_plots:
            plt.figure(num=num_x)
            if i == 0:
                plt.plot(y)
            plt.plot(baseline)
            plt.plot(func(y, lam=10**min_lam)[0], '--')

    fit = np.polynomial.polynomial.Polynomial.fit(np.log10(num_points), best_lams, 1)
    coeffs = fit.convert().coef
    print(f'{func_name:<16} {coeffs}')
    line = 10**fit(np.log10(num_points))

    handle_1 = ax.plot(num_points, line, label=func_name)[0]
    handle_2 = ax.plot(num_points, 10**best_lams, next(symbols))[0]
    legend[0].append((handle_1, handle_2))
    legend[1].append(func_name)

ax.loglog()
ax.legend(*legend)
ax.set_xlabel('Input Array Size, N')
ax.set_ylabel('Optimal lam Value')

plt.show()

# %%
# The results shown above demonstrate that the slope of the `lam` vs data
# size best fit line is much smaller for the P-spline based version of arpls.
# This means that once the number of knots is fixed for a particular baseline,
# the required `lam` value should be much less affected by a change in the
# number of data points (assuming the curvature of the data does not change).
#
# The above results are particularly useful when processing very large datasets.
# A `lam` value greater than ~1e14 typically causes numerical issues that can cause
# the solver to fail. Most Whittaker-smoothing-based algorithms reach that `lam`
# cutoff when the number of points is around ~20,000-500,000 (depends on the exact
# algorithm). Since the P-spline versions do not experience such a large increase in
# the required `lam`, they are more suited to fit those larger datasets. Additionally,
# the required `lam` value for the P-spline versions can be lowered simply by reducing
# the number of knots.
#
# It should be addressed that a similar result could be obtained using the regular
# Whittaker-smoothing-based version by truncating the number of points to a fixed
# value. That, however, would require additional processing steps to smooth out the
# resulting baseline after interpolating back to the original data size. Thus, the
# P-spline versions require less user-intervention to achieve the same result.