text:
# by Saeed Shaker

# This is a direct translation of Yuval Tassa's Matlab code into Python:
# https://benjaminmoll.com/wp-content/uploads/2020/06/LCP.m
# It solves LCP using a Newton type method
# To be consistent across platforms and with Yuval Tassa's code,
# I have tried to make as minimal changes as I could,
# so this code can be followed the same way as the original Matlab code does.

import numpy as np
import scipy.sparse as sparse
from scipy.sparse.linalg import spsolve

def LCP_python(M,q,l=[],u=[],x0=[],display=False):
    tol            = 1.0e-6;
    mu             = 1e-3;
    mu_step        = 5;
    mu_min         = 1e-5;
    max_iter       = 25;
    b_tol          = 1e-6;

    n              = M.shape[0]

    if l == []:
        l = np.zeros((n,1))
    if u == []:
        u = np.ones((n,1))*np.inf
    if x0 == []:
        x00 = np.maximum(np.ones((n,1)),l)
        x0 = np.minimum(x00,u)

    M = sparse.csc_matrix(M)
    q = q.reshape((-1, 1))
    l = l.reshape((-1, 1))
    u = u.reshape((-1, 1))
    x0 = x0.reshape((-1, 1))

    lu             = np.column_stack((l , u));
    x              = x0.copy();
    psi,phi,J    = FB(x,q,M,l,u);
    new_x          = True

    for iter1 in range(0,max_iter):
        if new_x:
            mlu      = np.min(np.column_stack((np.abs(x-l),np.abs(u-x))),1).reshape((-1, 1));
            ilu      = np.argmin(np.column_stack((np.abs(x-l),np.abs(u-x))),1).reshape((-1, 1));
            bad            = np.maximum(np.abs(phi),mlu) < b_tol;
            psi            = psi - 0.5*np.dot(phi[bad] , phi[bad])
            notbad = bad == False
            Jind = np.dot(notbad , notbad.T)
            notbad_trues = np.sum(notbad*1)
            J              = sparse.csc_matrix(np.reshape(J[Jind] , (notbad_trues,notbad_trues) ))
            phi            = phi[notbad];
            new_x          = False;
            nx             = x.copy();
            nx[bad]        = lu.flatten()[(bad[bad])*1+(ilu[bad]-1)*n]
        H              = np.dot(J.T , J) + mu*sparse.eye(notbad_trues);
        Jphi           = sparse.csc_matrix.dot(J.T,phi)
        d              = -spsolve(sparse.csc_matrix(H) , Jphi)
        nx[notbad]       = x[notbad] + d;
        npsi,nphi,nJ = FB(nx,q,M,l,u);

        r   = (psi - npsi)/ -(np.dot(Jphi.T,d) + 0.5*np.dot(sparse.csc_matrix.dot(d.T,H),d) );  # actual reduction / expected reduction

        if r < 0.3:
            mu = np.maximum(mu*mu_step,mu_min);
        if r > 0:
            x     = nx.copy();
            psi   = npsi.copy();
            phi   = nphi.copy();
            J     = nJ.copy();
            new_x = True;
            if r > 0.8:
                mu = mu/mu_step * (mu > mu_min);
        if display:
            print('iter = ', iter1 , ' --- psi = ' , psi ,' --- r = ' , r ,' --- mu = ' , mu);
        if psi < tol:
            break;
    x = np.minimum(np.maximum(x,l),u);
    return x

#----------------------------------------------------------

def FB(x,q,M,l,u):
    n     = x.size;
    Zl    = ((l >-np.inf) & (u==np.inf))
    Zu    = (l==-np.inf) & (u <np.inf);
    Zlu   = (l >-np.inf) & (u <np.inf);
    Zf    = (l==-np.inf) & (u==np.inf);

    a     = x.copy();
    b     = sparse.csc_matrix.dot(M,x)+q;
    a[Zl] = x[Zl]-l[Zl];
    a[Zu] = u[Zu]-x[Zu];
    b[Zu] = -b[Zu];

    if any(Zlu):
        nt     = np.sum(Zlu);
        at     = u[Zlu]-x[Zlu];
        bt     = -b[Zlu];
        st     = np.sqrt(np.power(at,2) + np.power(bt,2));
        a[Zlu] = x[Zlu]-l[Zlu];
        b[Zlu] = st -at -bt;

    s        = np.sqrt(np.power(a,2) + np.power(b,2));
    phi      = s - a - b;
    phi[Zu]  = -phi[Zu];
    phi[Zf]  = -b[Zf];

    psi      = 0.5*np.dot(phi.T , phi);

    if any(Zlu):
        M[Zlu,:] = -sparse.csc_matrix((at/st-np.ones((nt,1)),(np.arange(nt),Zlu[Zlu != 0])),nt,n , dtype=np.float) - sparse.csc_matrix.dot(sparse.csc_matrix((bt/st-np.ones((nt,1)),(np.arange(nt),np.arange(nt))) , dtype=np.float), M[Zlu,:]);

    da       = (a/s-np.ones((n,1))).reshape((-1, 1));
    db       = (b/s-np.ones((n,1))).reshape((-1, 1));
    da[Zf]   = 0;
    db[Zf]   = -1;
    J        =  sparse.csc_matrix((np.array(da[:,0]),(np.arange(n),np.arange(n))), dtype=np.float) + sparse.csc_matrix.dot(sparse.csc_matrix((np.array(db[:,0]),(np.arange(n),np.arange(n))) , dtype=np.float) , M);

    return psi,phi,J

max_stars_repo_name: madsankern/DynamicProgramming

max_issues_repo_name: madsankern/DynamicProgramming

max_forks_repo_name: madsankern/DynamicProgramming


text:
import unittest
import numpy as np
from scipy.special import logsumexp
from scipy.special import binom as binom_coef
from scipy.stats import binom as binom_dist
from pyascore import PyLogMath, PyBinomialDist, PyPowerSetSum

class TestLogMath(unittest.TestCase):
    lmath = PyLogMath()

    def test_log_sum(self):
        summands = [(-np.inf, 0),
                    (0, -np.inf)]

        n_to_add = 100
        np.random.seed(2345)
        for a, b in zip(np.random.randn(n_to_add), np.random.randn(n_to_add)):
            summands.append((a, b))

        for s in summands:
            self.assertTrue(np.isclose(self.lmath.log_sum(*s),
                                       logsumexp(s),
                                       rtol=0, atol=1e-6))

    def test_log_bin_coef(self):
        MAX_N = 50

        for n in range(1, MAX_N + 1):
            for k in range(1, n + 1):
                #print(k, n, self.lmath.log_bin_coef(k, n), np.log(binom_coef(n, k)))
                self.assertTrue(np.isclose(self.lmath.log_bin_coef(k, n),
                                           np.log(binom_coef(n, k)),
                                           rtol=0, atol=5e-5))

class TestBinomDist(unittest.TestCase):
    def test_log_pdf(self):
        MAX_N = 50

        probs = [.1, .25, .5, .75, .9]
        dists = [PyBinomialDist(p) for p in probs]

        for p, d in zip(probs, dists):
            for n in range(MAX_N):
                for k in range(1, n + 1):
                    #print(k, n, d.log_pmf(k, n), binom_dist.logpmf(k, n, p))
                    self.assertTrue(np.isclose(d.log_pmf(k, n),
                                               binom_dist.logpmf(k, n, p),
                                               rtol=0, atol=5e-5))
    def test_log_pvalue(self):
        MAX_N = 50

        probs = [.1, .25, .5, .75, .9]
        dists = [PyBinomialDist(p) for p in probs]

        for p, d in zip(probs, dists):
            for n in range(MAX_N):
                for k in range(1, n + 1):
                    scipy_pvalue = logsumexp([binom_dist.logpmf(k, n, p),
                                              binom_dist.logsf(k, n, p)])
                    #print(k, n, d.log_pvalue(k, n), scipy_pvalue)
                    self.assertTrue(np.isclose(d.log_pvalue(k, n),
                                               scipy_pvalue,
                                               rtol=0, atol=5e-5))
                    self.assertTrue(np.isclose(d.log10_pvalue(k, n),
                                               np.log10(np.exp(scipy_pvalue)),
                                               rtol=0, atol=5e-5))

class TestPowerSetSum(unittest.TestCase):
    def test_empty_set(self):
        pss = PyPowerSetSum()

        self.assertFalse( pss.has_next() )
        self.assertTrue( pss.get_sum() == 0.)

    def test_simple_sums(self):
        target = np.array([1., 2., 3.], dtype=np.float32)
        pss = PyPowerSetSum(target, 2)

        self.assertTrue( pss.has_next() )
        self.assertTrue( pss.get_sum() == 0.)

        true_sums = iter([1., 2., 3., 4., 5.])
        while pss.has_next():
            pss.next()
            self.assertTrue( pss.get_sum() == next(true_sums) )

        new_target = np.array([4., 5., 6.], dtype=np.float32)
        pss.reset(new_target, 2)

        self.assertTrue( pss.has_next() )
        self.assertTrue( pss.get_sum() == 0.)

        true_sums = iter([4., 5., 6., 9., 10, 11.])
        while pss.has_next():
            pss.next()
            self.assertTrue( pss.get_sum() == next(true_sums) )

