text:
# 3p
import numpy as np
from scipy.sparse import csr_matrix, diags
from scipy.sparse.linalg import eigsh

def cotangent(p):
    return np.cos(p) / np.sin(p)

# source: https://github.com/JM-data/PyFuncMap/blob/master/utils/MeshProcess.py
def cotangent_laplacian(verts, faces):
    '''
    Compute the cotangent matrix and the weight matrix of a shape
    :param S: a shape with VERT and TRIV
    :return: cotangent matrix W, and the weight matrix A
    '''
    T1 = faces[:, 0]
    T2 = faces[:, 1]
    T3 = faces[:, 2]

    V1 = verts[T1, :]
    V2 = verts[T2, :]
    V3 = verts[T3, :]

    nv = verts.shape[0]

    L1 = np.linalg.norm(V2 - V3, axis=1)
    L2 = np.linalg.norm(V1 - V3, axis=1)
    L3 = np.linalg.norm(V1 - V2, axis=1)
    L = np.column_stack((L1, L2, L3))  # Edges of each triangle

    Cos1 = (L2 ** 2 + L3 ** 2 - L1 ** 2) / (2 * L2 * L3)
    Cos2 = (L1 ** 2 + L3 ** 2 - L2 ** 2) / (2 * L1 * L3)
    Cos3 = (L1 ** 2 + L2 ** 2 - L3 ** 2) / (2 * L1 * L2)
    Cos = np.column_stack((Cos1, Cos2, Cos3))  # Cosines of opposite edges for each triangle
    Ang = np.arccos(Cos)  # Angles

    I = np.concatenate((T1, T2, T3))
    J = np.concatenate((T2, T3, T1))
    w = 0.5 * cotangent(np.concatenate((Ang[:, 2], Ang[:, 0], Ang[:, 1]))).astype(float)
    In = np.concatenate((I, J, I, J))
    Jn = np.concatenate((J, I, I, J))
    wn = np.concatenate((-w, -w, w, w))
    W = csr_matrix((wn, (In, Jn)), [nv, nv])  # Sparse Cotangent Weight Matrix

    cA = cotangent(Ang) / 2  # Half cotangent of all angles
    At = 1 / 4 * (L[:, [1, 2, 0]] ** 2 * cA[:, [1, 2, 0]] + L[:, [2, 0, 1]] ** 2 * cA[:, [2, 0, 1]]).astype(float)  # Voronoi Area

    Ar = np.linalg.norm(np.cross(V1 - V2, V1 - V3), axis=1)

    # Use Ar is ever cot is negative instead of At
    locs = cA[:, 0] < 0
    At[locs, 0] = Ar[locs] / 4
    At[locs, 1] = Ar[locs] / 8
    At[locs, 2] = Ar[locs] / 8

    locs = cA[:, 1] < 0
    At[locs, 0] = Ar[locs] / 8
    At[locs, 1] = Ar[locs] / 4
    At[locs, 2] = Ar[locs] / 8

    locs = cA[:, 2] < 0
    At[locs, 0] = Ar[locs] / 8
    At[locs, 1] = Ar[locs] / 8
    At[locs, 2] = Ar[locs] / 4

    Jn = np.zeros(I.shape[0])
    An = np.concatenate((At[:, 0], At[:, 1], At[:, 2]))
    Area = csr_matrix((An, (I, Jn)), [nv, 1])  # Sparse Vector of Area Weights

    In = np.arange(nv)
    A = csr_matrix((np.squeeze(np.array(Area.todense())), (In, In)),
                   [nv, nv])  # Sparse Matrix of Area Weights
    return W, A

def laplace_decomposition(verts, faces, neig=50):
    W, A = cotangent_laplacian(verts, faces)

    # sqrt_area = np.sqrt(A.diagonal().sum())

    try:
        # compute basis
        evals, evecs = eigsh(W, neig, A, 1e-6)
    except Exception as e:
        print(e)
        # add small diagonal matrix to W in case there is some isolated vertices
        c = np.ones(verts.shape[0]) * 1e-8
        damping = diags(c)
        W += damping
        # compute basis
        evals, evecs = eigsh(W, neig, A, 1e-6)

    evecs = np.array(evecs, ndmin=2)
    evecs_trans = evecs.T @ A
    evals = np.array(evals)
    return evals, evecs, evecs_trans, np.sqrt(A.diagonal().sum())

max_stars_repo_name: pvnieo/FMNet-pytorch

max_issues_repo_name: pvnieo/SURFMNet-pytorch

max_forks_repo_name: pvnieo/FMNet-pytorch

reason: import numpy,from scipy


text:
import os
import utils
import numpy as np
import matplotlib.pyplot as plt
from functools import partial
from scipy import io

'''
This is a script that simulates the movement of variable-length pendulum.
equation of motion:

   ..           .    .
  theta = [-2 * L * phi - g * sin(theta)] / l

The stabilization of this variable length pendulum is realized via the simple feedback control of the pendulum length.
The feedback control rule has the following form:

                             .
L = L0 * (1 + delta * phi * phi)

References lead to:
Bewley(2020) Stabilization of low-altitude balloon systems, Part 1: rigging with a single taut ground tether,
with analysis as a variable-length pendulum

=====================================
Author  :  Muhan Zhao
Date    :  Oct. 16, 2019
Location:  UC San Diego, La Jolla, CA
=====================================
'''

