text:
"""
Description: Calculates the Fisher information for a gravity measurement
Input: config.yaml

"""

from qutip import *
import yaml
import gc #Garbage collection
from numpy import *
import matplotlib.pyplot as plt
from scipy import constants as cp
import os
from matplotlib import rc
import datetime
import time
from scipy.special import factorial
from math import isnan as mathisnan
#from RKF45 import RKF45
import types
import tables

class Solver(object):
    """
    Description: Class that includes the option run and more.
    Input: config.yaml file.
    """
    def __init__(self, config):
        # Extract arguments from the config file
        self.args = {}
        if type(config) == str:
            with open(config) as cfile:
                self.args.update(yaml.load(cfile))
        elif type(config) == dict:
            self.args.update(config)
        else:
            print("Failed to load config arguments")

        # Assign values to constants
        self.function = str(self.args['function'])
        self.state = str(self.args['state'])
        self.squeezing = float(self.args['squeezing'])
        self.displace = float(self.args['displace'])
        self.start_time = float(self.args['start_time'])
        self.N = int(self.args['N'])
        self.M = int(self.args['M'])
        self.alpha = float(self.args['alpha'])
        self.beta = float(self.args['beta'])
        self.chi = float(self.args['chi'])
        if self.args['time'] == 6.24:
            self.time = 2.*pi
        else:
            self.time = float(self.args['time'])
        self.gbar = float(self.args['gbar'])
        self.k = float(self.args['k'])
        self.tolerance = float(self.args['tolerance'])
        self.folder = str(self.args['folder'])
        self.h = float(self.args['h'])
        self.steps = int(self.time/self.h)
        self.hmax = float(self.args['hmax'])
        self.hmin = float(self.args['hmin'])
        self.gamma = float(self.args['gamma'])

        print("Hilbert space is: " + str([self.N, self.M]))

        # Select the function to run
        if self.function == 'state_evolution':
            print("Starting function: state_evolution")
            self.state_evolution()
        if self.function == 'RK4F':
            print("Starting function: RK4F")
            self.run_RKF45()
        if self.function == 'triple_state_evolution':
            print("Starting function: triple_state_evolution")
            self.triple_state_evolution()

    def initialise_coherent(self, arg, objtype = 'matrix'):
        """
        Parameters:
            arg: string, determines type 'ket' or 'dm' of outputted state
        Ouput:  two-system state at t = self.start_time
        Note: State is initialised by creating an array of fock states with various n. The array is then summed and the state outputted.
        """
        terms = []
        eta = 1. - exp(-1j*self.start_time)
        for n in range(0, self.N):
            terms.append(exp(-absolute(self.alpha)**2/2.)*self.alpha**n/(sqrt(factorial(n)))*exp(1j*(self.start_time-sin(self.start_time))*(self.k*n - self.gbar)**2)*exp(((self.k*n - self.gbar)*eta*exp(1j*self.start_time)*self.beta - (self.k*n - self.gbar)*(1 - exp(1j*self.start_time))*exp(-1j*self.start_time)*self.beta)/2.)*tensor(fock(self.N,n), coherent(self.N, self.beta*exp(-1j*self.start_time) + (self.k*n- self.gbar)*(1 - exp(-1j*self.start_time)))))
        state = sum(terms)
        if arg =='ket':
            if objtype == 'numpy':
                return state/sqrt(state.overlap(state)).full()
            if objtype == 'qobj':
                return state/sqrt(state.overlap(state))
        if arg == 'dm':
            if objtype == 'numpy':
                return (state*state.dag()/state.overlap(state)).full() # turn into dm and normalise.
            if objtype == 'qobj':
                return (state*state.dag()/state.overlap(state)) # turn into dm and normalise.

    def initialise_d_start_time(self, arg):
        """
        Initialise the derivative of the state with respect to gbar at t = start_time
        """
        if arg == 'dm':
            alpha_terms = []
            beta_terms = []
            for n in range(0, (self.N)):
                for m in range(0, (self.N)):
                    alpha_terms.append(
                    -exp(- absolute(self.alpha)**2)
                    * ((conj(self.alpha)**m *self.alpha**n)/(sqrt(factorial(n)*factorial(m)))
                    * exp((self.start_time - sin(self.start_time))*1j*(self.k**2*(n**2 - m**2) - 2.*self.k*self.gbar*(n-m)))
                    * (2.*pi*(self.start_time - sin(self.start_time))*1j*self.k*(n-m))*fock(self.N, n)*fock(self.N, m).dag()))
                    beta_terms.append(exp(-absolute(self.beta)**2)*(self.beta**n*conj(self.beta)**m/(sqrt(factorial(n)*factorial(m)))*fock(self.N,n)*fock(self.M, m).dag()))
            drho = tensor(sum(alpha_terms), sum(beta_terms))
            return drho.full()
        if arg == 'ket':
            alpha_terms = []
            beta_terms = []
            for n in range(0, (self.N)):
                alpha_terms.append(exp(-absolute(self.alpha)**2/2.)*conj(self.alpha)**n/(sqrt(factorial(n)))*exp((self.start_time - sin(self.start_time))*1j*(self.k**2 * n**2 - 2.*self.k*self.gbar*n))*((self.start_time - sin(self.start_time))*2.*1j*self.k*n)*fock(self.N,n))
                beta_terms.append(exp(-absolute(self.beta)**2/2.)*self.beta**n/(sqrt(factorial(n)))*fock(self.N,n))
            return tensor(sum(alpha_terms), sum(beta_terms)).full()

    def initialise_squeezed(self, s, arg, objtype = 'matrix'):
        if arg == 'ket':
            if objtype == 'qobj':
                d = displace(self.N, self.displace)
                sq = squeeze(self.N, self.squeezing)
                print(sq)
                state = tensor(d*s*basis(self.N, 0), d*s*basis(self.N, 0))
                state = state/sqrt(state.overlap(state))
                print(state.overlap(state))
                return state

            if objtype == 'matrix':
                return tensor(squeeze(self.N, self.squeezing), squeeze(self.M, self.squeezing)).full()

    def initialise_squeezed2(self, z, arg, objtype = 'qobj'):
        terms = []
        for n in range(0, self.N/2):
            terms.append(((z-1)/(2*(z+1)))**n *sqrt(factorial(2*n))/(factorial(n))*basis(self.N, n))
        state = sum(terms)
        state = state/(sqrt(state.overlap(state)))
        return tensor(state, state)

    def run_RKF45(self):
        # Create the vector
        states = []
        dstates = []
        statevec = array([self.initialise_coherent('dm'), self.initialise_d_start_time('dm')])
        print("Starting simulation.")
        # Set decoherence
        self.Lind = sqrt(self.chi)*tensor(destroy(self.N), qeye(self.M))
        self.LinddagLind = self.Lind.dag()*self.Lind
        self.LinddagLind = self.LinddagLind.full()
        self.Linddag = self.Lind.dag().full()
        self.Lind = self.Lind.full()

        tic = time.clock()

        # Create folder destination, check if exists if not create
        st = datetime.datetime.fromtimestamp(time.time()).strftime('%Y-%m-%d-%H.%M.%S')
        foldername = self.folder + "/simulation" + st
        if not os.path.exists(foldername):
            os.makedirs(foldername)

        # Add the current config file into the simulation folder
        with open(foldername + "/config", 'w') as outfile:
            outfile.write(yaml.dump(self.args, default_flow_style = False))

        results = RKF45(self.Lindblad, self.start_time, self.time, statevec, self.tolerance, self.hmax, self.hmin, self.N)
        print("Simulation completed.")
        toc = time.clock()
        print("Time required: " + str((toc-tic)/60) + " minutes")

        # Save results to files
        save(foldername + "/times", results[0])
        save(foldername + "/fock_fisher", results[1])
        save(foldername + "/homodyne_fisher", results[2])
        save(foldername + "/heterodyne_fisher", results[3])
        save(foldername + "/x_quads", results[4])
        save(foldername + "/p_quads", results[5])

        plt.figure(figsize=(10,7.5))
        plt.plot(results[0], results[3])
        plt.show()

    def analytic_U(self, time):
        """
        Parameters: config file and time
        Output: Time evolution operator U generated at time.
        Note: Used for making sure numerical evolution was giving the correct value.
        """
        a = tensor(destroy(self.N), qeye(self.N))
        b = tensor(qeye(self.N), destroy(self.N))

        eta = 1 - exp(- 1j*time)
        U = (1j*(self.k*a.dag()*a - self.gbar)**2. *(time - sin(time))).expm()*((self.k*a.dag()*a - self.gbar)*(eta*b.dag() - conj(eta)*b)).expm()*(- 1j*time*b.dag()*b).expm()
        return U

    def state_evolution(self):
        """
        Description: Evolves a state five times and feeds them to a Fisher information function
        Parameters: None
        Output: Saves the resulting Fisher informatio, times and the config to path + self.foldername
        """
        st = datetime.datetime.fromtimestamp(time.time()).strftime('%Y-%m-%d-%H.%M.%S')

        a = tensor(destroy(self.N), qeye(self.N))
        b = tensor(qeye(self.N), destroy(self.N))

        Hamiltonian0 =   b.dag()*b + (b.dag() + b)*(self.gbar - self.k*a.dag()*a)
        times = linspace(self.start_time, self.time, self.steps)

        #state0 = tensor(coherent(self.N, self.alpha), coherent(self.N, self.beta))

        if self.state == 'coherent':
            state0 = self.initialise_coherent('ket', 'qobj')
        if self.state == 'squeezed':
            state0 = self.initialise_squeezed2(self.squeezing, 'ket', 'qobj')
        #state0 = tensor(coherent(self.N, 1.), coherent(self.N, 1.))

        g0 = self.gbar

        self.gbar = g0 - 2.*self.h
        Hamiltonianm2 = b.dag()*b + (b.dag() + b)*(self.gbar - self.k*a.dag()*a)
        if self.state == 'coherent':
            statem2 = self.initialise_coherent('ket', 'qobj')
        if self.state == 'squeezed':
            statem2 = self.initialise_squeezed2(self.squeezing, 'ket', 'qobj')

        self.gbar = g0 - self.h
        Hamiltonianm1 = b.dag()*b + (b.dag() + b)*(self.gbar - self.k*a.dag()*a)
        if self.state == 'coherent':
            statem1 = self.initialise_coherent('ket', 'qobj')
        if self.state == 'squeezed':
            statem1 = self.initialise_squeezed2(self.squeezing, 'ket', 'qobj')

        self.gbar = g0 + self.h
        Hamiltonian1 =  b.dag()*b + (b.dag() + b)*(self.gbar - self.k*a.dag()*a)
        if self.state == 'coherent':
            state1 = self.initialise_coherent('ket', 'qobj')
        if self.state == 'squeezed':
            state1 = self.initialise_squeezed2(self.squeezing, 'ket', 'qobj')

        self.gbar = g0 + 2.*self.h
        Hamiltonian2 =  b.dag()*b + (b.dag() + b)*(self.gbar - self.k*a.dag()*a)
        if self.state == 'coherent':
            print("Initialising coherent state")
            state2 = self.initialise_coherent('ket', 'qobj')
        if self.state == 'squeezed':
            print("Initialising squeezed state")
            state2 = self.initialise_squeezed2(self.squeezing, 'ket', 'qobj')

        # Set options
        options = Options()

        # Set decoherence
        deco = sqrt(self.chi)*tensor(destroy(self.N), qeye(self.N))

        psi0 = []
        psi1 = []
        psi2 = []
        psim1 = []
        psim2 = []

        # Perform simulation, trace out the states then delete the results. This is done to preserve memory.

        print("Starting simulation")
        results0 = mesolve(Hamiltonian0, state0, times, c_ops = deco, e_ops = [], args = [], progress_bar = True)
        for state in results0.states:
            psi0.append(state.ptrace(0))
        del results0
        gc.collect()

        results1 = mesolve(Hamiltonian1, state1, times, c_ops = deco, e_ops = [], args = [], progress_bar = True)
        for state in results1.states:
            psi1.append(state.ptrace(0))
        del results1
        gc.collect()

        results2 = mesolve(Hamiltonian2, state2, times, c_ops = deco, e_ops = [], args = [], progress_bar = True)
        for state in results2.states:
            psi2.append(state.ptrace(0))
        del results2
        gc.collect()

        resultsm1 = mesolve(Hamiltonianm1, statem1, times, c_ops = deco, e_ops = [], args = [], progress_bar = True)
        for state in resultsm1.states:
            psim1.append(state.ptrace(0))
        del resultsm1
        gc.collect()

        resultsm2 = mesolve(Hamiltonianm2, statem2, times, c_ops = deco, e_ops = [], args = [], progress_bar = True)
        for state in resultsm2.states:
            psim2.append(state.ptrace(0))
        del resultsm2
        gc.collect()

        # Check quadratures for clarity

        x_exp = []
        p_exp = []
        x = (destroy(self.N) + destroy(self.N).dag())/sqrt(2.)
        p = 1.j * (destroy(self.N).dag() - destroy(self.N))/sqrt(2.)

        for state in psi0:
            x_exp.append(expect(x, state))
            p_exp.append(expect(p, state))
        #plt.figure(figsize=(10,7.5))
        #plt.plot(x_exp, p_exp)
        #plt.show()

        tic = time.clock()

        print("Calculating Fisher information")
        [fisher_result_pos, fisher_result_mom] = self.triple_Fisher(psim2, psim1, psi0, psi1, psi2)

        toc = time.clock()
        print("Time required to do that: " + str((toc-tic)/60) + " minutes, or " + str((toc-tic)/(60*60)) + "hours.")

        # Create folder if it does not exist
        foldername = self.folder + "/simulation" + st
        if not os.path.exists(foldername):
            os.makedirs(foldername)

        with open(foldername + '/config', 'w') as outfile:
            outfile.write(yaml.dump(self.args, default_flow_style=False))
        qsave(psi0, foldername + "/psi0")
        qsave(psi1, foldername + "/psi1")
        qsave(psim1, foldername + "/psim1")
        save(foldername + "/fisher_position", fisher_result_pos)
        save(foldername + "/fisher_momentum", fisher_result_mom)
        save(foldername + "/times", times)

        plt.figure(figsize=(10,7.5))
        plt.plot(times, fisher_result_pos)
        plt.show()
        plt.plot(times, fisher_result_mom)
        plt.show()

    def triple_state_evolution(self):
        """
        In this section, we let a three-mode quantum state evolve. The last state is a quantum vacuum state that we will measure on.
        The question is whether we should let the qutip solver handle this one. The problem is that we need the same time steps.
        So: qutip objects: good solver, can trace out things, can easily take the dagger
        Numpy arrays: Fast solver, ensures same time steps.
        So the only problem is the times. As long as the solver gives that to us, it's worth keeping the objects.
        """
        # Take a time-stamp
        st = datetime.datetime.fromtimestamp(time.time()).strftime('%Y-%m-%d-%H.%M.%S')

        a = tensor(destroy(self.N), qeye(self.N), qeye(self.N))
        b = tensor(qeye(self.N), destroy(self.N), qeye(self.N))
        c = tensor(qeye(self.N), qeye(self.N), destroy(self.N))

        times = linspace(self.start_time, self.time, self.steps)

        # Create Hamiltonian and state for gbar = starting value
        # Add a vacuum state as system C
        Hamiltonian0 =  b.dag()*b + (b.dag() + b)*(self.gbar - self.k*a.dag()*a) + self.gamma*(a.dag()*c + a*c.dag())
        state0 = tensor(self.initialise_start_time('ket','qobj'), fock(self.N, 0))

        g0 = self.gbar
        self.gbar = g0 - 2*self.h
        Hamiltonianm2 = b.dag()*b + (b.dag() + b)*(self.gbar - self.k*a.dag()*a) + self.gamma*(a.dag()*c + a*c.dag())
        statem2 = tensor(self.initialise_start_time('ket', 'qobj'), fock(self.N,0))

        self.gbar = g0 - self.h
        Hamiltonianm1 = b.dag()*b + (b.dag() + b)*(self.gbar - self.k*a.dag()*a) + self.gamma*(a.dag()*c + a*c.dag())
        statem1 = tensor(self.initialise_start_time('ket', 'qobj'), fock(self.N,0))

        self.gbar = g0 + self.h
        Hamiltonian1 =  b.dag()*b + (b.dag() + b)*(self.gbar - self.k*a.dag()*a) + self.gamma*(a.dag()*c + a*c.dag())
        state1 = tensor(self.initialise_start_time('ket','qobj'), fock(self.N,0))

        self.gbar = g0 + 2*self.h
        Hamiltonian2 =  b.dag()*b + (b.dag() + b)*(self.gbar - self.k*a.dag()*a) + self.gamma*(a.dag()*c + a*c.dag())
        state2 = tensor(self.initialise_start_time('ket','qobj'), fock(self.N,0))

        psi0 = []
        psi1 = []
        psi2 = []
        psim1 = []
        psim2 = []

        # Add zero decoherence
        deco = []

        # Solve each system for a psi with different g
        results0 = mesolve(Hamiltonian0, state0, times, c_ops = deco, e_ops = [], args = [], progress_bar = True)
        print("Tracing out states")
        for state in results0.states:
            psi0.append(state.ptrace(2))
        del results0
        gc.collect()

        results1 = mesolve(Hamiltonian1, state1, times, c_ops = deco, e_ops = [], args = [], progress_bar = True)
        print("Tracing out states")
        for state in results1.states:
            psi1.append(state.ptrace(2))
        del results1
        gc.collect()

        results2 = mesolve(Hamiltonian2, state2, times, c_ops = deco, e_ops = [], args = [], progress_bar = True)
        print("Tracing out states")
        for state in results2.states:
            psi2.append(state.ptrace(2))
        del results2
        gc.collect()

        resultsm1 = mesolve(Hamiltonianm1, statem1, times, c_ops = deco, e_ops = [], args = [], progress_bar = True)
        print("Tracing out states")
        for state in resultsm1.states:
            psim1.append(state.ptrace(2))
        del resultsm1
        gc.collect()

        resultsm2 = mesolve(Hamiltonianm2, statem2, times, c_ops = deco, e_ops = [], args = [], progress_bar = True)
        print("Tracing out states")
        for state in resultsm2.states:
            psim2.append(state.ptrace(2))
        del resultsm2
        gc.collect()

        # Check that we get the correct quadratures
        x = (destroy(self.N) + destroy(self.N).dag())/sqrt(2.)
        p = 1.j * (destroy(self.N).dag() - destroy(self.N))/sqrt(2.)
        x_exp = []
        p_exp = []
        for state in psi0:
            x_exp.append(expect(x, state))
            p_exp.append(expect(p, state))
        plt.figure(figsize=(10,7.5))
        plt.plot(x_exp, p_exp)
        plt.show()

        print("Calculating Fisher information by tracing out states")
        tic = time.clock()

        # Here we trace out the states. I wonder however if it wouldn't be faster to just calculate the expectation values.
        print("Tracing out states")
        # Trace out system AB, to be left with the light state

        fisher_result_pos, fisher_result_mom = self.triple_Fisher(psim2, psim1, psi0, psi1, psi2)

        toc = time.clock()
        print("Time required to do that: " + str((toc-tic)/60) + " minutes, or " + str((toc-tic)/(60*60)) + "hours.")

        plt.figure(figsize=(10,7.5))
        plt.plot(times, fisher_result_pos)
        plt.show()
        plt.plot(times, fisher_result_mom)
        plt.show()

        # Create folder if it does not exist
        foldername = self.folder + "/simulation" + st
        if not os.path.exists(foldername):
            os.makedirs(foldername)

        with open(foldername + '/config', 'w') as outfile:
            outfile.write(yaml.dump(self.args, default_flow_style=False))
        qsave(psi0, foldername + "/psi0")
        qsave(psi1, foldername + "/psi1")
        qsave(psim1, foldername + "/psim1")
        save(foldername + "/fisher_leaky_position", fisher_result_pos)
        save(foldername + "/fisher_leaky_momentum", fisher_result_mom)
        save(foldername + "/times", times)

    def triple_Fisher(self, psim2, psim1, psi0, psi1, psi2):
        """
        Description: Calculates the Fisher information for the specified measurement.
        Parameters:
            psim2: States for a run with self.g - 2*h
            psim1: States for a run with self.g - h
            psi0: States for a run with self.g
            psi1: States for a run with self.g + h
            psi2: States for a run with self.g + 2*h
        Output: An array containig the Fisher information for each timestep
        """
        c_position = (destroy(self.N) + destroy(self.N).dag())/sqrt(2.)
        c_pos_eigen = c_position.eigenstates()

        c_momentum = 1j*(destroy(self.N).dag() - destroy(self.N))/sqrt(2.)
        c_mom_eigen = c_momentum.eigenstates()

        probs0_pos = []
        probs0_mom = []
        probs1_pos = []
        probs1_mom = []
        probs2_pos = []
        probs2_mom = []
        probsm1_pos = []
        probsm1_mom = []
        probsm2_pos = []
        probsm2_mom = []

        for state in psim2:
            probsm2_pos.append([expect(ket2dm(x), state) for x in c_pos_eigen[1]])
            probsm2_mom.append([expect(ket2dm(x), state) for x in c_mom_eigen[1]])
        for state in psim1:
            probsm1_pos.append([expect(ket2dm(x), state) for x in c_pos_eigen[1]])
            probsm1_mom.append([expect(ket2dm(x), state) for x in c_mom_eigen[1]])
        for state in psi0:
            probs0_pos.append([expect(ket2dm(x), state) for x in c_pos_eigen[1]])
            probs0_mom.append([expect(ket2dm(x), state) for x in c_mom_eigen[1]])
        for state in psi1:
            probs1_pos.append([expect(ket2dm(x), state) for x in c_pos_eigen[1]])
            probs1_mom.append([expect(ket2dm(x), state) for x in c_mom_eigen[1]])
        for state in psi2:
            probs2_pos.append([expect(ket2dm(x), state) for x in c_pos_eigen[1]])
            probs2_mom.append([expect(ket2dm(x), state) for x in c_mom_eigen[1]])

        # Check that probabilities are normalised (they should sum to 1 or 50 depending on the form)

        diffs_pos = []
        diffs_mom = []
        fisher_info_pos = array([])
        fisher_info_mom = array([])

        print("Calculating differentials")
        # Cycle through the probs and work out the differentials
        # Use a 2nd order central difference theorem
        for timestep in range(0, len(probs0_pos)):
            diffs_pos.append([(-probs2_pos[timestep][xvalue] + 8*probs1_pos[timestep][xvalue] - 8*probsm1_pos[timestep][xvalue] + probsm2_pos[timestep][xvalue])/(12*self.h) for xvalue in range(0, len(probs0_pos[0]))])
            diffs_mom.append([(-probs2_mom[timestep][xvalue] + 8*probs1_mom[timestep][xvalue] - 8*probsm1_mom[timestep][xvalue] + probsm2_mom[timestep][xvalue])/(12*self.h) for xvalue in range(0, len(probs0_mom[0]))])

        # Make into numpy arrays to simplify future operations
        diffs_pos = asarray(diffs_pos, dtype = float64)
        diffs_mom = asarray(diffs_mom, dtype = float64)
        probs0_pos = asarray(probs0_pos, dtype = float64)
        probs0_mom = asarray(probs0_mom, dtype = float64)

        for timestep in range(0, len(diffs_pos)):
            elements_pos = diffs_pos[timestep]*diffs_pos[timestep]/probs0_pos[timestep] # Use array multiplication to calculate each
            elements_mom = diffs_mom[timestep]*diffs_mom[timestep]/probs0_mom[timestep] # Use array multiplication to calculate each
            elements_pos = elements_pos[logical_not(isnan(elements_pos))] #Remove nan from array
            elements_mom = elements_mom[logical_not(isnan(elements_mom))] #Remove nan from array
            fisher_info_pos = append(fisher_info_pos, sum(elements_pos))
            fisher_info_mom = append(fisher_info_mom, sum(elements_mom))

        return [fisher_info_pos, fisher_info_mom]

    def RungeKutta4_Lindblad(self, statevector):
        k1 = self.Lindblad(statevector)
        k2 = self.Lindblad(statevector + (self.h/2.)*k1)
        k3 = self.Lindblad(statevector + (self.h/2.)*k2)
        k4 = self.Lindblad(statevector + self.h*k3)
        return (self.h/6.)*(k1 + 2.*k2 + 2.*k3 + k4)

    def Lindblad(self, statevec): # This function takes a vector consisting of the state and its derivative with respect to gbar.
        return array([- 1j*(self.H.dot(statevec[0]) - statevec[0].dot(self.H)) + self.Lind.dot(statevec[0].dot(self.Linddag)) - (1/2.)*(statevec[0].dot(self.LinddagLind) + self.LinddagLind.dot(statevec[0])),  1j*(self.dH.dot(statevec[0]) - statevec[0].dot(self.dH)) - 1j*(self.H.dot(statevec[1]) - statevec[1].dot(self.H)) + self.Lind.dot(statevec[1].dot(self.Linddag)) - (1/2.)*(statevec[1].dot(self.LinddagLind) + self.LinddagLind.dot(statevec[1]))])

