text:
# Copyright 2019 Allan Haldane.
#
# This file is part of Mi3-GPU.
#
# Mi3-GPU is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, version 3 of the License.
#
# Mi3-GPU is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Mi3-GPU.  If not, see <http://www.gnu.org/licenses/>.
#
#Contact: allan.haldane _AT_ gmail.com

import scipy
import numpy as np
from numpy.random import RandomState
import pyopencl as cl
import pyopencl.array as cl_array
import os, time, warnings
import textwrap
from utils import printsome
import collections

cf = cl.mem_flags

rng_buf_mul = 1024

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

os.environ['PYOPENCL_COMPILER_OUTPUT'] = '0'
os.environ['PYOPENCL_NO_CACHE'] = '1'
os.environ["CUDA_CACHE_DISABLE"] = '1'

def printPlatform(log, p, n=0):
    log("Platform {} '{}':".format(n, p.name))
    log("    Vendor: {}".format(p.vendor))
    log("    Version: {}".format(p.version))
    exts = ("\n" + " "*16).join(textwrap.wrap(p.extensions, 80-16))
    log("    Extensions: {}".format(exts))

def printDevice(log, d):
    log("  Device '{}':".format(d.name))
    log("    Vendor: {}".format(d.vendor))
    log("    Version: {}".format(d.version))
    log("    Driver Version: {}".format(d.driver_version))
    log("    Max Clock Frequency: {}".format(d.max_clock_frequency))
    log("    Max Compute Units: {}".format(d.max_compute_units))
    log("    Max Work Group Size: {}".format(d.max_work_group_size))
    log("    Global Mem Size: {}".format(d.global_mem_size))
    log("    Global Mem Cache Size: {}".format(d.global_mem_cache_size))
    log("    Local Mem Size: {}".format(d.local_mem_size))
    log("    Max Constant Buffer Size: {}".format(d.max_constant_buffer_size))

def printGPUs(log):
    for n,p in enumerate(cl.get_platforms()):
        printPlatform(log, p, n)
        for d in p.get_devices():
            printDevice(log, d)
        log("")

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

# The GPU performs two main types of computation: MCMC runs, and perturbed
# coupling updates. MCMCGPU methods are asynchronous on the host.
# Functions that return data do not return the data directly, but return a
# FutureBuf object. The data may be obtained by FutureBuf.read(), which is
# blocking.

# The gpu has two sequence buffers: A "small" buffer for MCMC gpu generation,
# and an optional "large buffer" for combined sequence sets.

# The opencl queue is created as an out-of-order queue, and so kernel order is
# managed by the MCMCGPU class itself. By default, it makes all opencl
# commands wait until the last command is finished, but all methods also
# have a wait_for argument to override this. `None` means wait until the last
# command is done, or it can be a list of opencl events to wait for. Set it
# to the empty list [] to run immediately.

# Note that in openCL implementations there is generally a limit on the number
# of queued commands allowed in a context. If you reach the limit, all queues
# will block until a kernel finishes. So we must be careful not to fill up a
# single queue before others, so do `for i in range(100): for g in gpus:
# g.command()` instead of `for g in gpus: for i in range(100): g.command()` as
# the latter may fill the first gpu's queue, blocking the rest.
# See  CL_DEVICE_QUEUE_ON_DEVICE_MAX_SIZE  and  CL_DEVICE_MAX_ON_DEVICE_EVENTS

# Note that on some systems there is a watchdog timer that kills any kernel
# that runs for too long. You will get a CL_OUT_OF_RESOURCES error if this
# happens, which occurs when the *following* kernel is run.

class FutureBuf:
    def __init__(self, buffer, event, postprocess=None):
        self.buffer = buffer
        self.event = event
        self.postfunc = postprocess

        self.shape = buffer.shape
        self.dtype = buffer.dtype

    def read(self):
        self.event.wait()
        if self.postfunc != None:
            return self.postfunc(self.buffer)
        return self.buffer