ext: py

max_stars_repo_name: RalfG/pyAscore

max_issues_repo_name: RalfG/pyAscore

max_forks_repo_name: RalfG/pyAscore

reason: import numpy,from scipy


text:
"""
This code implements the streaming version of scikit learn's IPCA algorithm -
https://github.com/scikit-learn/scikit-learn/blob/7813f7efb/sklearn/decomposition/incremental_pca.py
"""

import numpy as np
import random
import pickle as p
from scipy import linalg
from sklearn.utils.extmath import svd_flip

import sys
sys.path.append("../")
from IoTPy.core.stream import Stream, StreamArray, run
from IoTPy.agent_types.op import map_window
from IoTPy.agent_types.source import source_list, source_list_to_stream
from IoTPy.agent_types.sink import sink_window
from IoTPy.helper_functions.recent_values import recent_values
from IoTPy.helper_functions.print_stream import print_stream
from IoTPy.concurrency.multicore import get_processes_and_procs
from IoTPy.concurrency.multicore import terminate_stream, extend_stream
from examples.PCA.IPCAutils import _incremental_mean_and_var

source_list = [[1,5,0.1],[3,1,-0.1],[3,7,0.2],[-4,4,-0.3],
               [-1,-5,0.1],[2,0,0.2],[3,3,-0.2],[-4,8,0],
               [-7,-2,0.1],'eof']
#source_list = np.array(source_list)

def source(out_stream):
    return source_list_to_stream(source_list, out_stream)

## def make_and_run_process(compute_func):
##     proc = shared_memory_process(
##         compute_func=compute_func,
##         in_stream_names=['in'],
##         out_stream_names=[],
##         connect_sources=[('in', source)],
##         connect_actuators=[],
##         name='proc')
##     mp = Multiprocess(processes=[proc], connections=[])
##     mp.run()

def IPCA(k = 5,batches = 25):
    class IPCAValues(object):
        def __init__(self,k):
                self.n_components_ = k
                self.n_samples_seen_ = 0
                self.components_ = None
                self.mean_ = .0
                self.var_ = .0
                self.singular_values_ = None
                self.explained_variance_ = None
                self.explained_variance_ratio_ = None
                self.singular_values_ = None
                self.noise_variance_ = None
                self.count = 0

    state = IPCAValues(k)

    def gh(in_stream):

        def partial_fit(X, state):
            """Incremental fit with X. All of X is processed as a single batch.
            Parameters
            ----------
            X : array-like, shape (n_samples, n_features)
                Training data, where n_samples is the number of samples and
                n_features is the number of features.

            Returns
            -------
            self : object
                Returns the instance itself.
            """

            end = False
            print(X[0])
            if X[-1] == 'eof':
                end = True
                X= X[:-1]

            X = np.array(X,dtype = 'float64')
            n_samples, n_features = X.shape
            state.count += n_samples

            # Update stats - they are 0 if this is the fisrt step

            col_mean, col_var, n_total_samples = \
                _incremental_mean_and_var(
                    X, last_mean=state.mean_, last_variance=state.var_,
                    last_sample_count=np.repeat(state.n_samples_seen_, X.shape[1]))
            n_total_samples = n_total_samples[0]

            # Whitening
            if state.n_samples_seen_ == 0:
                # If it is the first step, simply whiten X
                X -= col_mean
            else:
                col_batch_mean = np.mean(X, axis=0)
                X -= col_batch_mean
                # Build matrix of combined previous basis and new data
                mean_correction = \
                    np.sqrt((state.n_samples_seen_ * n_samples) /
                            n_total_samples) * (state.mean_ - col_batch_mean)
                X = np.vstack((state.singular_values_.reshape((-1, 1)) *
                              state.components_, X, mean_correction))

            U, S, V = linalg.svd(X, full_matrices=False)
            U, V = svd_flip(U, V, u_based_decision=False)
            explained_variance = S ** 2 / (n_total_samples - 1)
            explained_variance_ratio = S ** 2 / np.sum(col_var * n_total_samples)

            state.n_samples_seen_ = n_total_samples
            state.components_ = V[:state.n_components_]
            state.singular_values_ = S[:state.n_components_]
            state.mean_ = col_mean
            state.var_ = col_var
            state.explained_variance_ = explained_variance[:state.n_components_]
            state.explained_variance_ratio_ = \
                explained_variance_ratio[:state.n_components_]
            if state.n_components_ < n_features:
                state.noise_variance_ = \
                    explained_variance[state.n_components_:].mean()
            else:
                state.noise_variance_ = 0.

            if state.count >= 999:
                    state.count = 0
                    print (state.components_)

            if end:
                np.savetxt("IPCA.dat", state.components_)

            return state

        sink_window(func = partial_fit, in_stream = in_stream,
                    window_size = batches, step_size=batches, state = state)

    s = Stream('s')
    gh(s)
    s.extend(source_list)
    run()

if __name__ == '__main__':

    print('Running the reverberation test ... ')
    IPCA(k = 2,batches = 5)
    print('reverb done')

max_stars_repo_name: AssembleSoftware/IoTPy

max_issues_repo_name: AssembleSoftware/IoTPy

max_forks_repo_name: AssembleSoftware/IoTPy


text:
from ase.io import read
from flosic_os import xyz_to_nuclei_fod,ase2pyscf
from pyscf import dft,gto
from flosic_scf import FLOSIC
import time
import numpy as np

# This example visualizes how a second order self-consistency solver can be used to drastically speed up and stabilize calculations.
# Second order here refers to the second derivatives of the energy w.r.t. the density.
# Drawing on these, the PySCF intrinsic functions .newton() and .scanner() allow the usage of a Newton optimization algorithm.
# Such an algorithm constructs the Hessian from these derivatives and uses it to perform anoptimization that is faster than the regular SCF solver.
# Unfortunately, it can ONLY BE USED WITH LDA OR GGA.
# PySCF currently DOES NOT SUPPORT MGGAs with this second order solver.
# To show the speedup, the calculations below will be timed.

# 1) Standard example WITHOUT usage of fixed VSIC properties

# To use this solver, we first have to simply define a regular calculator.
# The testing system is a Li atom.
molecule = read('Li.xyz')
geo,nuclei,fod1,fod2,included =  xyz_to_nuclei_fod(molecule)
spin = 1
charge = 0
b = 'sto3g'
xc = 'LDA,PW'
mol = gto.M(atom=ase2pyscf(nuclei), basis=b,spin=spin,charge=charge)
mf = FLOSIC(mol,xc=xc,fod1=fod1,fod2=fod2)

# Perform a calculation with the regular sovler.

t1_start = time.time()
e1 = mf.kernel()
t1_end = time.time()
delta_t1 = t1_end-t1_start

# 2) Standard example WITH usage of fixed VSIC properties

# To use this solver, we first have to simply define a regular calculator.
# The testing system is a Li atom.
molecule = read('Li.xyz')
geo,nuclei,fod1,fod2,included =  xyz_to_nuclei_fod(molecule)
spin = 1
charge = 0
b = 'sto3g'
xc = 'LDA,PW'
mol = gto.M(atom=ase2pyscf(nuclei), basis=b,spin=spin,charge=charge)
mf = FLOSIC(mol,xc=xc,fod1=fod1,fod2=fod2,vsic_every=8)

# Perform a calculation with the regular sovler.

t2_start = time.time()
e2 = mf.kernel()
t2_end = time.time()
delta_t2 = t2_end-t2_start

print('Etot with regular SCF solver (Hartree) = ',e1)
print('Duration for SCF cycle with regular SCF solver (in s):',delta_t1)
print('Etot with fixed VSIC (Hartree) = ',e2)
print('Duration for SCF cycle with fixed VSIC (in s):',delta_t2)
print('The magnitudes of these values depend on the power of your CPU and the number of processes that are running on your system simultaneously. Ideally, the second value should be significantly smaller.')

max_stars_repo_name: pyflosic/pyflosic

max_issues_repo_name: pyflosic/pyflosic

max_forks_repo_name: pyflosic/pyflosic

reason: import numpy


text:
# ------------------------------------------------------------------------------
#   @brief:
#       The optimization agent is responsible for doing the updates.
#   @author:
# ------------------------------------------------------------------------------
from .base_trainer import base_trainer
from mbbl.util.common import logger
import numpy as np
from collections import OrderedDict