class Pendulum:
    """
    Oscillating Pendulum Simulator: including a novel feedback rule to stabilize the pendulum regardless the
    parameters {m, g, L0} and the exact expression of the oscillating angle.

    Parameters
    ------------

    :param wave       : dict, including parameters:
    :keys frequency   : float, the frequency of the oscillating angle
    :keys phi         : array, the valeus of oscillating angle w.r.t. time
    :keys dphi        : array, the valeus of oscillating angle's 1st-order derivative w.r.t. time
    :keys ddphi       : array, the valeus of oscillating angle's 2nd-order derivative w.r.t. time

    :param attributes            : dict, including parameters:
    :keys max_t                  : float, the maximum time length, default 30
    :keys dt                     : float, the marching time step, default 0.001
    :keys delta_mode_adaptive    : string, adaptive mode is on or no
    :keys delta_adaptive_const   : float, the parameter C for adaptive delta, default 0.15
    :keys delta_mode_asymptotic  : string, asymptotic mode is on or no
    :keys delta_asymptotic_const : float, the value of delta for feedback rule, default 0.1
    :keys l0                     : float, initial length of pendulum, default 1
    :keys Lmax                   : float, maximum length of pendulum, default l0*1.2
    :keys Lmin                   : float, minimum length of pendulum, default l0*.8
    :keys Ldotmax                : float, maximum derivative of pendulum length, default 1.2
    :keys Ldotmin                : float, minimum derivative of pendulum length, default .8
    :keys save_fig               : bool, save fig or no
    :keys show_fig               : bool, display fig or no
    :keys format_fig             : string, fig format to be saved

    ------------

    Methods
    ------------

    :method free_pendulum_oscillation
    :method control_pendulum_oscillation
    :method adaptive_control_pendulum_oscillation
    ------------
    """
    def __init__(self, wave, attributes):
        self.max_t = attributes.get('max_t', 30)
        self.dt = attributes.get('dt', 0.001)

        # steps: the time slides of control inputs
        # Problem here: sometimes, even the result of max_t/dt is an integer, but it could be
        # not exactly the input you want, because how python stores the number. Sometimes there will be
        # a small error such like 1e-10. So thats why need round here.
        self.steps = int(round(self.max_t / self.dt))
        self.time = np.linspace(0, self.max_t, self.steps)

        self.l0 = attributes.get('l0', 1)
        self.Ldotmax = attributes.get('Ldotmax', 5)
        self.Ldotmin = attributes.get('Ldotmin', -5)

        self.Lmax = attributes.get('Lmax', 1.2 * self.l0)
        self.Lmin = attributes.get('Lmin', .8 * self.l0)

        self.g = attributes.get('g', 9.8)
        self.m = attributes.get('m', 1)

        self.wave_phi = np.atleast_1d(wave['phi'])
        self.wave_dphi = np.atleast_1d(wave['dphi'])

        self.frequency = wave.get('frequency', np.pi/2)
        self.control_start_time = self.dt * self.wave_phi.shape[0]
        self.prev_length = self.wave_phi.shape[0]
        self.entire_t = np.arange(0, self.dt * self.prev_length + self.max_t, self.dt)

        self.pose_constrain_L = attributes.get('constrain_L', False)

        # Below is an untested idea:
        self.delta_shrinkage = .8
        # initiate the Asymptotic control mode
        self.asymptotic_control_on = attributes.get('asymptotic_mode', True)
        if self.asymptotic_control_on:
            self.delta = attributes.get('delta_asymptotic_const', 0.2)

            # asymptotic control parameters sequence
            self.asym_control_phi = np.zeros(self.steps)
            self.asym_control_dphi = np.zeros(self.steps)

            # self.asym_control_ddphi = np.zeros(self.steps)
            self.asym_control_L = np.zeros(self.steps)
            self.asym_control_L[0] = self.l0
            self.asym_control_dL = np.zeros(self.steps)

        else:
            pass

        # initiate the Adaptive control mode
        self.adaptive_control_on = attributes.get('adaptive_mode', False)
        if self.adaptive_control_on:
            self.delta_adaptive_const = attributes.get('delta_adaptive_const', 0.15)
            self.delta = attributes.get('delta_adaptive_const', 0.2)

            # adaptive control parameters sequence
            self.adap_control_phi = np.zeros(self.steps)
            self.adap_control_dphi = np.zeros(self.steps)

            # self.adap_control_ddphi = np.zeros(self.steps)
            self.adap_control_length = np.zeros(self.steps)
            self.adap_control_dlength = np.zeros(self.steps)

            self.adap_delta_sequence = np.zeros(self.steps)
        else:
            pass

        self.no_control_on = attributes.get('no_control_mode', False)
        # oscillating control sequence
        if self.no_control_on:
            self.oscillating_phi = np.zeros(self.steps)
            self.oscillating_dphi = np.zeros(self.steps)
            self.oscillating_ddphi = np.zeros(self.steps)
        else:
            pass

        # The pendulum length and length derivative at the current time step
        self.L = self.l0
        self.dL = 0

        # The time marching scheme could be used, 1 = Forward Euler or 2 = RK4
        self.time_marching_method = attributes.get('time-marching scheme', 2)

        # The root folder path and images folder path
        self.ROOT_PATH = os.getcwd()
        self.IMG_PATH = os.path.join(self.ROOT_PATH, 'images')

        # plot setting
        self.plot_trigger = attributes.get('plot', False)
        self.ylim = np.max(self.wave_phi) * 1.1
        # save fig and show fig indicator
        self.save_fig = attributes.get('save_fig', False)
        self.show_fig = attributes.get('show_fig', False)
        # saving picture in designated format
        self.format_fig = attributes.get('format_fig', '.png')
        # save data
        self.save_data = attributes.get('save_data', True)

    # ====================================   EQUATION OF MOTION   ====================================
    def fixed_length_eom(self, x):
        '''
        ====================================
         ..
        phi = - g * sin(phi) / L;
        ====================================
        phi^n+1 - phi^n     .
        ---------------- = phi^n
               dt
         .         .
        phi^n+1 - phi^n
        ---------------- = -w^2 * phi^n
               dt
        ====================================
                                .
        phi^n+1 = phi^n + dt * phi^n
         .          .
        phi^n+1 = phi^n - dt * w^2 * phi^n
        ====================================
        :return:
        '''
        x_dot = np.empty((2, ))
        x_dot[0] = x[1]
        x_dot[1] = -self.frequency**2 * x[0]
        return x_dot

    def variable_length_eom(self, x):
        '''
        First compute the length for the states x;
        Second compute the time-change of the equation;
         ..         .    .
        phi = [-2 * L * phi - g * sin(phi)] / L
        states = [x0, x1]' = [phi, dphi]'
        :return:
        '''
        # Here we can't directly compute Ldot, otherwise it would fall into infinite loop.
        # Use Ldot = L0*delta*([d(phi)/dt]^2 + phi * d^2(phi)/dt^2), plug it back into the system dynamic
        L = self.compute_length(x)

        f = np.empty((2, ))
        f[0] = x[1]
        if abs(L + self.l0 * (1 + 3 * self.delta * x[0] * x[1])) > 1e-4:
            # Regarding the expression to compute f[1], use L instead of plugging L=L0[1+delta*phi*d(phi/dt)] directly,
            # easier to add constrain on L(t)
            f[1] = (-2 * self.l0 * self.delta * x[1]**3 - self.g * np.sin(x[0])) / \
                   (L + self.l0 * (1 + 3 * self.delta * x[0] * x[1]))
        else:
            f[1] = 0
            # f[1] = (-2 * self.l0 * self.delta * x[1]**3 - self.g * np.sin(x[0])) / \
            #        (L + self.l0 * (1 + 3 * self.delta_shrinkage * self.delta * x[0] * x[1]))
            print('Invalid values for denominator encountered in the computation of variable-length equation of motion')
            print('Set it to be 0 temporarily.')
        return f

    def compute_length(self, x):
        L = self.l0 * (1 + self.delta * x[0] * x[1])
        if self.pose_constrain_L:
            L = np.clip(L, self.Lmin, self.Lmax)
        else:
            pass
        return L

    def compute_length_dot(self, x):
        f = self.variable_length_eom(x)
        Ldot = self.l0 * self.delta * (x[1]**2 + x[0]*f[1])
        if self.pose_constrain_L:
            Ldot = np.clip(Ldot, self.Ldotmin, self.Ldotmax)
        else:
            pass
        return Ldot

    def delta_update(self, t):
        """
        Adaptive control; To accelerate the convergence of the pendulum, set the amplitude of control inputs consisted of
        pendulum length to be large enough
        :param t: time instance
        :return:
        """
        # combine controlled dphi with the dphi before control starts
        dphi_sequence = np.hstack((self.wave_dphi, self.adap_control_dphi[:t]))
        # 1st find the positions that dphi <= 1e-2/2
        dphi_zero_list = np.where(np.abs(dphi_sequence) <= 1e-2/2)[0]
        if len(dphi_zero_list) > 0:
            # 2nd find the last group of dphi <= 1e-2/2
            last_dphi_zero = utils.last_consecutives(dphi_zero_list)
            # 3rd find the largest amplitude of the angle in the last cycle
            phi_sequence = np.hstack((self.wave_phi, self.adap_control_phi[:t]))
            self.delta_adaptive = self.delta_adaptive_const / np.abs(np.max(phi_sequence[last_dphi_zero]))
        else:
            # dphi have not finished one peak yet
            self.delta_adaptive = self.delta_adaptive_const / 1
        self.adap_delta_sequence[t] = np.copy(self.delta_adaptive)

    # ====================================   MAIN SCRIPT   ====================================
    def free_pendulum_oscillation(self):
        '''

        Reference link:http://hplgit.github.io/Programming-for-Computations/pub/p4c/._p4c-solarized-Python022.html
        :return:
        '''
        for step, _ in enumerate(self.time):
            if step == 0:
                state = np.hstack((self.wave_phi[-1], self.wave_dphi[-1]))
            else:
                state = np.hstack((self.oscillating_phi[step - 1], self.oscillating_dphi[step - 1]))
            # time marching ODE
            self.oscillating_phi[step], self.oscillating_dphi[step] = self.time_marching(state, self.fixed_length_eom)

    def asymptotic_control_pendulum_oscillation(self):
        for i in range(self.steps):
            if i == 0:
                state = np.hstack((self.wave_phi[-1], self.wave_dphi[-1]))
            else:
                state = np.hstack((self.asym_control_phi[i-1], self.asym_control_dphi[i-1]))

            # update the length of the pendulum
            self.asym_control_L[i] = self.compute_length(state)
            self.asym_control_dL[i] = self.compute_length_dot(state)

            # time marching the system dynamic
            self.asym_control_phi[i], self.asym_control_dphi[i] = self.time_marching(state, self.variable_length_eom)

    def adaptive_control_pendulum_oscillation(self):
        for step, _ in enumerate(self.time):
            if step == 0:  # assemble state -> ODE, full_state -> length
                state = np.hstack((self.wave_phi[-1], self.wave_dphi[-1]))
            else:
                state = np.hstack((self.adap_control_phi[step - 1], self.adap_control_dphi[step - 1]))
                full_state = np.hstack(
                    (self.adap_control_phi[step - 1], self.adap_control_dphi[step - 1], self.adap_control_ddphi[step - 1]))
            # update delta
            self.delta_update(step)

            # update length <- phi, phi_dot, phi_double_dot
            self.adap_control_length[step], self.adap_control_dlength[step] = self.length_update(full_state)

            # assemble length vector
            length = np.hstack((self.adap_control_length[step], self.adap_control_dlength[step]))

            # update phi_double_dot, this is for updating the length at the next time step
            self.adap_control_ddphi[step] = self.update_ddphi(state, length)

            # assemble the ODE
            func = partial(self.variable_length_eom, input=length)

            # time marching ODE
            self.adap_control_phi[step], self.adap_control_dphi[step] = self.time_marching(state, func)

    def main(self):
        if self.no_control_on:
            self.free_pendulum_oscillation()
        else:
            pass

        # invoke the asymptotic control
        if self.asymptotic_control_on:
            self.asymptotic_control_pendulum_oscillation()
        else:
            pass

        # invoke the adaptive control
        if self.adaptive_control_on:
            self.adaptive_control_pendulum_oscillation()
        else:
            pass

        # plot trigger
        if self.plot_trigger:
            self.plot()
        else:
            pass

        # save data
        if self.save_data:
            self.data_saver()

    # ====================================   TIME MARCHING SCHEME   ====================================

    def forward_euler(self, x, func):
        x_dot = func(x)
        x_new = x + x_dot * self.dt
        return x_new

    def runge_kutta4(self, x, func):
        f1 = func(x)
        f2 = func(x + self.dt/2 * f1)
        f3 = func(x + self.dt/2 * f2)
        f4 = func(x + self.dt * f3)
        x_new = x + self.dt/6 * (f1 + 2*f2 + 2*f3 + f4)
        return x_new

    def time_marching(self, x, func):
        if self.time_marching_method == 1:
            # Forward Euler
            x_new = self.forward_euler(x, func)
        elif self.time_marching_method == 2:
            # RK4
            x_new = self.runge_kutta4(x, func)
        else:
            # new scheme needs to be claimed
            x_new = np.zeros(2)
            raise ValueError('New time-marching scheme need to specified. Otherwise set "time-marching scheme"=2.')
        return x_new[0], x_new[1]

    # ====================================   PLOT   ====================================
    def plot(self):
        self.phi_plot()
        self.phase_space_plot()
        self.length_plot()
        # self.delta_plot()
        # self.phi_dphi_time_plot()
        self.energy_plot()
        # self.frequency_plot()

    def phi_plot(self):
        plt.figure(figsize=[16, 9])
        plt.grid()
        self.control_indicator_plot()
        if self.no_control_on:
            self.free_pendulum_plot()
        self.control_pendulum_plot()
        plt.legend()
        self.fig_save_and_show('phi')

    def control_indicator_plot(self):
        y = np.linspace(-1 * self.ylim, self.ylim, 200)
        t = self.control_start_time * np.ones(y.shape[0])
        plt.plot(t, y, 'g-.', label=r'$Control\ starts$')

    def free_pendulum_plot(self):
        entire_free_phi = np.hstack((self.wave_phi, self.oscillating_phi))
        plt.plot(self.entire_t, entire_free_phi, 'b-', label=r'$\phi(t) - No Control$')
        plt.xlabel(r'$t$', fontsize=21)
        plt.ylabel(r'$\phi(t)$', fontsize=21, rotation=0)

    def control_pendulum_plot(self):
        control_t = np.arange(0, self.max_t, self.dt) + self.control_start_time
        if self.asymptotic_control_on:
            plt.plot(control_t, self.asym_control_phi, 'r--', label=r'$\phi(t)-Controlled$')

        if self.adaptive_control_on:
            plt.plot(control_t, self.adap_control_phi, 'k-.', label=r'$\phi(t)-Adaptive\ \delta$')

    def phase_space_plot(self):
        # 2D phase plot
        plt.figure(figsize=[16, 9])
        plt.gca().set_aspect('equal', adjustable='box')
        plt.grid()
        plt.scatter(self.asym_control_phi[0], self.asym_control_dphi[0], c='b', marker='s', label='initial')
        asym_x_range = asym_y_range = adap_x_range = adap_y_range = np.ones(1)

        if self.asymptotic_control_on:
            # plot the asymptotic convergence of phi(t)
            plt.plot(self.asym_control_phi, self.asym_control_dphi, 'r--', label='Asymptotic')
            self.phase_space_arrow_plot(self.asym_control_phi, self.asym_control_dphi)
            asym_x_range = np.ptp(self.asym_control_phi, axis=0)
            asym_y_range = np.ptp(self.asym_control_dphi, axis=0)

        if self.adaptive_control_on:
            # plot the exponential convergence of phi(t)
            plt.plot(self.adap_control_phi, self.adap_control_dphi, 'k--', label='Adaptive')
            self.phase_space_arrow_plot(self.asym_control_phi, self.asym_control_dphi)
            adap_x_range = np.ptp(self.adap_control_phi, axis=0)
            adap_y_range = np.ptp(self.adap_control_dphi, axis=0)

        # r = np.max(np.hstack((adap_y_range, asym_y_range, adap_x_range, asym_x_range))) * 1.2
        # plt.ylim(0 - r / 2, 0 + r / 2)
        # plt.xlim(0 - r / 2, 0 + r / 2)
        r = 2.0
        plt.ylim(0 - r, 0 + r)
        plt.xlim(0 - r, 0 + r)

        plt.xlabel(r'$\phi(t)$', fontsize=21)
        plt.ylabel(r'$\dot{\phi}(t)$', fontsize=21, rotation=0)
        plt.legend()
        self.fig_save_and_show('phi_dphi')

    def phase_space_arrow_plot(self, phi, dphi):
        # plot 10 arrows for every 100/1000 steps; if the steps < 100, plot just one arrow
        if self.steps < 100:
            num_arrows = 1
            arrow_position_sequence = np.array([int(round(self.steps/2))])
        else:
            num_arrows = 5
            arrow_position_sequence = np.array([int(step) for step in np.linspace(0, self.steps - 2, num_arrows)])

        for num, pos in enumerate(arrow_position_sequence[1:]):
            plt.annotate('', xy=(phi[pos+1], dphi[pos+1]), xytext=(phi[pos], dphi[pos]),
                         arrowprops={'arrowstyle': '->, head_width=.5, head_length=1.5', 'color': 'r'}, va='center')

    def phi_dphi_time_plot(self):
        # Debug use
        if self.asymptotic_control_on:
            plt.figure(figsize=[16, 9])
            plt.grid()
            plt.plot(self.time, self.asym_control_phi, 'r', label=r'$\phi(t)$')
            plt.plot(self.time, self.asym_control_dphi, 'b', label=r'$\dot{\phi}(t)$')
            plt.xlabel(r'$t$', size=21)
            plt.legend()
            self.fig_save_and_show('phi_dphi_time_asym')

        if self.adaptive_control_on:
            plt.figure(figsize=[16, 9])
            plt.grid()
            plt.plot(self.time, self.adap_control_phi, 'r', label=r'$\phi(t)$')
            plt.plot(self.time, self.adap_control_dphi, 'b', label=r'$\dot{\phi}(t)$')
            plt.xlabel(r'$t$', size=21)
            plt.legend()
            self.fig_save_and_show('phi_dphi_time_adap')

    def energy_plot(self):
        plt.figure(figsize=[16, 9])
        plt.grid()
        if self.asymptotic_control_on:
            energy_asym = 1/2 * self.m * ((self.asym_control_L * self.asym_control_dphi)**2 + self.asym_control_dL**2) \
                          + self.m * self.g * (-np.cos(self.asym_control_phi) * self.asym_control_L + self.l0)
            plt.plot(self.time, energy_asym, 'r', label=r'$V$ - Asymptotic')

        if self.adaptive_control_on:
            energy_adap = 1/2 * ((self.adap_control_length * self.adap_control_dphi)**2 + self.adap_control_dlength**2) \
                          + self.g*self.adap_control_length*(1-np.cos(self.adap_control_phi))
            plt.plot(self.time, energy_adap, 'k--', label=r'$V$ - Adaptive')
        plt.xlabel(r'$t$', size=21)
        plt.ylabel(r'$V(t)$', size=21, rotation=0, labelpad=22)
        plt.legend()
        self.fig_save_and_show('energy')

    def frequency_plot(self):
        plt.figure(figsize=[16, 9])
        plt.grid()
        if self.asymptotic_control_on:
            omega_asym = np.sqrt(self.g / self.asym_control_L - (self.asym_control_dL/self.asym_control_L)**2)
            plt.plot(self.time, omega_asym, 'r', label=r'$\omega$ - Asymptotic')

        if self.adaptive_control_on:
            omega_adap = np.sqrt(self.g / self.adap_control_length - (self.adap_control_dlength/self.adap_control_length)**2)
            plt.plot(self.time, omega_adap, 'k--', label=r'$\omega$- Adaptive')
        plt.xlabel(r'$t$', size=21)
        plt.ylabel(r'$\omega(t)$', size=21, rotation=0)
        plt.legend()
        self.fig_save_and_show('frequency')

    def length_plot(self):
        # length plot
        plt.figure(figsize=[16, 9])
        plt.grid()

        # L plot
        if self.asymptotic_control_on:
            # plot the asymptotic pendulum length
            entire_asym_length = np.hstack((self.l0 * np.ones(self.prev_length), self.asym_control_L))
            plt.plot(self.entire_t, entire_asym_length, 'r--', label='Asymptotic')

        if self.adaptive_control_on:
            # plot the adaptive pendulum length
            entire_adap_length = np.hstack((self.l0 * np.ones(self.prev_length), self.adap_control_length))
            plt.plot(self.entire_t, entire_adap_length, 'k-.', label='Adaptive')
        plt.xlabel(r'$t$', fontsize=21)
        plt.ylabel(r'$\frac{L(t)}{L_0}$', fontsize=21, rotation=0, labelpad=18)
        plt.legend()
        self.fig_save_and_show('length')

        # dL plot
        plt.figure(figsize=[16, 9])
        plt.grid()

        if self.asymptotic_control_on:
            # plot the asymptotic pendulum dlength
            entire_asym_dlength = np.hstack((np.zeros(self.prev_length), self.asym_control_dL))
            plt.plot(self.entire_t, entire_asym_dlength, 'r--', label='Asymptotic')
        if self.adaptive_control_on:
            # plot the adaptive pendulum length
            entire_adap_dlength = np.hstack((np.zeros(self.prev_length), self.adap_control_dlength))
            plt.plot(self.entire_t, entire_adap_dlength, 'k-.', label='Adaptive')
        plt.xlabel(r'$t$', fontsize=21)
        plt.ylabel(r'$\dot{L}(t)$', fontsize=21, rotation=0, labelpad=18)
        plt.legend()
        self.fig_save_and_show('dlength')

    def delta_plot(self):
        plt.figure(figsize=[16, 9])
        plt.grid()

        if self.asymptotic_control_on:
            # plot the asymptotic delta (const)
            entire_asym_delta = self.delta_asymptotic * np.ones(self.entire_t.shape[0])
            plt.plot(self.entire_t, entire_asym_delta, 'r--', label='Asymptotic')

        if self.adaptive_control_on:
            entire_adap_delta = np.hstack((np.zeros(self.prev_length), self.adap_delta_sequence))
            plt.plot(self.entire_t, entire_adap_delta, 'k-.', label='Adaptive')
        plt.legend()
        plt.xlabel(r'$t$', fontsize=21)
        plt.ylabel(r'$\delta(t)$', fontsize=21, rotation=0)
        self.fig_save_and_show('delta')

    def fig_save_and_show(self, name):
        if self.save_fig:
            name = os.path.join(self.IMG_PATH, name) + self.format_fig
            plt.savefig(name, format=self.format_fig[1:], dpi=300)
        if self.show_fig:
            plt.show()

    def data_saver(self):
        data = {
            'asym_phi': self.asym_control_phi,
            'asym_dphi': self.asym_control_dphi,
            'asym_L': self.asym_control_L,
            'asym_dL': self.asym_control_dL
        }
        io.savemat('pendulum_data.mat', data)