class MCMCGPU:
    def __init__(self, gpuinfo, L, q, nseq, wgsize, outdir,
                 vsize, seed, profile=False):
        if nseq%512 != 0:
            raise ValueError("nwalkers/ngpus must be a multiple of 512")
            # this guarantees that all kernel access to seqmem is coalesced and
            # simplifies the histogram kernels

        self.L = L
        self.q = q
        self.nPairs = L*(L-1)//2
        self.events = collections.deque()
        self.SWORDS = ((L-1)//4+1)    #num words needed to store a sequence
        self.SBYTES = (4*self.SWORDS) #num bytes needed to store a sequence
        self.nseq = {'main': nseq}
        self.nwalkers = nseq

        device, gpunum, ctx, prg = gpuinfo
        self.gpunum = gpunum
        self.ctx = ctx
        self.prg = prg
        self.device = device

        self.wgsize = wgsize
        self.nhist, self.histws = histogram_heuristic(q)

        # sanity checks (should be checked elsewhere before this)
        if nseq%wgsize != 0:
            raise Exception("nseq per GPU must be a multiple of wgsize")
        if wgsize < q*q:
            raise Exception("wgsize cannot be less than q*q")

        self.logfn = os.path.join(outdir, 'gpu-{}.log'.format(gpunum))
        with open(self.logfn, "wt") as f:
            printDevice(f.write, device)

        self.mcmcprg = prg.metropolis

        self.rngstate = RandomState(seed)

        #setup opencl for this device
        self.log("Getting CL Queue")

        qprop = cl.command_queue_properties.OUT_OF_ORDER_EXEC_MODE_ENABLE
        self.profile = profile
        if profile:
            qprop |= cl.command_queue_properties.PROFILING_ENABLE
        self.queue = cl.CommandQueue(ctx, device=device, properties=qprop)

        self.log("\nOpenCL Device Compilation Log:")
        self.log(self.prg.get_build_info(device, cl.program_build_info.LOG))
        maxwgs = self.mcmcprg.get_work_group_info(
                 cl.kernel_work_group_info.WORK_GROUP_SIZE, device)
        self.log("Max MCMC WGSIZE: {}".format(maxwgs))

        self.initted = []

        self.bufs = {}
        self.buf_spec = {}
        self.seqbufs = {}
        self.Ebufs = {}
        self.largebufs = []

        # setup essential buffers
        nPairs, SWORDS = self.nPairs, self.SWORDS
        j_pad = 3*self.wgsize
        self._setupBuffer(        'J', '<f4', (nPairs, q*q), pad=j_pad)
        self._setupBuffer('Junpacked', '<f4', (L*L, q*q), pad=j_pad)
        self._setupBuffer(       'bi', '<f4', (nPairs, q*q)),
        self._setupBuffer(  'bicount', '<u4', (nPairs, q*q)),
        self._setupBuffer( 'seq main', '<u4', (SWORDS, self.nseq['main'])),
        self._setupBuffer('seqL main', '<u4', (L, self.nseq['main']//4)),
        self._setupBuffer(   'E main', '<f4', (self.nseq['main'],)),
        self.unpackedJ = False #use to keep track of whether J is unpacked
        self.repackedSeqT = {'main': False}

        self.lastevt = None

    def log(self, msg):
        #logs are rare, so just open the file every time
        with open(self.logfn, "at") as f:
            print("{: 10.3f}".format(time.process_time()), msg, file=f)

    def logevt(self, name, evt, nbytes=None):
        self.lastevt = evt

        # don't save events if not profiling.
        # note that saved events use up memory - free it using logprofile
        if self.profile:
            if len(self.events)%1000 == 0 and len(self.events) != 0:
                self.log("Warning: Over {} profiling events are not flushed "
                    "(using up memory)".format(len(self.events)))
            if nbytes:
                self.events.append((evt, name, nbytes))
            else:
                self.events.append((evt, name))

        return evt

    def _evtlist(self, evts):
        if evts is None:
            return []
        elif isinstance(evts, cl.Event):
            return [evts]
        else:
            return evts

    def _waitevt(self, evts=None):
        if evts is None and self.lastevt is not None:
            return [self.lastevt]
        return self._evtlist(evts)

    def logProfile(self):
        if not self.profile:
            return

        def isComplete(e):
            return (e.command_execution_status ==
                    cl.command_execution_status.COMPLETE)

        with open(self.logfn, "at") as f:
            while len(self.events) != 0 and isComplete(self.events[0][0]):
                dat = self.events.popleft()
                evt, name, size = dat[0],dat[1],(dat[2] if len(dat)==3 else '')
                print("EVT", name, evt.profile.start, evt.profile.end,
                      size, file=f)

    def _setupBuffer(self, bufname, buftype, bufshape, pad=None,
                     flags=cf.READ_WRITE):
        flags = flags | cf.ALLOC_HOST_PTR

        nelem = int(np.product(bufshape))
        if pad:
            nelem = nelem + pad

        size = np.dtype(buftype).itemsize * nelem

        buf = cl.Buffer(self.ctx, flags, size=size)

        self.bufs[bufname] = buf
        self.buf_spec[bufname] = (buftype, bufshape, flags)

        # add it to convenience dicts if applicable
        names = bufname.split()
        if len(names) > 1:
            bufs = {'seq': self.seqbufs, 'E': self.Ebufs}
            if names[0] in bufs:
                bufs[names[0]][names[1]] = buf

    def require(self, *reqs):
        for r in reqs:
            if r not in self.initted:
                raise Exception("{} not initialized".format(r))

    def _initcomponent(self, cmp):
        if cmp in self.initted:
            raise Exception("Already initialized {}".format(cmp))
        self.initted.append(cmp)

    def initMCMC(self, nsteps, rng_offset, rng_span):
        self._initcomponent('MCMC')

        # rngstates should be size of mwc64xvec2_state_t
        self.nsteps = nsteps
        self._setupBuffer('rngstates', '<2u8', (self.nseq['main'],)),
        self._setupBuffer(       'Bs', '<f4',  (self.nseq['main'],)),
        self._setupBuffer(  'randpos', '<u4',  (self.nsteps*rng_buf_mul,))
        self.randpos_offset = rng_buf_mul*self.nsteps

        self.setBuf('Bs', np.ones(self.nseq['main'], dtype='<f4'))
        self._initMCMC_RNG(rng_offset, rng_span)
        self.nsteps = int(nsteps)

    def initLargeBufs(self, nseq_large):
        self._initcomponent('Large')

        self.nseq['large'] = nseq_large
        self._setupBuffer(    'seq large', '<u4', (self.SWORDS, nseq_large))
        self._setupBuffer(   'seqL large', '<u4', (self.L, nseq_large//4)),
        self._setupBuffer(      'E large', '<f4', (nseq_large,))
        self._setupBuffer('weights large', '<f4', (nseq_large,))

        self.largebufs.extend(['seq large', 'seqL large', 'E large',
                               'weights large'])
        self.nstoredseqs = 0

        # it is important to zero out the large seq buffer, because
        # if it is partially full we may need to compute energy
        # over the padding sequences at the end to get a full wg.
        buf = self.bufs['seq large']
        self.fillBuf('seq large', 0)
        self.repackedSeqT['large'] = False

    def initSubseq(self):
        self.require('Large')
        self._initcomponent('Subseq')
        self._setupBuffer('markpos', '<u1',  (self.SBYTES,), flags=cf.READ_ONLY)
        self.markPos(np.zeros(self.SBYTES, '<u1'))

    ## we may want to select replicas at a particular temperature
    #def initMarkSeq(self):
    #    self._initcomponent('Markseq')

    #    self._setupBuffer( 'markseqs', '<i4',  (self.nseq['main'],))
    #    self.setBuf('markseqs', np.arange(self.nseq['main'], dtype='<i4'))
    #    self.nmarks = self.nseq['main']

    def initJstep(self):
        self._initcomponent('Jstep')

        nPairs, q = self.nPairs, self.q
        self._setupBuffer(       'dJ', '<f4',  (nPairs, q*q))
        self._setupBuffer('bi target', '<f4',  (nPairs, q*q))
        self._setupBuffer(     'Creg', '<f4',  (nPairs, q*q))
        self._setupBuffer( 'Xlambdas', '<f4',  (nPairs,))
        self._setupBuffer(     'neff', '<f4',  (1,))
        self._setupBuffer(  'weights', '<f4',  (self.nseq['main'],))

    def packSeqs_4(self, seqs):
        """
        Converts seqs to 4-byte uint format on CPU, padded to 32bits, assumes
        little endian. Each row's bytes are
            a0 a1 a2 a3 b0 b1 b2 b3 c0 c1 c2 c3 ...
        for sequences a, b, c, so each uint32 correaponds to 4 seq bytes.
        """
        if seqs.dtype != np.dtype('<u1'):
            raise Exception("seqs must have u1 dtype")
        bseqs = np.zeros((seqs.shape[0], self.SBYTES), dtype='<u1', order='C')
        bseqs[:,:self.L] = seqs
        mem = np.zeros((self.SWORDS, seqs.shape[0]), dtype='<u4', order='C')
        for i in range(self.SWORDS):
            mem[i,:] = bseqs.view(np.uint32)[:,i]
        return mem

    def unpackSeqs_4(self, mem):
        """ reverses packSeqs_4 (on CPU)"""
        bseqs = np.zeros((mem.shape[1], self.SBYTES), dtype='<u1', order='C')
        for i in range(self.SWORDS): #undo memory rearrangement
            bseqs.view(np.uint32)[:,i] = mem[i,:]
        return bseqs[:,:self.L]

    def repackseqs_T(self, bufname, wait_for=None):
        """
        On GPU, copies the seq buffer (in 4-byte format) to a seqL buffer
        in "transpose" format, which is just the usual CPU sequence buffer
        but transposed.
        """
        self.log("repackseqs_T")

        nseq = self.nseq[bufname]
        inseq_dev = self.bufs['seq ' + bufname]
        outseq_dev = self.bufs['seqL ' + bufname]

        self.repackedSeqT[bufname] = True
        return self.logevt('repackseqs_T',
            self.prg.unpackseqs1(self.queue, (self.SWORDS*256,), (256,),
                            inseq_dev, np.uint32(nseq),
                            outseq_dev, np.uint32(nseq//4),
                            wait_for=self._waitevt(wait_for)))

    def unpackJ(self, wait_for=None):
        """convert J from format where every row is a unique ij pair (L choose 2
        rows) to format with every pair, all orders (L^2 rows)."""

        # quit if J already loaded/unpacked
        if self.unpackedJ:
            return wait_for

        self.log("unpackJ")

        q, nPairs = self.q, self.nPairs
        self.unpackedJ = True
        return self.logevt('unpackJ',
            self.prg.unpackfV(self.queue, (nPairs*q*q,), (q*q,),
                            self.bufs['J'], self.bufs['Junpacked'],
                            wait_for=self._waitevt(wait_for)))

    def _initMCMC_RNG(self, rng_offset, rng_span, wait_for=None):
        self.require('MCMC')
        self.log("initMCMC_RNG")

        # Idea is we want to divide the rng stream into non-overlapping chunks
        # for each walker. This GPU was given a span of rng_span, so divide it
        # by number of walkers (times 2 since each rng call advances by 2 for
        # vec2).

        # Read mwc64 docs for more info. mwc64 doc says each walker's stream
        # offset is (rng_offset + walker_span*(get_global_id(0)*vectorSize +
        # vecind)) where vectorSize is 2. Note that rng_offset is not
        # multiplied by 2!

        rng_offset = np.uint64(rng_offset)

        nwalkers = np.uint64(self.nseq['main'])
        v2 = np.uint64(2)
        # walker span is the # of rng calls assigned per walker
        walker_span = np.uint64(rng_span)//(v2*nwalkers) # factor of 2 for vec2
        self.log("RNG offset: {}  walker-span: {}  nwalkers {}".format(
                  rng_offset, walker_span, nwalkers))

        # Warning: It is very important that the walker rng stream offsets
        # across gpus are all distinct, or else some walkers will be highly
        # correlated. Touch this code with care.
        assert(walker_span*v2*nwalkers <= rng_span)

        wgsize = self.wgsize
        while wgsize > nwalkers:
            wgsize = wgsize//2
        return self.logevt('initMCMC_RNG',
            self.prg.initRNG2(self.queue, (nwalkers,), (wgsize,),
                         self.bufs['rngstates'],
                         np.uint64(rng_offset), walker_span,
                         wait_for=self._waitevt(wait_for)))

    def updateRngPos(self, wait_evt=None):
        self.randpos_offset = self.randpos_offset + self.nsteps
        rng_evt = None

        bufsize = rng_buf_mul*self.nsteps
        if self.randpos_offset >= bufsize:
            # all gpus use same position-rng series. This way there is no
            # difference between running on one gpu vs splitting on multiple
            rng = self.rngstate.randint(0, self.L, size=bufsize).astype('u4')
            rng_evt = self.setBuf('randpos', rng, wait_for=wait_evt)
            self.randpos_offset = 0
        return np.uint32(self.randpos_offset), rng_evt

    def runMCMC(self, wait_for=None):
        """Performs a single round of mcmc sampling (nsteps MC steps)"""
        t1 = time.time()
        self.require('MCMC')
        self.log("runMCMC")

        wait_evt = self._waitevt(wait_for)

        nseq = self.nseq['main']
        nsteps = self.nsteps
        wait_unpack = self.unpackJ(wait_for=wait_evt)
        rngoffset, wait_rng = self.updateRngPos(wait_evt)

        wait = self._evtlist(wait_unpack) + self._evtlist(wait_rng)

        self.repackedSeqT['main'] = False
        return self.logevt('mcmc',
            self.mcmcprg(self.queue, (nseq,), (self.wgsize,),
                         self.bufs['Junpacked'], self.bufs['rngstates'],
                         rngoffset, self.bufs['randpos'], np.uint32(nsteps),
                         self.Ebufs['main'], self.bufs['Bs'],
                         self.seqbufs['main'],
                         wait_for=wait))

    def measureFPerror(self, log, nloops=3):
        log("Measuring FP Error")
        for n in range(nloops):
            self.runMCMC()
            e1 = self.getBuf('E main').read()
            self.calcEnergies('main')
            e2 = self.getBuf('E main').read()
            log("Run", n, "Error:", np.mean((e1-e2)**2))
            log('    Final E MC', printsome(e1), '...')
            log("    Final E rc", printsome(e2), '...')

            seqs = self.getBuf('seq main').read()
            J = self.getBuf('J').read()
            e3 = getEnergies(seqs, J)
            log("    Exact E", e3[:5])
            log("    Error:", np.mean([float((a-b)**2) for a,b in zip(e1, e3)]))

    def calcBicounts(self, seqbufname, wait_for=None):
        self.log("calcBicounts " + seqbufname)
        L, q, nPairs, nhist = self.L, self.q, self.nPairs, self.nhist

        if seqbufname == 'main':
            nseq = self.nseq[seqbufname]
            buflen = nseq
        else:
            nseq = self.nstoredseqs
            buflen = self.nseq[seqbufname]
        seq_dev = self.seqbufs[seqbufname]

        localhist = cl.LocalMemory(nhist*q*q*np.dtype(np.uint32).itemsize)
        return self.logevt('calcBicounts',
            self.prg.countBivariate(self.queue, (nPairs*nhist,), (nhist,),
                     self.bufs['bicount'],
                     np.uint32(nseq), seq_dev, np.uint32(buflen), localhist,
                     wait_for=self._waitevt(wait_for)))

    def bicounts_to_bimarg(self, seqbufname='main', wait_for=None):
        self.log("bicounts_to_bimarg ")
        q, nPairs = self.q, self.nPairs

        if seqbufname == 'main':
            nseq = self.nseq['main']
        else:
            nseq = self.nstoredseqs

        nworkunits = self.wgsize*((nPairs*q*q-1)//self.wgsize+1)
        return self.logevt('bicounts_to_bimarg',
            self.prg.bicounts_to_bimarg(self.queue,
                     (nworkunits,), (self.wgsize,),
                     self.bufs['bicount'], self.bufs['bi'], np.uint32(nseq),
                     wait_for=self._waitevt(wait_for)))

    def calcEnergies(self, seqbufname, Jbufname='J', wait_for=None):
        self.log("calcEnergies " + seqbufname)

        energies_dev = self.Ebufs[seqbufname]
        seq_dev = self.seqbufs[seqbufname]
        buflen = self.nseq[seqbufname]

        if seqbufname == 'main':
            nseq = self.nseq[seqbufname]
        else:
            nseq = self.nstoredseqs
            # pad to be a multiple of wgsize (uses dummy seqs at end)
            nseq = nseq + ((self.wgsize - nseq) % self.wgsize)

        return self.logevt('getEnergies',
            self.prg.getEnergies(self.queue, (nseq,), (self.wgsize,),
                             self.bufs[Jbufname], seq_dev, np.uint32(buflen),
                             energies_dev, wait_for=self._waitevt(wait_for)))

    def weightedMarg(self, seqbufname='main', wait_for=None):
        self.require('Jstep')
        self.log("weightedMarg")
        q, L, nPairs = self.q, self.L, self.nPairs
        nhist, histws = self.nhist, self.histws

        if seqbufname == 'main':
            nseq = self.nseq[seqbufname]
            buflen = nseq//4
            weights_dev = self.bufs['weights']
        else:
            nseq = self.nstoredseqs
            buflen = self.nseq[seqbufname]//4
            weights_dev = self.bufs['weights large']
            # pad to be a multiple of 512 (uses dummy seqs at end)
            nseq = nseq + ((512 - nseq) % 512)
            # XXX find correct padding here

        if not self.repackedSeqT[seqbufname]:
            wait_for = self.repackseqs_T(seqbufname,
                                      wait_for=self._waitevt(wait_for))
        seq_dev = self.bufs['seqL ' + seqbufname]

        return self.logevt('weightedMarg',
            self.prg.weightedMarg(self.queue, (nPairs*histws,), (histws,),
                        self.bufs['bi'], weights_dev,
                        np.uint32(nseq), seq_dev, np.uint32(buflen),
                        wait_for=self._waitevt(wait_for)))

    def renormalize_bimarg(self, wait_for=None):
        self.log("renormalize_bimarg")
        q, nPairs = self.q, self.nPairs
        return self.logevt('renormalize_bimarg',
            self.prg.renormalize_bimarg(self.queue, (nPairs*q*q,), (q*q,),
                         self.bufs['bi'], wait_for=self._waitevt(wait_for)))

    def addBiBuffer(self, bufname, otherbuf, wait_for=None):
        # used for combining results from different gpus, where  otherbuf is a
        # buffer "belonging" to another gpu
        self.log("addbuf")

        selfbuf = self.bufs[bufname]
        if selfbuf.size != otherbuf.size:
            raise Exception('Tried to add bufs of different sizes')

        q, nPairs = self.q, self.nPairs
        nworkunits = self.wgsize*((nPairs*q*q-1)//self.wgsize+1)

        return self.logevt('addbuf',
            self.prg.addBiBufs(self.queue, (nworkunits,), (self.wgsize,),
                       selfbuf, otherbuf, wait_for=self._waitevt(wait_for)))

    def updateJ(self, gamma, pc, Jbuf='dJ', wait_for=None):
        self.require('Jstep')
        self.log("updateJ")
        q, nPairs = self.q, self.nPairs
        #find next highest multiple of wgsize, for num work units
        nworkunits = self.wgsize*((nPairs*q*q-1)//self.wgsize+1)

        bibuf = self.bufs['bi']
        Jin = Jout = self.bufs[Jbuf]
        self.unpackedJ = False
        return self.logevt('updateJ',
            self.prg.updatedJ(self.queue, (nworkunits,), (self.wgsize,),
                                self.bufs['bi target'], bibuf,
                                np.float32(gamma), np.float32(pc), Jin, Jout,
                                wait_for=self._waitevt(wait_for)))

    def reg_l1z(self, gamma, pc, lJ, wait_for=None):
        self.require('Jstep')
        self.log("reg_l1z")
        q, nPairs = self.q, self.nPairs

        bibuf = self.bufs['bi']
        self.unpackedJ = None
        return self.logevt('reg_l1z',
            self.prg.reg_l1z(self.queue, (nPairs*q*q,), (q*q,),
                            bibuf, np.float32(gamma), np.float32(pc),
                            np.float32(lJ), self.bufs['J'], self.bufs['dJ'],
                            wait_for=self._waitevt(wait_for)))

    def reg_l2z(self, gamma, pc, lJ, wait_for=None):
        self.require('Jstep')
        self.log("reg_l2z")
        q, nPairs = self.q, self.nPairs

        bibuf = self.bufs['bi']
        self.unpackedJ = None
        return self.logevt('reg_l2z',
            self.prg.reg_l2z(self.queue, (nPairs*q*q,), (q*q,),
                            bibuf, np.float32(gamma), np.float32(pc),
                            np.float32(lJ),
                            self.bufs['J'], self.bufs['dJ'],
                            wait_for=self._waitevt(wait_for)))

    def reg_X(self, gamma, pc, wait_for=None):
        self.require('Jstep')
        self.log("reg X")
        q, nPairs = self.q, self.nPairs

        bibuf = self.bufs['bi']
        self.unpackedJ = None
        return self.logevt('reg_X',
            self.prg.reg_X(self.queue, (nPairs*q*q,), (q*q,),
                                bibuf, self.bufs['Creg'],
                                np.float32(gamma), np.float32(pc),
                                self.bufs['J'], self.bufs['dJ'],
                                wait_for=self._waitevt(wait_for)))

    def reg_ddE(self, gamma, pc, lam, wait_for=None):
        self.require('Jstep')
        self.log("reg ddE")
        q, nPairs = self.q, self.nPairs

        bibuf = self.bufs['bi']
        self.unpackedJ = None
        return self.logevt('reg_ddE',
            self.prg.reg_ddE(self.queue, (nPairs*q*q,), (q*q,),
                                bibuf, np.float32(gamma), np.float32(pc),
                                np.float32(lam),
                                self.bufs['J'], self.bufs['dJ'],
                                wait_for=self._waitevt(wait_for)))

    def getBuf(self, bufname, truncateLarge=True, wait_for=None):
        """get buffer data. truncateLarge means only return the
        computed part of the large buffer (rest may be uninitialized)"""

        self.log("getBuf " + bufname)
        bufspec = self.buf_spec[bufname]
        buftype, bufshape = bufspec[0], bufspec[1]
        mem = np.zeros(bufshape, dtype=buftype)
        evt = cl.enqueue_copy(self.queue, mem, self.bufs[bufname],
                          is_blocking=False, wait_for=self._waitevt(wait_for))
        self.logevt('getBuf', evt, mem.nbytes)
        if bufname.split()[0] == 'seq':
            if bufname in self.largebufs and truncateLarge:
                nret = self.nstoredseqs
                return FutureBuf(mem, evt,
                                 lambda b: self.unpackSeqs_4(b)[:nret,:])
            return FutureBuf(mem, evt, self.unpackSeqs_4)

        if bufname in self.largebufs and truncateLarge:
            nret = self.nstoredseqs
            return FutureBuf(mem, evt, lambda b: b[:nret])

        return FutureBuf(mem, evt)

    def setBuf(self, bufname, buf, wait_for=None):
        self.log("setBuf " + bufname)

        # device-to-device copies skip all the checks
        if isinstance(buf, cl.Buffer):
            evt = cl.enqueue_copy(self.queue, self.bufs[bufname], buf,
                                  wait_for=self._waitevt(wait_for))
            self.logevt('setBuf', evt, buf.size)
            if bufname.split()[0] == 'J':
                self.unpackedJ = None
            return  evt

        if bufname.split()[0] == 'seq':
            buf = self.packSeqs_4(buf)

        bufspec = self.buf_spec[bufname]
        buftype, bufshape = bufspec[0], bufspec[1]
        if not isinstance(buf, np.ndarray):
            buf = array(buf, dtype=buftype)

        if np.dtype(buftype) != buf.dtype:
            raise ValueError("Buffer dtype mismatch.Expected {}, got {}".format(
                             np.dtype(buftype), buf.dtype))
        if bufshape != buf.shape and not (bufshape == (1,) or buf.size == 1):
            raise ValueError("Buffer size mismatch. Expected {}, got {}".format(
                            bufshape, buf.shape))

        evt = cl.enqueue_copy(self.queue, self.bufs[bufname], buf,
                            is_blocking=False, wait_for=self._waitevt(wait_for))
        self.logevt('setBuf', evt, buf.nbytes)
        #unset packedJ flag if we modified that J buf
        if bufname.split()[0] == 'J':
            self.unpackedJ = None
        if bufname == 'seq large':
            self.nstoredseqs = bufshape[1]
        if bufname.split()[0] == 'seq':
            self.repackedSeqT[bufname.split()[1]] = False

        return evt

    def fillBuf(self, bufname, val, wait_for=None):
        self.log("fillBuf " + bufname)

        buf = self.bufs[bufname]
        buftype = np.dtype(self.buf_spec[bufname][0]).type

        self.logevt('fill_buffer',
            cl.enqueue_fill_buffer(self.queue, buf, buftype(val), 0, buf.size,
                                   wait_for=self._waitevt()))

    def markPos(self, marks, wait_for=None):
        self.require('Subseq')

        marks = marks.astype('<u1')
        if len(marks) == self.L:
            marks.resize(self.SBYTES)
        return self.setBuf('markpos', marks, wait_for=wait_for)

    def fillSeqs(self, startseq, seqbufname='main', wait_for=None):
        #write a kernel function for this?
        self.log("fillSeqs " + seqbufname)
        nseq = self.nseq[seqbufname]
        self.setBuf('seq '+seqbufname, np.tile(startseq, (nseq,1)),
                    wait_for=wait_for)

    def storeSeqs(self, seqs=None, wait_for=None):
        """
        If seqs is None, stores main to large seq buffer. Otherwise
        stores seqs to large buffer
        """

        self.require('Large')

        offset = self.nstoredseqs
        self.log("storeSeqs " + str(offset))

        if seqs is not None:
            nseq, L = seqs.shape
            if L != self.L:
                raise Exception(
                    "Sequences have wrong length: {} vs {}".format(L, self.L))
            if offset + nseq > self.nseq['large']:
                raise Exception("cannot store seqs past end of large buffer")
            assert(seqs.dtype == np.dtype('u1'))
            buf = self.packSeqs_4(seqs)

            w, h = self.buf_spec['seq large'][1] # L/4, nseq
            # for some reason, rectangular copies in pyOpencl use opposite axis
            # order from numpy, and need indices in bytes not elements, so we
            # have to switch all this around. buf is uint32, or 4 bytes.
            evt = cl.enqueue_copy(self.queue, self.seqbufs['large'], buf,
                                  buffer_origin=(4*offset, 0),
                                  host_origin=(0, 0),
                                  region=(4*buf.shape[1], buf.shape[0]),
                                  buffer_pitches=(4*h, w),
                                  host_pitches=(4*buf.shape[1], buf.shape[0]),
                                  is_blocking=False,
                                  wait_for=self._waitevt(wait_for))
        else:
            nseq = self.nseq['main']
            if offset + nseq > self.nseq['large']:
                raise Exception("cannot store seqs past end of large buffer")
            evt = self.prg.storeSeqs(self.queue, (nseq,), (self.wgsize,),
                               self.seqbufs['main'], self.seqbufs['large'],
                               np.uint32(self.nseq['large']), np.uint32(offset),
                               wait_for=self._waitevt(wait_for))

        self.nstoredseqs += nseq
        self.repackedSeqT['large'] = False
        return self.logevt('storeSeqs', evt)

    def clearLargeSeqs(self):
        self.require('Large')
        self.nstoredseqs = 0
        self.repackedSeqT['large'] = False

    def restoreSeqs(self, wait_for=None):
        self.require('Large')
        self.log("restoreSeqs " + str(offset))
        nseq = self.nseq['main']

        if offset + nseq > self.nseq['large']:
            raise Exception("cannot get seqs past end of large buffer")
        if self.nstoredseqs < nseq:
            raise Exception("not enough seqs stored in large buffer")

        self.repackedSeqT['main'] = False
        return self.logevt('restoreSeqs',
            self.prg.restoreSeqs(self.queue, (nseq,), (self.wgsize,),
                           self.seqbufs['main'], self.seqbufs['large'],
                           np.uint32(self.nseq['large']), np.uint32(offset),
                           wait_for=self._waitevt(wait_for)))

    def copySubseq(self, seqind, wait_for=None):
        self.require('Subseq')
        self.log("copySubseq " + str(seqind))
        nseq = self.nseq['large']
        if seqind >= self.nseq['main']:
            raise Exception("given index is past end of main seq buffer")
        self.repackedSeqT['large'] = False
        return self.logevt('copySubseq',
            self.prg.copySubseq(self.queue, (nseq,), (self.wgsize,),
                            self.seqbufs['main'], self.seqbufs['large'],
                            np.uint32(self.nseq['main']), np.uint32(seqind),
                            self.bufs['markpos'],
                            wait_for=self._waitevt(wait_for)))

    def wait(self):
        self.log("wait")
        self.queue.finish()

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

def unpackJ_CPU(self, couplings):
    """convert from format where every row is a unique ij pair (L choose 2
    rows) to format with every pair, all orders (L^2 rows). Note that the
    GPU kernel unpackfV does the same thing faster"""

    L, q = seqsize_from_param_shape(couplings.shape)
    fullcouplings = np.zeros((L*L,q*q), dtype='<f4', order='C')
    pairs = [(i,j) for i in range(L-1) for j in range(i+1,L)]
    for n,(i,j) in enumerate(pairs):
        c = couplings[n,:]
        fullcouplings[L*i + j,:] = c
        fullcouplings[L*j + i,:] = c.reshape((q,q)).T.flatten()
    return fullcouplings

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

def setup_GPU_context(scriptpath, scriptfile, param, log):
    outdir = param.outdir
    L, q = param.L, param.q
    gpuspec = param.gpuspec
    measureFPerror = param.fperror

    with open(scriptfile) as f:
        src = f.read()

    #figure out which gpus to use
    gpudevices = []
    platforms = [(p, p.get_devices()) for p in cl.get_platforms()]
    if gpuspec is not None:
        try:
            dev = [tuple(int(x) for x in a.split(':'))
                                for a in gpuspec.split(',')]
        except:
            raise Exception("Error: GPU specification must be comma separated "
                            " list of platform:gpus, eg '0:0,0:1'")

        for p,d in dev:
            try:
                plat, devices = platforms[p]
                gpu = devices[d]
            except:
                raise Exception("No GPU with specification {}".format(d))
            gpudevices.append(gpu)
    else:
        #use gpus in first platform
        plat = platforms[0]
        for gpu in plat[1]:
            gpudevices.append(gpu)

    if len(gpudevices) == 0:
        raise Exception("Error: No GPUs found")

    #set up OpenCL. Assumes all gpus are identical
    log("Getting CL Context...")
    cl_ctx = cl.Context(gpudevices)

    nhist, histws = histogram_heuristic(q)

    #compile CL program
    options = [('q', q), ('L', L), ('NHIST', nhist), ('HISTWS', histws),
               ('WGSIZE', param.wgsize)]
    if measureFPerror:
        options.append(('MEASURE_FP_ERROR', 1))
    optstr = " ".join(["-D {}={}".format(opt,val) for opt,val in options])
    log("Compilation Options: ", optstr)
    extraopt = " -cl-nv-verbose -Werror -I {}".format(scriptpath)
    log("Compiling CL...")

    with warnings.catch_warnings():
        warnings.simplefilter("ignore", cl.CompilerWarning)
        cl_prg = cl.Program(cl_ctx, src).build(optstr + extraopt)

    # dump compiled program
    ptx = cl_prg.get_info(cl.program_info.BINARIES)[0].decode('utf-8')
    # get compile log (Nvidia truncates this at 4096 bytes.. annoying)
    compile_log = cl_prg.get_build_info(gpudevices[0], cl.program_build_info.LOG)

    return (cl_ctx, cl_prg), gpudevices, (ptx, compile_log)

def initGPU(devnum, cldat, device, nwalkers, param, log):
    cl_ctx, cl_prg = cldat
    outdir = param.outdir
    L, q = param.L, param.q
    profile = param.profile
    wgsize = param.wgsize
    seed = param.rngseed

    # wgsize = OpenCL work group size for MCMC kernel.
    # (also for other kernels, although would be nice to uncouple them)
    if wgsize not in [1<<n for n in range(32)]:
        raise Exception("wgsize must be a power of two")

    vsize = 1024 #power of 2. Work group size for 1d vector operations.

    gpu = MCMCGPU((device, devnum, cl_ctx, cl_prg), L, q,
                  nwalkers, wgsize, outdir, vsize, seed, profile=profile)
    return gpu

def wgsize_heuristic(q, wgsize='auto'):
    if wgsize == 'auto':
        wgsize = 256
        while wgsize < q*q:
            wgsize *= 2
        return wgsize

    wgsize = int(wgsize)
    if wgsize < q*q:
        raise Exception("Must have wgsize >= q*q, but "
                        "got {} < {}".format(wgsize, q*q))
    return wgsize

def histogram_heuristic(q):
    """
    Choose histogram size parameters (NHIST, HISTWS) for the bimarg GPU
    calculations.

    Each histogram is q*q float32, and we want to squeeze as many as possible
    into local memory (eg, 96k), such that multiple workgroups can run at
    once. Strategy below figues out # of histograms which can fit into 16k.
    Then we figure out the optimal wg size histws that does not waste too many
    work units, since in the worst part of the kernel only nhist wu are
    running. We also want nhist, histws to be powers of 2, and histws > nhist.
    """
    nhist = 4096//(q*q)
    if nhist == 0:
        raise Exception("alphabet size too large to make histogram on gpu")
    nhist = 2**int(np.log2(nhist)) # closest power of two

    # this seems like a roughly good heuristic on Titan X.
    if q <= 12:
        hist_ws = 512
    elif q <= 16:
        hist_ws = 256
    else:
        hist_ws = 128

    if nhist > hist_ws:
        nhist = hist_ws

    return nhist, hist_ws

reason: import numpy,from numpy,import scipy


text:
"""
On top of the :py:class:`ActorCritic` implementation, this module
provides a couple algorithms to solve a problem stated in terms of
Reinforcement Learning. All algorithms follow the same approach, namely
(action dependent) Heuristic Dynamic Programming. The baseline
algorithm is implemented in :py:class:`ADHDP`.

If a new specialisation of an :py:class:`ActorCritic` is created,
typically its :py:meth:`ActorCritic._step` method is adapted (this is
for example the case in :py:class:`ADHDP`). If so, the two methods
:py:meth:`ActorCritic._pre_increment_hook` (before returning) and
:py:meth:`ActorCritic._next_action_hook` (after computation of the
next action) should be called as other structures may rely on those.

Some variations of the baseline algorithm are implemented as well in
:py:class:`ActionGradient`, :py:class:`ActionRecomputation` and
:py:class:`ActionBruteForce`. They fulfill the same purpose but approach
it differently (specifically, the actor is implemented differently). The
details are given in details of the respective class.

"""
import numpy as np
import PuPy
from rl import ActorCritic
import warnings

class ADHDP(ActorCritic):
    """Action dependent Heuristic Dynamic Programming structure and
    basic algorithm implementation.

    In the :py:meth:`_step` method, this class provides the
    implementation of a baseline algorithm. By default, the behaviour
    is online, i.e. the critic is trained and the actor in effect.
    Note that the actor can be modified through the
    :py:meth:`_next_action_hook` routine.

    ``reservoir``
        A reservoir instance compliant with the interface of
        :py:class:`ReservoirNode`. Specifically, must provide
        a *reset* method and *reset_states* must be :py:const:`False`.
        The input dimension must be compliant with the specification
        of the ``action``.

    ``readout``
        The reservoir readout function. Usually an online linear
        regression (RLS) instance, like :py:class:`StabilizedRLS`.
        Note that the readout must include a bias. The dimensions
        of ``reservoir`` and ``readout``  must match and the
        output of the latter must be single dimensional.

    """
    def __init__(self, reservoir, readout, *args, **kwargs):
        self.reservoir = reservoir
        self.readout = readout
        super(ADHDP, self).__init__(*args, **kwargs)

        # Check assumptions
        assert self.reservoir.reset_states == False
        assert self.reservoir.get_input_dim() == self.child.action_space_dim() + self.plant.state_space_dim()
        assert self.reservoir.get_input_dim() >= self.child.initial_action().shape[0]

    def new_episode(self):
        """Start a new episode of the same experiment. This method can
        also be used to initialize the *ActorCritic*, for example when
        it is loaded from a file.
        """
        self.reservoir.reset()
        self.plant.reset()
        super(ADHDP, self).new_episode()

    def _critic_eval(self, state, action, simulate, action_name='a_curr'):
        """Evaluate the critic at ``state`` and ``action``."""
        in_state = self.plant.state_input(state)
        action_nrm = self.normalizer.normalize_value(action_name, action)
        r_input = np.vstack((in_state, action_nrm)).T
        #r_input += np.random.normal(scale=0.001, size=r_input.shape)
        r_state = self.reservoir(r_input, simulate=simulate)
        #o_state = r_state # TODO: Direct ESN Model
        o_state = np.hstack((r_state, r_input)) # TODO: Input/Output ESN Model
        j_curr = self.readout(o_state)
        return r_input, o_state, j_curr

    def _critic_deriv_io_model(self, r_state):
        """Return the critic's derivative at ``r_state``."""
        direct_input_size = self.plant.state_space_dim()+self.child.action_space_dim() # Input/Output ESN Model
        r_state = r_state[:, :-direct_input_size] # this is because _critic_eval appends the input to the state
        dtanh = (np.ones(r_state.shape) - r_state**2).T # Nx1
        dstate = dtanh * self.reservoir.w_in[:, -self._motor_action_dim:].toarray() # Nx1 .* NxA => NxA
        deriv = self.readout.beta[1:-direct_input_size].T.dot(dstate) # Input/Output ESN Model
        deriv += self.readout.beta[-self._motor_action_dim:].T # Input/Output ESN Model
        deriv = deriv.T # AxL
        scale = self.normalizer.get('a_curr')[1]
        deriv *= scale # Derivative denormalization
        return deriv

    def _critic_deriv_direct_model(self, r_state):
        """Return the critic's derivative at ``r_state``."""
        dtanh = (np.ones(r_state.shape) - r_state**2).T # Nx1
        dstate = dtanh * self.reservoir.w_in[:, -self._motor_action_dim:].toarray() # Nx1 .* NxA => NxA
        deriv = self.readout.beta[1:].T.dot(dstate) #  LxA # Direct ESN Model
        deriv = deriv.T # AxL
        scale = self.normalizer.get('a_curr')[1]
        deriv *= scale # Derivative denormalization
        return deriv

    def _critic_deriv(self, r_state):
        """Return the critic's derivative at ``r_state``."""
        return self._critic_deriv_io_model(r_state)

    def _step(self, s_curr, epoch, a_curr, reward):
        """Execute one step of the actor and return the next action.

        This is the baseline ADHDP algorithm. The next action is
        computed as

        .. math::
            a_{t+1} = m a_t + (1-m) \left( a_t + \\alpha \\frac{\\partial J(s_t, a_t)}{\\partial a} \\right)

        with :math:`m` the momentum and :math:`\\alpha` the step size.
        The critic trained on the TD error with discount rate
        :math:`\gamma`:

        .. math::
            err = r_{t+1} + \gamma J(s_{t+1}, a_{t+1}) - J(s_t, a_t)

        ``epoch``
            Latest observed state (s_next). :py:keyword:`dict`, same as ``epoch``
            of the :py:meth:`__call__`.

        ``s_curr``
            Previous observed state. :py:keyword:`dict`, same as ``s_curr``
            of the :py:meth:`__call__`.

        ``reward``
            Reward of ``epoch``

        """
        # ESN-critic, first instance: in(k) => J(k)
        i_curr, x_curr, j_curr = self._critic_eval(s_curr, a_curr, simulate=False, action_name='a_curr')

        # Next action
        deriv = self._critic_deriv(x_curr)

        # gradient training of action (acc. to eq. 10)
        a_next = a_curr + self.alpha(self.num_episode, self.num_step) * deriv
        a_next = self.momentum(a_curr, a_next, self.num_episode, self.num_step)
        a_next = self._next_action_hook(a_next)

        # ESN-critic, second instance: in(k+1) => J(k+1)
        i_next, x_next, j_next = self._critic_eval(epoch, a_next, simulate=True, action_name='a_next')

        # TD_error(k) = J(k) - U(k) - gamma * J(k+1)
        err = reward + self.gamma(self.num_episode, self.num_step) * j_next - j_curr

        # One-step RLS training => Trained ESN
        self.readout.train(x_curr, err=err)

        # update epoch with locals:
        epoch['reward'] = np.atleast_2d([reward]).T
        epoch['deriv'] = deriv.T
        epoch['err'] = err.T
        epoch['readout'] = self.readout.beta.T
        epoch['gamma'] = np.atleast_2d([self.gamma(self.num_episode, self.num_step)]).T
        epoch['i_curr'] = i_curr
        epoch['x_curr'] = x_curr
        epoch['j_curr'] = j_curr
        epoch['a_curr'] = a_curr.T
        epoch['i_next'] = i_next
        epoch['x_next'] = x_next
        epoch['j_next'] = j_next
        epoch['a_next'] = a_next.T

        # increment
        return epoch

    def init_episode(self, epoch, time_start_ms, time_end_ms, step_size_ms):
        """Initial behaviour (after reset)

        .. note::
            Assuming identical initial trajectories, the initial state
            is the same - and thus doesn't matter.
            Non-identical initial trajectories will result in
            non-identical behaviour, therefore the initial state should
            be different (initial state w.r.t. start of learning).
            Due to this, the critic is already updated in the initial
            trajectory.
        """
        if self.num_step > 1:
            i_curr, x_curr, j_curr = self._critic_eval(self.s_curr, self.a_curr, False, 'a_curr')
            epoch['x_curr'] = x_curr
            epoch['i_curr'] = i_curr
            epoch['a_next'] = self.a_curr.T
            epoch['a_curr'] = self.a_curr.T

        self.s_curr = epoch
        return self.child(epoch, time_start_ms, time_end_ms, step_size_ms)

## HDP VARIATIONS ##

class ReservoirActorADHDP(ADHDP):
    """The next action is determined by a separate readout of the reservoir.
    The action-readout weights are trained by gradient descending

    """
    def __init__(self, *args, **kwargs):
        self.gd_tol = kwargs.pop('gd_tolerance', 1e-15)
        self.gd_max_iter = kwargs.pop('gd_max_iter', 500)
        self.rho, self.beta = 0.001, 0.1
        self.alpha_max = 20.0
        self.ls_max_iter = 1000
        super(ActionGradient, self).__init__(*args, **kwargs)

    def _step(self, s_curr, epoch, a_curr, reward):
        """Execute one step of the actor and return the next action."""
        # ESN-critic, first instance: in(k) => J(k)
        in_state = self.plant.state_input(s_curr)
        action_nrm = self.normalizer.normalize_value('a_curr', a_curr)
        i_curr = np.vstack((in_state, action_nrm)).T
        r_prev = self.reservoir.states
        r_state = self.reservoir(i_curr, simulate=False)
        x_curr = np.hstack((r_state, i_curr)) # TODO: Input/Output ESN Model
        j_curr = self.readout(x_curr)

        # Gradient ascent of J(a|s_{t+1})
        direct_input_size = self.plant.state_space_dim()+self.child.action_space_dim() # Input/Output ESN Model
        dtanh = (np.ones(r_state.shape) - r_state**2).T # Nx1
        dstate = dtanh * self.reservoir.w_in[:, -self._motor_action_dim:].toarray() # Nx1 .* NxA => NxA
        deriv = self.readout.beta[1:-direct_input_size].T.dot(dstate) # Input/Output ESN Model
        deriv += self.readout.beta[-self._motor_action_dim:].T # Input/Output ESN Model
        deriv = deriv.T # AxL
        scale = self.normalizer.get('a_curr')[1]
        deriv *= scale # Derivative denormalization
        #return deriv

        a_next = self.gradient_descent(epoch, a_curr)
        deriv = (a_next - a_curr) / self.alpha(self.num_episode, self.num_step)

        # Next action
        a_next = self.momentum(a_curr, a_next, self.num_episode, self.num_step)
        a_next = self._next_action_hook(a_next)

        # ESN-critic, second instance: in(k+1) => J(k+1)
        i_next, x_next, j_next = self._critic_eval(epoch, a_next, simulate=True, action_name='a_next')

        # TD_error(k) = J(k) - U(k) - gamma * J(k+1)
        err = reward + self.gamma(self.num_episode, self.num_step) * j_next - j_curr

        # One-step RLS training => Trained ESN
        self.readout.train(x_curr, err=err)

        # fill epoch with locals:
        epoch['reward'] = np.atleast_2d([reward]).T
        epoch['deriv'] = deriv.T
        epoch['err'] = err.T
        epoch['readout'] = self.readout.beta.T
        epoch['gamma'] = np.atleast_2d([self.gamma(self.num_episode, self.num_step)]).T
        epoch['i_curr'] = i_curr
        epoch['x_curr'] = x_curr
        epoch['j_curr'] = j_curr
        epoch['a_curr'] = a_curr.T
        epoch['i_next'] = i_next
        epoch['x_next'] = x_next
        epoch['j_next'] = j_next
        epoch['a_next'] = a_next.T

        # increment
        return epoch

class ActionGradient(ADHDP):
    """Determine the next action by gradient ascent search.
    The gradient ascent computes the action which maximizes the
    predicted return for a fixed state.

    Additional keyword arguments:

    ``gd_tolerance``
        Stop gradient descent if gradient below this threshold.

    ``gd_max_iter``
        Maximum number of gradient descent steps

    """
    def __init__(self, *args, **kwargs):
        self.gd_tol = kwargs.pop('gd_tolerance', 1e-15)
        self.gd_max_iter = kwargs.pop('gd_max_iter', 500)
        self.rho, self.beta = 0.001, 0.1
        self.alpha_max = 20.0
        self.ls_max_iter = 1000
        super(ActionGradient, self).__init__(*args, **kwargs)

    def _step(self, s_curr, epoch, a_curr, reward):
        """Execute one step of the actor and return the next action."""
        # ESN-critic, first instance: in(k) => J(k)
        i_curr, x_curr, j_curr = self._critic_eval(s_curr, a_curr, simulate=False, action_name='a_curr')

        # Gradient ascent of J(a|s_{t+1})
        a_next = self.gradient_descent(epoch, a_curr)
        deriv = (a_next - a_curr) / self.alpha(self.num_episode, self.num_step)

        # Next action
        a_next = self.momentum(a_curr, a_next, self.num_episode, self.num_step)
        a_next = self._next_action_hook(a_next)

        # ESN-critic, second instance: in(k+1) => J(k+1)
        i_next, x_next, j_next = self._critic_eval(epoch, a_next, simulate=True, action_name='a_next')

        # TD_error(k) = J(k) - U(k) - gamma * J(k+1)
        err = reward + self.gamma(self.num_episode, self.num_step) * j_next - j_curr

        # One-step RLS training => Trained ESN
        self.readout.train(x_curr, err=err)

        # fill epoch with locals:
        epoch['reward'] = np.atleast_2d([reward]).T
        epoch['deriv'] = deriv.T
        epoch['err'] = err.T
        epoch['readout'] = self.readout.beta.T
        epoch['gamma'] = np.atleast_2d([self.gamma(self.num_episode, self.num_step)]).T
        epoch['i_curr'] = i_curr
        epoch['x_curr'] = x_curr
        epoch['j_curr'] = j_curr
        epoch['a_curr'] = a_curr.T
        epoch['i_next'] = i_next
        epoch['x_next'] = x_next
        epoch['j_next'] = j_next
        epoch['a_next'] = a_next.T

        # increment
        return epoch

    def _line_search(self, state, action, gradient):
        """Line search algorithm to compute the step size parameter for
        a gradient step at ``action`` into direction ``gradient``.
        The ``state`` is still fixed, but must be provided.

        .. todo::
            broken for some cases (no reduction, infinite loop)

        """
        #0 < rho  < 0.5
        #rho < beta < 1.0
        warnings.warn('This code is unreliable. Use a fixed step size instead.')

        def phi(alpha):
            """step-size centric gradient step representation."""
            action_query = action + alpha * gradient
            query_result = self._critic_eval(state, action_query, simulate=True, action_name='a_curr') # TODO: Check reservoir state!
            return - query_result[-1]

        def dphi(alpha):
            """step-size centric gradient step representation, first
            derivative."""
            # gradient * -derivative(action + alpha * gradient)
            action_query = action + alpha * gradient
            in_state = self.plant.state_input(state)
            action_nrm = self.normalizer.normalize_value('a_next', action_query)
            i_inter = np.vstack((in_state, action_nrm)).T
            x_inter = self.reservoir(i_inter, simulate=True) # TODO: Check reservoir state!
            deriv = self._critic_deriv(x_inter)
            return gradient * -deriv

        dphi_zero, phi_zero = dphi(0.0), phi(0.0)
        lambda_ = lambda al: phi_zero + self.rho * al * dphi_zero

        if (dphi_zero >= 0.0).any():
            #warnings.warn('step_size cannot increase the function')
            print 'WARNING: step_size cannot increase the function (dphi(0.0) = %f)' % dphi_zero
            return 0.0

        # enforce lower bound
        gamma = self.beta * dphi_zero
        a = 0
        b = min(1.0, self.alpha_max)
        num_iter = self.ls_max_iter
        while b < self.alpha_max and (phi(b) <= lambda_(b)).any() and (dphi(b) <= gamma).any() and num_iter > 0:
            #print a,b
            a = b
            b = min(2*b, self.alpha_max)
            num_iter -= 1

        # set alpha
        alpha = b

        # enforce both constraints
        # infinite loop possible! how & where??
        num_iter = self.ls_max_iter
        while ((phi(alpha) > lambda_(alpha)).any() or (dphi(alpha) < gamma).any()) and num_iter > 0:
            alpha, a, b = self._refine( (a, b), phi, dphi, lambda_)
            if abs(a - b) < 1e-14:
                break
            num_iter -= 1

        if (phi(alpha) >= phi_zero).any():
            #warnings.warn('step size does not increase the function')
            print 'WARNING: step size does not increase the function (phi(alpha)=%f)' % phi(alpha)
            alpha = 0.0

        return alpha

    def _refine(self, (a, b), phi, dphi, lambda_):
        """Refinement part of the line search algorithm. Adjusts
        alpha and narrows the interval [a,b].
        """
        D = b - a
        if D < 1e-15:
            print "WARNING: a ~= b (D=", D
            #return a, a, b

        c = (phi(b) - phi(a) - D * dphi(a) ) / D**2

        # compute alpha
        if (c > 0).any():
            alpha = a - dphi(a) / (2*c)
            alpha = min(max(alpha, a + 0.1 * D), b - 0.1 * D)
        else:
            alpha = (a+b)/2

        # adjust interval
        if (phi(alpha) < lambda_(alpha)).any():
            a = alpha
        else:
            b = alpha

        return alpha, a, b

    def gradient_descent(self, state, action):
        """Gradient ascent search to find the action which maximizes
        the predicted return given a ``state``. The search starts at
        ``action``.
        """
        in_state = self.plant.state_input(state)
        num_iter = self.gd_max_iter
        while True:
            # Compute the gradient
            action_nrm = self.normalizer.normalize_value('a_next', action)
            i_inter = np.vstack((in_state, action_nrm)).T
            x_inter = self.reservoir(i_inter, simulate=True)
            x_inter = np.hstack((x_inter, i_inter)) # FIXME: Input/Output ESN Model
            gradient = self._critic_deriv(x_inter)

            # Do line search and update the action
            #step_size = self._line_search(state, action, gradient)
            step_size = self.alpha(self.num_episode, self.num_step)
            action = action + step_size * gradient
            action = self._next_action_hook(action)

            # Exit condition
            num_iter -= 1
            if np.linalg.norm(gradient) < self.gd_tol or num_iter <= 0:
                break

        #print num_iter

        return action

class ActionRecomputation(ADHDP):
    """Determine the next action the same way as the baseline algorithm
    for critic training, then recompute it based on the updated critic
    and with the latest state information.

    """
    def _step(self, s_curr, epoch, a_curr, reward):
        """Execute one step of the actor and return the next action."""
        # ESN-critic, first instance: in(k) => J(k)
        i_curr, x_curr, j_curr = self._critic_eval(s_curr, a_curr, False, 'a_curr')

        # Next action
        deriv = self._critic_deriv(x_curr)

        # gradient training of action (acc. to eq. 10)
        a_next = a_curr + self.alpha(self.num_episode, self.num_step) * deriv
        a_next = self.momentum(a_curr, a_next, self.num_episode, self.num_step)
        a_next = self._next_action_hook(a_next)

        # ESN-critic, second instance: in(k+1) => J(k+1)
        i_next, x_next, j_next = self._critic_eval(epoch, a_next, True, 'a_next')

        # TD_error(k) = J(k) - U(k) - gamma * J(k+1)
        err = reward + self.gamma(self.num_episode, self.num_step) * j_next - j_curr

        # One-step RLS training => Trained ESN
        self.readout.train(x_curr, err=err)

        # Action re-computation
        deriv2 = self._critic_deriv(x_next)
        a_next = a_curr + self.alpha(self.num_episode, self.num_step) * deriv2
        a_next = self.momentum(a_curr, a_next, self.num_episode, self.num_step)
        a_next = self._next_action_hook(a_next)

        # fill epoch with locals:
        epoch['reward'] = np.atleast_2d([reward]).T
        epoch['deriv2'] = deriv2.T
        epoch['deriv'] = deriv.T
        epoch['err'] = err.T
        epoch['readout'] = self.readout.beta.T
        epoch['gamma'] = np.atleast_2d([self.gamma(self.num_episode, self.num_step)]).T
        epoch['i_curr'] = i_curr
        epoch['x_curr'] = x_curr
        epoch['j_curr'] = j_curr
        epoch['a_curr'] = a_curr.T
        epoch['i_next'] = i_next
        epoch['x_next'] = x_next
        epoch['j_next'] = j_next
        epoch['a_next'] = a_next.T

        # increment
        return epoch

class ActionBruteForce(ADHDP):
    """Find the optimal action by computing the expected return at
    different sampled locations and picking the action which yields the
    highest one.

    ``candidates``
        Action samples. Must be a list of valid actions.

    .. todo::
        Breaks old code

    """
    def __init__(self, candidates, *args, **kwargs):
        super(ActionBruteForce, self).__init__(*args, **kwargs)
        self.candidates = candidates

    def _step(self, s_curr, epoch, a_curr, reward):
        """Execute one step of the actor and return the next action."""
        # ESN-critic, first instance: in(k) => J(k)
        i_curr, x_curr, j_curr = self._critic_eval(s_curr, a_curr, False, 'a_curr')

        # Next action
        a_next = a_curr
        j_best = float('-inf')
        in_state = self.plant.state_input(epoch)
        for candidate in self.candidates:
            candidate_nrm = self.normalizer.normalize_value('a_next', candidate)
            i_cand = np.vstack((in_state, candidate_nrm)).T
            x_cand = self.reservoir(i_cand, simulate=True)
            j_cand = self.readout(x_cand)
            if j_cand > j_best:
                j_best = j_cand
                a_next = np.atleast_2d(candidate)

        a_next = a_curr + self.alpha(self.num_episode, self.num_step) * (a_next - a_curr)
        a_next = self._next_action_hook(a_next)

        # ESN-critic, second instance: in(k+1) => J(k+1)
        i_next, x_next, j_next = self._critic_eval(epoch, a_next, True, 'a_next')

        # TD_error(k) = J(k) - U(k) - gamma * J(k+1)
        err = reward + self.gamma(self.num_episode, self.num_step) * j_next - j_curr

        # One-step RLS training => Trained ESN
        self.readout.train(x_curr, err=err)

        # fill epoch with locals:
        epoch['reward'] = np.atleast_2d([reward]).T
        epoch['err'] = err.T
        epoch['deriv'] = a_next-a_curr
        epoch['readout'] = self.readout.beta.T
        epoch['gamma'] = np.atleast_2d([self.gamma(self.num_episode, self.num_step)]).T
        epoch['i_curr'] = i_curr
        epoch['x_curr'] = x_curr
        epoch['j_curr'] = j_curr
        epoch['a_curr'] = a_curr.T
        epoch['i_next'] = i_next
        epoch['x_next'] = x_next
        epoch['j_next'] = j_next
        epoch['a_next'] = a_next.T

        # increment
        return epoch

ext: py

max_stars_repo_name: igsor/HDPy

max_issues_repo_name: igsor/HDPy

max_forks_repo_name: igsor/HDPy

reason: import numpy


text:
import pytest
import numpy as np
import theano.tensor as T
import lasagne

@pytest.mark.parametrize('impl', ['conv1d_sc', 'conv1d_mc0',
                                  'conv1d_mc1', 'conv1d_unstrided',
                                  'conv1d_sd', 'conv1d_md'])
@pytest.mark.parametrize('stride', [1, 2])
def test_conv(impl, stride):
    import lasagne.theano_extensions.conv
    conv = getattr(lasagne.theano_extensions.conv, impl)

    X = T.tensor3()
    W = T.tensor3()
    input = lasagne.utils.floatX(np.ones((1, 1, 10)))
    kernel = lasagne.utils.floatX(np.random.uniform(-1, 1, (2, 1, 6)))

    conv_theano = conv(X, W, input.shape, kernel.shape, subsample=(stride,)
                       ).eval({X: input, W: kernel})

    output = []
    for b in input:
        temp = []
        for c in kernel:
            temp.append(
                np.convolve(b[0, :], c[0, :], mode='valid'))
        output.append(temp)
    conv_np = np.array(output)[:, :, ::stride]

    assert np.allclose(conv_theano, conv_np)

@pytest.mark.parametrize('impl', ['conv1d_sc', 'conv1d_mc0', 'conv1d_mc1'])
def test_conv_nones(impl):
    import lasagne.theano_extensions.conv
    conv = getattr(lasagne.theano_extensions.conv, impl)

    X = T.tensor3()
    W = T.tensor3()
    input = lasagne.utils.floatX(np.ones((1, 1, 12)))
    kernel = lasagne.utils.floatX(np.random.uniform(-1, 1, (2, 1, 3)))

    conv_theano = conv(X, W, None, None).eval({
        X: input, W: kernel
        })

    output = []
    for b in input:
        temp = []
        for c in kernel:
            temp.append(
                np.convolve(b[0, :], c[0, :], mode='valid'))
        output.append(temp)
    conv_np = np.array(output)

    assert np.allclose(conv_theano, conv_np)

@pytest.mark.parametrize('impl', ['conv1d_sc', 'conv1d_mc0',
                                  'conv1d_mc1', 'conv1d_unstrided',
                                  'conv1d_sd', 'conv1d_md'])
def test_conv_border_mode(impl):
    import lasagne.theano_extensions.conv
    conv = getattr(lasagne.theano_extensions.conv, impl)

    X = T.tensor3()
    W = T.tensor3()

    with pytest.raises(Exception):
        conv(X, W, (1, 1, 10), (2, 1, 3), border_mode=None)

@pytest.mark.parametrize('impl', ['conv1d_unstrided', 'conv1d_sd',
                                  'conv1d_md'])
def test_conv_stride(impl):
    import lasagne.theano_extensions.conv
    conv = getattr(lasagne.theano_extensions.conv, impl)

    X = T.tensor3()
    W = T.tensor3()

    with pytest.raises(Exception):
        conv(X, W, (1, 1, 10), (2, 1, 3), subsample=(2,))

@pytest.mark.parametrize('val', [0, 7])
@pytest.mark.parametrize('batch_ndim', [1, 2])
def test_pad(batch_ndim, val, width=3):
    from lasagne.theano_extensions.padding import pad

    X = T.tensor4()
    X0 = lasagne.utils.floatX(np.ones((2, 3, 4, 5)))
    X_pad_theano = pad(X, width, val, batch_ndim).eval({X: X0})

    pads = tuple((width, width) if i >= batch_ndim else (0, 0)
                 for i, _ in enumerate(X0.shape))
    X_pad_np = np.pad(X0, pads, mode='constant', constant_values=val)

    assert (X_pad_theano == X_pad_np).all()

@pytest.mark.parametrize('batch_ndim', [1, 2])
def test_pad_width_per_axis(batch_ndim, val=0):
    from lasagne.theano_extensions.padding import pad

    width = (1, 2, 3, 4)

    X = T.tensor4()
    X0 = lasagne.utils.floatX(np.ones((2, 3, 4, 5)))
    X_pad_theano = pad(X, width[batch_ndim:], val, batch_ndim).eval({X: X0})

    pads = tuple((w, w) if i >= batch_ndim else (0, 0)
                 for i, w in enumerate(width))
    X_pad_np = np.pad(X0, pads, mode='constant', constant_values=val)

    assert (X_pad_theano == X_pad_np).all()

@pytest.mark.parametrize('batch_ndim', [1, 2])
def test_pad_width_per_border(batch_ndim, val=0):
    from lasagne.theano_extensions.padding import pad

    width = [(1, 2), (3, 4), (1, 2), (3, 4)]

    X = T.tensor4()
    X0 = lasagne.utils.floatX(np.ones((2, 3, 4, 5)))
    X_pad_theano = pad(X, width[batch_ndim:], val, batch_ndim).eval({X: X0})

    pads = tuple(w if i >= batch_ndim else (0, 0)
                 for i, w in enumerate(width))
    X_pad_np = np.pad(X0, pads, mode='constant', constant_values=val)

    assert (X_pad_theano == X_pad_np).all()

max_stars_repo_name: MihailoIsakov/fishfish

max_issues_repo_name: MihailoIsakov/fishfish

max_forks_repo_name: MihailoIsakov/fishfish

reason: import numpy,import theano


text:
"""
This script will train a new network on a dataset and save the model as a checkpoint.
It prints out training loss, validation loss, and validation accuracy as the network trains.

Required Inputs:    - data_dir: ImageNet Dataset folder path (default="./flowers/")

Optional inputs:    - arch: choose architecture (vgg or densenet) - default = vgg
                    - hidden_units: customize hidden layers (default=512)
                    - learning_rate: customize lr (default=0.0005)
                    - droput: customize droput (default=0.4)
                    - epochs: customize epochs (default=5)
                    - gpu: gpu ON (default = cpu)
                    - save_dir: customize checkpoint.pth destination folder (default = "./")
"""

# Imports here
import time
import torch
from torch import nn
from torch import optim
#import torch.nn.functional as F
#from torchvision import datasets, models
#import json
#from collections import OrderedDict
import argparse
#import numpy as np
import myutils
import nn_model

# Command line arguments

parser = argparse.ArgumentParser(description='Train network arguments:')
parser.add_argument('data_dir', default="./flowers/", help='Path to dataset')
parser.add_argument('--arch', dest="arch", action="store", default="vgg", type = str, help='Choose architecture {vgg or densenet}')
parser.add_argument('--hidden_units', type=int, nargs='+', dest="hidden_units", default=512, help='If more than one layer put a space (i.e. 512 250)')
parser.add_argument('--learning_rate', type=float, dest="learning_rate", action="store", default=0.0005, help='Learning Rate')
parser.add_argument('--dropout', type=float, dest = "dropout", action = "store", default = 0.4, help='DropOut')
parser.add_argument('--epochs', dest="epochs", action="store", type=int, default=5, help='Epochs')
#parser.add_argument('--gpu', type=str, dest="gpu", action="store", default="gpu", help='Gpu')
parser.add_argument('--gpu', help='Use GPU for training instead of CPU', action='store_true')
parser.add_argument('--save_dir', type=str, dest="save_dir", action="store", default="./", help='Checkpoint destination folder')

# Hyperparameters
parcy = parser.parse_args()
data_dir = parcy.data_dir
arch = parcy.arch
hidden_layers = myutils.check_h(parcy.hidden_units)
learning_rate = parcy.learning_rate
drop = parcy.dropout
epochs = parcy.epochs
isgpu = parcy.gpu
save_dir = parcy.save_dir

output_size = 102

# Print Hyperparameters
print("{}data_dir:    {} {} ".format('\n','\t', data_dir))
print("arch:          {} {} ".format('\t', arch))
print("hidden_units:  {} {} ".format('\t',hidden_layers))
print("learning_rate: {} {} ".format('\t',learning_rate))
print("dropout:       {} {} ".format('\t',drop))
print("epochs:        {} {} ".format('\t',epochs))
print("save_dir:      {} {} ".format('\t',save_dir))

if isgpu:
    print("GPU:       {} {} {}".format('\t', 'ON', '\n'))
else:
    print("GPU:       {} {} {}".format('\t', 'OFF', '\n'))

# load datasets, trainloader and validloader
image_datasets, dataloaders = myutils.load_data(data_dir)

# load a pre-trained network
model, input_size = nn_model.load_pretrained_model(arch)

## Global use
# Set the device (using GPU or not)
device = torch.device("cuda:0" if isgpu is True else "cpu")

#create my feedforward classifier
classifier = nn_model.Network(input_size, output_size, hidden_layers, drop)

#substitute the model's classifier with this new classifier
#transfer learning connected here
model.classifier = classifier
#print(model)

#define criteria and optimizer
criterion = nn.NLLLoss()
optimizer = optim.Adam(model.classifier.parameters(), lr=learning_rate)

# Move the model to the device available (cpu or CUDA)
model.to(device)

## Train It ##

# Define some variable
#epochs = 5
steps = 0
running_loss = 0
print_every = 102

since = time.time()

print("START TRAINING ...")
# get images and labels batch (64)
for epoch in range(epochs):
    model.train()
    for images, labels in dataloaders['train']:
        steps += 1
        # move them to CPU or CUDA (It depend on the device)
        images, labels = images.to(device), labels.to(device)
        # training loop
        optimizer.zero_grad() # reset to zero the gradients
        logps = model.forward(images) # log probabilities (output)
        loss = criterion(logps, labels) # error calculation (cost function)
        loss.backward() # backpropagation pass
        optimizer.step() # updating weights
        running_loss += loss.item() # keeping track of the training loss

        ## Validate It ##

        if steps % print_every == 0:
            valid_loss = 0
            accuracy = 0
            model.eval() # evaluation inference mode which turns dropout off
            with torch.no_grad():

                for images, labels in dataloaders['valid']:
                    # move them to GPU or CUDA (It depends on the device)
                    images, labels = images.to(device), labels.to(device)
                    logps = model.forward(images) # log probabilities (output)
                    batch_loss = criterion(logps, labels) # error calculation (cost function)
                    valid_loss += batch_loss.item() # keeping track of the test loss

                    # calculate the accuracy
                    ps = torch.exp(logps) # transform the LogSoftmax output to get the actual probabilities
                    top_ps, top_class = ps.topk(1,dim=1) # first largest value in our ps along the columns
                    # comparing the one element in each row of top_class with each element in labels
                    # which returns 64 True/False boolean values for each row.
                    equality = top_class == labels.view(*top_class.shape)
                    accuracy += torch.mean(equality.type(torch.FloatTensor)).item()

            print(f"Epoch {epoch+1}/{epochs}.. "
                  f"Train loss: {running_loss/print_every:.3f}.. "
                  f"Valid loss: {valid_loss/len(dataloaders['valid']):.3f}.. "
                  f"Validation accuracy: {accuracy/len(dataloaders['valid']):.3f}")

            running_loss = 0
            model.train()
time_elapsed = time.time() - since
print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))

# set class_to_idx mapping for model
model.class_to_idx = image_datasets['train'].class_to_idx

nn_model.save_model_checkpoint(model, input_size, epochs, save_dir, arch, hidden_layers, learning_rate, drop, optimizer, output_size)
print("New network trained on a dataset and the model saved as a checkpoint")


text:
# !/usr/bin/env python3
# -*- coding:utf-8 -*-
__author__ = 'Zhiquan Wang'
__date__ = '2018/10/18'

from scipy.stats import t
from scipy.stats import f
import math

class SimpleLinearModel(object):
    def __init__(self, *, x_data, y_data):
        self.__x_data = x_data
        self.__y_data = y_data
        self.__b1 = 0
        self.__b0 = 0
        self.__mean_x = 0
        self.__mean_y = 0
        self.__MSE = 0
        self.__expectation_sigma_hat_square = 0
        self.__estimator_variance_b1 = 0
        self.__estimator_variance_b0 = 0
        self.__std_x = 0
        self.__std_y = 0
        self.__sxx = 0
        self.__MSE = 0
        self.__estimator_variance_b1 = 0
        self.__estimator_variance_b0 = 0
        self.__r_square = 0
        # init parameters
        self.__regression_analysis()

    def __regression_analysis(self):
        self.__mean_x = sum(self.__x_data) / len(self.__x_data)
        self.__mean_y = sum(self.__y_data) / len(self.__y_data)
        self.__sxx = sum(
            [pow((xh - self.__mean_x), 2) for xh in self.__x_data])
        self.__b1 = sum([(self.__x_data[i] - self.__mean_x) *
                         (self.__y_data[i] - self.__mean_y)
                         for i in range(len(self.__x_data))]) / self.__sxx
        self.__b0 = self.__mean_y - self.__b1 * self.__mean_x
        self.__cal_MSE()
        self.__estimator_variance_b0_b1()
        self.__cal_r_square()

    @property
    def x_data(self):
        return self.__x_data

    @property
    def y_data(self):
        return self.__y_data

    @property
    def dataset(self):
        return [self.__x_data, self.__y_data]

    @dataset.setter
    def dataset(self, *, x_data, y_data):
        self.__x_data = x_data
        self.__y_data = y_data
        self.__regression_analysis()

    @property
    def b1(self):
        return self.__b1

    @property
    def b0(self):
        return self.__b0

    @property
    def mean_x(self):
        return self.__mean_x

    @property
    def mean_y(self):
        return self.__mean_y

    @property
    def MSE(self):
        return self.__MSE

    @property
    def sxx(self):
        return self.__sxx

    @property
    def estimator_variance_b0(self):
        return self.__estimator_variance_b0

    @property
    def estimator_variance_b1(self):
        return self.__estimator_variance_b1

    @property
    def r_square(self):
        return self.__r_square

    def y_hat(self, *, x):
        return self.__b0 + self.__b1 * x

    def __estimator_variance_b0_b1(self):
        self.__estimator_variance_b1 = self.__MSE / self.__sxx
        self.__estimator_variance_b0 = self.__MSE * (
                1 / len(self.__x_data) + pow(self.__mean_x, 2) / self.__sxx)

    def __cal_MSE(self):
        self.__MSE = sum([
            pow((self.__y_data[i] - self.y_hat(x=self.__x_data[i])), 2)
            for i in range(len(self.__y_data))
        ]) / (len(self.__y_data) - 2)

    def __cal_r_square(self):
        tmp_sum_y_hat = 0
        tmp_sum_y = 0
        for i in range(len(self.__x_data)):
            tmp_sum_y_hat += pow((self.__b0 + self.__b1 * self.__x_data[i] - self.__mean_y), 2)
            tmp_sum_y += pow((self.__y_data[i] - self.__mean_y), 2)
        self.__r_square = tmp_sum_y_hat / tmp_sum_y

    def estimator_variance_yh(self, *, xh):
        return self.__MSE * (
                1 / len(self.__x_data) + pow(xh - self.__mean_x, 2) / self.__sxx)

    def confidence_interval_b0(self, *, alpha):
        tmp_t_value = t.isf(alpha * 2, len(self.__x_data) - 2)
        return [
            self.__b0 - tmp_t_value * math.sqrt(self.__estimator_variance_b0),
            self.__b0 + tmp_t_value * math.sqrt(self.__estimator_variance_b0)
        ]

    def confidence_interval_b1(self, *, alpha):
        tmp_t_value = t.isf(alpha / 2, len(self.__x_data) - 2)
        print(tmp_t_value)
        return [
            self.__b1 - tmp_t_value * math.sqrt(self.__estimator_variance_b1),
            self.__b1 + tmp_t_value * math.sqrt(self.__estimator_variance_b1)
        ]

    def confidence_interval_yh(self, *, xh, alpha):
        tmp_t_value = t.isf(alpha / 2, len(self.__x_data) - 2)
        print('2a - t value : ', tmp_t_value)
        print("2a - Y hat : ", self.y_hat(x=xh))
        print('2a - S^2 : ', self.estimator_variance_yh(xh=xh))
        return [
            self.y_hat(x=xh) -
            tmp_t_value * math.sqrt(self.estimator_variance_yh(xh=xh)),
            self.y_hat(x=xh) +
            tmp_t_value * math.sqrt(self.estimator_variance_yh(xh=xh))
        ]

    def multi_confidence_interval_yh(self, *, x_list, alpha):
        input_num = len(x_list)
        tmp_t_value = t.isf(alpha / (2 * input_num), len(self.__x_data) - 2)
        print('2b - t value : ', tmp_t_value)
        return_list = []
        for i in range(input_num):
            tmp_yh = self.y_hat(x=x_list[i])
            print('2b - CI Yh : ', tmp_yh)
            tmp_value = math.sqrt(self.estimator_variance_yh(xh=x_list[i]))
            print('2b - CI se_yh : ', tmp_value)
            return_list.append([
                tmp_yh - tmp_t_value * tmp_value,
                tmp_yh + tmp_t_value * tmp_value
            ])
        return return_list

    def multi_confidence_interval_yh_wh(self, *, x_list, alpha):
        input_num = len(x_list)
        tmp_f_value = math.sqrt(2 * f.isf(alpha, 2, len(self.__x_data) - 2))
        print('2c - f', tmp_f_value)
        return_list = []
        for i in range(input_num):
            tmp_yh = self.y_hat(x=x_list[i])
            print('Yh = ', tmp_yh)
            tmp_value = math.sqrt(self.estimator_variance_yh(xh=x_list[i]))
            print('tmp_v = ', tmp_value)

            return_list.append([
                tmp_yh - tmp_f_value * tmp_value,
                tmp_yh + tmp_f_value * tmp_value
            ])
        return return_list

    def prefiction_interval_yh(self, *, xh, alpha):
        tmp_t_value = t.isf(alpha / 2, len(self.__x_data) - 2)
        tmp_value = math.sqrt(self.__MSE * (1 + 1 / len(self.__x_data) + pow(
            (xh - self.__mean_x), 2) / self.__sxx))

        return [
            self.y_hat(x=xh) - tmp_t_value * tmp_value,
            self.y_hat(x=xh) + tmp_t_value * tmp_value
        ]

    def multi_preficition_interval_yh(self, *, x_list, alpha):
        input_num = len(x_list)
        tmp_t_value = t.isf(alpha / (2 * input_num), len(self.__x_data) - 2)
        return_list = []
        for i in range(input_num):
            tmp_yh = self.y_hat(x=x_list[i])
            print('2b - PI Yh : ', tmp_yh)
            tmp_value = math.sqrt(
                self.__MSE * (1 + 1 / len(self.__x_data) + pow(
                    (x_list[i] - self.__mean_x), 2) / self.__sxx))
            print('2b - PI tmp_value : ', tmp_value)
            return_list.append([
                tmp_yh - tmp_t_value * tmp_value,
                tmp_yh + tmp_t_value * tmp_value
            ])
        return return_list

reason: from scipy


text:
import numpy as np

from abmarl.sim.components.state import GridPositionState, LifeState
from abmarl.sim.components.observer import GridPositionBasedObserver, \
    GridPositionTeamBasedObserver, RelativePositionObserver
from abmarl.sim.components.agent import PositionObservingAgent, AgentObservingAgent, \
    ComponentAgent

class PositionTestAgent(PositionObservingAgent, AgentObservingAgent): pass
class PositionTeamTestAgent(PositionObservingAgent, AgentObservingAgent): pass
class PositionTeamNoViewTestAgent(ComponentAgent): pass

def test_grid_position_observer():
    agents = {
        'agent0': PositionTestAgent(id='agent0', initial_position=np.array([0, 0]), agent_view=1),
        'agent1': PositionTestAgent(id='agent1', initial_position=np.array([2, 2]), agent_view=2),
        'agent2': PositionTestAgent(id='agent2', initial_position=np.array([3, 2]), agent_view=3),
        'agent3': PositionTestAgent(id='agent3', initial_position=np.array([1, 4]), agent_view=4),
        'agent4': ComponentAgent(id='agent4', initial_position=np.array([1, 4])),
    }

    state = GridPositionState(agents=agents, region=5)
    life = LifeState(agents=agents)
    observer = GridPositionBasedObserver(position_state=state, agents=agents)
    state.reset()
    life.reset()

    np.testing.assert_array_equal(observer.get_obs(agents['agent0'])['position'], np.array([
        [-1., -1., -1.],
        [-1.,  0.,  0.],
        [-1.,  0.,  0.],
    ]))
    np.testing.assert_array_equal(observer.get_obs(agents['agent1'])['position'], np.array([
        [1., 0., 0., 0., 0.],
        [0., 0., 0., 0., 1.],
        [0., 0., 0., 0., 0.],
        [0., 0., 1., 0., 0.],
        [0., 0., 0., 0., 0.],
    ]))
    np.testing.assert_array_equal(observer.get_obs(agents['agent2'])['position'], np.array([
        [-1.,  1.,  0.,  0.,  0.,  0., -1.],
        [-1.,  0.,  0.,  0.,  0.,  1., -1.],
        [-1.,  0.,  0.,  1.,  0.,  0., -1.],
        [-1.,  0.,  0.,  0.,  0.,  0., -1.],
        [-1.,  0.,  0.,  0.,  0.,  0., -1.],
        [-1., -1., -1., -1., -1., -1., -1.],
        [-1., -1., -1., -1., -1., -1., -1.],
    ]))
    np.testing.assert_array_equal(observer.get_obs(agents['agent3'])['position'], np.array([
        [-1., -1., -1., -1., -1., -1., -1., -1., -1.],
        [-1., -1., -1., -1., -1., -1., -1., -1., -1.],
        [-1., -1., -1., -1., -1., -1., -1., -1., -1.],
        [ 1.,  0.,  0.,  0.,  0., -1., -1., -1., -1.],
        [ 0.,  0.,  0.,  0.,  1., -1., -1., -1., -1.],
        [ 0.,  0.,  1.,  0.,  0., -1., -1., -1., -1.],
        [ 0.,  0.,  1.,  0.,  0., -1., -1., -1., -1.],
        [ 0.,  0.,  0.,  0.,  0., -1., -1., -1., -1.],
        [-1., -1., -1., -1., -1., -1., -1., -1., -1.],
    ]))
    assert observer.get_obs(agents['agent4']) == {}

def test_grid_team_position_observer():
    agents = {
        'agent0': PositionTeamTestAgent(
            id='agent0', team=1, initial_position=np.array([0, 0]), agent_view=1
        ),
        'agent1': PositionTeamNoViewTestAgent(
            id='agent1', team=1, initial_position=np.array([0, 0])
        ),
        'agent2': PositionTeamTestAgent(
            id='agent2', team=1, initial_position=np.array([2, 2]), agent_view=2
        ),
        'agent3': PositionTeamTestAgent(
            id='agent3', team=2, initial_position=np.array([3, 2]), agent_view=3
        ),
        'agent4': PositionTeamTestAgent(
            id='agent4', team=2, initial_position=np.array([1, 4]), agent_view=4
        ),
        'agent5': PositionTeamNoViewTestAgent(
            id='agent5', team=2, initial_position=np.array([1, 4])
        ),
        'agent6': PositionTeamNoViewTestAgent(
            id='agent6', team=2, initial_position=np.array([1, 4])
        ),
        'agent7': PositionTeamTestAgent(
            id='agent7', team=3, initial_position=np.array([1, 4]), agent_view=2
        ),
    }
    for agent in agents.values():
        agent.position = agent.initial_position

    state = GridPositionState(agents=agents, region=5)
    life = LifeState(agents=agents)
    observer = GridPositionTeamBasedObserver(
        position_state=state, number_of_teams=3, agents=agents
    )
    state.reset()
    life.reset()

    np.testing.assert_array_equal(observer.get_obs(agents['agent0'])['position'][:,:,1], np.array([
        [-1., -1., -1.],
        [-1.,  1.,  0.],
        [-1.,  0.,  0.],
    ]))
    np.testing.assert_array_equal(observer.get_obs(agents['agent0'])['position'][:,:,2], np.array([
        [-1., -1., -1.],
        [-1.,  0.,  0.],
        [-1.,  0.,  0.],
    ]))
    np.testing.assert_array_equal(observer.get_obs(agents['agent0'])['position'][:,:,3], np.array([
        [-1., -1., -1.],
        [-1.,  0.,  0.],
        [-1.,  0.,  0.],
    ]))

    np.testing.assert_array_equal(observer.get_obs(agents['agent2'])['position'][:,:,1], np.array([
        [2., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.],
    ]))
    np.testing.assert_array_equal(observer.get_obs(agents['agent2'])['position'][:,:,2], np.array([
        [0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 3.],
        [0., 0., 0., 0., 0.],
        [0., 0., 1., 0., 0.],
        [0., 0., 0., 0., 0.],
    ]))
    np.testing.assert_array_equal(observer.get_obs(agents['agent2'])['position'][:,:,3], np.array([
        [0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 1.],
        [0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.],
    ]))

    np.testing.assert_array_equal(observer.get_obs(agents['agent3'])['position'][:,:,1], np.array([
        [-1.,  2.,  0.,  0.,  0.,  0., -1.],
        [-1.,  0.,  0.,  0.,  0.,  0., -1.],
        [-1.,  0.,  0.,  1.,  0.,  0., -1.],
        [-1.,  0.,  0.,  0.,  0.,  0., -1.],
        [-1.,  0.,  0.,  0.,  0.,  0., -1.],
        [-1., -1., -1., -1., -1., -1., -1.],
        [-1., -1., -1., -1., -1., -1., -1.],
    ]))
    np.testing.assert_array_equal(observer.get_obs(agents['agent3'])['position'][:,:,2], np.array([
        [-1.,  0.,  0.,  0.,  0.,  0., -1.],
        [-1.,  0.,  0.,  0.,  0.,  3., -1.],
        [-1.,  0.,  0.,  0.,  0.,  0., -1.],
        [-1.,  0.,  0.,  0.,  0.,  0., -1.],
        [-1.,  0.,  0.,  0.,  0.,  0., -1.],
        [-1., -1., -1., -1., -1., -1., -1.],
        [-1., -1., -1., -1., -1., -1., -1.],
    ]))
    np.testing.assert_array_equal(observer.get_obs(agents['agent3'])['position'][:,:,3], np.array([
        [-1.,  0.,  0.,  0.,  0.,  0., -1.],
        [-1.,  0.,  0.,  0.,  0.,  1., -1.],
        [-1.,  0.,  0.,  0.,  0.,  0., -1.],
        [-1.,  0.,  0.,  0.,  0.,  0., -1.],
        [-1.,  0.,  0.,  0.,  0.,  0., -1.],
        [-1., -1., -1., -1., -1., -1., -1.],
        [-1., -1., -1., -1., -1., -1., -1.],
    ]))

    np.testing.assert_array_equal(observer.get_obs(agents['agent4'])['position'][:,:,1], np.array([
        [-1., -1., -1., -1., -1., -1., -1., -1., -1.],
        [-1., -1., -1., -1., -1., -1., -1., -1., -1.],
        [-1., -1., -1., -1., -1., -1., -1., -1., -1.],
        [ 2.,  0.,  0.,  0.,  0., -1., -1., -1., -1.],
        [ 0.,  0.,  0.,  0.,  0., -1., -1., -1., -1.],
        [ 0.,  0.,  1.,  0.,  0., -1., -1., -1., -1.],
        [ 0.,  0.,  0.,  0.,  0., -1., -1., -1., -1.],
        [ 0.,  0.,  0.,  0.,  0., -1., -1., -1., -1.],
        [-1., -1., -1., -1., -1., -1., -1., -1., -1.],
    ]))
    np.testing.assert_array_equal(observer.get_obs(agents['agent4'])['position'][:,:,2], np.array([
        [-1., -1., -1., -1., -1., -1., -1., -1., -1.],
        [-1., -1., -1., -1., -1., -1., -1., -1., -1.],
        [-1., -1., -1., -1., -1., -1., -1., -1., -1.],
        [ 0.,  0.,  0.,  0.,  0., -1., -1., -1., -1.],
        [ 0.,  0.,  0.,  0.,  2., -1., -1., -1., -1.],
        [ 0.,  0.,  0.,  0.,  0., -1., -1., -1., -1.],
        [ 0.,  0.,  1.,  0.,  0., -1., -1., -1., -1.],
        [ 0.,  0.,  0.,  0.,  0., -1., -1., -1., -1.],
        [-1., -1., -1., -1., -1., -1., -1., -1., -1.],
    ]))
    np.testing.assert_array_equal(observer.get_obs(agents['agent4'])['position'][:,:,3], np.array([
        [-1., -1., -1., -1., -1., -1., -1., -1., -1.],
        [-1., -1., -1., -1., -1., -1., -1., -1., -1.],
        [-1., -1., -1., -1., -1., -1., -1., -1., -1.],
        [ 0.,  0.,  0.,  0.,  0., -1., -1., -1., -1.],
        [ 0.,  0.,  0.,  0.,  1., -1., -1., -1., -1.],
        [ 0.,  0.,  0.,  0.,  0., -1., -1., -1., -1.],
        [ 0.,  0.,  0.,  0.,  0., -1., -1., -1., -1.],
        [ 0.,  0.,  0.,  0.,  0., -1., -1., -1., -1.],
        [-1., -1., -1., -1., -1., -1., -1., -1., -1.],
    ]))

    np.testing.assert_array_equal(observer.get_obs(agents['agent7'])['position'][:,:,1], np.array([
        [-1., -1., -1., -1., -1.],
        [ 0.,  0.,  0., -1., -1.],
        [ 0.,  0.,  0., -1., -1.],
        [ 1.,  0.,  0., -1., -1.],
        [ 0.,  0.,  0., -1., -1.],
    ]))
    np.testing.assert_array_equal(observer.get_obs(agents['agent7'])['position'][:,:,2], np.array([
        [-1., -1., -1., -1., -1.],
        [ 0.,  0.,  0., -1., -1.],
        [ 0.,  0.,  3., -1., -1.],
        [ 0.,  0.,  0., -1., -1.],
        [ 1.,  0.,  0., -1., -1.],
    ]))
    np.testing.assert_array_equal(observer.get_obs(agents['agent7'])['position'][:,:,3], np.array([
        [-1., -1., -1., -1., -1.],
        [ 0.,  0.,  0., -1., -1.],
        [ 0.,  0.,  0., -1., -1.],
        [ 0.,  0.,  0., -1., -1.],
        [ 0.,  0.,  0., -1., -1.],
    ]))

def test_relative_position_observer():
    agents = {
        'agent0': PositionTestAgent(id='agent0', initial_position=np.array([0, 0]), agent_view=1),
        'agent1': PositionTestAgent(id='agent1', initial_position=np.array([2, 2]), agent_view=2),
        'agent2': PositionTestAgent(id='agent2', initial_position=np.array([3, 2]), agent_view=3),
        'agent3': PositionTestAgent(id='agent3', initial_position=np.array([1, 4]), agent_view=4),
        'agent4': ComponentAgent(id='agent4', initial_position=np.array([1, 4])),
    }

    state = GridPositionState(agents=agents, region=5)
    observer = RelativePositionObserver(position_state=state, agents=agents)
    state.reset()

    assert observer.get_obs(agents['agent0'])['relative_position']['agent1'][0] == 2
    assert observer.get_obs(agents['agent0'])['relative_position']['agent1'][1] == 2
    assert observer.get_obs(agents['agent0'])['relative_position']['agent2'][0] == 3
    assert observer.get_obs(agents['agent0'])['relative_position']['agent2'][1] == 2
    assert observer.get_obs(agents['agent0'])['relative_position']['agent3'][0] == 1
    assert observer.get_obs(agents['agent0'])['relative_position']['agent3'][1] == 4
    assert observer.get_obs(agents['agent0'])['relative_position']['agent4'][0] == 1
    assert observer.get_obs(agents['agent0'])['relative_position']['agent4'][1] == 4

    assert observer.get_obs(agents['agent1'])['relative_position']['agent0'][0] == -2
    assert observer.get_obs(agents['agent1'])['relative_position']['agent0'][1] == -2
    assert observer.get_obs(agents['agent1'])['relative_position']['agent2'][0] == 1
    assert observer.get_obs(agents['agent1'])['relative_position']['agent2'][1] == 0
    assert observer.get_obs(agents['agent1'])['relative_position']['agent3'][0] == -1
    assert observer.get_obs(agents['agent1'])['relative_position']['agent3'][1] == 2
    assert observer.get_obs(agents['agent1'])['relative_position']['agent4'][0] == -1
    assert observer.get_obs(agents['agent1'])['relative_position']['agent4'][1] == 2

    assert observer.get_obs(agents['agent2'])['relative_position']['agent0'][0] == -3
    assert observer.get_obs(agents['agent2'])['relative_position']['agent0'][1] == -2
    assert observer.get_obs(agents['agent2'])['relative_position']['agent1'][0] == -1
    assert observer.get_obs(agents['agent2'])['relative_position']['agent1'][1] == 0
    assert observer.get_obs(agents['agent2'])['relative_position']['agent3'][0] == -2
    assert observer.get_obs(agents['agent2'])['relative_position']['agent3'][1] == 2
    assert observer.get_obs(agents['agent2'])['relative_position']['agent4'][0] == -2
    assert observer.get_obs(agents['agent2'])['relative_position']['agent4'][1] == 2

    assert observer.get_obs(agents['agent3'])['relative_position']['agent0'][0] == -1
    assert observer.get_obs(agents['agent3'])['relative_position']['agent0'][1] == -4
    assert observer.get_obs(agents['agent3'])['relative_position']['agent1'][0] == 1
    assert observer.get_obs(agents['agent3'])['relative_position']['agent1'][1] == -2
    assert observer.get_obs(agents['agent3'])['relative_position']['agent2'][0] == 2
    assert observer.get_obs(agents['agent3'])['relative_position']['agent2'][1] == -2
    assert observer.get_obs(agents['agent3'])['relative_position']['agent4'][0] == 0
    assert observer.get_obs(agents['agent3'])['relative_position']['agent4'][1] == 0