class trainer(base_trainer):

    def __init__(self, args, network_type, task_queue, result_queue,
                 name_scope='trainer'):
        # the base agent
        super(trainer, self).__init__(
            args=args, network_type=network_type,
            task_queue=task_queue, result_queue=result_queue,
            name_scope=name_scope
        )

    def _update_parameters(self, rollout_data, training_info):
        # get the observation list
        self._update_whitening_stats(rollout_data)

        # generate the reward from discriminator
        rollout_data = self._network['reward'][0].generate_rewards(rollout_data)

        training_data = self._preprocess_data(rollout_data)
        training_stats = OrderedDict()
        training_stats['avg_reward'] = training_data['avg_reward']
        training_stats['avg_reward_std'] = training_data['avg_reward_std']

        assert 'reward' in training_info['network_to_train']
        # train the policy
        for key in training_info['network_to_train']:
            for i_network in range(self._num_model_ensemble[key]):
                i_stats = self._network[key][i_network].train(
                    training_data, self._replay_buffer, training_info={}
                )
                if i_stats is not None:
                    training_stats.update(i_stats)
        self._replay_buffer.add_data(training_data)

        # record the actual reward (not from the discriminator)
        self._get_groundtruth_reward(rollout_data, training_stats)
        return training_stats

    def _get_groundtruth_reward(self, rollout_data, training_stats):

        for i_episode in rollout_data:
            i_episode['raw_episodic_reward'] = sum(i_episode['raw_rewards'])
        avg_reward = np.mean([i_episode['raw_episodic_reward']
                              for i_episode in rollout_data])
        logger.info('Raw reward: {}'.format(avg_reward))
        training_stats['RAW_reward'] = avg_reward


text:
import warnings

import numpy as np

from starfish.core.intensity_table.decoded_intensity_table import DecodedIntensityTable
from starfish.core.morphology.binary_mask import BinaryMaskCollection
from starfish.core.types import Axes, Features
from ._base import AssignTargetsAlgorithm

class Label(AssignTargetsAlgorithm):
    """
    Extract cell ids for features in IntensityTable from a set of segmentation masks.
    """

    def __init__(self, **kwargs) -> None:
        pass

    @classmethod
    def _add_arguments(cls, parser) -> None:
        pass

    @staticmethod
    def _assign(
        masks: BinaryMaskCollection,
        decoded_intensities: DecodedIntensityTable,
        in_place: bool,
    ) -> DecodedIntensityTable:

        cell_ids = (
            Features.AXIS,
            np.full(decoded_intensities.sizes[Features.AXIS], fill_value='nan', dtype='<U8')
        )

        decoded_intensities[Features.CELL_ID] = cell_ids

        # it's 3D data.
        for _, mask in masks:
            has_z_data = Axes.ZPLANE.value in mask.coords
            if has_z_data:
                z_min, z_max = float(mask.z.min()), float(mask.z.max())
            else:
                warnings.warn(
                    "AssignTargets will require 3D masks in the future.", DeprecationWarning)
                z_min, z_max = np.NINF, np.inf
            y_min, y_max = float(mask.y.min()), float(mask.y.max())
            x_min, x_max = float(mask.x.min()), float(mask.x.max())

            in_bbox = decoded_intensities.where(
                (decoded_intensities.z >= z_min)
                & (decoded_intensities.z <= z_max)
                & (decoded_intensities.y >= y_min)
                & (decoded_intensities.y <= y_max)
                & (decoded_intensities.x >= x_min)
                & (decoded_intensities.x <= x_max),
                drop=True
            )

            selectors = {'y': in_bbox.y, 'x': in_bbox.x}
            if has_z_data:
                selectors['z'] = in_bbox.z
            in_mask = mask.sel(**selectors)
            spot_ids = in_bbox[Features.SPOT_ID][in_mask.values]
            decoded_intensities[Features.CELL_ID].loc[
                decoded_intensities[Features.SPOT_ID].isin(spot_ids)] = mask.name

        return decoded_intensities

    def run(
            self,
            masks: BinaryMaskCollection,
            decoded_intensity_table: DecodedIntensityTable,
            verbose: bool = False,
            in_place: bool = False,
    ) -> DecodedIntensityTable:
        """Extract cell ids for features in IntensityTable from a segmentation label image

        Parameters
        ----------
        masks : BinaryMaskCollection
            binary masks segmenting each cell
        decoded_intensity_table : IntensityTable
            spot information
        in_place : bool
            if True, process ImageStack in-place, otherwise return a new stack
        verbose : bool
            if True, report on the percentage completed during processing (default = False)

        Returns
        -------
        IntensityTable :
            IntensityTable with added features variable containing cell ids. Points outside of
            cells will be assigned `nan`.

        """
        return self._assign(masks, decoded_intensity_table, in_place=in_place)

max_stars_repo_name: haoxusci/starfish

max_issues_repo_name: lbgbox/starfish

max_forks_repo_name: lbgbox/starfish


text:
# Copyright 2021 The Distla Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# =============================================================================
"""
Backend class for bundling operations needed for distributed implementations of
various algorithms.
"""
from jax import lax
import jax.numpy as jnp

from distla_core.blas.summa import summa
from distla_core.utils import config
from distla_core.utils import pops

class DistributedBackend:
  """
  Bundles helper functions for distributed computations.
  """

  def __init__(self, p_sz: int, precision=lax.Precision.HIGHEST):
    self.p_sz = p_sz
    self.grid = config.GRID
    self.name = "DistributedBackend"
    self.precision = precision

  def add_to_diagonal(self, matrix, value, k=0, unpadded_dim=None):
    return pops.add_to_diagonal(matrix, value, k=k, unpadded_dim=unpadded_dim)

  def matmul(
      self,
      A,
      B,
      transpose_A=False,
      transpose_B=False,
      precision=None,
  ):
    if precision is None:
      precision = self.precision
    return summa.summa(
        A,
        B,
        self.p_sz,
        transpose_A,
        transpose_B,
        precision=precision,
    )

  def similarity_transform(self, A, B):
    AB = self.matmul(A, B)
    return self.matmul(B.conj(), AB, transpose_A=True)

  def trace(self, A):
    return pops.trace(A)

  def frobnorm(self, A):
    return pops.frobnorm(A)

  def vdot(self, A, B, precision=None):
    if precision is None:
      precision = self.precision
    return lax.psum(
        jnp.vdot(A, B, precision=precision),
        axis_name=pops.AXIS_NAME,
    )

  def gershgorin(self, A):
    return pops.gershgorin(A)

  def sum(self, A):
    return lax.psum(jnp.sum(A), axis_name=pops.AXIS_NAME)

  def shape(self, A):
    if len(A.shape) == 2:
      M = A.shape[0] * self.grid[0]
      N = A.shape[1] * self.grid[1]
    elif len(A.shape) == 3:
      M = A.shape[1] * self.grid[0]
      N = A.shape[2] * self.grid[1]
    else:
      raise TypeError(f"A had unexpected non-matrix shape {A.shape}.")
    return M, N

  def eye_like(self, A, unpadded_dim=None):
    return pops.eye(A.shape, A.dtype, unpadded_dim=unpadded_dim)

  def _can_transpose(self, A):
    # See documentation for pops.transpose to understand these restrictions.
    if pops.NCOLS == pops.NROWS * 2:
      return A.shape[0] % 2 == 0
    elif pops.NROWS == pops.NCOLS * 2:
      return A.shape[1] % 2 == 0
    else:
      return False

  def transpose(self, A):
    if self._can_transpose(A):
      return pops.transpose(A)
    else:
      msg = ("WARNING: Can't efficiently transpose a distributed matrix of "
             f"shape {A.shape} on a device grid {pops.GRID}. "
             "Falling back on matmul-based transpose.")
      # TODO use logging.warn
      print(msg)
      eye = self.eye_like(A)
      # Note that Precision.HIGHEST is used instead of self.precision.
      return self.matmul(
          A,
          eye,
          transpose_A=True,
          precision=lax.Precision.HIGHEST,
      )

  def __hash__(self):
    return hash((self.p_sz, self.grid, self.precision))

  def __eq__(self, other):
    return (self.p_sz == other.p_sz and self.grid == other.grid and
            self.precision == other.precision)

reason: import jax,from jax


text:
# -*- coding: utf-8 -*-
from Point import Point
import numpy as np
# import matplotlib.pyplot as plt
# from matplotlib import animation
# from matplotlib import style
from configs import SLICE_THRESHOLD, DIRECTION_THRESHOLD
from configs import COS_DIS_SAME_DIRECTION, COS_DIS_OPP_DIRECTION
from configs import OUTLIER_DIS
from configs import MIN_DIS_THRESHOLD
from configs import SPEED_THRESHOLD

def cos_vector(vec_1, vec_2):
    '''
    numerator = vec_1[0] * vec_2[0] + vec_1[1] * vec_2[1]
    denominator = np.sqrt(np.square(vec_1[0]) + np.square(vec_1[1])) *  np.sqrt(np.square(vec_2[0]) + np.square(vec_2[1]))
    if denominator == 0:
        print(vec_2)
    return numerator/denominator
    '''
    if np.linalg.norm(np.array(vec_1)) * np.linalg.norm(np.array(vec_2)) == 0:  # 分母判断
        return -1
    return np.dot(np.array(vec_1), np.array(vec_2)) / (
                np.linalg.norm(np.array(vec_1)) * (np.linalg.norm(np.array(vec_2))))