if __name__ == "__main__":
    # T: total control time
    T = 10
    dt = 0.02
    g = 9.8
    m = 1
    l = 1

    # t_length: time length before control starts
    t_length = dt
    t = np.arange(0, t_length, dt)

    # 1st: simple sin wave
    # signal = {
    #     'phi': a * np.sin(w0 * t),
    #     'dphi': a * w0 * np.cos(w0 * t),
    # }

    signal = {
        'phi': -2 * np.ones(1),
        'dphi': 3 * np.ones(1),
    }

    properties = {
        'm': m,
        'max_t': T,
        'dt': dt,
        'adaptive_mode': False,
        'delta_adaptive_const': .15,
        'asymptotic_mode': True,
        'delta_asymptotic_const': .05,
        'l0': l,
        'constrain_L': True,
        'Ldotmax': 5,
        'Ldotmin': -5,
        'Lmax': 1.5,
        'Lmin': 0.5,
        'g': g,
        'plot': True,
        'save_fig': True,
        'show_fig': False,
        'save_data': False,
    }
    # ================
    pendu = Pendulum(signal, properties)
    pendu.main()

ext: py


text:
import numpy as np
import pandas as pd
from . import data_transform
from sklearn.pipeline import Pipeline
from sklearn import feature_selection
from sklearn.model_selection import cross_val_score, cross_validate
import time
from sklearn import metrics
import pprint
from termcolor import colored

def single_mean_absolute_percentage_error(y_true, y_pred):
    """MAPE for single output"""
    y_true, y_pred = np.array(y_true), np.array(y_pred)
    return np.mean(np.abs((y_true - y_pred) / y_true)) * 100

def mean_absolute_percentage_error(y_true_arr, y_pred_arr):
    """MAPE for multi-output"""
    # print(y_true_arr.shape, y_pred_arr.shape)
    mapes = np.empty([y_true_arr.shape[1]])
    y_pred_arr = pd.DataFrame(y_pred_arr, columns=y_true_arr.columns)
    for ind, target in enumerate(y_true_arr):
        y_true = y_true_arr[target].to_numpy()
        y_pred = y_pred_arr[target].to_numpy()
        mapes[ind] = np.mean(np.abs((y_true - y_pred) / y_true)) * 100
    return mapes

class RegrModel:
    """
    This class provides methods and attributes for regression.
    It supports multi-output, target transform, and feature scale/transform.

    It can accept a model in a simple form or as:

    ('Name', Pipeline(steps=
        [(`'transform_X'`, transformer),
         ('scale', scaler),
         (`'regress'`, data_transform.transform_regressor(
             Estimator(),
             trans_y='log'))
        ])
    ),

    `_model`: can be an sklearn model or a `Pipeline` as above
    """

    __slots__ = [
        "_name",
        "_model",
        "_trans_y",
        "_n_features",
        "_train_time",
        "_pred_times",
        "_mean_pred_time",
        "_std_pred_time",
        "_r2s",
        "_avg_r2",
        "_median_r2",
        "_std_r2",
        "_maes",
        "_avg_mae",
        "_median_mae",
        "_std_mae",
        "_mapes",
        "_avg_mape",
        "_median_mape",
        "_std_mape",
    ]

    def __init__(self, name="regr_model", model=None):
        self._name = name
        self._model = model
        self._n_features = 0
        if isinstance(self._model, Pipeline):  # isinstance is True for subclasses to
            self._trans_y = model.steps[-1][1].trans_y
        else:
            self._trans_y = model.trans_y
        self._r2s, self._avg_r2, self._median_r2, self._std_r2 = (None, 0, 0, 0)
        self._maes, self._avg_mae, self._median_mae, self._std_mae = (None, 0, 0, 0)
        self._mapes, self._avg_mape, self._median_mape, self._std_mape = (None, 0, 0, 0)
        self._train_time, self._mean_pred_time, self._std_pred_time = 0, 0, 0
        self._pred_times = []

    # @property turns method into a “getter” for an attribute with the same name
    # it'd also make it easier to change what we return without changing the interface
    @property
    def name(self):
        return self._name

    @property
    def model(self):
        return self._model

    @model.setter
    def model(self, val):
        self._model = val

    @property
    def trans_y(self):
        return self._trans_y

    @property
    def n_features(self):
        return self._n_features

    @n_features.setter
    def n_features(self, val):
        self._n_features = val

    @property
    def train_time(self):
        return self._train_time

    @train_time.setter
    def train_time(self, val):
        self._train_time = val

    @property
    def pred_times(self):
        return self._pred_times

    # assign an empty list in `predict()`
    @pred_times.setter
    def pred_times(self, val_list):
        self._pred_times = val_list

    @property
    def mean_pred_time(self):
        return self._mean_pred_time

    @mean_pred_time.setter
    def mean_pred_time(self, val):
        self._mean_pred_time = val

    @property
    def std_pred_time(self):
        return self._std_pred_time

    @std_pred_time.setter
    def std_pred_time(self, val):
        self._std_pred_time = val

    # r2
    @property
    def r2s(self):
        return self._r2s

    @r2s.setter
    def r2s(self, val):
        self._r2s = val

    @property
    def avg_r2(self):
        return self._avg_r2

    @avg_r2.setter
    def avg_r2(self, val):
        self._avg_r2 = val

    @property
    def median_r2(self):
        return self._median_r2

    @median_r2.setter
    def median_r2(self, val):
        self._median_r2 = val

    @property
    def std_r2(self):
        return self._std_r2

    @std_r2.setter
    def std_r2(self, val):
        self._std_r2 = val

    # mae
    @property
    def maes(self):
        return self._maes

    @maes.setter
    def maes(self, val):
        self._maes = val

    @property
    def avg_mae(self):
        return self._avg_mae

    @avg_mae.setter
    def avg_mae(self, val):
        self._avg_mae = val

    @property
    def median_mae(self):
        return self._median_mae

    @median_mae.setter
    def median_mae(self, val):
        self._median_mae = val

    @property
    def std_mae(self):
        return self._std_mae

    @std_mae.setter
    def std_mae(self, val):
        self._std_mae = val

    # mape
    @property
    def mapes(self):
        return self._mapes

    @mapes.setter
    def mapes(self, val):
        self._mapes = val

    @property
    def avg_mape(self):
        return self._avg_mape

    @avg_mape.setter
    def avg_mape(self, val):
        self._avg_mape = val

    @property
    def median_mape(self):
        return self._median_mape

    @median_mape.setter
    def median_mape(self, val):
        self._median_mape = val

    @property
    def std_mape(self):
        return self._std_mape

    @std_mape.setter
    def std_mape(self, val):
        self._std_mape = val

    # regression
    def fit(self, X_df, y_df, verbose=False):
        # we do a separate fit to get _n_features so that the timing remains accurate
        if isinstance(self.model, Pipeline):
            self.model.steps.insert(-1, ("_Debug", data_transform.Debug()))
            t0 = time.time()
            fitted = self.model.fit(X_df, y_df)
            self.train_time = round(time.time() - t0, 3)

            self.n_features = self.model.steps[-2][1].X_head.shape[1]
            self.model.steps.pop(-2)
        else:
            t0 = time.time()
            fitted = self.model.fit(X_df, y_df)
            self.train_time = round(time.time() - t0, 3)

            self.n_features = X_df.shape[1]

        if verbose:
            print(
                colored("NOTE:", "red"),
                "due to a current bug in VSCode,"
                " `func` and `inverse_func` might not get printed in the notebook",
            )
            print(colored("Fitted:\n", "magenta"), str(fitted))
            print(colored("Model Parameters:\n", "magenta"))
            pprint.pprint(self.model.get_params(), width=1)

        return fitted

    def predict(self, test_X_df, test_y_df, pred_iter=10, verbose=False):
        self.pred_times = []
        for _ in range(0, pred_iter):
            t1 = time.time()
            ypred = self.model.predict(test_X_df)
            self.pred_times.append(round(time.time() - t1, 6))
        self.mean_pred_time = round(np.mean(self.pred_times), 4)
        self.std_pred_time = round(np.std(self.pred_times), 4)
        # print(self.name, "pred and, train times: ",
        #     self.pred_time, self.train_time)
        self.r2s = metrics.r2_score(test_y_df, ypred, multioutput="raw_values")
        # take the average manually,
        # since multioutput can take other values than 'raw_values'
        self.avg_r2 = np.mean(self.r2s)
        self.median_r2 = np.median(self.r2s)
        self.std_r2 = np.std(self.r2s)
        self.maes = metrics.mean_absolute_error(
            test_y_df, ypred, multioutput="raw_values"
        )
        self.avg_mae = np.mean(self.maes)
        self.median_mae = np.median(self.maes)
        self.std_mae = np.std(self.maes)
        self.mapes = mean_absolute_percentage_error(test_y_df, ypred)
        self.avg_mape = np.mean(self.mapes)
        self.median_mape = np.median(self.mapes)
        self.std_mape = np.std(self.mapes)

        if verbose:
            print(self.name, ", model.score:", self.model.score(test_X_df, test_y_df))
            print(self.trans_y, "y-transform, raw scores: ", self.r2s)
            print("avg. r2 score ('uniform_average'): ", self.avg_r2)
            print("pred time iters:", self.pred_times)

        return ypred

    def get_metrics(self, verbose=False):
        res = self.__repr__()

        if verbose:
            res = (
                res
                + "\n   raw r2s: "
                + "\033[0m"
                + str(self.r2s)
                + "\n   raw maes: "
                + str(self.maes)
                + "\n   raw mapes: "
                + str(self.mapes)
            )

        return res + "\n"

    # print
    def __repr__(self):
        """Print object.
        """
        return (
            colored(self.name + "\n", "cyan")
            + (
                str(self.model.steps[-1][1])
                if isinstance(self.model, Pipeline)
                else str((self.model))
            )
            + colored("\navg r2: ", "red")
            + colored("%.4f" % self.avg_r2, "green")
            + colored(", avg mae: ", "red")
            + colored("%.4f" % self.avg_mae, "green")
            + colored(", avg mape: ", "red")
            + colored("%.2f" % self.avg_mape, "green")
            + ", median mape: %.2f" % self.median_mape
            + ", train time: %.4f" % self.train_time
            + ", mean_pred time: %.4f" % self.mean_pred_time
            + colored("\nn_features: ", "red")
            + str(self.n_features)
            + (
                "\nPipeline Steps: " + str(self.model.named_steps.keys())
                if isinstance(self.model, Pipeline)
                else ""
            )
            + "\n----------\n"
        )