if __name__ == "__main__":
    system = Solver('config.yaml')
    system.state_evolution()

reason: from numpy,from scipy


text:
import torch
import math
from torch.optim.lr_scheduler import _LRScheduler

class PolyLR(_LRScheduler):
    def __init__(self, optimizer, max_iter, decay_iter=1, power=0.9, last_epoch=-1) -> None:
        self.decay_iter = decay_iter
        self.max_iter = max_iter
        self.power = power
        super().__init__(optimizer, last_epoch=last_epoch)

    def get_lr(self):
        factor = (1 - self.last_epoch / float(self.max_iter)) ** self.power
        return [factor*lr for lr in self.base_lrs]
        # if self.last_epoch % self.decay_iter or self.last_epoch % self.max_iter:
        #     return self.base_lrs
        # else:
        #     factor = (1 - self.last_epoch / float(self.max_iter)) ** self.power
        #     return [factor*lr for lr in self.base_lrs]

class WarmupLR(_LRScheduler):
    def __init__(self, optimizer, warmup_iter=500, warmup_ratio=5e-4, warmup='exp', last_epoch=-1) -> None:
        self.warmup_iter = warmup_iter
        self.warmup_ratio = warmup_ratio
        self.warmup = warmup
        super().__init__(optimizer, last_epoch)

    def get_lr(self):
        ratio = self.get_lr_ratio()
        return [ratio * lr for lr in self.base_lrs]

    def get_lr_ratio(self):
        return self.get_warmup_ratio() if self.last_epoch < self.warmup_iter else self.get_main_ratio()

    def get_main_ratio(self):
        raise NotImplementedError

    def get_warmup_ratio(self):
        assert self.warmup in ['linear', 'exp']
        alpha = self.last_epoch / self.warmup_iter

        return self.warmup_ratio + (1. - self.warmup_ratio) * alpha if self.warmup == 'linear' else self.warmup_ratio ** (1. - alpha)