def min_distance(middle_point_list_1, middle_point_list_2):
    len_1 = len(middle_point_list_1)
    len_2 = len(middle_point_list_2)
    if len_1 <= 1 or len_2 <= 1:
        return 0.
    dis_list = []
    for i in range(0, len_1):
        for j in range(0, len_2):
            x1 = middle_point_list_1[i][0] * 1e5
            y1 = middle_point_list_1[i][1] * 1e5
            x2 = middle_point_list_2[j][0] * 1e5
            y2 = middle_point_list_2[j][1] * 1e5
            # dis_list.append(np.sqrt(np.square(x1-x2)+np.square(y1-y2)))
            dis_list.append(((x1 - x2) * (x1 - x2) + (y1 - y2) * (y1 - y2)) ** 0.5)
    return min(dis_list)

def mean_vector(vector_list, begin, end):
    if end <= begin:
        return [0., 0.]
    length = end - begin
    mean_x = 0.
    mean_y = 0.
    for i in range(begin, end):
        mean_x += vector_list[i][0]  # 不scale attention!
        mean_y += vector_list[i][1]  # 不scale attention!
    return [mean_x / length, mean_y / length]

class Track(object):
    '''
    track:
        id,
        sequence(point类的序列),
        video_id 所在的video

    '''

    def __init__(self, id, sequence):
        self.id = id
        self.sequence = sequence
        # self.simple_sequence = simple_sequence
        self.length = 0
        self.valid_speed = 0
        self.move_vector = [0, 0]
        self.speed_list = []

    def add_point(self, point):
        if point.id == self.id:
            self.sequence.append(point)
            self.length += 1

    def print_sequence(self):
        for item in self.sequence:
            item.print_point()

    def sort_every_point_by_time(self):
        # 对sequence进行排序

        self.sequence.sort(key=lambda x: x.t)  # 正序，时间由早到晚
        # self.print_sequence()

    def sort_by_time_and_smooth(self):  # 删除离群点
        # 对sequence按照时间先后进行排序
        self.sequence.sort(key=lambda x: x.t)

        length = len(self.sequence)
        if length <= 2:
            return 0
        vector_list = []
        index_list = []
        point_num_list = []
        speed_list = []
        for i in range(1, length):
            det_x = self.sequence[i].x - self.sequence[i - 1].x
            det_y = self.sequence[i].y - self.sequence[i - 1].y
            vector_list.append([det_x * 1e5, det_y * 1e5])  # attention !!!!!
            speed = self.sequence[i].speed_between_point(self.sequence[i - 1])

            speed_list.append(speed)
            point_num_list.append(i)

            if speed > SPEED_THRESHOLD:  # 速度超出了限制
                if i not in index_list:
                    index_list.append(i)  # 起点的index
                if i == 1:
                    index_list.append(0)

            if i - 1 in index_list:  # todo 如果上一个点跟上上一个点的速度超出限制
                if i >= 2:
                    speed_with_prepre = self.sequence[i].speed_between_point(self.sequence[i - 2])
                    if speed_with_prepre > SPEED_THRESHOLD:
                        if i not in index_list:
                            index_list.append(i)

        # print(speed_list)
        over_speed_num = len(index_list)

        begin = 0
        end = 0
        max_gap = 0
        for i in range(0, len(index_list) - 1):
            if index_list[i + 1] - index_list[i] >= max_gap:
                begin = index_list[i]
                end = index_list[i + 1]
                max_gap = index_list[i + 1] - index_list[i]
        # 0到begin ,begin-end,end到尾,取最长的连续一段
        if len(index_list) != 0:
            gap_to_tail = len(self.sequence) - 1 - index_list[-1]
            gap_from_begin = index_list[0]
            if gap_to_tail > gap_from_begin and gap_to_tail > max_gap:
                begin = index_list[-1]
                end = len(self.sequence) - 1
            elif gap_from_begin > gap_to_tail and gap_from_begin > max_gap:
                end = index_list[0]
                begin = 0
            # test

        # 一定要先算运动向量和速度，再删除离群点
        self.move_vector = mean_vector(vector_list, begin, end)
        if begin == 0 and end == 0:  # 如果没有离群点，index_list里面只有尾巴
            self.valid_speed = self.sequence[0].speed_between_point(self.sequence[-1])
        else:
            self.valid_speed = self.sequence[begin + 1].speed_between_point(
                self.sequence[min(end + 1, len(self.sequence) - 1)])

        # 删除离群点
        # if 1 not in index_list:
        #    index_list.append(1)
        # if 2 not in index_list:
        #    index_list.append(2)  # todo 一开始就非常偏离的点 把前半部分删掉
        index_list.sort()

        if len(index_list) > len(self.sequence) / 3:
            return over_speed_num
        for i in range(len(index_list) - 1, -1, -1):  # todo attention!!!! 倒序删除
            self.sequence.pop(index_list[i])

        return over_speed_num

    def sort_by_time_and_deduplicate(self):  # , anchor)#todo 希望anchor是整个map的也就是整个场景之下的中间点
        # 对sequence进行排序
        # TODO计算排序时间
        self.sequence.sort(key=lambda x: x.t, reverse=True)  # 倒序
        # todo 如果时间不连续 插值
        # 如果相距太近，删除点（O(1) pop()，O(nlogn) sort())
        length = len(self.sequence)
        if length <= 5:
            return
        for i in range(length - 1, 0, -1):
            if self.sequence[i].similar_point_in_pos(self.sequence[i - 1]):
                self.sequence[-1], self.sequence[i - 1] = self.sequence[i - 1], self.sequence[-1]
                self.sequence.pop()

        self.sequence.sort(key=lambda x: x.t)

    # def visualization(self):
    #     #对一条track做可视化
    #     #每一帧加一个点
    #     style.use('fivethirtyeight')
    #
    #     fig = plt.figure(figsize=(20, 10), facecolor='white')
    #     ax = plt.gca()
    #     #color = np.ones((n, 4)) * (0, 0, 0, 1)
    #
    #     plt.xlim((6000, 9000))
    #     plt.ylim((2000, 5000))
    #     # 设置x轴、y轴名称
    #     ax.set_xlabel('gps-x')
    #     ax.set_ylabel('gps-y')
    #
    #     def animate(point_number):
    #         pos_list = [self.sequence[point_number].x, self.sequence[point_number].y]
    #         # scale
    #         x = int(pos_list[0] * 1e5 - 42490000)
    #         y = int(pos_list[1] * 1e5 - 90670000)
    #
    #         scat = ax.scatter(x, y, s=10, c='r', lw=0.5)
    #         # plt.pause(1)
    #         return scat,
    #
    #     ani = animation.FuncAnimation(fig, animate, interval=40, frames=len(self.sequence),repeat=False)
    #     plt.draw()
    #     plt.pause(30)
    #     plt.close()
    #     return plt

    def get_middle_point(self, begin, end, step=1):
        # 算一个点列表的中间坐标，代表这个点
        # 返回[x,y,t]
        middle_x = 0
        middle_y = 0
        middle_t = 0
        length = end - begin
        if length == 0:
            return [0., 0., 0.]
        number = 0
        for i in range(begin, end, step):  # todo 步长可以继续调
            if i >= end:
                break  # 因为有step
            middle_x += self.sequence[i].x
            middle_y += self.sequence[i].y
            middle_t += self.sequence[i].t
            number += 1
        if number == 0:
            return [0., 0., 0.]
        middle_point = [middle_x / number, middle_y / number, middle_t / number]
        # print(middle_point)
        return middle_point

    def get_middle_points_list(self, threshold):
        length = len(self.sequence)
        # SLICE_THRESHOLD = 10
        # DIRECTION_THRESHOLD = 50
        middle_points_list = []
        if length < 2 * threshold:  # 长度甚至小于两倍的阈值的，
            if length < threshold:
                return [self.get_middle_point(0, int(length / 2), 1), self.get_middle_point(int(length / 2), length, 1)]
            else:
                return [self.get_middle_point(0, int(length / 2), 1), self.get_middle_point(int(length / 2), length, 1),
                        self.get_middle_point(threshold, length, 1)]
        else:
            for i in range(0, length, threshold):
                middle_points_list.append(self.get_middle_point(i, min(i + threshold, length)))

            return middle_points_list

    def vector_list(self, middle_point_list):
        vector_list = []
        length = len(middle_point_list)
        if length <= 1:
            return vector_list
        for i in range(0, length - 1, 1):  # todo
            det_x = middle_point_list[i + 1][0] - middle_point_list[i][0]
            det_y = middle_point_list[i + 1][1] - middle_point_list[i][1]
            vector_list.append([det_x * 1e5, det_y * 1e5])  # attention !!!!!
        return vector_list

    def angle_with_another_track(self, track_2):  # track已经按时间进行了排序 #TODO
        middle_points_list_1 = self.get_middle_points_list(threshold=int(DIRECTION_THRESHOLD / 2))

        middle_points_list_2 = track_2.get_middle_points_list(threshold=int(DIRECTION_THRESHOLD / 2))

        vector_list_1 = self.vector_list(middle_points_list_1)
        # if len(vector_list_1) == 0:
        # print(middle_points_list_1)
        # print(self.id)
        vector_list_2 = track_2.vector_list(middle_points_list_2)
        # if len(vector_list_2) == 0:
        # print(middle_points_list_2)
        # print(track_2.id)
        angle_list = []
        len_1 = len(vector_list_1)
        len_2 = len(vector_list_2)
        if len_1 <= 0 or len_2 <= 0:
            return [0, 0, 0]
        else:
            for i in range(0, len_1):
                for j in range(0, len_2):
                    angle_list.append(cos_vector(vector_list_1[i], vector_list_2[j]))

        return [min(angle_list), max(angle_list), np.mean(angle_list)]
        # return angle_list

    def angle_and_mindis_with_another_track(self, track_2):
        middle_points_list_1 = self.get_middle_points_list(threshold=int(DIRECTION_THRESHOLD / 2))

        middle_points_list_2 = track_2.get_middle_points_list(threshold=int(DIRECTION_THRESHOLD / 2))

        vector_list_1 = self.vector_list(middle_points_list_1)
        vector_list_2 = track_2.vector_list(middle_points_list_2)

        min_dis = min_distance(middle_points_list_1, middle_points_list_2)  # todo new

        angle_list = []
        len_1 = len(vector_list_1)
        len_2 = len(vector_list_2)
        if len_1 <= 0 or len_2 <= 0:
            return [0, 0, 0]
        else:
            for i in range(0, len_1):
                for j in range(0, len_2):
                    angle_list.append(cos_vector(vector_list_1[i], vector_list_2[j]))

        return [min(angle_list), max(angle_list), np.mean(angle_list), min_dis]

    def can_not_reach(self, track_2):
        vector_between_track = [(track_2.sequence[0].x - self.sequence[0].x),
                                (track_2.sequence[0].y - self.sequence[0].y)]
        cos_dis = cos_vector(self.move_vector, vector_between_track)
        if cos_dis < COS_DIS_OPP_DIRECTION:
            return 0
        else:
            return 1

    def time_to_another_track(self, track_2):
        # return [5.0]
        # 方向不能完全不一致
        time = 0  # 为了跟angle_and_mindis_with_another_track的返回值合并，用list返回
        min_max_mean_mindis = self.angle_and_mindis_with_another_track(track_2)
        if min_max_mean_mindis[-1] < MIN_DIS_THRESHOLD:
            return [0]  # 方向完全不一致，认为不可能到达,  两条track
        if min_max_mean_mindis[1] < COS_DIS_OPP_DIRECTION or self.can_not_reach(track_2):
            return [-99]
        else:
            return [99]

        # x1 = self.sequence[-1].x * 1e5
        # y1 = self.sequence[-1].y * 1e5
        # if len(track_2.sequence) == 0:
        #     print(track_2.id)
        # x2 = track_2.sequence[0].x * 1e5
        # y2 = track_2.sequence[0].y * 1e5
        #
        # dis = ((x1 - x2) * (x1 - x2) + (y1 - y2) * (y1 - y2)) ** 0.5
        # if self.valid_speed != 0:
        #     time = dis / self.valid_speed
        #
        # #0430 time modified
        # vector_between_track = [(track_2.sequence[0].x - self.sequence[-1].x),
        #                         (track_2.sequence[0].y - self.sequence[-1].y)]
        # cos_dis = cos_vector(self.move_vector, vector_between_track)
        # if cos_dis < 0:
        #     time = -time
        #
        # return [time]

    def angle_with_another_track_coarse(self, track_2):  # 使用首尾的方向向量，用于验证
        if len(self.sequence) == 1 or len(track_2.sequence) == 1:
            return [0, 0, 0]
        vec_1 = [(self.sequence[1].x - self.sequence[0].x) * 1e5, (self.sequence[1].y - self.sequence[0].y) * 1e5]
        vec_2 = [(track_2.sequence[1].x - track_2.sequence[0].x) * 1e5,
                 (track_2.sequence[1].y - track_2.sequence[0].y) * 1e5]
        cos = cos_vector(vec_1, vec_2)
        return [cos, cos, cos]