def check_scores(predictor, train_X, train_y, test_X, test_y, cv=None, verbose=False):
    """`predictor`: can be a regressor or a `Pipeline` with a regressor at the last step

    `cv`: if `None`, cross-validation scores will not be computed
    """
    if cv is not None:
        cv_scoring = ["r2", "neg_mean_absolute_error"]
        cv_dict = cross_validate(
            predictor, train_X, train_y, scoring=cv_scoring, cv=cv, n_jobs=-1,
        )

    # fit on train
    # Pipeline
    if isinstance(predictor, Pipeline):
        predictor.steps.insert(-1, ("_Debug", data_transform.Debug()))
        # fit
        t_tr0 = time.time()
        fitted = predictor.fit(train_X, train_y)
        t_tr1 = time.time()
        n_features = predictor.steps[-2][1].X_head.shape[1]
        predictor.steps.pop(-2)
    else:
        # fit
        t_tr0 = time.time()
        fitted = predictor.fit(train_X, train_y)
        t_tr1 = time.time()
        # meta-estimator feature selector such as RFE or RFECV
        if hasattr(predictor, "n_features_"):
            n_features = predictor.n_features_
        # regressor
        else:
            n_features = train_X.shape[1]

    print("Fitting time:", round(t_tr1 - t_tr0, 4))
    print("Num of features:", n_features)
    if cv is not None:
        print(
            "Cross-validation MAE: {:.4f}".format(
                -np.mean(cv_dict["test_neg_mean_absolute_error"])
            )
        )
        print("Cross-validation R2: {:.4f}".format(np.mean(cv_dict["test_r2"])))
        if verbose:
            print("Cross-validation detailed results:")
            for res in cv_dict:
                print("\t" + res + ": ", cv_dict[res])

    # predict on test
    t_ts0 = time.time()
    y_pred = predictor.predict(test_X)
    t_ts1 = time.time()
    print("Prediction time:", round(t_ts1 - t_ts0, 4))

    print(
        "MAE on Test:: {:.4f}".format(
            metrics.mean_absolute_error(test_y, y_pred, multioutput="uniform_average")
        )
    )
    print(
        "MAPE on Test:: {:.2f}%".format(
            np.mean(mean_absolute_percentage_error(test_y, y_pred))
        )
    )
    print(
        "R2 on Test: {:.2f}".format(
            metrics.r2_score(test_y, y_pred, multioutput="uniform_average")
        )
    )
    print("\n")
    return fitted

reason: import numpy


text:
import os

import numpy as np

from . import msis2f, msis00f

def run(dates, lons, lats, alts, f107s, f107as, aps,
        options=None, version=2):
    """Call MSIS looping over all possible inputs.

    Parameters
    ----------
    dates : list of dates
        Dates and times of interest
    lons : list of floats
        Longitudes of interest
    lats : list of floats
        Latitudes of interest
    alts : list of floats
        Altitudes of interest
    f107s : list of floats
        F107 value for the given date(s)
    f107as : list of floats
        F107 running 100-day average for the given date(s)
    aps : list of floats
        Ap for the given date(s)
    options : list of floats (length 25) [optional]
        A list of options (switches) to the model
    version : int or string, default: 2
        MSIS version number, [00, 2]

    Returns
    -------
    ndarray (ndates, nlons, nlats, nalts, 11)
        | The data calculated at each grid point:
        | [Total mass density (kg/m3)
        | N2 # density (m-3),
        | O2 # density (m-3),
        | O # density (m-3),
        | He # density (m-3),
        | H # density (m-3),
        | Ar # density (m-3),
        | N # density (m-3),
        | Anomalous oxygen # density (m-3),
        | empty (will contain NO in future release),
        | Temperature (K)]

    Notes
    -----
    The 10.7 cm radio flux is at the Sun-Earth distance,
    not the radio flux at 1 AU.

    """
    if options is None:
        options = create_options()
    elif len(options) != 25:
        raise ValueError("options needs to be a list of length 25")

    input_shape, input_data = create_input(dates, lons, lats, alts,
                                           f107s, f107as, aps)

    if int(version) == 0:
        msis00f.pytselec(options)
        output = msis00f.pygtd7d(input_data[:, 0], input_data[:, 1],
                                 input_data[:, 2], input_data[:, 3],
                                 input_data[:, 4], input_data[:, 5],
                                 input_data[:, 6], input_data[:, 7:])

    elif int(version) == 2:
        # We need to point to the MSIS parameter file that was installed with
        # the Python package
        msis_path = os.path.dirname(os.path.realpath(__file__)) + "/"
        msis2f.pyinitswitch(options, parmpath=msis_path)
        output = msis2f.pymsiscalc(input_data[:, 0], input_data[:, 1],
                                   input_data[:, 2], input_data[:, 3],
                                   input_data[:, 4], input_data[:, 5],
                                   input_data[:, 6], input_data[:, 7:])

    else:
        raise ValueError(f"The MSIS version selected: {version} is not "
                         "a valid version number [0, 2]")

    # The Fortran code puts 9.9e-38 in as NaN
    # Have to make sure this doesn't overlap 0 due to really small values
    # so atol should be less than the comparison value
    output[np.isclose(output, 9.9e-38, atol=1e-38)] = np.nan

    return output.reshape(input_shape + (11,))

def create_options(f107=1, time_independent=1, symmetrical_annual=1,
                   symmetrical_semiannual=1, asymmetrical_annual=1,
                   asymmetrical_semiannual=1, diurnal=1, semidiurnal=1,
                   geomagnetic_activity=1, all_ut_effects=1, longitudinal=1,
                   mixed_ut_long=1, mixed_ap_ut_long=1, terdiurnal=1):
    """Creates the options list based on keyword argument choices.

    Defaults to all 1's for the input options.

    Parameters
    ----------
    f107 : float
        Account for F10.7 variations
    time_independent : float
        Account for time variations
    symmetrical_annual : float
        Account for symmetrical annual variations
    symmetrical_semiannual : float
        Account for symmetrical semiannual variations
    asymmetrical_annual : float
        Account for asymmetrical annual variations
    asymmetrical_semiannual : float
        Account for asymmetrical semiannual variations
    diurnal : float
        Account for diurnal variations
    semidiurnal : float
        Account for semidiurnal variations
    geomagnetic_activity : float
        Account for geomagnetic activity
        (1 = Daily Ap mode, -1 = Storm-time Ap mode)
    all_ut_effects : float
        Account for all UT/longitudinal effects
    longitudinal : float
        Account for longitudinal effects
    mixed_ut_long : float
        Account for UT and mixed UT/longitudinal effects
    mixed_ap_ut_long : float
        Account for mixed Ap, UT, and longitudinal effects
    terdiurnal : float
        Account for terdiurnal variations

    Returns
    -------
    list
        25 options as a list ready for msis2 input
    """
    options = [f107, time_independent, symmetrical_annual,
               symmetrical_semiannual, asymmetrical_annual,
               asymmetrical_semiannual, diurnal, semidiurnal,
               geomagnetic_activity, all_ut_effects, longitudinal,
               mixed_ut_long, mixed_ap_ut_long, terdiurnal] + [1]*11
    return options

def create_input(dates, lons, lats, alts, f107s, f107as, aps):
    """Combine all input values into a single flattened array.

    Parameters
    ----------
    dates : list of dates
        Dates and times of interest
    lons : list of floats
        Longitudes of interest
    lats : list of floats
        Latitudes of interest
    alts : list of floats
        Altitudes of interest
    f107s : list of floats
        F107 values for the given date(s)
    f107as : list of floats
        F107 running 100-day average for the given date(s)
    aps : list of floats
        Ap for the given date(s)

    Returns
    -------
    (shape, flattened_input)
        The shape of the data as a tuple (ndates, nlons, nlats, nalts) and
        the flattened version of the input data
        (ndates*nlons*nlats*nalts, 14).
    """
    # Turn everything into arrays
    dates = np.atleast_1d(np.array(dates, dtype='datetime64'))
    dyear = (dates.astype('datetime64[D]') -
             dates.astype('datetime64[Y]')).astype(float) + 1  # DOY 1-366
    dseconds = (dates.astype('datetime64[s]') -
                dates.astype('datetime64[D]')).astype(float)
    # TODO: Make it a continuous day of year?
    #       The new code mentions it should be and accepts float, but the
    #       regression tests indicate it should still be integer DOY
    # dyear += dseconds/86400
    lons = np.atleast_1d(lons)
    # If any longitudes were input as negatives, try to change them to
    # the (0, 360) range
    lons[lons < 0] += 360
    lats = np.atleast_1d(lats)
    alts = np.atleast_1d(alts)
    f107s = np.atleast_1d(f107s)
    f107as = np.atleast_1d(f107as)
    aps = np.atleast_1d(aps)

    ndates = len(dates)
    nlons = len(lons)
    nlats = len(lats)
    nalts = len(alts)
    shape = (ndates, nlons, nlats, nalts)

    if not (ndates == len(f107s) == len(f107as) == len(aps)):
        raise ValueError(f"The length of dates ({ndates}), f107s "
                         f"({len(f107s)}), f107as ({len(f107as)}), "
                         f"and aps ({len(aps)}) must all be equal")

    # Make a grid of indices
    indices = np.stack(np.meshgrid(np.arange(ndates),
                                   np.arange(nlons),
                                   np.arange(nlats),
                                   np.arange(nalts), indexing='ij'),
                       -1).reshape(-1, 4)

    # Now stack all of the arrays, indexing by the proper indices
    arr = np.stack([dyear[indices[:, 0]], dseconds[indices[:, 0]],
                    lons[indices[:, 1]], lats[indices[:, 2]],
                    alts[indices[:, 3]],
                    f107s[indices[:, 0]], f107as[indices[:, 0]]], -1)
    # ap has 7 components, so we need to concatenate it onto the
    # arrays rather than stack
    return shape, np.concatenate([arr, aps[indices[:, 0], :]], axis=1,
                                 dtype=np.float32)

max_stars_repo_name: SWxTREC/pymsis

max_issues_repo_name: SWxTREC/pymsis

max_forks_repo_name: SWxTREC/pymsis


text:
import numpy as np