class WarmupPolyLR(WarmupLR):
    def __init__(self, optimizer, power, max_iter, warmup_iter=500, warmup_ratio=5e-4, warmup='exp', last_epoch=-1) -> None:
        self.power = power
        self.max_iter = max_iter
        super().__init__(optimizer, warmup_iter, warmup_ratio, warmup, last_epoch)

    def get_main_ratio(self):
        real_iter = self.last_epoch - self.warmup_iter
        real_max_iter = self.max_iter - self.warmup_iter
        alpha = real_iter / real_max_iter

        return (1 - alpha) ** self.power

class WarmupExpLR(WarmupLR):
    def __init__(self, optimizer, gamma, interval=1, warmup_iter=500, warmup_ratio=5e-4, warmup='exp', last_epoch=-1) -> None:
        self.gamma = gamma
        self.interval = interval
        super().__init__(optimizer, warmup_iter, warmup_ratio, warmup, last_epoch)

    def get_main_ratio(self):
        real_iter = self.last_epoch - self.warmup_iter
        return self.gamma ** (real_iter // self.interval)

class WarmupCosineLR(WarmupLR):
    def __init__(self, optimizer, max_iter, eta_ratio=0, warmup_iter=500, warmup_ratio=5e-4, warmup='exp', last_epoch=-1) -> None:
        self.eta_ratio = eta_ratio
        self.max_iter = max_iter
        super().__init__(optimizer, warmup_iter, warmup_ratio, warmup, last_epoch)

    def get_main_ratio(self):
        real_iter = self.last_epoch - self.warmup_iter
        real_max_iter = self.max_iter - self.warmup_iter

        return self.eta_ratio + (1 - self.eta_ratio) * (1 + math.cos(math.pi * self.last_epoch / real_max_iter)) / 2

__all__ = ['polylr', 'warmuppolylr', 'warmupcosinelr', 'warmupsteplr']

def get_scheduler(scheduler_name: str, optimizer, max_iter: int, power: int, warmup_iter: int, warmup_ratio: float):
    assert scheduler_name in __all__, f"Unavailable scheduler name >> {scheduler_name}.\nAvailable schedulers: {__all__}"
    if scheduler_name == 'warmuppolylr':
        return WarmupPolyLR(optimizer, power, max_iter, warmup_iter, warmup_ratio, warmup='linear')
    elif scheduler_name == 'warmupcosinelr':
        return WarmupCosineLR(optimizer, max_iter, warmup_iter=warmup_iter, warmup_ratio=warmup_ratio)
    return PolyLR(optimizer, max_iter)

if __name__ == '__main__':
    model = torch.nn.Conv2d(3, 16, 3, 1, 1)
    optim = torch.optim.SGD(model.parameters(), lr=1e-3)

    max_iter = 20000
    sched = WarmupPolyLR(optim, power=0.9, max_iter=max_iter, warmup_iter=200, warmup_ratio=0.1, warmup='exp', last_epoch=-1)

    lrs = []

    for _ in range(max_iter):
        lr = sched.get_lr()[0]
        lrs.append(lr)
        optim.step()
        sched.step()

    import matplotlib.pyplot as plt
    import numpy as np

    plt.plot(np.arange(len(lrs)), np.array(lrs))
    plt.grid()
    plt.show()

ext: py

max_stars_repo_name: sunilhoho/semseg

max_issues_repo_name: sunilhoho/semseg

max_forks_repo_name: sunilhoho/semseg

reason: import numpy


text:
import pandas as pd
import os
from datetime import datetime, date, timedelta
import numpy as np

__author__ = 'Duy Cao'
__copyright__ = 'Duy Cao, 2020'
__license__ = 'MIT'
__status__ = 'release'
__url__ = 'https://github.com/caominhduy/bitcoin-indicated'
__version__ = '1.0'

NOMICS_API_KEY = 'e19a802cb19367a00e3d66638e9e974f'

def live_data():
    # Read historical exchange rate from Nomics
    nomics_url = f'https://api.nomics.com/v1/exchange-rates/history?key={NOMICS_API_KEY}&start=2020-01-01T00%3A00%3A00Z&currency=BTC'
    nomics_df = pd.read_json(nomics_url)
    nomics_df['timestamp'] = nomics_df['timestamp'].dt.date
    nomics_df = nomics_df.rename(columns={'timestamp':'date', 'rate':'nomics'})
    latest_date = nomics_df['date'].max()

    # Read historical exchange rate from Coindesk
    coindesk_url = f'https://api.coindesk.com/v1/bpi/historical/close.json?start=2020-01-01&end={latest_date}'
    df = pd.read_json(coindesk_url)
    df = df.drop(['disclaimer', 'time'], 1)
    df = df[:-2].reset_index().rename(columns={'index':'date', 'bpi':'coindesk'})
    df['date'] = pd.to_datetime(df['date']).dt.date

    df = df.merge(nomics_df, on=['date'], how='inner')

    return df

max_stars_repo_name: caominhduy/bitcoin-indicated

max_issues_repo_name: caominhduy/bitcoin-indicated

max_forks_repo_name: caominhduy/bitcoin-indicated


text:
#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Copyright 1999-2020 Alibaba Group Holding Ltd.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import numpy as np
import scipy.sparse as sps

from mars.config import options
from mars.core import get_tiled
from mars.tensor.datasource import tensor, arange, zeros
from mars.tensor.indexing import take, compress, extract, choose, \
    unravel_index, nonzero, flatnonzero, fill_diagonal
from mars.tensor import mod, stack, hstack
from mars.tests.core import ExecutorForTest, TestBase

class Test(TestBase):
    def setUp(self):
        self.executor = ExecutorForTest('numpy')
        self.old_chunk = options.chunk_size
        options.chunk_size = 10

    def tearDown(self):
        options.chunk_size = self.old_chunk

    def testBoolIndexingExecution(self):
        raw = np.random.random((11, 8, 12, 14))
        arr = tensor(raw, chunk_size=3)

        index = arr < .5
        arr2 = arr[index]
        size_res = self.executor.execute_tensor(arr2, mock=True)
        res = self.executor.execute_tensor(arr2)

        self.assertEqual(sum(s[0] for s in size_res), arr.nbytes)
        np.testing.assert_array_equal(np.sort(np.concatenate(res)), np.sort(raw[raw < .5]))

        index2 = tensor(raw[:, :, 0, 0], chunk_size=3) < .5
        arr3 = arr[index2]
        res = self.executor.execute_tensor(arr3, concat=True)[0]

        expected = raw[raw[:, :, 0, 0] < .5]
        self.assertEqual(sum(it.size for it in res), expected.size)
        self.assertEqual(res.shape, expected.shape)

        raw = np.asfortranarray(np.random.random((11, 8, 12, 14)))
        arr = tensor(raw, chunk_size=3)

        index = tensor(raw[:, :, 0, 0], chunk_size=3) < .5
        arr2 = arr[index]
        res = self.executor.execute_tensor(arr2, concat=True)[0]
        expected = raw[raw[:, :, 0, 0] < .5].copy('A')

        self.assertEqual(res.flags['C_CONTIGUOUS'], expected.flags['C_CONTIGUOUS'])
        self.assertEqual(res.flags['F_CONTIGUOUS'], expected.flags['F_CONTIGUOUS'])

    def testFancyIndexingNumpyExecution(self):
        # test fancy index of type numpy ndarray
        raw = np.random.random((11, 8, 12, 14))
        arr = tensor(raw, chunk_size=(2, 5, 7, 8))

        index = [9, 10, 3, 1, 8, 10]
        arr2 = arr[index]

        res = self.executor.execute_tensor(arr2, concat=True)[0]
        np.testing.assert_array_equal(res, raw[index])

        index = np.random.permutation(8)
        arr3 = arr[:2, ..., index]

        res = self.executor.execute_tensor(arr3, concat=True)[0]
        np.testing.assert_array_equal(res, raw[:2, ..., index])

        index = [1, 3, 9, 10]
        arr4 = arr[..., index, :5]

        res = self.executor.execute_tensor(arr4, concat=True)[0]
        np.testing.assert_array_equal(res, raw[..., index, :5])

        index1 = [8, 10, 3, 1, 9, 10]
        index2 = [1, 3, 9, 10, 2, 7]
        arr5 = arr[index1, :, index2]

        res = self.executor.execute_tensor(arr5, concat=True)[0]
        np.testing.assert_array_equal(res, raw[index1, :, index2])

        index1 = [1, 3, 5, 7, 9, 10]
        index2 = [1, 9, 9, 10, 2, 7]
        arr6 = arr[index1, :, index2]

        res = self.executor.execute_tensor(arr6, concat=True)[0]
        np.testing.assert_array_equal(res, raw[index1, :, index2])
        # fancy index is ordered, no concat required
        self.assertGreater(len(get_tiled(arr6).nsplits[0]), 1)

        index1 = [[8, 10, 3], [1, 9, 10]]
        index2 = [[1, 3, 9], [10, 2, 7]]
        arr7 = arr[index1, :, index2]

        res = self.executor.execute_tensor(arr7, concat=True)[0]
        np.testing.assert_array_equal(res, raw[index1, :, index2])

        index1 = [[1, 3], [3, 7], [7, 7]]
        index2 = [1, 9]
        arr8 = arr[0, index1, :, index2]

        res = self.executor.execute_tensor(arr8, concat=True)[0]
        np.testing.assert_array_equal(res, raw[0, index1, :, index2])

    def testFancyIndexingTensorExecution(self):
        # test fancy index of type tensor

        raw = np.random.random((11, 8, 12, 14))
        arr = tensor(raw, chunk_size=(6, 5, 7, 8))

        raw_index = [8, 10, 3, 1, 9, 10]
        index = tensor(raw_index, chunk_size=4)
        arr2 = arr[index]

        res = self.executor.execute_tensor(arr2, concat=True)[0]
        np.testing.assert_array_equal(res, raw[raw_index])

        raw_index = np.random.permutation(8)
        index = tensor(raw_index, chunk_size=3)
        arr3 = arr[:2, ..., index]

        res = self.executor.execute_tensor(arr3, concat=True)[0]
        np.testing.assert_array_equal(res, raw[:2, ..., raw_index])

        raw_index = [1, 3, 9, 10]
        index = tensor(raw_index)
        arr4 = arr[..., index, :5]

        res = self.executor.execute_tensor(arr4, concat=True)[0]
        np.testing.assert_array_equal(res, raw[..., raw_index, :5])

        raw_index1 = [8, 10, 3, 1, 9, 10]
        raw_index2 = [1, 3, 9, 10, 2, 7]
        index1 = tensor(raw_index1, chunk_size=4)
        index2 = tensor(raw_index2, chunk_size=3)
        arr5 = arr[index1, :, index2]

        res = self.executor.execute_tensor(arr5, concat=True)[0]
        np.testing.assert_array_equal(res, raw[raw_index1, :, raw_index2])

        raw_index1 = [1, 3, 5, 7, 9, 10]
        raw_index2 = [1, 9, 9, 10, 2, 7]
        index1 = tensor(raw_index1, chunk_size=3)
        index2 = tensor(raw_index2, chunk_size=4)
        arr6 = arr[index1, :, index2]

        res = self.executor.execute_tensor(arr6, concat=True)[0]
        np.testing.assert_array_equal(res, raw[raw_index1, :, raw_index2])

        raw_index1 = [[8, 10, 3], [1, 9, 10]]
        raw_index2 = [[1, 3, 9], [10, 2, 7]]
        index1 = tensor(raw_index1)
        index2 = tensor(raw_index2, chunk_size=2)
        arr7 = arr[index1, :, index2]

        res = self.executor.execute_tensor(arr7, concat=True)[0]
        np.testing.assert_array_equal(res, raw[raw_index1, :, raw_index2])

        raw_index1 = [[1, 3], [3, 7], [7, 7]]
        raw_index2 = [1, 9]
        index1 = tensor(raw_index1, chunk_size=(2, 1))
        index2 = tensor(raw_index2)
        arr8 = arr[0, index1, :, index2]

        res = self.executor.execute_tensor(arr8, concat=True)[0]
        np.testing.assert_array_equal(res, raw[0, raw_index1, :, raw_index2])

        raw_a = np.random.rand(30, 30)
        a = tensor(raw_a, chunk_size=(13, 17))
        b = a.argmax(axis=0)
        c = a[b, arange(30)]
        res = self.executor.execute_tensor(c, concat=True)[0]

        np.testing.assert_array_equal(res, raw_a[raw_a.argmax(axis=0), np.arange(30)])

        # test one chunk
        arr = tensor(raw, chunk_size=20)

        raw_index = [8, 10, 3, 1, 9, 10]
        index = tensor(raw_index, chunk_size=20)
        arr9 = arr[index]

        res = self.executor.execute_tensor(arr9, concat=True)[0]
        np.testing.assert_array_equal(res, raw[raw_index])

        raw_index1 = [[1, 3], [3, 7], [7, 7]]
        raw_index2 = [1, 9]
        index1 = tensor(raw_index1)
        index2 = tensor(raw_index2)
        arr10 = arr[0, index1, :, index2]

        res = self.executor.execute_tensor(arr10, concat=True)[0]
        np.testing.assert_array_equal(res, raw[0, raw_index1, :, raw_index2])

        # test order
        raw = np.asfortranarray(np.random.random((11, 8, 12, 14)))
        arr = tensor(raw, chunk_size=(2, 3, 2, 3))

        raw_index = [8, 10, 3, 1, 9, 10]
        index = tensor(raw_index, chunk_size=4)
        arr11 = arr[index]

        res = self.executor.execute_tensor(arr11, concat=True)[0]
        expected = raw[raw_index].copy('A')
        np.testing.assert_array_equal(res, expected)
        self.assertEqual(res.flags['C_CONTIGUOUS'], expected.flags['C_CONTIGUOUS'])
        self.assertEqual(res.flags['F_CONTIGUOUS'], expected.flags['F_CONTIGUOUS'])

    def testSliceExecution(self):
        raw = np.random.random((11, 8, 12, 14))
        arr = tensor(raw, chunk_size=3)

        arr2 = arr[2:9:2, 3:7, -1:-9:-2, 12:-11:-4]
        res = self.executor.execute_tensor(arr2, concat=True)[0]

        np.testing.assert_array_equal(res, raw[2:9:2, 3:7, -1:-9:-2, 12:-11:-4])

        arr3 = arr[-4, 2:]
        res = self.executor.execute_tensor(arr3, concat=True)[0]
        np.testing.assert_equal(res, raw[-4, 2:])

        raw = sps.random(12, 14, density=.1)
        arr = tensor(raw, chunk_size=3)

        arr2 = arr[-1:-9:-2, 12:-11:-4]
        res = self.executor.execute_tensor(arr2, concat=True)[0]

        np.testing.assert_equal(res.toarray(), raw.toarray()[-1:-9:-2, 12:-11:-4])

        # test order
        raw = np.asfortranarray(np.random.random((11, 8, 12, 14)))
        arr = tensor(raw, chunk_size=3)

        arr2 = arr[2:9:2, 3:7, -1:-9:-2, 12:-11:-4]
        res = self.executor.execute_tensor(arr2, concat=True)[0]
        expected = raw[2:9:2, 3:7, -1:-9:-2, 12:-11:-4].copy('A')

        np.testing.assert_array_equal(res, expected)
        self.assertEqual(res.flags['C_CONTIGUOUS'], expected.flags['C_CONTIGUOUS'])
        self.assertEqual(res.flags['F_CONTIGUOUS'], expected.flags['F_CONTIGUOUS'])

        arr3 = arr[0:13, :, None]
        res = self.executor.execute_tensor(arr3, concat=True)[0]
        expected = raw[0:13, :, None].copy('A')

        np.testing.assert_array_equal(res, expected)
        self.assertEqual(res.flags['C_CONTIGUOUS'], expected.flags['C_CONTIGUOUS'])
        self.assertEqual(res.flags['F_CONTIGUOUS'], expected.flags['F_CONTIGUOUS'])

    def testMixedIndexingExecution(self):
        rs = np.random.RandomState(0)
        raw = rs.random((11, 8, 12, 13))
        arr = tensor(raw, chunk_size=3)

        raw_cond = raw[0, :, 0, 0] < .5
        cond = tensor(raw[0, :, 0, 0], chunk_size=3) < .5
        arr2 = arr[10::-2, cond, None, ..., :5]
        size_res = self.executor.execute_tensor(arr2, mock=True)
        res = self.executor.execute_tensor(arr2, concat=True)[0]

        new_shape = list(arr2.shape)
        new_shape[1] = cond.shape[0]
        self.assertEqual(sum(s[0] for s in size_res), int(np.prod(new_shape) * arr2.dtype.itemsize))
        np.testing.assert_array_equal(res, raw[10::-2, raw_cond, None, ..., :5])

        b_raw = np.random.random(8)
        raw_cond = b_raw < .5
        conds = [raw_cond, tensor(b_raw, chunk_size=2) < .5]
        for cond in conds:
            arr3 = arr[-2::-3, cond, ...]
            res = self.executor.execute_tensor(arr3, concat=True)[0]

            np.testing.assert_array_equal(res, raw[-2::-3, raw_cond, ...])

        # test multiple bool index and fancy index
        cond1 = np.zeros(11, dtype=bool)
        cond1[rs.permutation(11)[:5]] = True
        cond2 = np.zeros(12, dtype=bool)
        cond2[rs.permutation(12)[:5]] = True
        f3 = np.random.randint(13, size=5)

        expected = raw[cond1, ..., cond2, f3]

        t = arr[cond1, ..., cond2, f3]
        res = self.executor.execute_tensor(t, concat=True)[0]
        np.testing.assert_array_equal(res, expected)

        ctx, executor = self._create_test_context(self.executor)
        with ctx:
            t = arr[tensor(cond1), ..., tensor(cond2), tensor(f3)]
            res = executor.execute_tensors([t])[0]
            np.testing.assert_array_equal(res, expected)

    def testSetItemExecution(self):
        rs = np.random.RandomState(0)

        raw = data = rs.randint(0, 10, size=(11, 8, 12, 13))
        arr = tensor(raw.copy(), chunk_size=3)
        raw = raw.copy()

        idx = slice(2, 9, 2), slice(3, 7), slice(-1, -9, -2), 2
        arr[idx] = 20
        res = self.executor.execute_tensor(arr, concat=True)[0]

        raw[idx] = 20
        np.testing.assert_array_equal(res, raw)
        self.assertEqual(res.flags['C_CONTIGUOUS'], raw.flags['C_CONTIGUOUS'])
        self.assertEqual(res.flags['F_CONTIGUOUS'], raw.flags['F_CONTIGUOUS'])

        raw = data
        shape = raw[idx].shape

        arr2 = tensor(raw.copy(), chunk_size=3)
        raw = raw.copy()

        replace = rs.randint(10, 20, size=shape[:-1] + (1,)).astype('f4')
        arr2[idx] = tensor(replace, chunk_size=4)
        res = self.executor.execute_tensor(arr2, concat=True)[0]

        raw[idx] = replace
        np.testing.assert_array_equal(res, raw)

        raw = np.asfortranarray(np.random.randint(0, 10, size=(11, 8, 12, 13)))
        arr = tensor(raw.copy('A'), chunk_size=3)
        raw = raw.copy('A')

        idx = slice(2, 9, 2), slice(3, 7), slice(-1, -9, -2), 2
        arr[idx] = 20
        res = self.executor.execute_tensor(arr, concat=True)[0]

        raw[idx] = 20
        np.testing.assert_array_equal(res, raw)
        self.assertEqual(res.flags['C_CONTIGUOUS'], raw.flags['C_CONTIGUOUS'])
        self.assertEqual(res.flags['F_CONTIGUOUS'], raw.flags['F_CONTIGUOUS'])

        # test bool indexing set
        raw = data

        arr = tensor(raw.copy(), chunk_size=3)
        raw1 = rs.rand(11)
        arr[tensor(raw1, chunk_size=4) < 0.6, 2: 7] = 3
        res = self.executor.execute_tileable(arr, concat=True)[0]

        raw[raw1 < 0.6, 2: 7] = 3
        np.testing.assert_array_equal(res, raw)

        raw = np.random.randint(3, size=10).astype(np.int64)
        raw2 = np.arange(3)

        arr = zeros((10, 3))
        arr[tensor(raw) == 1, tensor(raw2) == 1] = 1
        res = self.executor.execute_tileable(arr, concat=True)[0]

        expected = np.zeros((10, 3))
        expected[raw == 1, raw2 == 1] = 1
        np.testing.assert_array_equal(res, expected)

        ctx, executor = self._create_test_context(self.executor)
        with ctx:
            raw = data

            arr = tensor(raw.copy(), chunk_size=3)
            raw1 = rs.rand(11)
            set_data = rs.rand((raw1 < 0.8).sum(), 8, 12, 13)
            arr[tensor(raw1, chunk_size=4) < 0.8] = tensor(set_data)

            res = self.executor.execute_tileables([arr])[0]

            raw[raw1 < 0.8] = set_data
            np.testing.assert_array_equal(res, raw)

        # test error
        with self.assertRaises(ValueError):
            t = tensor(raw, chunk_size=3)
            t[0, 0, 0, 0] = zeros(2, chunk_size=10)
            _ = self.executor.execute_tensor(t)

    def testSetItemStructuredExecution(self):
        rec_type = np.dtype([('a', np.int32), ('b', np.double),
                             ('c', np.dtype([('a', np.int16), ('b', np.int64)]))])

        raw = np.zeros((4, 5), dtype=rec_type)
        arr = tensor(raw.copy(), chunk_size=3)

        arr[1:4, 1] = (3, 4., (5, 6))
        arr[1:4, 2] = 8
        arr[1:3] = np.arange(5)
        arr[2:4] = np.arange(10).reshape(2, 5)
        arr[0] = np.arange(5)

        raw[1:4, 1] = (3, 4., (5, 6))
        raw[1:4, 2] = 8
        raw[1:3] = np.arange(5)
        raw[2:4] = np.arange(10).reshape(2, 5)
        raw[0] = np.arange(5)

        res = self.executor.execute_tensor(arr, concat=True)[0]
        self.assertEqual(arr.dtype, raw.dtype)
        self.assertEqual(arr.shape, raw.shape)
        np.testing.assert_array_equal(res, raw)

    def testTakeExecution(self):
        data = np.random.rand(10, 20, 30)
        t = tensor(data, chunk_size=10)

        a = t.take([4, 1, 2, 6, 200])

        res = self.executor.execute_tensor(a, concat=True)[0]
        expected = np.take(data, [4, 1, 2, 6, 200])
        np.testing.assert_array_equal(res, expected)

        a = take(t, [5, 19, 2, 13], axis=1)

        res = self.executor.execute_tensor(a, concat=True)[0]
        expected = np.take(data, [5, 19, 2, 13], axis=1)
        np.testing.assert_array_equal(res, expected)

        with self.assertRaises(ValueError):
            take(t, [1, 3, 4], out=tensor(np.random.rand(4)))

        out = tensor([1, 2, 3, 4])
        a = take(t, [4, 19, 2, 8], out=out)

        res = self.executor.execute_tensor(out, concat=True)[0]
        expected = np.take(data, [4, 19, 2, 8])
        np.testing.assert_array_equal(res, expected)

    def testCompressExecution(self):
        data = np.array([[1, 2], [3, 4], [5, 6]])
        a = tensor(data, chunk_size=1)

        t = compress([0, 1], a, axis=0)

        res = self.executor.execute_tensor(t, concat=True)[0]
        expected = np.compress([0, 1], data, axis=0)
        np.testing.assert_array_equal(res, expected)

        t = compress([0, 1], a, axis=1)

        res = self.executor.execute_tensor(t, concat=True)[0]
        expected = np.compress([0, 1], data, axis=1)
        np.testing.assert_array_equal(res, expected)

        t = a.compress([0, 1, 1])

        res = self.executor.execute_tensor(t, concat=True)[0]
        expected = np.compress([0, 1, 1], data)
        np.testing.assert_array_equal(res, expected)

        t = compress([False, True, True], a, axis=0)

        res = self.executor.execute_tensor(t, concat=True)[0]
        expected = np.compress([False, True, True], data, axis=0)
        np.testing.assert_array_equal(res, expected)

        t = compress([False, True], a, axis=1)

        res = self.executor.execute_tensor(t, concat=True)[0]
        expected = np.compress([False, True], data, axis=1)
        np.testing.assert_array_equal(res, expected)

        with self.assertRaises(np.AxisError):
            compress([0, 1, 1], a, axis=1)

        # test order
        data = np.asfortranarray([[1, 2], [3, 4], [5, 6]])
        a = tensor(data, chunk_size=1)

        t = compress([0, 1, 1], a, axis=0)

        res = self.executor.execute_tensor(t, concat=True)[0]
        expected = np.compress([0, 1, 1], data, axis=0)
        np.testing.assert_array_equal(res, expected)
        self.assertEqual(res.flags['C_CONTIGUOUS'], expected.flags['C_CONTIGUOUS'])
        self.assertEqual(res.flags['F_CONTIGUOUS'], expected.flags['F_CONTIGUOUS'])

        t = compress([0, 1, 1], a, axis=0, out=tensor(np.empty((2, 2), order='F', dtype=int)))

        res = self.executor.execute_tensor(t, concat=True)[0]
        expected = np.compress([0, 1, 1], data, axis=0, out=np.empty((2, 2), order='F', dtype=int))
        np.testing.assert_array_equal(res, expected)
        self.assertEqual(res.flags['C_CONTIGUOUS'], expected.flags['C_CONTIGUOUS'])
        self.assertEqual(res.flags['F_CONTIGUOUS'], expected.flags['F_CONTIGUOUS'])

    def testExtractExecution(self):
        data = np.arange(12).reshape((3, 4))
        a = tensor(data, chunk_size=2)
        condition = mod(a, 3) == 0

        t = extract(condition, a)

        res = self.executor.execute_tensor(t, concat=True)[0]
        expected = np.extract(np.mod(data, 3) == 0, data)
        np.testing.assert_array_equal(res, expected)

    def testChooseExecution(self):
        options.chunk_size = 2

        choices = [[0, 1, 2, 3], [10, 11, 12, 13],
                   [20, 21, 22, 23], [30, 31, 32, 33]]
        a = choose([2, 3, 1, 0], choices)

        res = self.executor.execute_tensor(a, concat=True)[0]
        expected = np.choose([2, 3, 1, 0], choices)

        np.testing.assert_array_equal(res, expected)

        a = choose([2, 4, 1, 0], choices, mode='clip')  # 4 goes to 3 (4-1)
        expected = np.choose([2, 4, 1, 0], choices, mode='clip')

        res = self.executor.execute_tensor(a, concat=True)[0]
        np.testing.assert_array_equal(res, expected)

        a = choose([2, 4, 1, 0], choices, mode='wrap')  # 4 goes to (4 mod 4)
        expected = np.choose([2, 4, 1, 0], choices, mode='wrap')  # 4 goes to (4 mod 4)

        res = self.executor.execute_tensor(a, concat=True)[0]
        np.testing.assert_array_equal(res, expected)

        a = [[1, 0, 1], [0, 1, 0], [1, 0, 1]]
        choices = [-10, 10]

        b = choose(a, choices)
        expected = np.choose(a, choices)

        res = self.executor.execute_tensor(b, concat=True)[0]
        np.testing.assert_array_equal(res, expected)

        a = np.array([0, 1]).reshape((2, 1, 1))
        c1 = np.array([1, 2, 3]).reshape((1, 3, 1))
        c2 = np.array([-1, -2, -3, -4, -5]).reshape((1, 1, 5))

        b = choose(a, (c1, c2))
        expected = np.choose(a, (c1, c2))

        res = self.executor.execute_tensor(b, concat=True)[0]
        np.testing.assert_array_equal(res, expected)

        # test order
        a = np.array([0, 1]).reshape((2, 1, 1), order='F')
        c1 = np.array([1, 2, 3]).reshape((1, 3, 1), order='F')
        c2 = np.array([-1, -2, -3, -4, -5]).reshape((1, 1, 5), order='F')

        b = choose(a, (c1, c2))
        expected = np.choose(a, (c1, c2))

        res = self.executor.execute_tensor(b, concat=True)[0]
        np.testing.assert_array_equal(res, expected)
        self.assertEqual(res.flags['C_CONTIGUOUS'], expected.flags['C_CONTIGUOUS'])
        self.assertEqual(res.flags['F_CONTIGUOUS'], expected.flags['F_CONTIGUOUS'])

        b = choose(a, (c1, c2), out=tensor(np.empty(res.shape, order='F')))
        expected = np.choose(a, (c1, c2), out=np.empty(res.shape, order='F'))

        res = self.executor.execute_tensor(b, concat=True)[0]
        np.testing.assert_array_equal(res, expected)
        self.assertEqual(res.flags['C_CONTIGUOUS'], expected.flags['C_CONTIGUOUS'])
        self.assertEqual(res.flags['F_CONTIGUOUS'], expected.flags['F_CONTIGUOUS'])

    def testUnravelExecution(self):
        a = tensor([22, 41, 37], chunk_size=1)
        t = stack(unravel_index(a, (7, 6)))

        res = self.executor.execute_tensor(t, concat=True)[0]
        expected = np.stack(np.unravel_index([22, 41, 37], (7, 6)))

        np.testing.assert_array_equal(res, expected)

    def testNonzeroExecution(self):
        data = np.array([[1, 0, 0], [0, 2, 0], [1, 1, 0]])
        x = tensor(data, chunk_size=2)
        t = hstack(nonzero(x))

        res = self.executor.execute_tensor(t, concat=True)[0]
        expected = np.hstack(np.nonzero(data))

        np.testing.assert_array_equal(res, expected)

        t = hstack((x > 1).nonzero())

        res = self.executor.execute_tensor(t, concat=True)[0]
        expected = np.hstack(np.nonzero(data > 1))

        np.testing.assert_array_equal(res, expected)

    def testFlatnonzeroExecution(self):
        x = arange(-2, 3, chunk_size=2)

        t = flatnonzero(x)

        res = self.executor.execute_tensor(t, concat=True)[0]
        expected = np.flatnonzero(np.arange(-2, 3))

        np.testing.assert_equal(res, expected)

    def testFillDiagonalExecution(self):
        # 2-d
        raws = [np.random.rand(30, 11),
                np.random.rand(15, 15),
                np.random.rand(11, 30),
                sps.random(30, 11, density=0.1, format='csr')]

        def copy(x):
            if hasattr(x, 'nnz'):
                # sparse
                return x.A
            else:
                return x.copy()

        for raw in raws:
            # test 1 chunk, wrap=False
            t = tensor(raw, chunk_size=30)
            fill_diagonal(t, 1)

            res = self.executor.execute_tensor(t, concat=True)[0]
            expected = copy(raw)
            np.fill_diagonal(expected, 1)

            np.testing.assert_array_equal(np.asarray(res), expected)

            # test 1 chunk, wrap=True
            t = tensor(raw, chunk_size=30)
            fill_diagonal(t, 1, wrap=True)

            res = self.executor.execute_tensor(t, concat=True)[0]
            expected = copy(raw)
            np.fill_diagonal(expected, 1, wrap=True)

            np.testing.assert_array_equal(np.asarray(res), expected)

            # test multiple chunks, wrap=False
            t = tensor(raw, chunk_size=(12, 4))
            fill_diagonal(t, 1)

            res = self.executor.execute_tensor(t, concat=True)[0]
            expected = copy(raw)
            np.fill_diagonal(expected, 1)

            np.testing.assert_array_equal(np.asarray(res), expected)

            t = tensor(raw, chunk_size=(4, 12))
            fill_diagonal(t, 1)

            res = self.executor.execute_tensor(t, concat=True)[0]
            expected = copy(raw)
            np.fill_diagonal(expected, 1)

            np.testing.assert_array_equal(np.asarray(res), expected)

            # test multiple chunk, val with list type
            t = tensor(raw, chunk_size=(12, 4))
            fill_diagonal(t, [1, 2, 3])

            res = self.executor.execute_tensor(t, concat=True)[0]
            expected = copy(raw)
            np.fill_diagonal(expected, [1, 2, 3])

            np.testing.assert_array_equal(np.asarray(res), expected)

            # test multiple chunk, val with tensor type
            t = tensor(raw, chunk_size=(12, 4))
            fill_diagonal(t, tensor([1, 2, 3]))

            res = self.executor.execute_tensor(t, concat=True)[0]
            expected = copy(raw)
            np.fill_diagonal(expected, [1, 2, 3])

            np.testing.assert_array_equal(np.asarray(res), expected)

            # test multiple chunks, wrap=True
            t = tensor(raw, chunk_size=(12, 4))
            fill_diagonal(t, 1, wrap=True)

            res = self.executor.execute_tensor(t, concat=True)[0]
            expected = copy(raw)
            np.fill_diagonal(expected, 1, wrap=True)

            np.testing.assert_array_equal(np.asarray(res), expected)

            t = tensor(raw, chunk_size=(4, 12))
            fill_diagonal(t, 1, wrap=True)

            res = self.executor.execute_tensor(t, concat=True)[0]
            expected = copy(raw)
            np.fill_diagonal(expected, 1, wrap=True)

            np.testing.assert_array_equal(np.asarray(res), expected)

            # test multiple chunk, val with list type
            t = tensor(raw, chunk_size=(12, 4))
            fill_diagonal(t, [1, 2, 3], wrap=True)

            res = self.executor.execute_tensor(t, concat=True)[0]
            expected = copy(raw)
            np.fill_diagonal(expected, [1, 2, 3], wrap=True)

            np.testing.assert_array_equal(np.asarray(res), expected)

            # test multiple chunk, val with tensor type
            t = tensor(raw, chunk_size=(12, 4))
            fill_diagonal(t, tensor([[1, 2], [3, 4]]), wrap=True)

            res = self.executor.execute_tensor(t, concat=True)[0]
            expected = copy(raw)
            np.fill_diagonal(expected, [1, 2, 3, 4], wrap=True)

            np.testing.assert_array_equal(np.asarray(res), expected)

        # 3-d
        raw = np.random.rand(11, 11, 11)

        expected = raw.copy()
        np.fill_diagonal(expected, 1)
        expected2 = raw.copy()
        np.fill_diagonal(expected2, 1, wrap=True)
        np.testing.assert_array_equal(expected, expected2)

        # test 1 chunk
        t = tensor(raw, chunk_size=30)
        fill_diagonal(t, 1)

        res = self.executor.execute_tensor(t, concat=True)[0]

        np.testing.assert_array_equal(res, expected)

        t = tensor(raw, chunk_size=30)
        # wrap = True does not take effect when ndim > 2
        fill_diagonal(t, 1, wrap=True)

        res = self.executor.execute_tensor(t, concat=True)[0]

        np.testing.assert_array_equal(res, expected)

        # test multiple chunk
        t = tensor(raw, chunk_size=(3, 4, 5))
        fill_diagonal(t, 1)

        res = self.executor.execute_tensor(t, concat=True)[0]

        np.testing.assert_array_equal(res, expected)

        t = tensor(raw, chunk_size=(3, 4, 5))
        # wrap = True does not take effect when ndim > 2
        fill_diagonal(t, 1, wrap=True)

        res = self.executor.execute_tensor(t, concat=True)[0]

        np.testing.assert_array_equal(res, expected)

        # test val with list type
        t = tensor(raw, chunk_size=(3, 4, 5))
        fill_diagonal(t, [[1, 2], [3, 4]])

        res = self.executor.execute_tensor(t, concat=True)[0]
        expected = raw.copy()
        np.fill_diagonal(expected, [1, 2, 3, 4])

        np.testing.assert_array_equal(res, expected)

        # test val with tensor type
        t = tensor(raw, chunk_size=(3, 4, 5))
        fill_diagonal(t, tensor([1, 2, 3]))

        res = self.executor.execute_tensor(t, concat=True)[0]
        expected = raw.copy()
        np.fill_diagonal(expected, [1, 2, 3])

        np.testing.assert_array_equal(res, expected)

        # test val with tensor type which ndim == 0
        t = tensor(raw, chunk_size=(3, 4, 5))
        fill_diagonal(t, tensor([1, 2, 3]).sum())

        res = self.executor.execute_tensor(t, concat=True)[0]
        expected = raw.copy()
        np.fill_diagonal(expected, 6)

        np.testing.assert_array_equal(res, expected)

        # test val with ndarray type which size is too long
        t = tensor(raw, chunk_size=(3, 4, 5))
        fill_diagonal(t, np.arange(20))

        res = self.executor.execute_tensor(t, concat=True)[0]
        expected = raw.copy()
        np.fill_diagonal(expected, np.arange(20))

        np.testing.assert_array_equal(res, expected)

max_stars_repo_name: snsnlou/mars

max_issues_repo_name: snsnlou/mars

max_forks_repo_name: snsnlou/mars

reason: import numpy,import scipy


text:
#__BEGIN_LICENSE__
# Copyright (c) 2017, United States Government, as represented by the
# Administrator of the National Aeronautics and Space Administration.
# All rights reserved.
#
# The GeoRef platform is licensed under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# http://www.apache.org/licenses/LICENSE-2.0.
#
# Unless required by applicable law or agreed to in writing, software distributed
# under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
# CONDITIONS OF ANY KIND, either express or implied. See the License for the
# specific language governing permissions and limitations under the License.
#__END_LICENSE__

import os
import sys
import math
import subprocess
import traceback
import json
import numpy
import shutil
import piexif
import datetime
import tempfile

import IrgGeoFunctions
import offline_config

basepath    = os.path.abspath(sys.path[0]) # Scott debug
sys.path.insert(0, basepath + '/../geocamTiePoint')
sys.path.insert(0, basepath + '/../geocamUtilWeb')

from geocamTiePoint import transform
from django.conf import settings

# These codes are used to define the confidence in the detected image registration
CONFIDENCE_NONE = 0
CONFIDENCE_LOW  = 1
CONFIDENCE_HIGH = 2

CONFIDENCE_STRINGS = ['NONE', 'LOW', 'HIGH']

def confidenceFromString(s):
    if s == 'LOW':
        return CONFIDENCE_LOW
    if s == 'HIGH':
        return CONFIDENCE_HIGH
    return CONFIDENCE_NONE

# Geographig projection used to write output files
OUTPUT_PROJECTION = '+proj=longlat +datum=WGS84'

# TODO: Split up this file!

# TODO: Make sure that information entered via the GUI gets handled properly!

class TemporaryDirectory(object):
    """Context manager for tempfile.mkdtemp() so it's usable with "with" statement."""
    def __enter__(self):
        self.name = tempfile.mkdtemp()
        return self.name

    def __exit__(self, exc_type, exc_value, traceback):
        if '/tmp' in self.name:
            shutil.rmtree(self.name)

def safeMakeDir(folder):
    '''Make sure a folder exists and ignore any errors.'''
    try:
        os.mkdir(folder)
    except:
        pass

def cropImageLabel(jpegPath, outputPath):
    '''Create a copy of a jpeg file with any label cropped off'''

    # Check if there is a label using a simple command line tool
    cmdPath = settings.PROJ_ROOT + '/apps/georef_imageregistration/build/detectImageTag'
    cmd    = [cmdPath, jpegPath]
    print cmd
    p = subprocess.Popen(cmd, stdout=subprocess.PIPE)
    textOutput, err = p.communicate()

    if 'NO_LABEL' in textOutput:
        # The file is fine, just copy it.
        print 'Copy ' + jpegPath +' --> '+ outputPath
        try:
            shutil.copy(jpegPath, outputPath)
        except:
            print 'Copy failed, try again!'
#            shutil.copy(jpegPath, outputPath)
            os.system('cp ' + jpegPath +' '+ outputPath)
            if not os.path.exists(outputPath):
                raise Exception('Still failed!')
            print 'Retry successful!'
    else:
        lines = textOutput.strip().split('\n') # Get the parts of the last line
        parts = lines[-1].split()
        if len(parts) != 3:
            raise Exception('Error running detectImageTag, got response: ' + textOutput)
        side     = parts[1]
        labelPos = int(parts[2])
        print 'Detected image label: ' + side + ' at index ' + str(labelPos)
        # Trim the label off of the bottom of the image
        imageSize = IrgGeoFunctions.getImageSize(jpegPath)
        x = 0
        y = 0
        width  = imageSize[0]
        height = imageSize[1]
        if side == 'LEFT':
            x = labelPos
            width = width - labelPos
        if side == 'RIGHT':
            width = labelPos
        if side == 'TOP':
            y = labelPos
            height = height - labelPos
        if side == 'BOTTOM':
            height = labelPos
        cmd = ('gdal_translate -of jpeg -srcwin %d %d %d %d %s %s'
                % (x, y, width, height, jpegPath, outputPath))
        print cmd
        os.system(cmd)

def updateExif(exifSourcePath, geotiffFilePath):
    '''Copy EXIF info from the source file to the geotiff file'''

    # These two files contain a bunch of arguments that are read by exiftool
    creationArgsFile = settings.STATIC_ROOT + '/georef_imageregistration/creation-args.txt'
    extrasArgsFile = settings.STATIC_ROOT + '/georef_imageregistration/extras-args.txt'

    outputFileName = geotiffFilePath
    # rename the geotiff input to "temp" so that we can generate a new geotiffFilePath geotiff file with updated exif.
    #     tempFileName = geotiffFilePath + ".temp"
    filename, file_extension = os.path.splitext(outputFileName)
    tempFileName = os.path.dirname(outputFileName) + "/temp-%s%s" % (datetime.datetime.utcnow().strftime('%Y-%m-%d_%H:%M:%S%Z'), file_extension)

    os.rename(outputFileName, tempFileName)
    print "Exif Source Path: %s"

    try:
        exifCmd = 'exiftool -tagsFromFile %s -@ %s -@ %s -ModifyDate="`date \'+%%Y:%%m:%%d %%H:%%M:%%S\'`" -EXIF:Software="%s" -o %s %s' \
                    % (exifSourcePath, creationArgsFile, extrasArgsFile, "GeoRef", outputFileName, tempFileName)
        os.system(exifCmd)
        os.remove(tempFileName)
    except Exception as e:
        os.rename(tempFileName, outputFileName)
        print "Failed to copy over the exif information. %s" % e

def getIdentityTransform():
    '''Return an identity transform from the transform.py file'''
    return transform.ProjectiveTransform(numpy.matrix([[1,0,0],[0,1,0],[0,0,1]],dtype='float64'))

def isPixelValid(pixel, size):
    '''Simple pixel bounds check'''
    return ((pixel[0] >= 0      ) and (pixel[1] >= 0      ) and
            (pixel[0] <  size[0]) and (pixel[1] <  size[1])    )

def estimateGroundResolution(focalLength, width, height, sensorWidth, sensorHeight,
                             stationLon, stationLat, stationAlt, centerLon, centerLat, tilt=0.0):
    '''Estimates a ground resolution in meters per pixel using the focal length.'''
    # TODO: Use the angle to get a more accurate computation!
    # Divide by four since it is half the distance squared
    sensorDiag = math.sqrt(sensorWidth*sensorWidth/4 + sensorHeight*sensorHeight/4)
    pixelDiag  = math.sqrt(width*width/4 + height*height/4)

    angle = math.atan2(sensorDiag, focalLength)

    NAUTICAL_MILES_TO_METERS = 1852.0
    DEGREES_TO_RADIANS = 3.14159 / 180.0
    distance   = (stationAlt * NAUTICAL_MILES_TO_METERS) / math.cos(tilt*DEGREES_TO_RADIANS)
    groundDiag = math.tan(angle) * distance

    #print str(math.cos(tilt*DEGREES_TO_RADIANS))
    #print 'sensorDiag = ' + str(sensorDiag)
    #print 'angle      = ' + str(angle)
    #print 'distance   = ' + str(distance)
    #print 'groundDiag = ' + str(groundDiag)
    #print 'pixelDiag  = ' + str(pixelDiag)

    pixelSize = groundDiag / pixelDiag # Meters / pixels
    return pixelSize

def getWorkingDir(mission, roll, frame):
    # Break up frames so that there are 1000 per folder
    FRAMES_PER_FOLDER = 1000
    FRAME_DIGITS = 6
    frameFolderNum = (int(frame) // FRAMES_PER_FOLDER)*1000
    frameFolder    = str(frameFolderNum).rjust(FRAME_DIGITS, '0')

    # Store data in /mission/roll/frameF/file, skip roll if E.
    safeMakeDir(offline_config.OUTPUT_IMAGE_FOLDER)

    # make mission directory
    subFolder = os.path.join(offline_config.OUTPUT_IMAGE_FOLDER, mission)
    safeMakeDir(subFolder)

    # make roll directory
    if not (roll.lower() == 'e'):
        subFolder = os.path.join(subFolder, roll)
        safeMakeDir(subFolder)

    # make frame directory
    subFolder = os.path.join(subFolder, frameFolder)
    safeMakeDir(subFolder)

    return subFolder

def getZipFilePath(mission, roll, frame):
    '''Get the full path for an output zip file'''

    # Store data in /mission/mission-roll-frame/file
    issIdFolder = mission + '-' + roll + '-' + frame
    safeMakeDir(offline_config.OUTPUT_ZIP_FOLDER)

    subFolder = os.path.join(offline_config.OUTPUT_ZIP_FOLDER, mission)
    safeMakeDir(subFolder)

    subFolder = os.path.join(subFolder, issIdFolder)
    safeMakeDir(subFolder)

    return subFolder

def getWorkingPath(mission, roll, frame):
    '''Get a good location to process this image.'''

    if offline_config.USE_RAW:
        ext = '.tif'
    else:
        ext = '.jpg'

    # Generate a file name similar to the RAW storage scheme
    FRAME_DIGITS = 6
    zFrame   = frame.rjust(FRAME_DIGITS, '0')
    filename = mission.lower() + roll.lower() + zFrame + ext

    subFolder = getWorkingDir(mission, roll, frame)

    return os.path.join(subFolder, filename)

# TODO: Move this to the transform.py file?
def getFitError(imageInliers, gdcInliers):
    '''Computes the RMS error of the transform fit of the provided points.'''

    # The error is computed in pixels, but meters might be better.
    lonlatToPixels = transform.ProjectiveTransform.fit(numpy.asarray(imageInliers),
                                                       numpy.asarray(gdcInliers))

    numPoints = float(len(imageInliers))
    rms = 0.0
    for (pixel, lonlat) in zip(imageInliers, gdcInliers):
        tformPixel = lonlatToPixels.forward(lonlat)
        dx         = pixel[0] - tformPixel[0]
        dy         = pixel[1] - tformPixel[1]
        errSq      = dx*dx + dy*dy
        #print 'pixel = ' + str(pixel) +', tform = ' + str(tformPixel)
        rms += errSq / numPoints

    return math.sqrt(rms)

def getPixelToGdcTransform(imagePath, pixelToProjectedTransform=None):
    '''Returns a pixel to GDC transform.
       The input image must either be a nicely georegistered image from Earth Engine
       or a pixel to projected coordinates transform must be provided.'''

    if pixelToProjectedTransform:
        # Have image to projected transform, convert it to an image to GDC transform.

        # Use the simple file info call (the input file may not have geo information)
        (width, height) = IrgGeoFunctions.getImageSize(imagePath)

        imagePoints = []
        gdcPoints   = []

        # Loop through a spaced out grid of pixels in the image
        pointPixelSpacing = (width + height) / 20 # Results in about 100 points
        for r in range(0, width, pointPixelSpacing):
            for c in range(0, height, pointPixelSpacing):
                # This pixel --> projected coords --> lonlat coord
                thisPixel           = numpy.array([float(c), float(r)])
                projectedCoordinate = pixelToProjectedTransform.forward(thisPixel)
                gdcCoordinate       = transform.metersToLatLon(projectedCoordinate)

                imagePoints.append(thisPixel)
                gdcPoints.append(gdcCoordinate)
        # Solve for a transform with all of these point pairs
        pixelToGdcTransform = transform.getTransform(numpy.asarray(gdcPoints),
                                                     numpy.asarray(imagePoints))

    else: # Using a reference image from EE which will have nice bounds.

        # Use the more thorough file info call
        stats  = IrgGeoFunctions.getImageGeoInfo(imagePath, False)
        (width, height) = stats['image_size']
        (minLon, maxLon, minLat, maxLat) = stats['lonlat_bounds']

        # Make a transform from ref pixel to GDC using metadata on disk
        xScale = (maxLon - minLon) / width
        yScale = (maxLat - minLat) / height
        transformMatrix = numpy.array([[xScale,  0,      minLon],
                                       [0,      -yScale, maxLat],
                                       [0 ,      0,      1     ]])
        pixelToGdcTransform = transform.LinearTransform(transformMatrix)

    return pixelToGdcTransform

def getGdcTransformFromPixelTransform(imageSize, pixelTransform, refImageGdcTransform):
    '''Converts an image-to-image transform chained with a GDC transform
       to a pixel to GDC transform for this image.'''

    # Convert the image-to-image transform parameters to a class
    temp = numpy.array([tform[0:3], tform[3:6], tform[6:9]] )
    imageToRefTransform = transform.ProjectiveTransform(temp)

    newImageSize = ImageFetcher.miscUtilities.getImageSize(newImagePath)

    # Generate a list of point pairs
    imagePoints = []
    gdcPoints   = []

    # Loop through an evenly spaced grid of pixels in the new image
    # - For each pixel, compute the desired output coordinate
    pointPixelSpacing = (newImageSize[0] + newImageSize[1]) / 20 # Results in about 100 points
    for r in range(0, newImageSize[0], pointPixelSpacing):
        for c in range(0, newImageSize[1], pointPixelSpacing):
            # Get pixel in new image and matching pixel in the reference image,
            #  then pass that into the GDC transform.
            thisPixel       = numpy.array([float(c), float(r)])
            pixelInRefImage = pixelTransform.forward(thisPixel)
            gdcCoordinate   = refImageGdcTransform.forward(pixelInRefImage)

            imagePoints.append(thisPixel)
            gdcPoints.append(gdcCoordinate)

    # Compute a transform object that converts from the new image to projected coordinates
    imageToGdcTransform = transform.getTransform(numpy.asarray(worldPoints),
                                                 numpy.asarray(gdcPoints))

    return imageToGdcTransform

def alignImages(testImagePath, refImagePath, workPrefix, force, debug=False, slowMethod=False):
    '''Call the C++ code to find the image alignment'''

    transformPath = workPrefix + '-transform.txt'

    # The computed transform is from testImage to refImage

    # Run the C++ command if we need to generate the transform
    if (not os.path.exists(transformPath) or force):
        if os.path.exists(transformPath):
            os.remove(transformPath) # Clear out any old results

        print 'Running C++ image alignment tool...'
        cmdPath = settings.PROJ_ROOT + '/apps/georef_imageregistration/build/registerGeocamImage'

        cmd = [cmdPath, refImagePath, testImagePath, transformPath]
        if debug: cmd.append('y')
        else:     cmd.append('n')
        if slowMethod: cmd.append('y')
        else:          cmd.append('n')

        if debug:
            print cmd
        #os.system('build/registerGeocamImage '+ refImagePath+' '+testImagePath+' '+transformPath+' --debug')
        p = subprocess.Popen(cmd, stdout=subprocess.PIPE)
        textOutput, err = p.communicate()
        print textOutput

    if not os.path.exists(transformPath):
        raise Exception('Failed to compute transform!')
        tform = [1, 0, 0, 0, 1, 0, 0, 0, 1]
        confidence = CONFIDENCE_NONE
        return (tform, confidence)

    # Load the computed transform, confidence, and inliers.
    handle   = open(transformPath, 'r')
    fileText = handle.read()
    handle.close()
    lines = fileText.split('\n')
    confidence = CONFIDENCE_NONE
    if 'CONFIDENCE_LOW' in lines[0]:
        confidence = CONFIDENCE_LOW
    if 'CONFIDENCE_HIGH' in lines[0]:
        confidence = CONFIDENCE_HIGH
    tform = [float(f) for f in lines[2].split(',')] +  \
            [float(f) for f in lines[3].split(',')] +  \
            [float(f) for f in lines[4].split(',')]
    refInliers  = []
    testInliers = []
    for line in lines[6:]:
        if len(line) < 2:
            break
        numbers = [float(f) for f in line.split(',')]
        refInliers.append( (numbers[0], numbers[1]))
        testInliers.append((numbers[2], numbers[3]))

    #print tform
    #print refInliers
    #print testInliers

    return (tform, confidence, testInliers, refInliers)

def alignScaledImages(testImagePath, refImagePath, testImageScaling, workPrefix, force, debug=False, slowMethod=False):
    '''Align a possibly higher resolution input image with a reference image.
       This call handles the fact that registration should be performed at the same resolution.'''

    # If the scale is within this amount, don't bother rescaling.
    SCALE_TOLERANCE = 0.10
    if abs(testImageScaling - 1.0) < SCALE_TOLERANCE:
        # In this case just use the lower level function
        return alignImages(testImagePath, refImagePath, workPrefix, force, debug, slowMethod)

    # Generate a scaled version of the input image
    scaledImagePath = workPrefix + '-scaledInputImage.tif'
    outPercentage = str(testImageScaling*100.0)
    cmd = 'gdal_translate -outsize ' + outPercentage +'% ' + outPercentage +'% '+ testImagePath +' '+ scaledImagePath
    print cmd
    os.system(cmd)
    if not os.path.exists(scaledImagePath):
        raise Exception('Failed to rescale image with command:\n' + cmd)

    # Call alignment with the scaled version
    (scaledTform, confidence, scaledImageInliers, refInliers) = \
            alignImages(scaledImagePath, refImagePath, workPrefix, force, debug, slowMethod)

    # De-scale the output transform so that it applies to the input sized image.
    testInliers = []
    for pixel in scaledImageInliers:
        testInliers.append( (pixel[0]/testImageScaling, pixel[1]/testImageScaling) )
    print 'scaled tform = \n' + str(scaledTform)
    tform = scaledTform
    for i in [0, 1, 3, 4, 6, 7]: # Scale the six coefficient values
        tform[i] = tform[i] * testImageScaling
    print 'tform = \n' + str(tform)

    if not debug: # Clean up the scaled image
        os.remove(scaledImagePath)

    return (tform, confidence, testInliers, refInliers)

def logRegistrationResults(outputPath, pixelTransform, confidence,
                           refImagePath, imageToGdcTransform=None):
    '''Log the registration results so they can be read back in later.
       Provides enough data so that the image can be '''

    # TODO: Delete this function?  What was it for?
    raise Exception('BROKEN FUNCTION!')

    # Generate a scaled version of the input image
    scaledImagePath = workPrefix + '-scaledInputImage.tif'
    outPercentage = str(testImageScaling*100.0)
    cmd = 'gdal_translate -outsize ' + outPercentage +'% ' + outPercentage +'% '+ testImagePath +' '+ scaledImagePath
    print cmd
    os.system(cmd)
    if not os.path.exists(scaledImagePath):
        raise Exception('Failed to rescale image with command:\n' + cmd)

    # Call alignment with the scaled version
    (scaledTform, confidence, scaledImageInliers, refInliers) = \
            alignImages(scaledImagePath, refImagePath, workPrefix, force, debug, slowMethod)

    # De-scale the output transform so that it applies to the input sized image.
    testInliers = []
    for pixel in scaledImageInliers:
        testInliers.append( (pixel[0]/testImageScaling, pixel[1]/testImageScaling) )
    print 'scaled tform = \n' + str(scaledTform)
    tform = scaledTform
    for i in [0, 1, 3, 4, 6, 7]: # Scale the six coefficient values
        tform[i] = tform[i] * testImageScaling
    print 'tform = \n' + str(tform)

    if not debug: # Clean up the scaled image
        os.remove(scaledImagePath)

    return (tform, confidence, testInliers, refInliers)

def convertGcps(inputGdcCoords, imageToProjectedTransform, width, height):
    '''Given a set of GDC coordinates and image registration info,
       produces a set of GCPs for that image.'''

    size = (width, height)

    imageCoords = []
    gdcCoords   = []

    for inputCoord in inputGdcCoords:
        # Convert from GDC to Google projected coordinate
        coordMeters = transform.lonLatToMeters(inputCoord)

        # Get image coordinate
        pixel = imageToProjectedTransform.reverse(coordMeters)

        # Only keep GCPs that actually fall within the image
        if isPixelValid(pixel, size):
            gdcCoords.append(inputCoord)
            imageCoords.append(pixel)

    return (imageCoords, gdcCoords)

def recordOutputImages(sourceImagePath, exifSourcePath, outputPrefix, imageInliers,
                       gdcInliers, minUncertaintyMeters, centerPointSource,
                       isManualRegistration=False, overwrite=True):
    '''Generates all the output image files that we create for each successfully processed image.'''

    # We generate two pairs of images, one containing the image data
    #  and another with the same format but containing the uncertainty distances.
    outputPrefix = outputPrefix + '-' + centerPointSource
    uncertaintyOutputPrefix = outputPrefix + '-uncertainty'
    rawUncertaintyPath      = outputPrefix + '-uncertainty_raw.tif'

    # Create the raw uncertainty image
    (width, height) = IrgGeoFunctions.getImageSize(sourceImagePath)
    posError = generateUncertaintyImage(width, height, imageInliers,
                                        minUncertaintyMeters, rawUncertaintyPath)

    # Get a measure of the fit error
    fitError = getFitError(imageInliers, gdcInliers)

    # Generate the two pairs of images in the same manner
    try:
        (noWarpOutputPath, warpOutputPath) = \
            generateGeotiff(sourceImagePath, outputPrefix, imageInliers, gdcInliers,
                                            posError, fitError, isManualRegistration,
                                            exifSourcePath,
                                            writeHeaders=True, overwrite=True)
    except Exception as e:
        print str(e)

    try:
        (noWarpOutputPath, warpOutputPath) = \
            generateGeotiff(rawUncertaintyPath, uncertaintyOutputPrefix, imageInliers, gdcInliers,
                                                posError, fitError, isManualRegistration,
                                                exifSourcePath,
                                                writeHeaders=False, overwrite=True)
    except Exception as e:
        print str(e)

    # Clean up the raw uncertainty image and any extraneous files
    rawXmlPath = rawUncertaintyPath + '.aux.xml'
    os.remove(rawUncertaintyPath)
    if os.path.exists(rawXmlPath):
        os.remove(rawXmlPath)

def generateUncertaintyImage(width, height, imageInliers, minUncertainty, outputPath):
    '''Given a list of GCPs in an image, generate a distance image containing the
       distance from each pixel to the nearest GCP location.
       Returns the RMS error.'''

    # Create a white image with black dots at each GCP coordinate
    drawLine = ''
    for point in imageInliers:
        drawLine += (" point "+ str(point[0]) +","+ str(point[1]))

    tempPath1 = outputPath + '-tempDot.tif'
    tempPath2 = outputPath + '-tempDist.tif'
    cmd = ("convert +depth -size "+str(width)+"x"+str(height)+
           " xc:white  -fill black -draw '"+drawLine+"' " +tempPath1)
    #print cmd
    os.system(cmd)
    if not os.path.exists(tempPath1):
        print "tempPath1 is %s" % tempPath1
        raise Exception('Failed to generate GCP point image!')

    # Get the distance from each
    cmd2 = ('convert '+tempPath1+' -morphology Distance Euclidean:1,1 '+ tempPath2)
    #print cmd2
    os.system(cmd2)
    if not os.path.exists(tempPath2):
        raise Exception('Failed to generate GCP distance image!')

    # TODO: Improve this calculation!

    # For each pixel away from a GCP, the uncertainty increases by this
    #  fraction of the minimum uncertainty.
    UNCERTAINTY_STEP_FRACTION = 0.03
    uncertaintyStep = minUncertainty * UNCERTAINTY_STEP_FRACTION

    # Figure out the uncertainty range in the output image and generate a scale string.
    UINT16_MAX = 65535
    maxUncertainty    = UINT16_MAX * UNCERTAINTY_STEP_FRACTION + minUncertainty
    uncertaintyString = ((' 0 %d %f %f ') % (UINT16_MAX, minUncertainty, maxUncertainty))

    # Use gdal_translate to convert from Uint16 to a scaled 32 bit floating point image
    #  with the final error numbers.
    cmd3 = ('gdal_translate -b 1 -ot Float32 -scale '+uncertaintyString+ tempPath2 +' '+ outputPath)
    #print cmd3
    os.system(cmd3)
    if not os.path.exists(outputPath):
        raise Exception('Failed to generate uncertainty image!')

    # Compute the RMS error using a simple command line tool
    cmdPath = settings.PROJ_ROOT + '/apps/georef_imageregistration/build/computeImageRms'
    cmd4    = [cmdPath, outputPath]
    #print cmd4
    p = subprocess.Popen(cmd4, stdout=subprocess.PIPE)
    textOutput, err = p.communicate()

    # Parse the line "RMS: 123"
    parts    = textOutput.split(':')
    rmsError = float(parts[1])

    # Clean up
    os.remove(tempPath1)
    os.remove(tempPath2)

    return rmsError

def qualityGdalwarp(imagePath, outputPath, imagePoints, gdcPoints):
    '''Use some workarounds to get a higher quality gdalwarp output than is normally possible.'''

    # Generate a high resolution grid of fake GCPs based on a transform we compute,
    # then call gdalwarp using a high order polynomial to accurately match our transform.

    #trans = transform.ProjectiveTransform.fit(numpy.asarray(gdcPoints),numpy.asarray(imagePoints))ls
    trans = transform.getTransform(numpy.asarray(gdcPoints),numpy.asarray(imagePoints))
    transformName = trans.getJsonDict()['type']

    tempPath = outputPath + '-temp.tif'

    # Generate a temporary image containing the grid of fake GCPs
    cmd = ('gdal_translate -co "COMPRESS=LZW" -co "tiled=yes"  -co "predictor=2" -a_srs "'
           + OUTPUT_PROJECTION +'" '+ imagePath +' '+ tempPath)

    # Generate the GCPs in a grid, keeping the total under about 500 points so
    # that GDAL does not complain.
    (width, height) = IrgGeoFunctions.getImageSize(imagePath)
    xStep = width /22
    yStep = height/22
    MAX_DEG_SIZE = 20
    minLon = 999 # Keep track of the lonlat size and don't write if it is too big.
    minLat = 999 # - This would work better if it was in pixels, but how to get that size?
    maxLon = -999
    maxLat = -999
    for r in range(0,height,yStep):
        for c in range(0,width,xStep):
            pixel  = (c,r)
            lonlat = trans.forward(pixel)
            cmd += ' -gcp '+ str(c) +' '+str(r) +' '+str(lonlat[0]) +' '+str(lonlat[1])
            if lonlat[0] < minLon:
                minLon = lonlat[0]
            if lonlat[1] < minLat:
                minLat = lonlat[1]
            if lonlat[0] > maxLon:
                maxLon = lonlat[0]
            if lonlat[1] > maxLat:
                maxLat = lonlat[1]
    #print cmd
    os.system(cmd)
    if max((maxLon - minLon), (maxLat - minLat)) > MAX_DEG_SIZE:
        raise Exception('Warped image is too large to generate!\n'
                        '-> LonLat bounds: ' + str((minLon, minLat, maxLon, maxLat)))

    # Now generate a warped geotiff.
    # - "order 2" looks terrible with fewer GCPs, but "order 1" does not accurately
    #   capture the footprint of higher tilt images.
    # - tps seems to work well with the evenly spaced grid of virtual GCPs.
    cmd = ('gdalwarp -co "COMPRESS=LZW" -co "tiled=yes"  -co "predictor=2"'
               + ' -dstalpha -overwrite -tps -multi -r cubic -t_srs "'
           + OUTPUT_PROJECTION +'" ' + tempPath +' '+ outputPath)
    print cmd
    os.system(cmd)

    # Check output and cleanup
    os.remove(tempPath)
    if not os.path.exists(outputPath):
        raise Exception('Failed to create warped geotiff file: ' + outputPath)

    return transformName

def generateGeotiff(imagePath, outputPrefix, imagePoints, gdcPoints, posError, fitError,
                    isManualRegistration, exifSourcePath, writeHeaders, overwrite=False):
    '''Converts a plain tiff to a geotiff using the provided geo information.'''

    # Check inputs
    if len(imagePoints) != len(gdcPoints):
        raise Exception('Unequal length correspondence points passed to generateGeoTiff!')

    noWarpOutputPath = outputPrefix + '-no_warp.tif'
    warpOutputPath   = outputPrefix + '-warp.tif'

    if isManualRegistration:
        registrationMethodString = 'Manual'
    else:
        registrationMethodString = 'Automated'

    exifData = piexif.load(exifSourcePath)
    acquisitionTime = exifData['Exif'][piexif.ExifIFD.DateTimeOriginal]

    # TODO: Do the manual registrations not use Landsat?
    extraMetadataString = ('-mo POSITION_UNCERTAINTY_RMS_METERS=' + str(posError)
                         + ' -mo FIT_ERROR_RMS_PIXELS=' + str(fitError)
                         + ' -mo REGISTRATION_METHOD=' + registrationMethodString
                         + ' -mo REGISTRATION_REFERENCE=Landsat'
                         + ' -mo ACQUISITION_DATETIME="' + acquisitionTime + '"'
                         + ' -mo ACQUISITION_DATETIME_TIMEZONE="GMT"')

    # First generate a geotiff that adds metadata but does not change the image data.
    # TODO - This may not be useful unless we can duplicate how they processed their RAW data!
    if (not os.path.exists(noWarpOutputPath)) or overwrite:
        print 'Generating UNWARPED output tiff'
        cmd = ('gdal_translate ' + extraMetadataString
               + ' -co "COMPRESS=LZW" -co "tiled=yes"  -co "predictor=2" -a_srs "'
               + OUTPUT_PROJECTION +'" '+ imagePath +' '+ noWarpOutputPath)

        ## Include an arbitrary tag with our estimated error amount
        #if errorMeters:
        #    cmd += ' -mo "REGISTRATION_ERROR=+/-'+str(errorMeters)+' meters" '

        # Include the actual GCPs that we matched to our Landsat data.
        MAX_NUM_GCPS = 500 # Too many GCPs breaks gdal!
        count = 0
        for (imagePoint, gdcPoint) in zip(imagePoints, gdcPoints):
            s = (' -gcp %f %f %f %f' % (imagePoint[0], imagePoint[1], gdcPoint[0], gdcPoint[1]))
            cmd += s
            count += 1
            if count == MAX_NUM_GCPS:
                break

        # Generate the file using gdal_translate
        print cmd
        os.system(cmd)
        if not os.path.exists(noWarpOutputPath):
            raise Exception('Failed to create geotiff file: ' + noWarpOutputPath)

        if writeHeaders:
            generateStandaloneMetadataFile(noWarpOutputPath)

        updateExif(exifSourcePath, noWarpOutputPath)

    # Now generate a warped geotiff.
    if (not os.path.exists(warpOutputPath)) or overwrite:
        print 'Generating WARPED output tiff'

        transformName = qualityGdalwarp(imagePath, warpOutputPath, imagePoints, gdcPoints)

        # Add some extra metadata fields.
        cmd = ('gdal_edit.py -mo TIFFTAG_DOCUMENTNAME= ' + extraMetadataString
                + ' -mo RESAMPLING_METHOD=cubic -mo WARP_TRANSFORM='+transformName+' ' + warpOutputPath)
        #print cmd
        os.system(cmd)

        if not os.path.exists(warpOutputPath):
            raise Exception('Failed to create warped geotiff file: ' + warpOutputPath)

        if writeHeaders:
            generateStandaloneMetadataFile(warpOutputPath)

        updateExif(exifSourcePath, warpOutputPath)

    return (noWarpOutputPath, warpOutputPath)

def generateStandaloneMetadataFile(inputImagePath):
    '''Convert geotiff metadata into a nicely formatted external text file.'''

    # Silently quit when the input image does not exist, that error should already have
    #  been handled.
    if not os.path.exists(inputImagePath):
        return

    outputPath = os.path.splitext(inputImagePath)[0] + '_metadata.txt'

    print 'Generating metadata file ' + outputPath

    cmd = ['gdalinfo', inputImagePath]
    p = subprocess.Popen(cmd, stdout=subprocess.PIPE)
    textOutput, err = p.communicate()

    lines = textOutput.split('\n')

    imageStartLine = metadataStartLine = coordSystemStartLine = gcpStartLine = imageSizeLine = 0
    index = 0
    for line in lines:
        if 'Image Structure Metadata' in line:
            imageStartLine = index
        if line == 'Metadata:':
            metadataStartLine = index
        if 'Coordinate System is' in line:
            coordSystemStartLine = index
        if 'GCP[  0]' in line:
            gcpStartLine = index
        if 'Size is ' in line:
            imageSizeLine = line

        index += 1

    headerText = 'File Type: GeoTiff\n' + imageSizeLine + '\n'
    ipHeader = '''[Tie-Points Used For Georeferencing] {Point format is: (pixel column, pixel row) -> (longitude, latitude, 0)}\n'''

    if gcpStartLine > 0:
        ipText = ipHeader + '\n'.join(lines[gcpStartLine:metadataStartLine])
    else:
        ipText = ''

    geoText = '[Geographic Coordinate Information]\n' + '\n'.join(lines[coordSystemStartLine+1:coordSystemStartLine+9])

    imageText = '[Image Structure Metadata]\n' + '\n'.join(lines[imageStartLine+1:])

    METADATA_SKIP_LIST = ['TIFFTAG', 'AREA_OR_POINT']

    metadataText = lines[metadataStartLine+1:imageStartLine]

    accuracyText = '[Accuracy Measures For Georeferencing Result]\n'
    cameraText = '[Camera Metadata]\n'
    for line in metadataText:
        # Ignore certain lines
        skip = False
        for item in METADATA_SKIP_LIST:
            if item in line:
                skip = True
        if skip:
            continue
        # Otherwise send the line to the correct section
        if 'EXIF' in line:
            cameraText += line + '\n'
        else:
            accuracyText += line + '\n'

    # Generate the output file
    f = open(outputPath, 'w')
    f.write(headerText + '\n')
    f.write(imageText + '\n')
    f.write(geoText      + '\n\n')
    if len(ipText) > 20:
        f.write(ipText       + '\n\n')
    f.write(accuracyText)
    if len(cameraText) > 20: # The warped image does not have this information
        f.write('\n' + cameraText)
    f.close()

    print 'Finished writing header file.'