'''

    def get_mean_vector_last_10(self):
        #todo 如果最后几秒钟停下不运动怎么办
        mean_vector = [0,0]
        length = len(self.sequence)#可以优化的吧
        if length <= 1:
            return mean_vector#少于一个点的情况，方向向量为（0，0）
        mean_x = 0
        mean_y = 0
        for i in range(max(0, length-10),length-1,1):#最后10个点的平均方向向量，如果少于十个点，取全部点的平均方向向量
            det_x = self.sequence[i+1].x - self.sequence[i].x
            det_y = self.sequence[i+1].x - self.sequence[i].y
            mean_x += det_x*100000
            mean_y += det_y*100000#scale 但是似乎不是必需的
        mean_vector[0] = mean_x/min(length,10)
        mean_vector[1] = mean_y/min(length,10)

        return mean_vector#平均x的增量或减量，平均y的增量或减量

    def near_field_of_track_head(self):
        # TODO 一条track开头部分的一个周边范围，返回x,y,r(半径）

        length = len(self.sequence)
        if length <= 1:
            return [0, 0, 0]
        mean_point_head = self.get_middle_point(0, max(length, 20))  # todo 阈值是猜的
        r = 1e-4

        return [mean_point_head[0], mean_point_head[1], r]
        '''
'''
 #todo track 画一条速度曲线
        id_to_check_1 = ['011_322145','011_322128', '013_321503','014_321601', '015_321067']  # 011_322145','011_322128', '013_321503',
        id_to_check_2 = ['029_324834',
                         '029_325577','029_324964','032_322801','036_324009','036_324120','036_323997','037_323592','037_323581','037_323582',
                         '037_323583','037_323565','037_323568','037_323559','037_323560','037_323609','037_323601','037_323598','038_322551']
        # ['016_323518','016_323547','016_323460','017_323800','020_322959','020_322973','020_322956','020_323051','020_323000',
        # '020_323018','020_323052','025_322897','025_322924','029_325103','029_325323','029_325424','029_324974',
        if self.id in id_to_check_2:
            print(index_list)
        #if True:
            plt.figure('Line fig')
            ax = plt.gca()
            # 设置x轴、y轴名称
            ax.set_xlabel('number')
            ax.set_ylabel('speed')
            ax.plot(point_num_list, speed_list, color='b', linewidth=1, alpha=0.6)
            plt.draw()
            plt.pause(3)
            plt.close()

        ###################attention!##################
        new_speed_list = []
        point_num_list = []
        new_length = len(self.sequence)
        for i in range(1, new_length):
            det_x = self.sequence[i].x - self.sequence[i - 1].x
            det_y = self.sequence[i].y - self.sequence[i - 1].y
            speed = self.sequence[i].speed_between_point(self.sequence[i - 1])
            new_speed_list.append(speed)
            point_num_list.append(i)

        if self.id in id_to_check_2:

        #if True:
            plt.figure('Line fig')
            ax = plt.gca()
            # 设置x轴、y轴名称
            ax.set_xlabel('number')
            ax.set_ylabel('speed')
            ax.plot(point_num_list, new_speed_list, color='r', linewidth=1, alpha=0.6)
            plt.draw()
            plt.pause(3)

            plt.close()

         new_speed_list = []
        point_num_list = []
        new_length = len(self.sequence)
        for i in range(1, new_length):
            det_x = self.sequence[i].x - self.sequence[i - 1].x
            det_y = self.sequence[i].y - self.sequence[i - 1].y
            speed = self.sequence[i].speed_between_point(self.sequence[i - 1])
            new_speed_list.append(speed)
            point_num_list.append(i)

        #if self.id in id_to_check_2:
        if True:
            plt.figure('Line fig')
            ax = plt.gca()
            # 设置x轴、y轴名称
            ax.set_xlabel('number')
            ax.set_ylabel('speed')
            ax.plot(point_num_list, new_speed_list, color='r', linewidth=1, alpha=0.6)
            print(self.id)
            plt.draw()
            plt.pause(0.001)

            plt.close()

'''

max_stars_repo_name: didi/mtmc-vt

max_issues_repo_name: didichuxing/mtmc-vt

max_forks_repo_name: didichuxing/mtmc-vt


text:
import copy as cp

import numpy as np

from ..multi_tasks_learning.GPREPS import *
from envs.mujoco.utils.quaternion import mat2Quat
from .plot_result import *

import time

from envs.abb.models import utils