###############
## Class for univariate gaussian
## p(x) = 1/sqrt(2*pi*simga^2) * e ^ - (x-miu)^2/2*sigma^2
## Where miu is the gaussian mean, and sigma^2 is the gaussian variance
################

class Gaussian():

    def __init__(self,mean,variance):
        self.mean = mean;
        self.variance = variance;

    def sample(self,points):
        return np.random.normal(self.mean,self.variance,points)

### Returns the mean and the variance of a data set of X points assuming that the points come from a gaussian distribution
### X
def estimate_gaussian(X):
    mean = np.mean(X,0)
    variance = np.var(X,0)
    return Gaussian(mean,variance)

max_stars_repo_name: SimonSuster/lxmls-toolkit

max_issues_repo_name: daviddao/LxMLS-labs-solution

max_forks_repo_name: daviddao/LxMLS-labs-solution


text:
# Copyright (c) Peking University 2018
#
# The software is released under the Open-Intelligence Open Source License V1.0.
# The copyright owner promises to follow "Open-Intelligence Open Source Platform
# Management Regulation V1.0", which is provided by The New Generation of
# Artificial Intelligence Technology Innovation Strategic Alliance (the AITISA).

import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.distributed as dist
import torch.optim
import torch.utils.data
import torch.utils.data.distributed
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import torchvision.models as models
import os
import time
from sklearn.cluster import KMeans
from sklearn.externals import joblib
import numpy as np
import tqdm
import json
import math
from tensorboardX import SummaryWriter
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
import matplotlib.pyplot as plt
from shrink import Model_WeightLoss, Part_WeightLoss

model = nn.DataParallel(models.__dict__['resnet18'](pretrained=True))
model.cuda()

data_batch_size = 128
train_num = 1281168
batches = math.ceil(train_num/data_batch_size)
interval_epoch = 2
in_epoch = 36

# imagenet train dir
traindir = os.path.join('/home/yzy/2018-work/pytorch-example', 'train')

normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                     std=[0.229, 0.224, 0.225])

train_dataset = datasets.ImageFolder(
        traindir,
        transforms.Compose([
            transforms.RandomResizedCrop(224),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            normalize,
        ]))

train_loader = torch.utils.data.DataLoader(
        train_dataset, batch_size=data_batch_size, shuffle=True,
        num_workers=4, pin_memory=True)

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

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

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

def accuracy(output, target, topk=(1,)):
    """Computes the precision@k for the specified values of k"""
    maxk = max(topk)
    batch_size = target.size(0)

    _, pred = output.topk(maxk, 1, True, True)
    pred = pred.t()
    correct = pred.eq(target.view(1, -1).expand_as(pred))

    res = []
    for k in topk:
        correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)
        res.append(correct_k.mul_(100.0 / batch_size))
    return res

def train(train_loader, model, criterion, optimizer, epoch, writer,w_cen_end=None):
    batch_time = AverageMeter()
    data_time = AverageMeter()
    losses = AverageMeter()
    loss_reg = AverageMeter()
    top1 = AverageMeter()
    top5 = AverageMeter()

    reg_alpha = 0.00001
    model.train()

    end = time.time()
    for i, (input, target) in (enumerate(train_loader)):
        data_time.update(time.time() - end)
        input = input.cuda(async=True)
        target = target.cuda(async=True)
        input_var = torch.tensor(input)
        target_var = torch.tensor(target)

        output = model(input_var)
        loss = criterion[0](output, target_var)
        loss_w_reg = criterion[1].forward()
        loss_all = loss + loss_w_reg

        prec1, prec5 = accuracy(output.data, target, topk=(1, 5))
        losses.update(loss_all.item(), input.size(0))
        loss_reg.update(loss_w_reg.item(), input.size(0))
        top1.update(prec1.item(), input.size(0))
        top5.update(prec5.item(), input.size(0))

        optimizer.zero_grad()
        loss_all.backward()

        criterion[1].shrink(i*(epoch%interval_epoch+1)/batches/interval_epoch if i*(epoch%interval_epoch+1)/batches/interval_epoch<1 else 1)
        criterion[1].freeze()
        optimizer.step()
        criterion[1].constrain()
        batch_time.update(time.time() - end)
        end = time.time()

        if i % 50 == 0:
            # criterion[1].update_center()
            writer.add_scalars('data/Prec', {
                                             "Prec@1": top1.avg,
                                             "Prec@5": top5.avg}, i)
            writer.add_scalars('data/Loss', {
                                             "loss_all":losses.avg,
                                             "loss_w_loss":loss_reg.avg}, i)
        if i % 1000==0:
            print('Epoch: [{0}][{1}/{2}]\t, batch_time: {3}'.format(epoch, i, len(train_loader), batch_time.avg))

# Weight centers
with open('compress_resnet18/res18_cen16.json', 'r') as f:
    modelinit = json.load(f)
nllloss = nn.CrossEntropyLoss().cuda()
class_wl = Model_WeightLoss(model.module, modelinit, loss_weight =0.1)
class_wl.convert2dict()
class_wl.compute_boundary()
class_wl.init_model_criterion()
criterion = [nllloss, class_wl]
optimizer = torch.optim.RMSprop(model.parameters(), lr= 8e-7, momentum=0.5)
writer = SummaryWriter()
up_hist = []
for i in range(in_epoch):
    train(train_loader, model, criterion, optimizer, i, writer, modelinit)
    torch.save(model.state_dict(), './compress_resnet18/'+str(i)+'.torch')
    if i%interval_epoch == (interval_epoch-1):
        class_wl.phase_shrink()
        up_hist.append(criterion[1].update_phase())
class_wl.counter = 0
class_wl.Traverse_layer(model.children(), cen_dict = class_wl.cen_dict_list, in_func=class_wl.last_prune)
torch.save(model.state_dict(), './compress_resnet18/compressed.torch')

max_stars_repo_name: Antsypc/qizhi

max_issues_repo_name: Antsypc/qizhi

max_forks_repo_name: Antsypc/qizhi


text:
#!/usr/bin/env python
# Copyright (c) 2011-2020, wradlib developers.
# Distributed under the MIT License. See LICENSE.txt for more info.

"""
Clutter Identification
^^^^^^^^^^^^^^^^^^^^^^

.. autosummary::
   :nosignatures:
   :toctree: generated/

   {}
"""
__all__ = [
    "filter_gabella",
    "filter_gabella_a",
    "filter_gabella_b",
    "filter_cloudtype",
    "filter_window_distance",
    "histo_cut",
    "classify_echo_fuzzy",
]
__doc__ = __doc__.format("\n   ".join(__all__))

import numpy as np
from scipy import ndimage

from wradlib import dp, util

def filter_gabella_a(img, wsize, tr1, cartesian=False, radial=False):
    """First part of the Gabella filter looking for large reflectivity \
    gradients.

    This function checks for each pixel in ``img`` how many pixels surrounding
    it in a window of ``wsize`` are by ``tr1`` smaller than the central pixel.

    Parameters
    ----------
    img : :py:class:`numpy:numpy.ndarray`
        radar image to which the filter is to be applied
    wsize : int
        Size of the window surrounding the central pixel
    tr1 : float
        Threshold value
    cartesian : bool
        Specify if the input grid is Cartesian or polar
    radial : bool
        Specify if only radial information should be used

    Returns
    -------
    output : :py:class:`numpy:numpy.ndarray`
        an array with the same shape as ``img``, containing the
        filter's results.

    See Also
    --------
    :func:`~wradlib.clutter.filter_gabella` - the complete filter

    :func:`~wradlib.clutter.filter_gabella_b` - the second part of the filter

    Examples
    --------

    See :ref:`/notebooks/classify/wradlib_clutter_gabella_example.ipynb`.

    """
    nn = wsize // 2
    count = -np.ones(img.shape, dtype=int)
    range_shift = range(-nn, nn + 1)
    azimuth_shift = range(-nn, nn + 1)
    if radial:
        azimuth_shift = [0]
    for sa in azimuth_shift:
        refa = np.roll(img, sa, axis=0)
        for sr in range_shift:
            refr = np.roll(refa, sr, axis=1)
            count += img - refr < tr1
    count[:, 0:nn] = wsize ** 2
    count[:, -nn:] = wsize ** 2
    if cartesian:
        count[0:nn, :] = wsize ** 2
        count[-nn:, :] = wsize ** 2
    return count

def filter_gabella_b(img, thrs=0.0):
    """Second part of the Gabella filter comparing area to circumference of \
    contiguous echo regions.

    Parameters
    ----------
    img : :py:class:`numpy:numpy.ndarray`
    thrs : float
        Threshold below which the field values will be considered as no rain

    Returns
    -------
    output : :py:class:`numpy:numpy.ndarray`
        contains in each pixel the ratio between area and circumference of the
        meteorological echo it is assigned to or 0 for non precipitation
        pixels.

    See Also
    --------
    :func:`~wradlib.clutter.filter_gabella` - the complete filter

    :func:`~wradlib.clutter.filter_gabella_a` - the first part of the filter

    Examples
    --------

    See :ref:`/notebooks/classify/wradlib_clutter_gabella_example.ipynb`.

    """
    conn = np.ones((3, 3))
    # create binary image of the rainfall field
    binimg = img > thrs
    # label objects (individual rain cells, so to say)
    labelimg, nlabels = ndimage.label(binimg, conn)
    # erode the image, thus removing the 'boundary pixels'
    binimg_erode = ndimage.binary_erosion(binimg, structure=conn)
    # determine the size of each object
    labelhist, edges = np.histogram(
        labelimg, bins=nlabels + 1, range=(-0.5, labelimg.max() + 0.5)
    )
    # determine the size of the eroded objects
    erodelabelhist, edges = np.histogram(
        np.where(binimg_erode, labelimg, 0),
        bins=nlabels + 1,
        range=(-0.5, labelimg.max() + 0.5),
    )
    # the boundary is the difference between these two
    boundarypixels = labelhist - erodelabelhist
    # now get the ratio between object size and boundary
    ratio = labelhist.astype(np.float32) / boundarypixels
    # assign it back to the objects
    # first get the indices
    indices = np.digitize(labelimg.ravel(), edges) - 1
    # then produce a new field with the ratios in the right place
    result = ratio[indices.ravel()].reshape(img.shape)

    return result

def filter_gabella(
    img,
    wsize=5,
    thrsnorain=0.0,
    tr1=6.0,
    n_p=6,
    tr2=1.3,
    rm_nans=True,
    radial=False,
    cartesian=False,
):
    """Clutter identification filter developed by :cite:`Gabella2002`.

    This is a two-part identification algorithm using echo continuity and
    minimum echo area to distinguish between meteorological (rain) and non-
    meteorological echos (ground clutter etc.)

    Parameters
    ----------
    img : :py:class:`numpy:numpy.ndarray`
    wsize : int
        Size of the window surrounding the central pixel
    thrsnorain : float
    tr1 : float
    n_p : int
    tr2 : float
    rm_nans : bool
        True replaces nans with Inf
        False takes nans into acount
    radial : bool
        True to use radial information only in
        :func:`~wradlib.clutter.filter_gabella_a`.
    cartesian : bool
        True if cartesian data are used, polar assumed if False.

    Returns
    -------
    output : :py:class:`numpy:numpy.ndarray`
        boolean array with pixels identified as clutter set to True.

    See Also
    --------
    :func:`~wradlib.clutter.filter_gabella_a` - the first part of the filter

    :func:`~wradlib.clutter.filter_gabella_b` - the second part of the filter

    Examples
    --------

    See :ref:`/notebooks/classify/wradlib_clutter_gabella_example.ipynb`.

    """
    bad = np.isnan(img)
    if rm_nans:
        img = img.copy()
        img[bad] = np.Inf
    ntr1 = filter_gabella_a(img, wsize, tr1, cartesian, radial)
    if not rm_nans:
        f_good = ndimage.filters.uniform_filter((~bad).astype(float), size=wsize)
        f_good[f_good == 0] = 1e-10
        ntr1 = ntr1 / f_good
        ntr1[bad] = n_p
    clutter1 = ntr1 < n_p
    ratio = filter_gabella_b(img, thrsnorain)
    clutter2 = np.abs(ratio) < tr2
    return clutter1 | clutter2