class Solver(object):
	def __init__(self, args, env, project_path):

		self.args = args

		# #####################  Training Parameters  ######################
		self.K = self.args.num_policy_update  # 上层policy训练循环总数
		self.N = self.args.num_real_episodes  # 在上层policy的一个训练周期中，下层RL训练，改变context参数的次数
		self.n = 1  # 下层RL训练，每改变一次context参数，执行RL的episode数
		self.d = 1  # 下层RL每d个step更新一次网络参数
		self.M = self.args.num_simulated_episodes  # 在上层policy的一个训练周期中，通过model获得artificial trajectory的过程中，改变context参数的次数
		self.L = 5  # 每改变一次context参数，通过model获得artificial trajectory的次数

		# #####################  Policy Parameters  ######################
		self.MAX_EP_STEPS = self.args.max_episode_steps  # RL的最大步数
		self.context_dim = 6  # including initial state (x, y, z) and terminal state (x, y, z)
		self.num_waypoints = 3  # depends on the task to solve ::: predefined
		self.eps_max = 0.25
		self.eps_min = 0.10

		# Impedance parameters ::: stiffness damping is derived from a constant relation Kd = beta * sqrt(Kp)
		self.contextual_impedance_dim = 12
		self.contextual_impedance_lower_bound = np.array([0, 0.2, 0.2, 1.0, 1.0, 0.5, 1.0, 1.0, 0.0, 0.5, 0.5, 2])
		self.contextual_impedance_upper_bound = np.array([0.0, 1.0, 1.0, 5.0, 5.0, 3.0, 5.0, 5.0, 2.0, 2.0, 2.0, 8])

		self.memory_dim = 500  # Context parameters
		# Give several initial expe
		# print("weights :::", weights)
		# print("stiffness pos :::", stiffness_pos)rt data ::: very important for policy search
		self.initial_a = np.array([1.0, 0.8, 0.1, 4, 4, 1, 3, 3, 1, 1, 1, 4])

		# attractor points
		self.pos_list = None
		self.quat_list = None

		# lower-level RL controller :::
		self.observation_dim = 12  # state
		self.action_dim = 6  # control command ::: action

		self.env = env
		self.done = False
		self.safe = True
		self.render = self.args.render

		self.file_name = ''
		self.project_path = project_path
		self.result_path = project_path + "results/runs/mujoco"

		self.evaluations = []
		self.eval_episodes_states = []

		# # Set seeds
		# torch.manual_seed(args.seed)
		# np.random.seed(args.seed)points

		# context parameter model
		self.replay_buffer = utils.ReplayBuffer()

		# lower-level memory
		self.replay_buffer_model = utils.ReplayBuffer(1e4)

		# Initialize GPREPS
		self.gpreps = GPREPS(self.context_dim, self.contextual_impedance_dim, self.memory_dim,
							 self.contextual_impedance_lower_bound, self.contextual_impedance_upper_bound,
							 self.initial_a, 0.25)

		# For model-based reps :::
		# self.r_model = R_MODEL(self.policy,
		# 					   self.env,
		# 					   self.context_dim,
		# 					   self.contextual_impedance_dim,
		# 					   self.observation_dim,
		# 					   self.action_dim,
		# 					   self.MAX_EP_STEPS)

		# self.s_model = S_MODEL()

		self.total_timesteps = 0
		self.episode_timesteps = 0
		self.episode_number = 0
		self.episode_reward = 0
		self.reward_scale = 0.001
		self.pre_num_steps = self.total_timesteps
		self.best_reward = 0.0
		self.timesteps_since_eval = 0

		""" training performance """
		self.training_reward = []
		self.training_time = []
		self.training_states = []
		self.training_im_actions = []

		""" evaluation performance """
		self.evaluations_reward = []
		self.evaluations_time = []
		self.evaluations_actions = []
		self.evaluations_im_actions = []
		self.evaluations_states = []
		self.evaluations_options = []

		self.episode_reward_average_list = []
		self.successful_rate_list = []

		self.episode_reward_average_list = []
		self.successful_rate_list = []

	def reset(self):
		"""
			Random offset value :::
		"""
		# initial point ::: add initial offset based on the top of hole
		initial_offset = np.zeros(3)
		initial_offset[:2] = np.random.uniform(-0.035, 0.035, size=2)
		initial_offset[2] = 0.07

		# final optimal point
		target_pos, target_rot = self.env._get_target_obs()
		quat_hole_top = mat2Quat(target_rot[2])
		quat_hole_base = mat2Quat(target_rot[1])

		# attractor points ::: could be constant
		self.pos_list = np.concatenate(([target_pos[1]], [target_pos[2]], [target_pos[2] - [0.0, 0.0, -0.05]]), axis=0)
		self.quat_list = np.concatenate(([quat_hole_base], [quat_hole_top], [quat_hole_top]), axis=0)
		# print("pos_list :::", self.pos_list)
		# print("quat_list :::", self.quat_list)

		self.obs, self.initial_state, self.target_state = \
			self.env.vic_reset(initial_offset=initial_offset, pos_list=self.pos_list, quat_list=self.quat_list)

		self.episode_reward = 0
		self.episode_timesteps = 0

		# successful episode number
		# self.episode_number = 0
		self.done = False
		self.safe = True
		self.state = self.initial_state

		# return context of one episode
		# self.context_state = np.concatenate((self.initial_state, self.target_state))
		self.context_state = self.get_context()
		return self.context_state, self.state, self.done, self.safe

	def get_context(self):
		# print('Get Context Parameter for one episode!!!')
		context = np.concatenate((self.initial_state, self.target_state))
		return context

	def train(self, w=None):
		"""
			execute reinforcement learning and store data for model training
		"""
		self.episode_number = 0
		self.episode_reward_average = 0
		for i in range(self.N):
			# print("Iteration ::::::::::::::::::::::", i)

			z, _, _, _ = self.reset()
			z = self.get_context().reshape(-1, self.context_dim)
			# print("context parameters :::", z)

			# choose impedance parameters ::: only stiffness and weights, damping is related to stiffness
			if w is None:
				w = self.gpreps.choose_action(z)[0]

			# print("Choose Impedance Parameters :::", w)

			# Extract Parameters for IMOGIC
			weights = np.array(w[:3])

			stiffness_pos = np.array([w[3:6], w[6:9], w[9:]])
			# print("weights :::", weights)
			# print("stiffness pos :::", stiffness_pos)

			stiffness_rotation = np.ones((3, 3))  # fixed in this experiment :::

			stiffness_list_test = np.concatenate((stiffness_pos, stiffness_rotation), axis=1)

			# stiffness_list = np.array([[4, 4, 0.1, 1, 1, 1],
			# 						   [0.5, 0.5, 8, 1, 1, 1]])
			# weights_list = np.array([0.9, 0.1])

			stiffness_list = stiffness_list_test
			beta = weights[0]
			damping_list = 2 * np.sqrt(stiffness_list)
			weights_list = weights
			weights_list[0] = 1

			# Start realistic simulation
			# while self.episode_number < self.n:
			# print("Episode Number :::", self.episode_number)
			# self.reset()
			while self.episode_timesteps < self.MAX_EP_STEPS:
				"""
					basic action from VIC
				"""
				obs, reward, self.done, info, self.state, action, self.safe = \
					self.env.step_im(stiffness=stiffness_list, damping=damping_list, weights=weights_list)

				new_obs = obs[:12]

				if self.render:
					self.env.render()
				# time.sleep(0.001)

				# store data into buffer ::: for RL training
				# self.replay_buffer.add((self.obs, new_obs, action, reward, self.done, 0))

				# here ::: self.obs and new_obs # for model training
				# self.replay_buffer_model.add((self.obs, new_obs, action))

				self.episode_reward += reward
				self.obs = new_obs
				self.episode_timesteps += 1

				if self.done or not self.safe or self.episode_timesteps == self.MAX_EP_STEPS - 1:
					print('RL episode :::', self.episode_number,
						  'step :::', self.episode_timesteps,
						  'done? :::', self.done,
						  'safe? :::', self.safe,
						  'episode reward :::', self.episode_reward)
					if self.done:
						# calculate successful rate
						self.episode_number += 1
					break

			# print("z :::", z)
			# print("w :::", w)
			# print("episode reward :::", self.episode_reward * 0.01)

			self.gpreps.store_realistic_data(z[0], w, self.episode_reward * self.reward_scale)
			self.episode_reward_average += self.episode_reward * self.reward_scale

		return self.episode_number / self.N, self.episode_reward_average / self.N

	def contextual_main(self):
		"""
			main function of GPRREPS for contextual policy training
		"""
		self.episode_reward_average_list = []
		self.successful_rate_list = []

		for k in range(self.K):
			print('Training Cycle...', k)
			successful_rate, episode_reward_average = self.train()

			# evaluate learning performance every cycle
			print('Successful Rate...', successful_rate)
			print("Average reward :::", episode_reward_average)
			self.episode_reward_average_list.append(cp.deepcopy(episode_reward_average))
			self.successful_rate_list.append(cp.deepcopy(successful_rate))

			# run RL and collect data
			# print('Running RL...')

			# train reward and context model
			# print('Training GMR models...')
			# self.r_model.train_reward_model(self.replay_buffer_model)

			# self.s_model.train_context_model()

			##################################
			# Predict Rewards and Store Data #
			##################################
			# print('Generate artificial trajectories !!!')
			# for j in range(self.M):
			# 	# Predict Rewards
			# 	R = 0.
			# 	Z = self.get_context()
			# 	W = self.gpreps.choose_action(Z)
			#
			# 	# Predict L Trajectories
			# 	for l in range(self.L):
			# 		R += self.r_model.trajectory(Z, W)
			# 	reward = R / self.L
			# 	print('Artificial running cycle', j + 1,
			# 		  'ra:', Z[0], 'rd:', Z[1],
			# 		  'reward:', reward)
			#
			# 	# Construct Artificial Dataset D
			# 	self.gpreps.store_simulated_data(Z, W, reward)

			# Sample and Update Policy
			# self.context_average = np.array([0.8, 0.8, 1.3, 0.8, 0.1, 1.29])
			#
			# print('memory :::', self.gpreps.memory_realistic)

			print('****************** Start Training *****************')
			eps_k = self.eps_min + k * (self.eps_max - self.eps_min) / self.K
			self.gpreps.learn(training_type='Realistic', eps=eps_k)

			# evaluation policy
			print('************ Training cycle finished **************')

		# Plot learning results
		plot_single_reward(data=self.successful_rate_list, font_size=18)
		plot_single_reward(self.episode_reward_average_list, font_size=18, y_label_list=['Episode Reward Average'])

		return self.episode_reward_average_list, self.successful_rate_list

	def model_contextual_main(self):
		"""
			main function of GPRREPS for contextual policy training
		"""
		self.episode_reward_average_list = []
		self.successful_rate_list = []

		for k in range(self.K):

			print('Training Cycle...', k)
			successful_rate, episode_reward_average = self.train()

			# evaluate learning performance every cycle
			print('Successful Rate...', successful_rate)
			print("Average reward :::", episode_reward_average)
			self.episode_reward_average_list.append(cp.deepcopy(episode_reward_average))
			self.successful_rate_list.append(cp.deepcopy(successful_rate))

			# run RL and collect data
			# print('Running RL...')

			# train reward and context model
			# print('Training GMR models...')
			# self.r_model.train_reward_model(self.replay_buffer_model)

			# self.s_model.train_context_model()

			##################################
			# Predict Rewards and Store Data #
			##################################
			# print('Generate artificial trajectories !!!')
			# for j in range(self.M):
			# 	# Predict Rewards
			# 	R = 0.
			# 	Z = self.get_context()
			# 	W = self.gpreps.choose_action(Z)
			#
			# 	# Predict L Trajectories
			# 	for l in range(self.L):
			# 		R += self.r_model.trajectory(Z, W)
			# 	reward = R / self.L
			# 	print('Artificial running cycle', j + 1,
			# 		  'ra:', Z[0], 'rd:', Z[1],
			# 		  'reward:', reward)
			#
			# 	# Construct Artificial Dataset D
			# 	self.gpreps.store_simulated_data(Z, W, reward)

			# Sample and Update Policy
			# self.context_average = np.array([0.8, 0.8, 1.3, 0.8, 0.1, 1.29])
			#
			# print('memory :::', self.gpreps.memory_realistic)

			if k > self.args.start_policy_update_idx:
				self.gpreps.train_reward_model(sample_number=100, type='GP')

				for m in range(self.M):
					z, _, _, _ = self.reset()
					z = self.get_context().reshape(-1, self.context_dim)
					w = self.gpreps.choose_action(z)[0]
					sample_reward = self.gpreps.generate_artifical_trajectory([np.array(z[0])], [w])
					self.gpreps.store_simulated_data(z[0], w, sample_reward)

			if k > self.args.start_policy_update_idx:
				print('****************** Start Training *****************')
				eps_k = self.eps_min + k * (self.eps_max - self.eps_min) / self.K
				self.gpreps.learn(training_type='Simulated', eps=eps_k)

			if k % 100 == 0:
				# Plot learning results
				plot_single_reward(data=self.successful_rate_list, font_size=18,
								   y_label_list=['Successful Rate'])
				plot_single_reward(data=self.episode_reward_average_list, font_size=18,
								   y_label_list=['Episode Reward Average'])

		# Plot learning results
		plot_single_reward(data=self.successful_rate_list, font_size=18)
		plot_single_reward(data=self.episode_reward_average_list, font_size=18, y_label_list=['Episode Reward Average'])

		return self.episode_reward_average_list, self.successful_rate_list

	def CGPUCB_main(self, num_rolluts=1000, beta=100):
		"""
			main function of GPUCB for contextual policy training
		"""
		# give initial parameters
		num_mesh_grid = 3
		mesh_grid_dist = np.array([0., 0.4, 0.4, 2.0, 2.0, 1.25, 2.0, 2.0, 1.0, 0.75, 0.75, 4.0])
		contextual_impedance_lower_bound = np.array([0.0, 0.2, 0.2, 1.0, 1.0, 0.5, 1.0, 1.0, 0.0, 0.5, 0.5, 2])
		contextual_impedance_upper_bound = np.array([0.0, 1.4, 1.4, 7.0, 7.0, 4.25, 7.0, 7.0, 3.0, 2.75, 2.75, 12])

		para_samples = [np.array([0., 0., 0.])]
		for i in range(1, self.contextual_impedance_dim):
			sample_list = np.arange(contextual_impedance_lower_bound[i], contextual_impedance_upper_bound[i],
									mesh_grid_dist[i])
			para_samples.append(sample_list)

		# print("para_samples :::", para_samples)

		self.meshgrid = np.array(
			np.meshgrid(para_samples[0], para_samples[1], para_samples[2], para_samples[3], para_samples[4],
						para_samples[5],
						para_samples[6], para_samples[7], para_samples[8], para_samples[9], para_samples[10],
						para_samples[11]))

		self.sample_grid = self.meshgrid.reshape(self.meshgrid.shape[0], -1).T

		self.mu = np.array([0. for _ in range(self.sample_grid.shape[0])])
		self.sigma = np.array([0.5 for _ in range(self.sample_grid.shape[0])])
		self.beta = beta
		# evaluation results

		self.realistic_sample_list = []
		self.episode_reward_average_list = []
		self.successful_rate_list = []

		gp = GaussianProcessRegressor()

		for k in range(num_rolluts):
			grid_idx = self.argmax_ucb(self.mu, self.sigma, self.beta)
			print("gird_idx", grid_idx)
			print("sample_gird", self.sample_grid)
			optimal_sample = self.sample_grid[grid_idx]
			print("optimal_sample", optimal_sample)
			successful_rate, episode_reward_average = self.train(w=optimal_sample)
			self.realistic_sample_list.append(optimal_sample)
			self.episode_reward_average_list.append(episode_reward_average)
			self.successful_rate_list.append(successful_rate)

			gp.fit(self.realistic_sample_list, self.episode_reward_average_list)
			self.mu, self.sigma = gp.predict(self.sample_grid, return_std=True)

		# Plot learning results
		plot_single_reward(data=self.successful_rate_list, font_size=18)
		plot_single_reward(self.episode_reward_average_list, font_size=18,
						   y_label_list=['Episode Reward Average'])

		return self.episode_reward_average_list, self.successful_rate_list

	def argmax_ucb(self, mu, sigma, beta):
		return np.argmax(mu + sigma * np.sqrt(beta))

	def eval_once(self):
		# self.pbar.update(self.total_timesteps - self.pre_num_steps)
		# self.pre_num_steps = self.total_timesteps
		# if self.timesteps_since_eval >= self.args.eval_freq:
		#     self.timesteps_since_eval %= self.args.eval_freq
		""" evaluate the policy for once """
		avg_reward, avg_time, eval_actions, eval_states, eval_im_actions, eval_options = \
			evaluate_assembly_policy(self.env, self.policy, self.args)
		self.evaluations_reward.append(cp.deepcopy(avg_reward))
		self.evaluations_time.append(cp.deepcopy(avg_time))
		self.evaluations_actions.append(cp.deepcopy(eval_actions))
		self.evaluations_states.append(cp.deepcopy(eval_states))
		self.evaluations_im_actions.append(cp.deepcopy(eval_im_actions))
		self.evaluations_options.append(cp.deepcopy(eval_options))
		print('evaluations_reward :::::::::::::::', self.evaluations_reward)
		print('evaluations_time :::::::::::::::::', self.evaluations_time)

		""" save test data numpy """
		np.save(self.log_dir + "/test_reward", self.evaluations_reward)
		np.save(self.log_dir + "/test_time", self.evaluations_time)
		np.save(self.log_dir + "/test_actions", self.evaluations_actions)
		np.save(self.log_dir + "/test_options", self.evaluations_options)
		np.save(self.log_dir + "/test_im_actions", self.evaluations_im_actions)
		np.save(self.log_dir + "/test_states", self.evaluations_states)

		utils.write_table(self.log_dir + "/test_reward", np.asarray(self.evaluations_reward))
		utils.write_table(self.log_dir + "/test_time", np.asarray(self.evaluations_time))

		if self.args.save_all_policy:
			self.policy.save(
				self.file_name + str(int(int(self.total_timesteps / self.args.eval_freq) * self.args.eval_freq)),
				directory=self.log_dir)
		else:
			self.policy.save(self.file_name, directory=self.log_dir)

# print('total_timesteps ::::::::::::::::::::::::::', self.total_timesteps)
# print('episode_reward :::::::::::::::::::::::::::', self.episode_reward)
# self.training_reward.append(cp.deepcopy(self.episode_reward))
# self.training_time.append(cp.deepcopy(self.episode_timesteps))
# self.training_states.append(cp.deepcopy(epi_states))
# self.training_im_actions.append(cp.deepcopy(epi_actions))
#
# np.save(self.log_dir + "/train_reward", self.training_reward)
# np.save(self.log_dir + "/train_time", self.training_time)
# np.save(self.log_dir + "/train_states", self.training_states)
# np.save(self.log_dir + "/train_im_actions", self.training_im_actions)
#
# utils.write_table(self.log_dir + "/train_reward", np.asarray(self.training_reward))
# utils.write_table(self.log_dir + "/train_time", np.asarray(self.training_time))

def evaluate_assembly_policy(env, policy, args):
	"""
		Runs policy for X episodes and returns average reward
	"""
	avg_reward = 0.
	eval_actions = []
	eval_im_actions = []
	eval_states = []
	eval_options = []
	start_time = time.time()
	for _ in range(args.num_eval_episodes):
		obs, state, done = env.reset()
		done = False
		episode_step = 0
		epi_actions = []
		epi_im_actions = []
		epi_options = []
		epi_states = []
		while not done and episode_step < args.max_episode_steps:
			if 'HRLACOP' in args.policy_name:
				action, option = policy.select_evaluate_action([np.array(obs)])
				epi_options.append(cp.deepcopy(option))
			else:
				action = policy.select_action(np.array(obs))

			epi_states.append(cp.deepcopy(state))
			obs, state, reward, done, _, execute_action = env.step(action)
			epi_actions.append(cp.deepcopy(action))
			epi_im_actions.append(cp.deepcopy(execute_action))
			avg_reward += reward
			episode_step += 1
		eval_states.append(cp.deepcopy(epi_states))
		eval_actions.append(cp.deepcopy(epi_actions))
		eval_im_actions.append(cp.deepcopy(epi_im_actions))
		eval_options.append(cp.deepcopy(epi_options))
	avg_time = (time.time() - start_time) / args.num_eval_episodes
	avg_reward /= args.eval_episodes
	return avg_reward, avg_time, eval_actions, eval_states, eval_im_actions, eval_options