def histo_cut(prec_accum):
    """Histogram based clutter identification.

    This identification algorithm uses the histogram of temporal accumulated
    rainfall. It iteratively detects classes whose frequency falls below a
    specified percentage (1% by default) of the frequency of the class with the
    biggest frequency and remove the values from the dataset until the changes
    from iteration to iteration falls below a threshold. This algorithm is able
    to detect static clutter as well as shadings. It is suggested to choose a
    representative time periode for the input precipitation accumulation. The
    recommended time period should cover one year.

    Parameters
    ----------
    prec_accum : :py:class:`numpy:numpy.ndarray`
        spatial array containing rain accumulation

    Returns
    -------
    output : :py:class:`numpy:numpy.ndarray`
        boolean array with pixels identified as clutter/shadings set to True.

    Examples
    --------

    See :ref:`/notebooks/classify/wradlib_histo_cut_example.ipynb`.

    """

    prec_accum = np.array(prec_accum)

    # initialization of data bounds for clutter and shade definition
    lower_bound = 0
    upper_bound = prec_accum.max()

    # predefinitions for the first iteration
    lower_bound_before = -51
    upper_bound_before = -51

    # iterate as long as the difference between current and
    # last iteration doesn't fall below the stop criterion
    while (abs(lower_bound - lower_bound_before) > 1) or (
        abs(upper_bound - upper_bound_before) > 1
    ):

        # masks for bins with sums over/under the data bounds
        upper_mask = (prec_accum <= upper_bound).astype(int)
        lower_mask = (prec_accum >= lower_bound).astype(int)
        # NaNs in place of masked bins
        # Kopie der Datenmatrix mit Nans an Stellen,
        # wo der Threshold erreicht wird
        prec_accum_masked = np.where((upper_mask * lower_mask) == 0, np.nan, prec_accum)

        # generate a histogram of the valid bins with 50 classes
        (n, bins) = np.histogram(
            prec_accum_masked[np.isfinite(prec_accum_masked)].ravel(), bins=50
        )
        # get the class with biggest occurence
        index = np.where(n == n.max())
        index = index[0][0]
        # separated stop criterion check in case one of the bounds
        # is already robust
        if abs(lower_bound - lower_bound_before) > 1:
            # get the index of the class which underscores the occurence of
            # the biggest class by 1%, beginning from the class with the
            # biggest occurence to the first class
            for i in range(index, -1, -1):
                if n[i] < (n[index] * 0.01):
                    break
        if abs(upper_bound - upper_bound_before) > 1:
            # get the index of the class which underscores the occurence of
            # the biggest class by 1%, beginning from the class with the
            # biggest occurence to the last class
            for j in range(index, len(n)):
                if n[j] < (n[index] * 0.01):
                    break

        lower_bound_before = lower_bound
        upper_bound_before = upper_bound
        # update the new boundaries
        lower_bound = bins[i]
        upper_bound = bins[j + 1]

    return np.isnan(prec_accum_masked)

def classify_echo_fuzzy(dat, weights=None, trpz=None, thresh=0.5):
    """Fuzzy echo classification and clutter identification based on \
    polarimetric moments.

    The implementation is based on :cite:`Vulpiani2012`. At the
    moment, it only distinguishes between meteorological and non-meteorological
    echos.

    .. versionchanged:: 1.4.0
       The implementation was extended using depolarization ratio (dr)
       and clutter phase alignment (cpa).

    For Clutter Phase Alignment (CPA) see :cite:`Hubbert2009a` and
    :cite:`Hubbert2009b`

    For each decision variable and radar bin, the algorithm uses trapezoidal
    functions in order to define the membership to the non-meteorological
    echo class.
    Based on pre-defined weights, a linear combination of the different degrees
    of membership is computed. The echo is assumed to be non-meteorological
    in case the linear combination exceeds a threshold.

    At the moment, the following decision variables are considered:

        - Texture of differential reflectivity (zdr) (mandatory)

        - Texture of correlation coefficient (rho) (mandatory)

        - Texture of differential propagation phase (phidp) (mandatory)

        - Doppler velocity (dop) (mandatory)

        - Static clutter map (map) (mandatory)

        - Correlation coefficient (rho2) (additional)

        - Depolarization Ratio (dr), computed from
          correlation coefficient & differential reflectivity (additional)

        - clutter phase alignment (cpa) (additional)

    Parameters
    ----------
    dat : dict
        dictionary of arrays.
        Contains the data of the decision variables. The shapes of the arrays
        should be (..., number of beams, number of gates) and the shapes need
        to be identical or be broadcastable.
    weights : dict
        dictionary of floats.
        Defines the weights of the decision variables.
    trpz : dict
        dictionary of lists of floats.
        Contains the arguments of the trapezoidal membership functions for each
        decision variable
    thresh : float
       Threshold below which membership in non-meteorological membership class
       is assumed.

    Returns
    -------
    output : tuple
        a tuple of two boolean arrays of same shape as the input arrays
        The first array boolean array indicates non-meteorological echos based
        on the fuzzy classification.
        The second boolean array indicates where all the polarimetric moments
        had missing values which could be used as an additional information
        criterion.

    See Also
    --------
    :func:`~wradlib.dp.texture` - texture

    :func:`~wradlib.dp.depolarization` - depolarization ratio

    """
    # Check the inputs
    # mandatory data keys
    dkeys = ["zdr", "rho", "phi", "dop", "map"]
    # usable wkeys
    wkeys = ["zdr", "rho", "phi", "dop", "map", "rho2", "dr", "cpa"]
    # usable tkeys
    tkeys = ["zdr", "rho", "phi", "dop", "map", "rho2", "dr", "cpa"]

    # default weights
    weights_default = {
        "zdr": 0.4,
        "rho": 0.4,
        "phi": 0.1,
        "dop": 0.1,
        "map": 0.5,
        "rho2": 0.4,
        "dr": 0.4,
        "cpa": 0.4,
    }
    if weights is None:
        weights = weights_default
    else:
        weights = dict(list(weights_default.items()) + list(weights.items()))

    # default trapezoidal membership functions
    trpz_default = {
        "zdr": [0.7, 1.0, 9999, 9999],
        "rho": [0.1, 0.15, 9999, 9999],
        "phi": [15, 20, 10000, 10000],
        "dop": [-0.2, -0.1, 0.1, 0.2],
        "map": [1, 1, 9999, 9999],
        "rho2": [-9999, -9999, 0.95, 0.98],
        "dr": [-20, -12, 9999, 9999],
        "cpa": [0.6, 0.9, 9999, 9999],
    }
    if trpz is None:
        trpz = trpz_default
    else:
        trpz = dict(list(trpz_default.items()) + list(trpz.items()))

    # check data conformity
    assert np.all(np.in1d(dkeys, list(dat.keys()))), (
        "Argument dat of classify_echo_fuzzy must be a dictionary "
        f"with mandatory keywords {*dkeys,}."
    )
    assert np.all(np.in1d(wkeys, list(weights.keys()))), (
        "Argument weights of classify_echo_fuzzy must be a dictionary "
        f"with keywords {*wkeys,}."
    )
    assert np.all(np.in1d(tkeys, list(trpz.keys()))), (
        "Argument trpz of classify_echo_fuzzy must be a dictionary "
        "with keywords {*tkeys,}."
    )

    # copy rho to rho2
    dat["rho2"] = dat["rho"].copy()

    shape = None
    for key in dkeys:
        if not dat[key] is None:
            if shape is None:
                shape = dat[key].shape
            else:
                assert dat[key].shape[-2:] == shape[-2:], (
                    "Arrays of the decision variables have inconsistent "
                    f"shapes: {dat[key].shape} vs. {shape}"
                )
        else:
            print(f"WARNING: Missing decision variable: {key}")

    # If all dual-pol moments are NaN, can we assume that and echo is
    # non-meteorological?
    # Successively identify those bins where all moments are NaN
    nmom = ["rho", "zdr", "phi", "dr", "cpa"]  # 'dop'
    nan_mask = np.isnan(dat["rho"])
    for mom in nmom[1:]:
        try:
            nan_mask &= np.isnan(dat[mom])
        except KeyError:
            pass

    # Replace missing data by NaN
    dummy = np.zeros(shape) * np.nan
    for key in dat.keys():
        if dat[key] is None:
            dat[key] = dummy

    # membership in meteorological class for each variable
    qres = {}
    for key in dat.keys():
        if key not in tkeys:
            continue
        if key in ["zdr", "rho", "phi"]:
            d = dp.texture(dat[key])
        else:
            d = dat[key]
        qres[key] = 1.0 - util.trapezoid(
            d, trpz[key][0], trpz[key][1], trpz[key][2], trpz[key][3]
        )

    # create weight arrays which are zero where the data is NaN
    # This way, each pixel "adapts" to the local data availability
    wres = {}
    for key in dat.keys():
        if key not in wkeys:
            continue
        wres[key] = _weight_array(qres[key], weights[key])

    # Membership in meteorological class after combining all variables
    qsum = []
    wsum = []
    for key in dat.keys():
        if key not in wkeys:
            continue
        # weighted sum, also removing NaN from data
        qsum.append(np.nan_to_num(qres[key]) * wres[key])
        wsum.append(wres[key])

    q = np.array(qsum).sum(axis=0) / np.array(wsum).sum(axis=0)

    # flag low quality
    return np.where(q < thresh, True, False), nan_mask

def _weight_array(data, weight):
    """
    Generates weight array where valid values have the weight value
    and NaNs have 0 weight value.
    """
    w_array = weight * np.ones(np.shape(data))
    w_array[np.isnan(data)] = 0.0
    return w_array

def filter_cloudtype(
    img,
    cloud,
    thrs=0,
    snow=False,
    low=False,
    cirrus=False,
    smoothing=None,
    grid="polar",
    scale=None,
):
    """Identification of non-meteorological echoes based on cloud type.

    Parameters
    ----------
    img : :py:class:`numpy:numpy.ndarray`
        radar image to which the filter is to be applied
    cloud : :py:class:`numpy:numpy.ndarray`
        image with collocated cloud value from MSG SAFNWC PGE02 product
    thrs : float
        Threshold above which to identify clutter
    snow : bool
        Switch to use PGE02 class "land/sea snow" for clutter identification
    low : bool
        Switch to use PGE02 class very low stratus, very low cumulus and
        low cumulus for clutter identification
    cirrus : bool
        Switch to use PGE02 class "very thin cirrus" and "fractional clouds"
        for clutter identification
    smoothing : float
        Size [m] of the smoothing window used to take into account various
        localisation errors (e.g. advection, parallax)
    grid : str
        "polar" or "cartesian"
    scale : float or tuple
        float or tuple of 2 floats
        range [m] scale for polar grid
        x[m] and y[m] scale for cartesian grid

    Returns
    -------
    output : :py:class:`numpy:numpy.ndarray`
        a boolean array containing TRUE where clutter has been identified.

    """
    noprecip = (cloud == 1) | (cloud == 2)
    if snow:
        noprecip = noprecip | (cloud == 3) | (cloud == 4)
    if low:
        noprecip = noprecip | (cloud == 5) | (cloud == 6) | (cloud == 7)
    if cirrus:
        noprecip = noprecip | (cloud == 14) | (cloud == 18)
    if smoothing is not None:
        myfilter = getattr(util, f"filter_window_{grid}")
        noprecip = myfilter(noprecip, smoothing, "minimum", scale)
    clutter = noprecip & (img > thrs)
    return clutter