# class GPREPS(object):
# 	def __init__(self, w_dim, z_dim, memory_dim, w_bound):
# 		# initialize parameters
# 		self.memory = []
# 		self.pointer = 0
# 		self.w_dim, self.z_dim, self.memory_dim, self.w_bound = w_dim, z_dim, memory_dim, w_bound
#
# 		# build actor
# 		self.a = np.ones((w_dim, 1), dtype=np.float32)
# 		self.A = np.ones((w_dim, z_dim), dtype=np.float32)
# 		self.COV = np.ones((w_dim, w_dim), dtype=np.float32)
#
# 	def choose_action(self, z):
# 		z = np.array([z])
# 		u = self.a + np.dot(self.A, z.transpose())
# 		u = u.transpose()[0]
# 		return np.random.multivariate_normal(mean=u, cov=self.COV, size=1)
#
# 	def learn(self, z_):
# 		eta, theta = argmin(self.memory, z_, self.z_dim)
#
# 		p = 0.
# 		P_ = []
# 		Z = []
# 		B = []
# 		for i in range(self.memory_dim):
# 			z, w, r = self.memory[i]
# 			z = np.array([z])
# 			w = np.array(w)
# 			r = np.array([r])
# 			p = np.exp((r - np.dot(z, theta)) / eta)
# 			z_ = np.c_[np.array([1.]), z]
# 			Z.append(z_[0])
# 			B.append(w[0])
# 			P_.append(p[0])
# 		p, B, Z = np.array(p), np.array(B), np.array(Z)
# 		P = P_ * np.eye(self.memory_dim)
#
# 		# calculate mean action
# 		target1 = np.linalg.inv(np.dot(np.dot(Z.transpose(), P), Z))
# 		# print(np.shape(P), np.shape(Z.transpose()), np.shape(B))
# 		target2 = np.dot(np.dot(Z.transpose(), P), B)
# 		target = np.dot(target1, target2).transpose()
# 		self.a = target[:, :1]
# 		self.A = target[:, 1:]
#
# 		# calculate the COV
# 		Err = 0
# 		for i in range(self.memory_dim):
# 			z, w, r = self.memory[i]
# 			z = np.array([z])
# 			w = np.array([w])
# 			err = w - self.a - np.dot(self.A, z.transpose())
# 			Err += np.dot(err, err.transpose()) * P_[i]
# 		self.COV = Err / np.sum(P_)
# 		# COV 减维到2维，否则维度是(1, 12, 12, 1)
# 		COV = np.zeros([12, 12])
# 		for i in range(12):
# 			for j in range(12):
# 				COV[i][j] = self.COV[0][i][j][0]
# 		print(np.shape(self.COV))
# 		print(np.shape(COV))
# 		self.COV = COV
# 		print('Contextual policy search upper level parameters updated')
#
# 	def store_data(self, z, w, r):
# 		transition = [z, w, [r]]
# 		if len(self.memory) == self.memory:
# 			index = self.pointer % self.memory_dim  # replace the old memory with new memory
# 			self.memory[index] = transition
# 		else:
# 			self.memory.append(transition)
# 		self.pointer += 1

# **** 程序结构 ****
# 主函数：solver.contextual_main()
# 运行RL: solver.train()
# 上层policy：class GPREPS
# 获得环境变量z：solver.get_context()
# model：class R_MODEL
# 训练model：R_MODEL.train_reward_model()
# 生成trajectory：R_MODEL.trajectory()
# w=u(z): class S_MODEL (未完成），目前是直接读取solver.get_context()

# **** 环境变量w设定 ****
# w = np.array([ra, rd)]
# ra：速度系数（0.6~1)，乘在action上面
# rd：深度系数（0.6~1)，乘在depth上面

# **** 待定参数： ****
# solver.init() 参数 K N d n M L w_boundary ......
# env
# safe or not: 安全边界力大小 (solver.train(), model.trajectory())
# trajectory初始位置obs 及其方差大小 (model.trajectory())
# trajectory reward function (model.trajectory())


text:
# coding: utf-8
import pickle
import os
import numpy as np
import json
classes = ["line","stopline"]

def convert(size, box):
    dw = 1./(size[0])
    dh = 1./(size[1])
    x1 = box[0]*dw
    x2 = box[1]*dw
    y1 = box[2]*dh
    y2 = box[3]*dh
    return (x1,y1,x2,y2)

def convert_annotation(img_filename,xml_file_name):

    out_name = img_filename.replace(os.path.splitext(img_filename)[1], '.txt')

    if os.path.exists(xml_file_name):
        with open(xml_file_name,encoding='gbk') as f:
            data = json.load(f)
            w = int(data['imageWidth'])
            h = int(data['imageHeight'])
            out_file = open(out_name, 'w')
            if w == 0 or h == 0:
                out_file.close()
                print('invalid json file')
                return -1
            for shape in data['shapes']:
                if shape['shape_type'] != 'line':
                    print('Skipping shape: label={label}, shape_type={shape_type}'.format(**shape))
                    continue
                class_name = shape['label']
                class_id = classes.index(class_name)
                pt1, pt2 = shape['points']
                pt1 = np.array(pt1).astype('int')
                pt2 = np.array(pt2).astype('int')

                # 调换位置(左，右）
                if pt1[0] > pt2[0]:
                    pt1[0], pt2[0] = pt2[0], pt1[0]
                    pt1[1],pt2[1] = pt2[1],pt1[1]

                b = np.array([float(pt1[0]), float(pt2[0]), float(pt1[1]), float(pt2[1])]).astype('int')
                # if (b[1]-b[0] < 16) or (b[3] - b[2] < 16):#过滤过小的目标
                #     continue
                bb = convert((w,h), b)
                if class_id == 0:
                    out_file.write(str(class_id) + " " + " ".join([str(a) for a in bb]) + '\n')
            out_file.close()
    else:
        out_file = open(out_name, 'w')
        out_file.close()

if __name__ == '__main__':

    img_path = r'/media/hzh/work/workspace/yolov3-line_detect/data/coco'
    xml_path = r'/media/hzh/work/workspace/yolov3-line_detect/data/coco'
    suffix = '.jpg'

    img_file_list = os.listdir(img_path)
    xml_file_list = os.listdir(xml_path)
    img_file_list = [os.path.join(img_path,filename) for filename in img_file_list if filename.endswith('.jpg') or filename.endswith('.jpeg')]
    xml_file_list = [os.path.join(xml_path,filename) for filename in xml_file_list if filename.endswith('.json')]

    # #skip
    # skip = 85873
    # img_file_list = img_file_list[skip:]
    # xml_file_list = xml_file_list[skip:]

    idx = 1
    for img_filename in img_file_list:
        print('process ', img_filename,idx)
        idx += 1
        basename = os.path.split(img_filename)[1]
        xml_filename = os.path.join(xml_path, basename.replace('.jpg', '.json')).replace('.jpeg','.json')
        res = convert_annotation(img_filename,xml_filename)
        if res == -1:
            print(img_filename)

max_stars_repo_name: stwrd/ImageProcessTool

max_issues_repo_name: stwrd/ImageProcessTool

max_forks_repo_name: stwrd/ImageProcessTool