def filter_window_distance(img, rscale, fsize=1500, tr1=7):
    """2d filter looking for large reflectivity gradients.

    This function counts for each bin in ``img`` the percentage of surrounding
    bins in a window of half size ``fsize`` which are not ``tr1`` smaller than
    the central bin. The window is defined using geometrical distance.

    Parameters
    ----------
    img : :py:class:`numpy:numpy.ndarray`
        2d polar data to which the filter is to be applied
    rscale : float
        range [m] scale of the polar grid
    fsize : int
        Half-size [m] of the square window surrounding the central pixel
    tr1 : float
        Threshold value

    Returns
    -------
    output : :py:class:`numpy:numpy.ndarray`
        an array with the same shape as ``img``, containing the
        filter's results.

    See Also
    --------
    :func:`~wradlib.clutter.filter_gabella_a` - Original version of the filter

    :func:`~wradlib.clutter.filter_gabella_b` - filter using a echo area
    """
    ascale = 2 * np.pi / img.shape[0]
    count = np.ones(img.shape, dtype=int)
    similar = np.zeros(img.shape, dtype=float)
    good = np.ones(img.shape, dtype=float)
    valid = ~np.isnan(img)
    hole = np.sum(~valid) > 0
    nr = int(round(fsize / rscale))
    range_shift = range(-nr, nr + 1)
    r = np.arange(img.shape[1]) * rscale + rscale / 2
    adist = r * ascale
    na = np.around(fsize / adist).astype(int)
    max_na = img.shape[0] / 10
    sa = 0
    while sa < max_na:
        imax = np.where(na >= sa)[0][-1] + 1
        refa1 = util.roll2d_polar(img, sa, axis=0)
        refa2 = util.roll2d_polar(img, -sa, axis=0)
        for sr in range_shift:
            refr1 = util.roll2d_polar(refa1, sr, axis=1)
            similar[:, 0:imax] += img[:, 0:imax] - refr1[:, 0:imax] < tr1
            if sa > 0:
                refr2 = util.roll2d_polar(refa2, sr, axis=1)
                similar[:, 0:imax] += img[:, 0:imax] - refr2[:, 0:imax] < tr1
        count[:, 0:imax] = 2 * sa + 1
        sa += 1
    similar[~valid] = np.nan
    count[~valid] = -1
    count[:, nr:-nr] = count[:, nr:-nr] * (2 * nr + 1)
    for i in range(0, nr):
        count[:, i] = count[:, i] * (nr + 1 + i)
        count[:, -i - 1] = count[:, -i - 1] * (nr + 1 + i)
    if hole:
        good = util.filter_window_polar(valid.astype(float), fsize, "uniform", rscale)
        count = count * good
        count[count == 0] = 1
    similar -= 1
    count -= 1
    similar = similar / count
    return similar

if __name__ == "__main__":
    print("wradlib: Calling module <clutter> as main...")

max_stars_repo_name: syedhamidali/wradlib

max_issues_repo_name: syedhamidali/wradlib

max_forks_repo_name: syedhamidali/wradlib


text:
import os
import glob
import torch
import random
import numpy as np
from torch.utils.data import Dataset, DataLoader
import torchaudio
from utils.utils import read_wav_np
from utils.stft import TacotronSTFT
from scipy.io.wavfile import read

def create_dataloader(hp, args, train):
    dataset = MelFromDisk(hp, args, train)

    '''
    if train:
        return DataLoader(dataset=dataset, batch_size=hp.train.batch_size, shuffle=True,
            num_workers=0, pin_memory=False, drop_last=True) # hp.train.num_workers
    else:
        return DataLoader(dataset=dataset, batch_size=1, shuffle=False,
            num_workers=0, pin_memory=False, drop_last=False)
    '''

    if train:
        return DataLoader(dataset=dataset, batch_size=hp.train.batch_size, shuffle=True,
            num_workers=hp.train.num_workers, pin_memory=True, drop_last=True)
    else:
        return DataLoader(dataset=dataset, batch_size=1, shuffle=False,
            num_workers=0, pin_memory=False, drop_last=False)

class MelFromDisk(Dataset):
    def __init__(self, hp, args, train):
        self.hp = hp
        self.args = args
        self.train = train
        self.path = hp.data.train if train else hp.data.validation

        with open(self.path, encoding='utf-8') as f:
            filepaths = [line.strip().split("|")[0] for line in f]

        self.wav_list = filepaths
        self.stft = TacotronSTFT(filter_length=hp.audio.filter_length,
                        hop_length=hp.audio.hop_length,
                        win_length=hp.audio.win_length,
                        n_mel_channels=hp.audio.n_mel_channels,
                        sampling_rate=hp.audio.sampling_rate,
                        mel_fmin=hp.audio.mel_fmin,
                        mel_fmax=hp.audio.mel_fmax)
        self.mel_segment_length = hp.audio.segment_length // hp.audio.hop_length + 2
        self.mapping = [i for i in range(len(self.wav_list))]

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

    def __getitem__(self, idx):

        if self.train:
            idx1 = idx
            idx2 = self.mapping[idx1]
            return self.my_getitem(idx1), self.my_getitem(idx2)
        else:
            return self.my_getitem(idx)

    def shuffle_mapping(self):
        random.shuffle(self.mapping)

    def load_wav_to_torch(self, full_path):
        sampling_rate, data = read(full_path)
        return torch.FloatTensor(data.astype(np.float32)), sampling_rate

    def my_getitem(self, idx):
        wavpath = self.wav_list[idx]

        audio, sampling_rate = self.load_wav_to_torch(wavpath)
        audio_norm = audio / 32768.0
        audio_norm = audio_norm.unsqueeze(0)

        if len(audio) < self.hp.audio.segment_length + self.hp.audio.pad_short:
            audio_norm = torch.nn.functional.pad(audio_norm, (0, self.hp.audio.segment_length + self.hp.audio.pad_short - audio.size(0)), 'constant').data

        '''
        audio_norm = torch.autograd.Variable(audio_norm, requires_grad=False)
        mel = self.stft.mel_spectrogram(audio_norm)
        mel = torch.squeeze(mel, 0)
        '''

        mel = torch.from_numpy(np.load(wavpath.replace('.wav', '.npy').replace('data_22050', 'mel_22050')))

        if self.train:
            max_mel_start = mel.size(1) - self.mel_segment_length
            mel_start = random.randint(0, max_mel_start)
            mel_end = mel_start + self.mel_segment_length
            mel = mel[:, mel_start:mel_end]

            audio_start = mel_start * self.hp.audio.hop_length
            audio_norm = audio_norm[:, audio_start:audio_start+self.hp.audio.segment_length]

        audio_norm = audio_norm + (1.0/32768.0) * torch.randn_like(audio_norm)
        return mel, audio_norm

def collate_fn(batch):

    sr = 22050
    # perform padding and conversion to tensor
    mels_g = [x[0][0] for x in batch]
    audio_g = [x[0][1] for x in batch]

    mels_g = torch.stack(mels_g)
    audio_g = torch.stack(audio_g)

    sub_orig_1 = torchaudio.transforms.Resample(sr, (sr // 2))(audio_g)
    sub_orig_2 = torchaudio.transforms.Resample(sr, (sr // 4))(audio_g)
    sub_orig_3 = torchaudio.transforms.Resample(sr, (sr // 8))(audio_g)
    sub_orig_4 = torchaudio.transforms.Resample(sr, (sr // 16))(audio_g)

    mels_d = [x[1][0] for x in batch]
    audio_d = [x[1][1] for x in batch]
    mels_d = torch.stack(mels_d)
    audio_d = torch.stack(audio_d)
    sub_orig_1_d = torchaudio.transforms.Resample(sr, (sr // 2))(audio_d)
    sub_orig_2_d = torchaudio.transforms.Resample(sr, (sr // 4))(audio_d)
    sub_orig_3_d = torchaudio.transforms.Resample(sr, (sr // 8))(audio_d)
    sub_orig_4_d = torchaudio.transforms.Resample(sr, (sr // 16))(audio_d)

    return [mels_g, audio_g, sub_orig_1, sub_orig_2, sub_orig_3, sub_orig_4],\
           [mels_d, audio_d, sub_orig_1_d, sub_orig_2_d, sub_orig_3_d, sub_orig_4_d]


text:
import cv2
import os
import shutil
import numpy as np

# size filter
def size_filter(root, size=(512, 512), mode='move'):
    paths = os.listdir(root)
    length = len(paths)
    for i, path in enumerate(paths, start=1):
        img_path = os.path.join(root, path)
        img = cv2.imread(img_path)
        if img is None:
            continue
        h, w, c = img.shape
        if h < size[0] or w < size[1]:
            if mode == 'move':
                shutil.move(img_path, 'path/to/small_img_dir')
            elif mode == 'remove':
                os.remove(img_path)
        if i % 1000 == 0:
            print(f'{i}/{length} images ended')

# judge image color or black-white
def image_filter(root, mode='move', with_size_filter=False, size=(512, 512), threshold=0.1):
    paths = os.listdir(root)
    length = len(paths)
    for i, path in enumerate(paths, start=1):
        img_path = os.path.join(root, path)
        img = cv2.imread(img_path)
        if img is None:
            continue
        h, w, c = img.shape
        if with_size_filter:
            if h < size[0] or w < size[1]:
                if mode == 'move':
                    shutil.move(img_path, 'path/to/small_img_dir')
                elif mode == 'remove':
                    os.remove(img_path)
                continue
        b, g, r = cv2.split(img)

        r_g = np.count_nonzero(abs(r - g))
        r_b = np.count_nonzero(abs(r - b))
        g_b = np.count_nonzero(abs(g - b))
        diff_sum = float(r_g + r_b + g_b)

        ratio = diff_sum / img.size
        # ----------------------------------------------
        # FIXME
        # please change threshold to preferred value.
        # i use 0.1 because i want brilliant illustration.
        # i fill accuracy is about 80%.
        # ----------------------------------------------
        # if color
        if ratio > threshold:
            continue
        # if grey
        else:
            if mode == 'move':
                shutil.move(img_path, 'path/to/grey_img_dir')
            elif mode == 'remove':
                os.remove(img_path)
        if i % 10000 == 0:
            print(f'{i} / {length} ended')

if __name__ == '__main__':
    # FIXME fix root
    image_filter(root='path/to/ColorArtFolder',
                 mode='move',
                 with_size_filter=True,
                 threshold=0.1)


text:
import os
import cv2
import numpy as np
import json
from PIL import Image
import pickle
import matplotlib.pyplot as plt
import matplotlib.patches as patches

# set the path to the downloaded data:
data_path = '../RedLights2011_Medium'

# set a path for saving predictions:
preds_path = './hw01_preds'
os.makedirs(preds_path,exist_ok=True) # create directory if needed

# get sorted list of files:
file_names = sorted(os.listdir(data_path))

# remove any non-JPEG files:
file_names = [f for f in file_names if '.jpg' in f]

with open('filter.pkl', 'rb') as f:
    filt = pickle.load(f)

def read_im(fname, data_path=data_path):
    # read image using PIL:
    I = Image.open(os.path.join(data_path,fname))
    # convert to numpy array:
    return np.asarray(I)

def get_circles(file, filepath=data_path):
    '''
    This function returns a list of centers and radii
    for red circles in the image given by file
    '''
    font = cv2.FONT_HERSHEY_SIMPLEX
    img = cv2.imread(os.path.join(filepath,file))
    cimg = img
    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)

    # ranges we'll consider red
    lower_red1 = np.array([0,100,100])
    upper_red1 = np.array([10,255,255])
    lower_red2 = np.array([160,100,100])
    upper_red2 = np.array([180,255,255])
    mask1 = cv2.inRange(hsv, lower_red1, upper_red1)
    mask2 = cv2.inRange(hsv, lower_red2, upper_red2)
    maskr = cv2.add(mask1, mask2)

    m,n,_ = img.shape

    # hough circle detect
    r_circles = cv2.HoughCircles(maskr, cv2.HOUGH_GRADIENT, 1, 80,
                               param1=50, param2=10, minRadius=0, maxRadius=30)

    # needs to be in top bound proportion of image, since lights are high
    # usually
    bound = 4.0 / 10

    centers = []
    radii = []

    if r_circles is not None:
        r_circles = np.uint16(np.around(r_circles))

        for (x,y,r) in r_circles[0, :]:
            # Ensure circle we're finding is in the picture
            if x > n or y > m or y > m*bound:
                continue
            centers.append((x,y))
            radii.append(r)

    return centers, radii

def detect_red_light(fname, filt=filt, thresh = 0.85, filepath=data_path):
    '''
    This function takes a filename <fname> and returns a list <bounding_boxes>.
    The list <bounding_boxes> should have one element for each red light in the
    image. Each element of <bounding_boxes> should itself be a list, containing
    four integers that specify a bounding box: the row and column index of the
    top left corner and the row and column index of the bottom right corner (in
    that order). See the code below for an example.

    Note that PIL loads images in RGB order, so:
    I[:,:,0] is the red channel
    I[:,:,1] is the green channel
    I[:,:,2] is the blue channel
    '''

    bboxes = []
    centers, radii = get_circles(fname, filepath)
    I = read_im(fname, filepath)
    m, n, _ = I.shape

    for (x,y), r in zip(centers, radii):
        # [1, 2] is aspect ratio, 3 is scaling since width of
        # traffic light is about 3 times the radius of the
        # actual red circle
        dims = np.array([1, 2]) * r * 3
        anchor = (x - (dims[0] // 2), y - dims[1] // 4)
        bbox = [anchor[0], anchor[1], anchor[0] + dims[0], anchor[1] + dims[1]]
        bbox = [max(0, bbox[0]), max(0, bbox[1]),
                min(n, bbox[2]), min(m, bbox[3])]

        I_spliced = I[bbox[1]:bbox[3], bbox[0]:bbox[2], :]
        pil = Image.fromarray(I_spliced)
        I2 = np.asarray(pil.resize((filt.shape[1], filt.shape[0])))
        I3 = I2 / np.linalg.norm(I2)
        if (np.sum(I3 * filt) > thresh):
            bboxes.append(bbox)

    for i in range(len(bboxes)):
        assert len(bboxes[i]) == 4

    return bboxes

preds = {}
for i in range(len(file_names)):
    preds[file_names[i]] = [[int(n) for n in val] \
                            for val in detect_red_light(file_names[i])]

# save preds (overwrites any previous predictions!)
with open(os.path.join(preds_path,'preds.json'),'w') as f:
    json.dump(preds,f)


text:
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Oct 27 17:53:29 2020

@author: Isaac Ayala

See LICENSE for details.

Sample code to simulate a simple pendulum in Python

--------------
Explanation
--------------

1) Import libraries needed for the simulation.
2) Define custom functions.
3) Execute main script.

"""

from matplotlib import pyplot as plt
from  matplotlib import animation as animation
import numpy as np
from scipy import integrate

from celluloid import Camera

def system_dynamics(t, x, params,):
    """
    Parameters
    ----------
    x0 : State vector
    t : Current time step
    params : Simulation parameters

    Returns
    -------
    dx : State vector dynamics for time step integration

    """

    # Extract state variables and parameters
    # Python starts counting with 0 like any sane programming language!
    x1, x2, = x

    # Params is a dictionary with key and value pairs
    m = params['mass']
    g = params['gravity']
    l = params['length']
    k = params['friction']

    # Solve system dynamics for the current time step
    dot_x1 = x2
    dot_x2 = - (g/l) * np.sin(x1) - (k/m) * x2

    # Store each state variable into array
    dx = np.array([dot_x1, dot_x2])

    return dx

def plot_solution(sol, params):
    # Set figure size in inches
    fig = plt.figure(figsize=(12, 6))

    ax = fig.add_subplot(1, 2, 1)
    ax.set_xlabel(r'$x$')
    ax.set_ylabel(r'$y$')

    linear_plot = fig.add_subplot(2, 2, 2)
    linear_plot.set_xlabel(r'$t$')
    linear_plot.set_ylabel(r'$\chi$')

    angular_plot = fig.add_subplot(2, 2, 4)
    angular_plot.set_xlabel(r'$t$')
    angular_plot.set_ylabel(r'$\xi$')

    l = params['length']
    dt = params['step']
    x = l * np.sin(sol.y[0])
    y = -l * np.cos(sol.y[0])

    ax.plot(x, y, '-', label = r'Linear')
    ax.set_xlim([-1.1*l, 1.1*l])
    ax.set_ylim([-1.1*l, 1.1*l])
    ax.legend()

    linear_plot.plot(sol.t, x, '-', label = r'$x$')
    linear_plot.plot(sol.t, y, '-', label = r'$y$')
    linear_plot.legend()

    angular_plot.plot(sol.t, sol.y[0], '-', label = r'$\theta$')
    angular_plot.plot(sol.t, sol.y[1], '-', label = r'$\dot \theta$')
    angular_plot.legend()

    fig.tight_layout()

    # Animation

    fig2 = plt.figure()
    anim = fig2.add_subplot(111, autoscale_on=False, xlim=(-1.1*l, 1.1*l), ylim=(-1.1*l, 1.1*l))
    anim.grid()

    line, = anim.plot([], [], 'o-', lw=2)
    time_template = 'time = %.1fs'
    time_text = anim.text(0.05, 0.9, '', transform=anim.transAxes)

    def init():
        line.set_data([], [])
        time_text.set_text('')
        return line, time_text

    def animate(i):
        thisx = [0, x[i]]
        thisy = [0, y[i]]

        line.set_data(thisx, thisy)
        time_text.set_text(time_template % (i*dt))
        return line, time_text

    ani = animation.FuncAnimation(fig2, animate, np.arange(1, len(sol.t)),
                              interval=25, blit=True, init_func=init)

    plt.show()

if __name__ == '__main__':
    # Close figures
    plt.close("all")

    # Simulation parameters
    dt = 0.01
    max_time = 10
    t_span = [0, max_time]
    t = np.arange(0, max_time, dt)

    params = dict(
                    gravity = 9.81, # m/s^2
                    mass = 0.3, # kg
                    length = 0.15, # m
                    friction = 0.35,
                    step = dt,
                    )

    # Initial vector state
    x0 = np.array([np.pi/2, 0])

    sol = integrate.solve_ivp(
                                system_dynamics,
                                t_span,
                                x0,
                                method='RK45',
                                t_eval = t,
                                args = (params, )
                                )
    plot_solution(sol, params)


text:
import os
import rnnSMAP
from rnnSMAP import runTrainLSTM
import matplotlib.pyplot as plt
import numpy as np
import matplotlib

import imp
imp.reload(rnnSMAP)
rnnSMAP.reload()

#################################################
# intervals temporal test
doOpt = []
doOpt.append('test')
doOpt.append('plotMap')
doOpt.append('plotVS')

trainName = 'CONUSv2f1'
testName = 'CONUSv2f1'
strSigmaLst = ['sigma']
strErrLst = ['ubRMSE', 'RMSE']
saveFolder = os.path.join(
    rnnSMAP.kPath['dirResult'], 'Sigma', 'paper')
matplotlib.rcParams.update({'font.size': 16})
matplotlib.rcParams.update({'lines.linewidth': 2})
matplotlib.rcParams.update({'lines.markersize': 10})

#################################################
if 'test' in doOpt:
    rootOut = rnnSMAP.kPath['OutSigma_L3_NA']
    rootDB = rnnSMAP.kPath['DB_L3_NA']

    out = trainName+'_y15_Forcing'
    ds = rnnSMAP.classDB.DatasetPost(
        rootDB=rootDB, subsetName=testName, yrLst=[2016,2017])
    ds.readData(var='SMAP_AM', field='SMAP')
    ds.readPred(rootOut=rootOut, out=out, drMC=100, field='LSTM')
    statErr = ds.statCalError(predField='LSTM', targetField='SMAP')
    statSigma = ds.statCalSigma(field='LSTM')

#################################################
if 'plotMap' in doOpt:
    cRangeErr = [0, 0.08]
    for s in strErrLst:
        grid = ds.data2grid(data=getattr(statErr, s))
        saveFile = os.path.join(saveFolder, 'map_'+testName+'_'+s)
        titleStr = s
        fig = rnnSMAP.funPost.plotMap(
            grid, crd=ds.crdGrid, cRange=cRangeErr, title=titleStr, showFig=False)
        fig.savefig(saveFile)
    for s in strSigmaLst:
        grid = ds.data2grid(data=getattr(statSigma, s))
        saveFile = os.path.join(saveFolder, 'map_'+testName+'_'+s)
        titleStr = s
        if s == 'sigmaMC':
            cRangeSigma = [0, 0.04]
        else:
            cRangeSigma = [0, 0.08]
        fig = rnnSMAP.funPost.plotMap(
            grid, crd=ds.crdGrid, cRange=cRangeSigma, title=titleStr, showFig=False)
        fig.savefig(saveFile, dpi=1000)

#################################################
if 'plotVS' in doOpt:
    for iS in range(0, len(strSigmaLst)):
        for iE in range(0, len(strErrLst)):
            strE = strErrLst[iE]
            strS = strSigmaLst[iS]
            y = getattr(statErr, strE)
            x = getattr(statSigma, strS)
            # ub = np.percentile(y, 95)
            # lb = np.percentile(y, 5)
            # ind = np.logical_and(y >= lb, y <= ub)
            # x = x[ind]
            # y = y[ind]
            # ax = axes[iE]
            fig, ax = plt.subplots(figsize=(8, 6))
            rnnSMAP.funPost.plotVS(x, y, ax=ax, doRank=False)
            # rnnSMAP.funPost.plot121Line(ax)
            ax.set_ylabel(strE)
            ax.set_xlabel(strS)
            fig.suptitle(strS+' vs '+strE)
            saveFile = os.path.join(saveFolder, strS+'2'+strE+'_'+testName)
            fig.savefig(saveFile, dpi=1000)
            plt.close(fig)

max_stars_repo_name: fkwai/geolearn

max_issues_repo_name: fkwai/geolearn

max_forks_repo_name: fkwai/geolearn


text:
from skimage import io, color, measure, draw, img_as_bool
import numpy as np
from scipy import optimize
import matplotlib.pyplot as plt
from joblib import Parallel, delayed
import cv2
from matplotlib.patches import Ellipse, Circle, Rectangle
from scipy.linalg.basic import LinAlgError

from circle_fitting import *

def process_image(image, image_name):
    N = image['scores'].size
    result = []
    for i in range(N):
        Nroi = image['rois'][i].size
        s = list(image["rois"][i])
        narrow = np.zeros((1, 1))

        # Check the shape compared to the input image
        if s[2]-s[0] == 1:
            s[2] + =1
        if s[2] > image['masks'].shape[0]:
            s[2] -= 1
            s[0] -= 1
        if s[3]-s[1] == 1:
            s[3] += 1
        if s[3] > image['masks'].shape[1]:
            s[3] -= 1
            s[1] -= 1
        aspect = (s[3]-s[1])/(s[2]-s[0])
        athreshold = 3

        # If it's not close to a circle, ignore it
        if aspect > athreshold or aspect < 1/athreshold:
            continue

        shp = (s[3]-s[1], s[2]-s[0])
        area_BB = shp[0]*shp[1]
        res = [0.5*(s[3]+s[1]),
               0.5*(s[0]+s[2]),
               shp[0],
               shp[1],
               image['scores'][i],
               aspect]

        # Fit the centroid 'mass' weighted circle
        try:
            cx, cy, cd = centroid(image['masks'][..., i], s)
            area_centroid = np.pi*(cd/2)**2
            res.extend([cx, cy, cd, area_centroid / area_BB])
        except FitError as e:
            res.extend([None]*4)

        # Fit the least square error circle
        try:
            cx, cy, cd = proc_circle(image['masks'][..., i], s)
            area_circle = np.pi*(cd/2)**2
            res.extend([cx, cy, cd, area_centroid / area_BB])
        except FitError as e:
            res.extend([None]*4)

        # Fit the least square error direct fit ellipse
        try:
            cx, cy, cdx, cdy, cphi = proc_ellipse(image['masks'][..., i], s)
            area_circle = np.pi * (cdx / 2) * (cdy / 2)
            res.extend([cx, cy, cdx, cdy, cphi, area_centroid / area_BB])
        except FitError as e:
            res.extend([None]*6)

        result.append(res)

    columns = ['bbCx',  # Bounding box centre X
               'bbCy',   # Bounding box centre Y
               'bbNx',  # Bounding box width
               'bbNy',  # Bounding box height
               "Score",  # certainty of identification
               "Aspect",  # aspect ratio of the box
               'x_centroid',  # x location of the mass centroid
               'y_centroid',  # y location of the mass centroid
               'diameter_centroid',
               'fraction_centroid',  # centroid area over BBox area
               'x_circle',  # LMSE circle X
               'y_circle',  # LMSE circle Y
               'diameter_circle',  # LMSE circle diameter
               'fraction_circle',  # LMSE circle / BBox area
               'x_ellipse',  # direct ellipse X
               'y_ellipse',  # direct ellipse X
               'width_ellipse',  # direct ellipse width
               'height_ellipse',  # direct ellipse height
               'angle_ellipse',  # direct ellipse angle from horizontal
               'fraction_ellipse'  # ellipse / BBox area
               ]
    df = pd.DataFrame(result, columns=columns)
    df["image_name"] = image_name
    return df

max_stars_repo_name: eelsirhc/DeepMars

max_issues_repo_name: eelsirhc/DeepMars

max_forks_repo_name: eelsirhc/DeepMars