<filename>examples/dielectric/broad_dielectric_box_charging.py
# # Dielectric charging script development and testing
#
# **7/18/2017**
#
# This notebook will develop and troubleshoot scripts for testing dielectric charging from electron impact.
#
# 1. Simple geometry (fast) - 1 micron x 1 micron (z). 100 x 100 grid
# 2. Rectangular dielectric - 500 nm x 100 nm (z). Centered transversely, closer to anode along z.
# 3. Uniform parallel emission -> 20 particles, user-injected,no transverse velocity, clustered about central axis.
#
# <NAME>

from __future__ import division
import sys
del sys.argv[1:] #  Necessry to run 'from warp import *' in IPython notebook without conflict.
from warp import *
import numpy as np
import matplotlib as mpl
mpl.use('TkAgg')
import matplotlib.pyplot as plt
import os
import pickle
import h5py
from re import findall
from scipy.special import erfinv
from datetime import datetime

import rswarp
from warp.data_dumping.openpmd_diag import ParticleDiagnostic
from rswarp.diagnostics import FieldDiagnostic
from rswarp.utilities.file_utils import cleanupPrevious
from rswarp.utilities.file_utils import readparticles
from rswarp.utilities.file_utils import loadparticlefiles
from rswarp.cathode import sources
from rswarp.cathode import injectors
from warp.particles.singleparticle import TraceParticle

import matplotlib.lines as mlines
import matplotlib.patches as patches

# Constants imports
from scipy.constants import e, m_e, c, k
kb_eV = 8.6173324e-5 #Bolztmann constant in eV/K
kb_J = k #Boltzmann constant in J/K
m = m_e

diagDir = 'diags/xzsolver/hdf5/'
field_base_path = 'diags/fields/'
diagFDir = {'magnetic':'diags/fields/magnetic','electric':'diags/fields/electric'}

# Cleanup previous files
cleanupPrevious(diagDir,diagFDir)

### Grid parameters, Solver, and Boundaries
epsn = 7.

if comm_world.size != 1:
    synchronizeQueuedOutput_mpi4py(out=False, error=False)

print "rank:", comm_world.rank

top.inject = 0
top.npinject = 0

#Dimensions

PLATE_SPACING = 1.e-6 #plate spacing
CHANNEL_WIDTH = 1e-6 #width of simulation box

X_MAX = CHANNEL_WIDTH*0.5
X_MIN = -1.*X_MAX
Y_MAX = CHANNEL_WIDTH*0.5
Y_MIN = -1.*Y_MAX
Z_MIN = 0.
Z_MAX = PLATE_SPACING

#Grid parameters
NUM_X = 100 #256-1#64-1
NUM_Z = 100 #256-1#64-1

# # Solver Geometry

w3d.solvergeom = w3d.XZgeom

# Set boundary conditions
w3d.bound0  = dirichlet
w3d.boundnz = dirichlet
w3d.boundxy = periodic

# Set grid boundaries
w3d.xmmin = X_MIN
w3d.xmmax = X_MAX
w3d.zmmin = 0.
w3d.zmmax = Z_MAX

w3d.nx = NUM_X
w3d.nz = NUM_Z

w3d.dx = (w3d.xmmax-w3d.xmmin)/w3d.nx
w3d.dz = (w3d.zmmax-w3d.zmmin)/w3d.nz

zmesh = np.linspace(0,Z_MAX,NUM_Z+1) #holds the z-axis grid points in an array

ANODE_VOLTAGE = 10.
CATHODE_VOLTAGE = 0.
vacuum_level = ANODE_VOLTAGE - CATHODE_VOLTAGE
beam_beta = 5e-4
#Determine an appropriate time step based upon estimated final velocity
vzfinal = sqrt(2.*abs(vacuum_level)*np.abs(e)/m_e)+beam_beta*c
dt = w3d.dz/vzfinal #5e-15
top.dt = 0.1*dt

if vzfinal*top.dt > w3d.dz:
    print "Time step dt = {:.3e}s does not constrain motion to a single cell".format(top.dt)

#### Set up field solver

top.depos_order = 1
f3d.mgtol = 1e-6 # Multigrid solver convergence tolerance, in volts. 1 uV is default in Warp.
solverE = MultiGrid2DDielectric()
registersolver(solverE)

#### Define conductors and dielectrics using new wrapper

source = ZPlane(zcent=w3d.zmmin+0*w3d.dz,zsign=-1.,voltage=CATHODE_VOLTAGE)
solverE.installconductor(source, dfill=largepos)

plate = ZPlane(voltage=ANODE_VOLTAGE, zcent=Z_MAX-0.*w3d.dz)
solverE.installconductor(plate,dfill=largepos)

box = Box(xsize=0.5*(w3d.xmmax-w3d.xmmin),
          ysize=0.5*(w3d.ymmax-w3d.ymmin),
          zsize=0.1*(w3d.zmmax-w3d.zmmin),
          xcent=0.5*(w3d.xmmax+w3d.xmmin),
          ycent=0.5*(w3d.ymmax+w3d.ymmin),
          zcent=0.8*(w3d.zmmax+w3d.zmmin),
          permittivity=epsn)

solverE.installconductor(box,dfill=largepos)

### Diagnostics

particleperiod = 100
particle_diagnostic_0 = ParticleDiagnostic(period = particleperiod, top = top, w3d = w3d,
                                          species = {species.name: species for species in listofallspecies},
                                          comm_world=comm_world, lparallel_output=False, write_dir = diagDir[:-5])
fieldperiod = 100
efield_diagnostic_0 = FieldDiagnostic.ElectrostaticFields(solver=solverE, top=top, w3d=w3d, comm_world = comm_world,
                                                          period=fieldperiod)

installafterstep(particle_diagnostic_0.write)
installafterstep(efield_diagnostic_0.write)

### Generate and Run

#Generate PIC code and Run Simulation
solverE.mgmaxiters = 1

#prevent GIST from starting upon setup
top.lprntpara = false
top.lpsplots = false
top.verbosity = 0

solverE.mgmaxiters = 10000 #rough approximation needed for initial solve to converge
package("w3d")
generate()
solverE.mgmaxiters = 100

#Need to compute the fields first
epsilon_array = solverE.epsilon/eps0

#Now plot
fig = plt.figure(figsize=(12,6))

X_CELLS = NUM_X
Z_CELLS = NUM_Z

xl = 0
xu = NUM_X
zl = 0
zu = NUM_Z

plt.xlabel("z ($\mu$m)")
plt.ylabel("x ($\mu$m)")
plt.title(r"$\kappa$ across domain - with box")

pxmin = ((X_MAX - X_MIN) / X_CELLS * xl + X_MIN) * 1e6
pxmax = ((X_MAX - X_MIN) / X_CELLS * xu + X_MIN) * 1e6
pzmin = (Z_MIN + zl / Z_CELLS * Z_MAX) * 1e6
pzmax = (Z_MAX * zu / Z_CELLS) * 1e6

plt.xlim(pzmin, pzmax)
plt.ylim(pxmin, pxmax)

eps_plt = plt.imshow(epsilon_array[xl:xu,zl:zu],cmap='viridis',extent=[pzmin, pzmax, pxmin, pxmax],aspect='auto')

cbar = fig.colorbar(eps_plt)
cbar.ax.set_xlabel(r"$\kappa$")
cbar.ax.xaxis.set_label_position('top')

plt.savefig('eps_broad_box.png',bbox_inches='tight')

#### Specify emission

electrons_tracked_t0 = Species(type=Electron, weight=1.0)
ntrack = 20
Z_PART_MIN = w3d.dz/8 #Add a minimum z coordinate to prevent absorption

# Uniform velocity used for all particles
x_vals = np.arange(-0.25*CHANNEL_WIDTH,0.25*CHANNEL_WIDTH,0.5*CHANNEL_WIDTH / ntrack)
y_vals = CHANNEL_WIDTH*(np.random.rand(ntrack)-0.5)
z_vals = np.zeros(ntrack) + Z_PART_MIN #Add a minimum z coordinate to prevent absorption

vx_vals = np.zeros(ntrack)
vy_vals = np.zeros(ntrack)
vz_vals = beam_beta*clight*np.ones(ntrack) #beta = 0.0005

eptclArray = np.asarray([x_vals,vx_vals,y_vals,vy_vals,z_vals,vz_vals]).T

electron_tracker_0 = TraceParticle(js=electrons_tracked_t0.jslist[0],
                     x=x_vals,
                     y=y_vals,
                     z=z_vals,
                     vx=vx_vals,
                     vy=vy_vals,
                     vz=vz_vals)

num_steps = 2000
step(num_steps)

print solverE.getselfe().shape
print getselfe('z').shape
zfield = getselfe('z')
if comm_world.size > 1:
	if comm_world.rank == 0:
		np.save('diel_para.npy',zfield)
elif comm_world.size == 1:
	np.save('diel_ser.npy',zfield)

### Plot particle trajectories

def particle_trace(trace,ntrack):
    kept_electronsx = []
    kept_electronsz = []
    lost_electronsx = []
    lost_electronsz = []
    for electron in range(ntrack):
        for step in range(len(trace.getx(i=electron)) - 1):
            if abs(trace.getx(i=electron)[step] -
                   trace.getx(i=electron)[step + 1]) > (X_MAX - X_MIN) / 2.:
                lost_electronsx.append(trace.getx(i=electron)[0:step])
                lost_electronsz.append(trace.getz(i=electron)[0:step])
                break
            if step == (len(trace.getx(i=electron)) - 2):
                kept_electronsx.append(trace.getx(i=electron))
                kept_electronsz.append(trace.getz(i=electron))
    return [kept_electronsx,kept_electronsz], [lost_electronsx,lost_electronsz]

kept_electrons, lost_electrons = particle_trace(electron_tracker_0,ntrack)

cond_list = solverE.conductordatalist

fig = plt.figure(figsize=(12,6))
plt.title("Broad Dielectric Particle Trace")

scale = 1e6

ax2 = plt.subplot(111)

steps2cross = 1900 #computed # of steps to cross

cond_list = solverE.conductordatalist
for cond in cond_list[2:]: #ignore first two conductors - these are the plates
    co = cond[0]
    specs = co.getkwlist()
    xw = specs[0] #x-dimension (y in plot)
    xc = specs[3] #center
    xll = xc - xw/2. #lower left
    zw = specs[2] #z-dimension (x in plot)
    zc = specs[-1] #center
    zll = zc - zw/2. #lower left

    ax2.add_patch(
        patches.Rectangle(
            (zll * scale, xll * scale),
            zw * scale,
            xw * scale,
            facecolor="grey",
            edgecolor="grey"
        )
    )

ax1 = plt.subplot(111)
kept_electrons, lost_electrons = particle_trace(electron_tracker_0,ntrack)

for i in range(len(kept_electrons[1])):
    ax1.plot(kept_electrons[1][i][:steps2cross] * scale,kept_electrons[0][i][:steps2cross] * scale, c = '#1f77b4')

for i in range(len(lost_electrons[1])):
    ax1.plot(lost_electrons[1][i][:steps2cross] * scale,lost_electrons[0][i][:steps2cross] * scale, c = '#2ca02c')

kept = mlines.Line2D([], [], color='#1f77b4',label='Absorbed Particles')
lost = mlines.Line2D([], [], color='#2ca02c',label='Reflected Particles')

box = ax1.get_position()
ax1.set_position([box.x0, box.y0, box.width * 0.8, box.height])

plt.xlim(Z_MIN * scale,Z_MAX * scale)
plt.ylim(X_MIN * scale, X_MAX * scale)
plt.legend(handles=[kept, lost],loc='best', bbox_to_anchor=(1, 1))
plt.xlabel('z ($\mu$m)')
plt.ylabel('x ($\mu$m)')
plt.savefig('broad_dielectric_trace.png')
plt.show()

### Plot fields

#Need to compute the fields first
fieldEz = solverE.getez()

#Now plot
fig = plt.figure(figsize=(12,6))

X_CELLS = NUM_X
Z_CELLS = NUM_Z

xl = 0
xu = NUM_X
zl = 0
zu = NUM_Z

plt.xlabel("z ($\mu$m)")
plt.ylabel("x ($\mu$m)")
plt.title("$E_z$ across domain - with box")

pxmin = ((X_MAX - X_MIN) / X_CELLS * xl + X_MIN) * 1e6
pxmax = ((X_MAX - X_MIN) / X_CELLS * xu + X_MIN) * 1e6
pzmin = (Z_MIN + zl / Z_CELLS * Z_MAX) * 1e6
pzmax = (Z_MAX * zu / Z_CELLS) * 1e6

plt.xlim(pzmin, pzmax)
plt.ylim(pxmin, pxmax)

ez_plt = plt.imshow(fieldEz[xl:xu,zl:zu],cmap='viridis',extent=[pzmin, pzmax, pxmin, pxmax],aspect='auto')

cbar = fig.colorbar(ez_plt)
cbar.ax.set_xlabel("V/m")
cbar.ax.xaxis.set_label_position('top')

plt.savefig('Ez_broad_box.png',bbox_inches='tight')
plt.close()

#Need to compute the fields first
fieldEx = solverE.getex()

Exr = fieldEx[::-1]

#Now plot
fig = plt.figure(figsize=(12,6))

X_CELLS = NUM_X
Z_CELLS = NUM_Z

xl = 0
xu = NUM_X
zl = 0
zu = NUM_Z

plt.xlabel("z ($\mu$m)")
plt.ylabel("x ($\mu$m)")
plt.title("$E_x$ across domain - with box")

pxmin = ((X_MAX - X_MIN) / X_CELLS * xl + X_MIN) * 1e6
pxmax = ((X_MAX - X_MIN) / X_CELLS * xu + X_MIN) * 1e6
pzmin = (Z_MIN + zl / Z_CELLS * Z_MAX) * 1e6
pzmax = (Z_MAX * zu / Z_CELLS) * 1e6

plt.xlim(pzmin, pzmax)
plt.ylim(pxmin, pxmax)

ex_plt = plt.imshow(Exr[xl:xu,zl:zu],cmap='viridis',extent=[pzmin, pzmax, pxmin, pxmax],aspect='auto')

cbar = fig.colorbar(ex_plt)
cbar.ax.set_xlabel("V/m")
cbar.ax.xaxis.set_label_position('top')

plt.savefig('Ex_broad_box.png',bbox_inches='tight')
plt.close()

#### Plot potential

#Need to compute the potential first
potential = solverE.getphi()

#Now plot
fig = plt.figure(figsize=(12,6))

X_CELLS = NUM_X
Z_CELLS = NUM_Z

potential = solverE.getphi()

xl = 0
xu = NUM_X
zl = 0
zu = NUM_Z

plt.xlabel("z ($\mu$m)")
plt.ylabel("x ($\mu$m)")
plt.title("$\phi$ across domain -with box")

pxmin = ((X_MAX - X_MIN) / X_CELLS * xl + X_MIN) * 1e6
pxmax = ((X_MAX - X_MIN) / X_CELLS * xu + X_MIN) * 1e6
pzmin = (Z_MIN + zl / Z_CELLS * Z_MAX) * 1e6
pzmax = (Z_MAX * zu / Z_CELLS) * 1e6

plt.xlim(pzmin, pzmax)
plt.ylim(pxmin, pxmax)

phi_plt = plt.imshow(potential[xl:xu,zl:zu],cmap='RdBu',extent=[pzmin, pzmax, pxmin, pxmax],aspect='auto')

cbar = fig.colorbar(phi_plt)
cbar.ax.set_xlabel("Volts")
cbar.ax.xaxis.set_label_position('top')

plt.savefig('phi_broad_box.png',bbox_inches='tight')
plt.close()


import csv, json, argparse, copy, re, os, urllib2

import numpy as np
from scipy.spatial import distance
from sklearn import manifold, metrics, decomposition, preprocessing

import igraph

import jsmin

import rdkit
from rdkit import Chem, DataStructs, Geometry
from rdkit.DataStructs import cDataStructs
from rdkit.Chem import Draw, AllChem, Scaffolds, Lipinski, Crippen, rdMolDescriptors, TemplateAlign
from rdkit.Chem.Scaffolds import MurckoScaffold

PROPS_ORDER = ["mw", "hba", "hbd", "rb", "rc", "arc", "logp", "tpsa"]

PROP2FNC = {
    "mw": rdMolDescriptors.CalcExactMolWt,
    "hba": Lipinski.NumHAcceptors,
    "hbd": Lipinski.NumHDonors,
    "rb": Lipinski.NumRotatableBonds,
    "rc": Lipinski.RingCount,
    "arc": Lipinski.NumAromaticRings,
    "logp": Crippen.MolLogP,
    "tpsa": rdMolDescriptors.CalcTPSA,
}

PROP2LABEL = {
    "mw": "Molecular weight",
    "hba": "H-bond acceptors",
    "hbd": "H-bond donors",
    "rb": "Rotatable bonds",
    "rc": "Rings",
    "arc": "Aromatic rings",
    "logp": "cLogP",
    "tpsa": "TPSA"
}

FP2FNC = {
    "ecfp4": lambda rdmol: AllChem.GetMorganFingerprintAsBitVect(rdmol, radius=2, nBits=1024),
    "ecfp6": lambda rdmol: AllChem.GetMorganFingerprintAsBitVect(rdmol, radius=3, nBits=1024),
    "apfp": lambda rdmol: AllChem.GetHashedAtomPairFingerprintAsBitVect(rdmol, nBits=1024),
    "ttfp": lambda rdmol: AllChem.GetHashedTopologicalTorsionFingerprintAsBitVect(rdmol, nBits=1024),
    "maccs": lambda rdmol: AllChem.GetMACCSKeysFingerprint(rdmol),
}

AVAILABLE_METRICS = ["Tanimoto", "Dice", "Cosine", "Sokal", "Russel", "RogotGoldberg", "AllBit", "Kulczynski", "McConnaughey", "Asymmetric", "BraunBlanquet"]

class ChemSpace():

    def __init__(self):
        self.category_field = False
        self.category_field_delimiter = False
        self.label_field = False
        self.compound_structure_field = False
        self.sdf = False
        self.write_structures = True
        self.fp = "ecfp4"
        self.fingerprint_field = False
        self.metric = "Tanimoto"

        if self.metric not in AVAILABLE_METRICS:
            raise Exception("Metric '{}' not found in available similarity metrics: {}".format(self.metric, AVAILABLE_METRICS))

        self.index2rdmol = {}
        self.index2fpobj = {}

    def read_csv(self, filename, delimiter=",", header=False, missing_value=False, remove_columns=False):
        """Reads data from the CSV file"""
        print("Reading file: {}".format(filename))

        self.filename = filename
        with open(self.filename, "r") as input_file:
            reader = csv.reader(input_file, delimiter=delimiter)
            rows = [row for row in reader]

        self.read_data(rows, header, missing_value, remove_columns)

    def read_sdf(self, filename):
        """Reads data from a sdf file"""
        print("Reading file: {}".format(filename))
        self.sdf = True
        self.header = False
        self.data = []
        self.filename = filename
        self.index2rdmol = {}
        self.index2fpobj = {}
        self.index2props = {}
        self.index2category = {}
        self.index2label = {}
        self.index2id = {}

        molsupplier = Chem.SDMolSupplier(str(filename))
        not_parsed = []

        for index, m in enumerate(molsupplier):
            try:
                Chem.SanitizeMol(m)
                self.index2rdmol[index] = m
                self.index2fpobj[index] = FP2FNC[self.fp](m)
                self.index2props[index] = m.GetPropsAsDict()
                self.index2id[index] = index

            except Exception, e:
                print(e)
                not_parsed.append(index)

        self.index_order = self.index2rdmol.keys()
        self.index_order.sort()
        self.index2row = {i: [] for i in self.index_order}
        self.data = self.index2row.values()

        if self.label_field is not False and self.label_field in self.index2props[self.index_order[0]]:
            self.index2label = {i: self.index2props[i].get(self.label_field) for i in self.index_order}

        if self.category_field is not False and self.category_field in self.index2props[self.index_order[0]]:
            self.index2category = {i: self.index2props[i].get(self.category_field) for i in self.index_order}

        self.__create_chemspace_format__()

    def add_compounds_from_file(self, filename, delimiter=","):
        print("Reading compounds: {}".format(filename))
        self.filename = filename

        with open(self.filename, "r") as input_file:
            reader = csv.reader(input_file, delimiter=delimiter)
            rows = [row for row in reader]

        self.add_compounds(rows)

    def add_category(self, category):
        if not "categories" in self.chemical_space:
            self.chemical_space["categories"] = []

        self.chemical_space["categories"].append(category)

    def add_compounds(self, rows):
        """Reads data in a form of list of lists (tuples)"""
        self.compounds = {r[0]: r[1] for r in rows}
        self.chemical_space["compounds"] = {}
        self.__parse_compounds__()

        for key in self.chemical_space["points"]:
            if key in self.id2rdmol:
                self.chemical_space["compounds"][key] = {"structure": self.__get_compound__(key)}

    def read_data(self, rows, header=False, missing_value=False, remove_columns=False):
        """Reads data in a form of list of lists (tuples)"""
        self.header = header
        self.missing_value = missing_value
        data_start = 0

        self.data = rows
        self.index2id = {}
        self.index2row = {}
        self.index2compound = {}
        self.index2label = {}
        self.index2category = {}

        if self.header:
            self.header = self.data[0]
            self.data = self.data[1:]

        if self.header:
            if remove_columns is not False and len(remove_columns) > 0:
                for col in remove_columns:
                    self.__remove_field__(col)

            if self.compound_structure_field and self.compound_structure_field in self.header:
                self.index2compound = self.__extract_field__(self.compound_structure_field)
                self.__read_compounds__()

            if self.label_field and self.label_field in self.header:
                self.index2label = self.__extract_field__(self.label_field)

            if self.category_field and self.category_field in self.header:
                self.index2category = self.__extract_field__(self.category_field)

            if self.fingerprint_field and self.fingerprint_field in self.header:
                self.index2fp = self.__extract_field__(self.fingerprint_field)
                self.index2fpobj = {}

                for index, fp in self.index2fp.items():
                    self.index2fpobj[index] = self.__get_bitvect_for_fp__(fp)

            # remove ID field
            self.header.pop(0)

        self.index2id = {i: row[0] for i, row in enumerate(self.data)}
        self.index2row = {i: [round(float(v), 2) if v not in ["", None, "None", self.missing_value] else None for v in row[1:]] for i, row in enumerate(self.data)}
        self.index_order = [i for i, row in enumerate(self.data)]
        self.data = [self.index2row[i] for i in self.index_order]

        if self.missing_value is not False:
            self.data, self.missing_values_indexes = self.__impute_missing_values__(self.data)
            # self.original_data = self.__return_missing_values__(copy.deepcopy(self.data), self.missing_values_indexes)

        # self.original_data = copy.deepcopy(self.index2row)
        self.__create_chemspace_format__()

    def __read_compounds__(self):
        for i, smi in self.index2compound.items():
            try:
                self.index2rdmol[i] = Chem.MolFromSmiles(smi)
                self.index2fpobj[i] = FP2FNC[self.fp](self.index2rdmol[i])

            except Exception, e:
                print(e)
                self.index2rdmol[i] = None
                self.index2fpobj[i] = None

    def __remove_field__(self, field):
        if field in self.header:
            index = self.header.index(field)

            if index is not False:
                self.header.pop(index)

                for i, row in enumerate(self.data):
                    self.data[i].pop(index)

    def __extract_field__(self, field):
        index2value = {}

        if field in self.header:
            index = self.header.index(field)

            if index is not False:
                self.header.pop(index)

                for i, row in enumerate(self.data):
                    index2value[i] = row[index]
                    self.data[i].pop(index)

        return index2value

    def __impute_missing_values__(self, data):
        datatype2impute = {"numeric": {"strategy":"mean",
                                        "value": lambda x: round(float(value), 3)},
                           "binary": {"strategy":"most_frequent",
                                      "value": lambda x: int(value)}
                           }

        missing_values_indexes = []

        for i, row in enumerate(self.data):
            missing_values_indexes.append([j for j, v in enumerate(row) if v == self.missing_value])

            for j, value in enumerate(row):
                if value == self.missing_value:
                    data[i][j] = np.nan
        imputer = preprocessing.Imputer(missing_values="NaN", strategy=datatype2impute["numeric"]["strategy"])
        #error when using median strategy - minus one dimension in imputed data... omg
        imputed_data = [list(row) for row in imputer.fit_transform(self.data)]
        imputed_data = [[datatype2impute["numeric"]["value"](value) for value in row] for row in imputed_data]
        return imputed_data, missing_values_indexes

    def __return_missing_values__(self, data, missing_values_indexes):
        for i, indexes in enumerate(missing_values_indexes):
            if indexes:
                for index in indexes:
                    data[i][index] = None
        return data

    def __create_chemspace_format__(self):
        self.chemical_space = {"points": {}}

        for index in self.index_order:
            self.chemical_space["points"][index] = {"object_ids": [self.index2id[index]]}

        if len(self.index2category):
            self.__parse_categories__()

        if len(self.index2label):
            for index, label in self.index2label.items():
                self.chemical_space["points"][index]["label"] = label

        for index, row in self.index2row.items():
            self.chemical_space["points"][index]["features"] = copy.copy(row)

        if self.header:
            current_header = self.chemical_space.get("feature_names", [])
            current_header.extend(self.header)
            self.chemical_space["feature_names"] = current_header

        if len(self.index2rdmol) and self.write_structures:
            self.chemical_space["compounds"] = {}

            for index, rdmol in self.index2rdmol.items():
                # self.chemical_space["compounds"][index] = {"structure": self.__get_compound__(rdmol), "smiles": Chem.MolToSmiles(rdmol, True)}

                self.chemical_space["compounds"][index] = {"smiles": Chem.MolToSmiles(rdmol, True)}

    def __parse_categories__(self):
        category2ids = {}

        for index, category in self.index2category.items():
            categories = [category] if self.category_field_delimiter is False else [c.strip() for c in category.split(self.category_field_delimiter)]

            for c in categories:
                if c in category2ids:
                    category2ids[c].add(index)
                else:
                    category2ids[c] = {index}

        if not "categories" in self.chemical_space:
            self.chemical_space["categories"] = []

        for c, ids in category2ids.items():
            self.chemical_space["categories"].append({"label": c, "points": list(ids)})

    def add_paths(self, paths):
        if not self.chemical_space.get("paths", False):
            self.chemical_space["paths"] = []

        self.chemical_space["paths"].extend(paths)

    def add_paths_from_file(self):
        pass

    def add_physico_chemical_properties(self):
        print("Calculating physico-chemical properties: {} compounds".format(len(self.index2rdmol)))
        self.pcp = True
        if len(self.index2rdmol):
            count = len(self.index2rdmol)
            i = 0

            id2pcp = {}
            for index, rdmol in self.index2rdmol.items():
                if i%100 == 0 or i == count:
                    print("{}/{}".format(i, count))

                id2pcp[index] = self.__get_pcp_for_rdmol__(rdmol)
                i+=1

            empty = [None for x in PROP2LABEL]
            for i, index in enumerate(self.index_order):

                if id2pcp.get(index, False):
                    pcps = id2pcp[index]
                else:
                    pcps = empty

                self.chemical_space["points"][index]["features"].extend(pcps)
                self.data[i].extend(pcps)

            current_header = self.chemical_space.get("feature_names", [])
            current_header.extend([PROP2LABEL[prop] for prop in PROPS_ORDER])
            self.chemical_space["feature_names"] = current_header
            self.original_data = copy.deepcopy(self.data)

    def __get_pcp_for_rdmol__(self, rdmol):
        return [round(PROP2FNC[prop](rdmol), 2) for prop in PROPS_ORDER]

    def __get_compound__(self, rdmol):
        if rdmol is not None:
            Chem.Kekulize(rdmol)
            AllChem.Compute2DCoords(rdmol)
            compound = {"atoms": {}}
            atoms = [a for a in rdmol.GetAtoms()]
            bond_types = []
            for i, a in enumerate(atoms, 1):
                number = a.GetIdx()
                position = rdmol.GetConformer().GetAtomPosition(number)

                compound["atoms"][number] = {
                    "bonds": {b.GetEndAtomIdx():b.GetBondTypeAsDouble() for b in a.GetBonds() if b.GetEndAtomIdx() != number},
                    "symbol": a.GetSymbol(),
                    "charge": a.GetFormalCharge(),
                    "coordinates": [round(position.x, 3), round(position.y, 3)]
                }

                bond_types.extend(compound["atoms"][number]["bonds"].values())
        else:
            compound = None

        return compound

    def normalize_data(self, feature_range=(0,1)):
        """Normalizes data to a scale from 0 to 1."""
        print("Data normalization (scale): {}".format(feature_range))

        min_max_scaler = preprocessing.MinMaxScaler(feature_range)
        self.data = min_max_scaler.fit_transform(self.data)
        self.data = [[round(v, 3) for v in row] for row in self.data]

    def __calculate_distance_matrix__(self, similarity_threshold):
        print("\nCalculating distance matrix: {} compounds".format(len(self.index2fpobj)))

        self.dist_matrix = {x:[] for x in self.index_order}
        self.edges = []
        self.index2edges = {}

        fps_count = len(self.index_order)

        for i, index_1 in enumerate(self.index_order):
            self.index2edges[index_1] = []

            if i%100 == 0 or i == fps_count:
                print("{}/{}".format(i, fps_count))

            for j, index_2 in enumerate(self.index_order[i:], i):
                sim = DataStructs.FingerprintSimilarity(self.index2fpobj[index_1], self.index2fpobj[index_2], metric=getattr(DataStructs, "{}Similarity".format(self.metric)))
                self.dist_matrix[index_1].append(1-sim)

                if index_1 != index_2:
                    self.dist_matrix[index_2].append(1-sim)

                    if sim >= similarity_threshold:
                        self.edges.append((index_1, index_2))
                        self.index2edges[index_1].append([index_2])

    def __get_edges__(self, similarity_threshold=0.7, k=2):
        print("\nCalculating edges [similarity threshold={}]: {} compounds".format(similarity_threshold, len(self.index2fpobj)))
        self.edges = []
        self.index2edges = {}
        count = len(self.index_order)

        for i, index in enumerate(self.index_order):
            if (i+1)%100 == 0:
                print("{}/{}".format(i, count))

            values = [[idx, v] for idx, v in zip(self.index_order, self.dist_matrix[index]) if idx != index]
            values.sort(key=lambda x: x[1])

            if 1-values[1][1] >= similarity_threshold:
                self.index2edges[index] = []

                for v in values:
                    if 1-v[1] >= similarity_threshold:
                        self.edges.append((index, v[0]))
                        self.index2edges[index].append([v[0]])

                        if len(self.index2edges[index]) == k:
                            break
                    else:
                        break

        print("EDGES: {}".format(len(self.edges)))

    def __convert_fps_to_bitvects__(self, fps):
        converted = []

        for fp in fps:
            row = [fp[0], self.__get_bitvect_for_fp__(fp[1:])]
            converted.append(row)

        return converted

    def __get_bitvect_for_fp__(self, fp):
        if type(fp) is list and len(fp) == 1:
            fp = fp[0]
        bitvect = cDataStructs.ExplicitBitVect(len(fp))
        on_indexes = [i for i, b in enumerate(fp) if int(b) == 1]
        bitvect.SetBitsFromList(on_indexes)
        return bitvect

    def arrange(self, by="fps", fps=[], method="pca", similarity_threshold=0.7, add_edges=False, k=None):
        self.index2edges = False
        self.edges = False
        self.dist_matrix = False
        bitvects = False

        if type(method) is not list:
            methods = [method]
        else:
            methods = method

        for method in methods:
            if by == "scaffolds":
                self.__arrange_by_scaffolds__()

                g = igraph.Graph(len(self.index_order))
                print("\nCalculating Chemical Space Network...")
                feature_names = ["CSN1", "CSN2"]

                g.add_edges(self.edges)
                layout = g.layout_fruchterman_reingold()
                coords = layout.coords

            elif by == "dm" or method == "sas":
                if len(fps) == 0:
                    for index in self.index_order:
                        fps.append(self.index2fpobj[index])
                    bitvects = True

                elif type(fps[0][1]) in [unicode, int] and not bitvects:
                    fps = self.__convert_fps_to_bitvects__(fps)
                    bitvects = True

                if not self.dist_matrix:
                    self.__calculate_distance_matrix__(similarity_threshold)

                dist_matrix = np.matrix([np.array(self.dist_matrix[index]) for index in self.index_order])
                g = igraph.Graph(len(self.index_order))

                if method == "csn":
                    print("\nCalculating Chemical Space Network...")
                    feature_names = ["CSN1", "CSN2"]
                    if k is not None:
                        # k = len(self.index_order)
                        self.__get_edges__(similarity_threshold=similarity_threshold, k=k)

                    print("Fruchterman-Reingold Layout calculation...")
                    g.add_edges(self.edges)
                    layout = g.layout_fruchterman_reingold()
                    coords = layout.coords

                elif method == "mds":
                    print("\nCalculating MDS...")
                    feature_names = ["MDS1", "MDS2"]
                    # sklearn implementation
                    mds = manifold.MDS(n_components=2, dissimilarity='precomputed')
                    coords = mds.fit_transform(dist_matrix)

                    # igraph implementation
                    # layout = g.layout_mds(dist_matrix, 2, arpack_options=igraph.ARPACKOptions(iter=1000))
                    # coords = layout.coords

                elif method == "pca":
                    print("\nCalculating PCA...")
                    feature_names = ["PC1", "PC2"]
                    pca = decomposition.PCA(n_components=2)
                    coords = pca.fit_transform(dist_matrix)

                elif method == "fa":
                    print("\nCalculating Factor Analysis...")
                    feature_names = ["FA1", "FA2"]
                    fa = decomposition.FactorAnalysis(n_components=2)
                    coords = fa.fit_transform(dist_matrix)

                elif method == "isomap":
                    print("\nCalculating Isomap...")
                    feature_names = ["Isomap1", "Isomap2"]
                    isomap = manifold.Isomap(n_neighbors=200, n_components=2)
                    coords = isomap.fit_transform(dist_matrix)

                elif method == "tsne":
                    print("\nCalculating t-SNE...")
                    feature_names = ["t-SNE1", "t-SNE2"]
                    tsne = manifold.TSNE(n_components=2, metric='precomputed')
                    coords = tsne.fit_transform(dist_matrix)

                elif method == "sas":
                    print("\nCalculating SAS...")
                    feature_names = ["Similarity", "Activity difference"]
                    self.chemical_space = {"points": {}, "feature_names": ["SALI"]}
                    ai = self.header.index(self.activity_field)
                    ids = []
                    coords = []

                    for i, index_1 in enumerate(self.index_order[:-1]):
                        for j, index_2 in enumerate(self.index_order[i:], i):
                            if i != j:
                                activity_diff = round(abs(float(self.data[i][ai]) - float(self.data[j][ai])), 2)
                                distance = self.dist_matrix[i][j]
                                distance = distance if distance > 0 else 0.01

                                sali = round(activity_diff/distance, 2)
                                coord = [round(1 - self.dist_matrix[i][j], 2), activity_diff]
                                self.chemical_space["points"]["{}_{}".format(index_1, index_2)] = {"features": [sali]}
                                ids.append("{}_{}".format(index_1, index_2))
                                coords.append(coord)

                    self.index_order = ids

            elif by in ["data", "fps"]:
                if by == "fps":
                    if len(fps) == 0 and len(self.index2fpobj):
                        for index in self.index_order:
                            fps.append(self.index2fpobj[index])

                    data = [[int(b) for b in fp] for fp in fps]
                else:
                    data = self.data

                if method == "pca":
                    print("\nCalculating PCA...")
                    feature_names = ["PC1", "PC2"]
                    pca = decomposition.PCA(n_components=2)
                    coords = pca.fit_transform(data)

                elif method == "fa":
                    print("\nCalculating Factor Analysis...")
                    feature_names = ["FA1", "FA2"]
                    fa = decomposition.FactorAnalysis(n_components=2)
                    coords = fa.fit_transform(data)

            if method in ["csn", "nn"] or add_edges or by == "scaffolds":
                if self.dist_matrix is False and self.edges is False:
                    self.__calculate_distance_matrix__()

                if self.edges is False:
                    if k is None:
                        k = len(self.index_order)
                    self.__get_edges__(similarity_threshold=similarity_threshold, k=k)

                if by != "scaffolds":
                    for cid, es in self.index2edges.items():
                        if not self.chemical_space["points"][cid].get("links", False):
                            self.chemical_space["points"][cid]["links"] = []

                        for e in es:
                            self.chemical_space["points"][cid]["links"].extend(e)

            index2coords = {index:coords[i] for i, index in enumerate(self.index_order)}

            for index, values in self.chemical_space["points"].items():
                if index in index2coords:
                    point_features = self.chemical_space["points"][index]["features"]
                    features = [round(index2coords[index][0], 3), round(index2coords[index][1], 3)]
                    features.extend(point_features)
                    self.chemical_space["points"][index]["features"] = features

                else:
                    self.chemical_space["points"].pop(index, None)

            feature_names.extend(self.chemical_space.get("feature_names", []))
            self.chemical_space["feature_names"] = feature_names

    def __arrange_by_scaffolds__(self, align_by_scaffold=True):
        self.scaffold2indexes = {}
        self.scaffold2rdmol = {}
        self.index2scaffold = {}

        for index, rdmol in self.index2rdmol.items():
            AllChem.Compute2DCoords(rdmol)

            if align_by_scaffold:
                scaffold = Scaffolds.MurckoScaffold.GetScaffoldForMol(rdmol)
                AllChem.Compute2DCoords(scaffold)
                scaffold_smiles = Chem.MolToSmiles(scaffold)
                self.scaffold2rdmol[scaffold_smiles] = scaffold

                matched = rdmol.GetSubstructMatch(scaffold)
                coords = [rdmol.GetConformer().GetAtomPosition(x) for x in matched]
                coords2D = [Geometry.Point2D(pt.x,pt.y) for pt in coords]

                coordDict = {}
                for i,coord in enumerate(coords2D):
                    coordDict[matched[i]] = coord

                AllChem.Compute2DCoords(rdmol, coordMap=coordDict)

                if scaffold_smiles in self.scaffold2indexes:
                    self.scaffold2indexes[scaffold_smiles].append(index)
                else:
                    self.scaffold2indexes[scaffold_smiles] = [index]

        self.scaffold2indexes = {scaffold: indexes for scaffold, indexes in self.scaffold2indexes.items() if len(indexes) > 1}

        for index, scaffold in enumerate(self.scaffold2indexes.keys(), len(self.index_order)):
                self.index2scaffold[index] = scaffold
                self.index2rdmol[index] = self.scaffold2rdmol[scaffold]
                self.index_order.append(index)

        self.edges = []
        self.index2edges = {}

        for index_1, scaffold in self.index2scaffold.items():
            self.index2edges[index_1] = []

            for index_2 in self.scaffold2indexes[scaffold]:
                self.edges.append((index_1, index_2))
                self.index2edges[index_1].append(index_2)

        self.__add_scaffolds_to_chemical_space__()

    def __add_scaffolds_to_chemical_space__(self):
        for index, scaffold in self.index2scaffold.items():
            self.chemical_space["points"][index] = {
                # "features": self.__get_pcp_for_rdmol__(self.scaffold2rdmol[scaffold]),
                "object_ids": self.scaffold2indexes[scaffold],
                "links": self.scaffold2indexes[scaffold],
            }
            if self.pcp:
                self.chemical_space["points"][index]["features"] = [None for i in range(len(self.chemical_space["feature_names"]) - len(PROPS_ORDER))]
                self.chemical_space["points"][index]["features"].extend(self.__get_pcp_for_rdmol__(self.scaffold2rdmol[scaffold]))
            self.chemical_space["compounds"][index] = {"smiles": scaffold, "color": "red"}

    def get_chemspace_compound_from_smiles(self, smi):
        mol_obj = Chem.MolFromSmiles(smi)
        if mol_obj is not None:
            AllChem.Compute2DCoords(mol_obj)
            mol = self.__get_compound__(mol=mol_obj)
        else:
            mol = False
        return mol

    def export_chemical_space_as_html(self, htmldir=".", ):
        """Export a simple HTML page with embedded chemical space and dependencies into a given directory."""
        if not os.path.exists(htmldir):
            os.makedirs(htmldir)

        chemspace_json = self.export_chemical_space_as_json(minify=True, dump=True)

        libs = [
            ("chemspace-0.2.0.min.js", "https://openscreen.cz/software/chemspace/static/js/chemspace-0.2.0.min.js"),
            ("jquery-3.3.1.min.js", "https://code.jquery.com/jquery-3.3.1.min.js"),
            ("konva.min.js", "https://cdn.rawgit.com/konvajs/konva/1.7.6/konva.min.js")
        ]

        js_html = []
        for l in libs:
            js_html.append("<script src='{}'></script>".format(l[0]))

        settings = {
            "target": "chemspace"
        }

        template = """<html>
        <head>
            {}
            <script>
                $(document).ready(function() {{
                    var data = {};
                    var chemspace = new ChemSpace({});
                    chemspace.read_data(data);
                    chemspace.draw();
                }});
            </script>
        </head>

        <body>
            <div id="chemspace"></div>
        </body>
        </html>""".format('\n'.join(js_html), chemspace_json, json.dumps(settings))

        for l in libs:
            lib, url = l
            try:
                source = urllib2.urlopen(url)
                source_html = source.read()

                with open(os.path.join(htmldir, lib), "w") as output:
                    output.write(source_html)
            except urllib2.URLError, e:
                raise Exception("""
                        \nCan't download file {}.\nPlease check your internet connection and try again.\nIf the error persists there can be something wrong with the InCHlib server.\n""".format(url)
                    )

        with open(os.path.join(htmldir, "chemspace.html"), "w") as output:
            output.write(template)

    def export_chemical_space_as_json(self, filename=None, minify=False, dump=True):
        """Returns space in a JSON format or exports it to the file specified by the filename parameter."""
        space_json = self.chemical_space

        if minify:
            space_json = json.dumps(space_json)
            space_json = self.__minify_data(space_json)
        elif dump:
            space_json = json.dumps(space_json, indent=4)

        if filename:
            output = open(filename, "w")
            output.write(space_json)

        return space_json

    def __minify_data(self, data):
        return jsmin.jsmin(data)

def _process_(arguments):
    s = ChemSpace()
    s.sdf_file = False
    s.write_structures = False if arguments.dont_write_structures else True
    s.fp = arguments.fingerprint
    s.category_field = arguments.category_field
    s.category_field_delimiter = arguments.category_field_delimiter
    s.label_field = arguments.label_field
    s.activity_field = arguments.activity_field
    s.compound_structure_field = arguments.compound_structure_field
    s.fingerprint_field = arguments.fingerprint_field
    s.metric = arguments.similarity_metric

    if arguments.data_file.split(".")[-1].lower() == "sdf":
        s.read_sdf(arguments.data_file)
    else:
        s.read_csv(arguments.data_file, arguments.data_delimiter, arguments.data_header, arguments.missing_values, arguments.remove_columns)

    if s.compound_structure_field is not False or s.sdf == True:
        if arguments.physico_chemical_properties:
            s.add_physico_chemical_properties()

    if arguments.normalize:
        s.normalize_data()

    if arguments.arrange_by:
        s.arrange(
            method=arguments.dimensional_reduction_method,
            similarity_threshold=float(arguments.compound_similarity_threshold),
            add_edges=arguments.add_edges,
            by=arguments.arrange_by,
            k=arguments.knn
        )

    if arguments.html_dir:
        s.export_chemical_space_as_html(arguments.html_dir)
    elif arguments.output_file:
        s.export_chemical_space_as_json(arguments.output_file, minify=arguments.minify_output)
    else:
        print(s.export_chemical_space_as_json(minify=arguments.minify_output))

if __name__ == '__main__':
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    parser.add_argument("data_file", type=str, help="csv(text) data file with delimited values or a sdf file")
    parser.add_argument("-dh", "--data_header", default=False, help="whether the first row of data file is a header", action="store_true")
    parser.add_argument("-dd", "--data_delimiter", type=str, default=",", help="delimiter of values in data file")
    parser.add_argument("-o", "--output_file", type=str, help="the name of output file")
    parser.add_argument("-fpf", "--fingerprint_field", type=str, default=False, help="set a fingerprint field name in case it is in the data file")
    parser.add_argument("-cf", "--category_field", type=str, default=False, help="set a category field name in case it is in the data file")
    parser.add_argument("-cfd", "--category_field_delimiter", type=str, default=False, help="a category field delimiter")
    parser.add_argument("-lf", "--label_field", type=str, default=False, help="set a label field name in case it is in the data file")
    parser.add_argument("-af", "--activity_field", type=str, default=False, help="set an activity field name in case it is in the data file")
    parser.add_argument("-csf", "--compound_structure_field", type=str, default=False, help="the name of a column with a compound structure")
    parser.add_argument("-fp", "--fingerprint", type=str, default="ecfp4", help="fingerprint used for a compound representation (ecfp4, ecfp6, maccs, topological, atom_pairs)")
    parser.add_argument("-arr", "--arrange_by", default=False, help="arrange data by compound structures (distance matrix) or by input data (data/fps)", type=str)
    parser.add_argument("-cst", "--compound_similarity_threshold", default=0.7, help="compound similarity threshold")
    parser.add_argument("-drm", "--dimensional_reduction_method", nargs='+', type=str, default="pca", help="which method use for dimensional reduction (pca/isomap/csn)")
    parser.add_argument("-dws", "--dont_write_structures", default=False, help="dont write structures to output file", action="store_true")
    parser.add_argument("-min", "--minify_output", default=False, help="minify the JSON output format", action="store_true")
    parser.add_argument("-html", "--html_dir", type=str, default=False, help="the directory to store HTML page with dependencies")
    parser.add_argument("-pcp", "--physico_chemical_properties", default=False, help="calculate basic phyisico-chemical properties and add them ass features", action='store_true')
    parser.add_argument("-edges", "--add_edges", default=False, help="add edges based on compound similarity to the graph", action='store_true')
    parser.add_argument("-n", "--normalize", default=False, help="normalize data to [0, 1] range", action="store_true")
    parser.add_argument("-mv", "--missing_values", type=str, default=False, help="define the string representating missing values in the data")
    parser.add_argument("-k", "--knn", type=int, default=None, help="the number of neighbours (k) used for the construction of csn using the nn method")
    parser.add_argument("-sm", "--similarity_metric", type=str, default="Tanimoto", help="similarity metric")
    parser.add_argument('-rmc','--remove_columns', nargs='+', default=False, help='columns in data that should not be used')

    args = parser.parse_args()
    _process_(args)


"""This module implements numerical functions."""
from __future__ import annotations

import numpy as np
import scipy as sp

from bqskit.qis.pauli import PauliMatrices
from bqskit.qis.unitary.unitary import RealVector

def dexpmv(M: np.ndarray, dM: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    """
    Compute the Matrix exponential F = e^M and its derivative dF.

    User must provide M and its derivative dM. If the argument dM is a
    vector of partials then dF will be the respective partial vector.
    This is done using a Pade Approximat with scaling and squaring.

    Args:
        M (np.ndarray): Matrix to exponentiate.

        dM (np.ndarray): Derivative(s) of M.

    Returns:
        tuple: Tuple containing
            - F (np.ndarray): Exponentiated matrix, i.e. e^M.

            - dF (np.ndarray): Derivative(s) of F.

    References:
        Brančík, Lubomír. "Matlab programs for matrix exponential function
        derivative evaluation." Proc. of Technical Computing Prague 2008
        (2008): 17-24.
    """

    e = np.log2(np.linalg.norm(M, np.inf))
    r = int(max(0, e + 1))
    M = M / (2 ** r)
    dM = dM / (2 ** r)
    X = M
    Y = dM
    c = 0.5
    F = np.identity(M.shape[0]) + c * M
    D = np.identity(M.shape[0]) - c * M
    dF = c * dM
    dD = -c * dM
    q = 6
    p = True
    for k in range(2, q + 1):
        c = c * (q - k + 1) / (k * (2 * q - k + 1))
        Y = dM @ X + M @ Y
        X = M @ X
        cX = c * X
        cY = c * Y
        F = F + cX
        dF = dF + cY
        if p:
            D = D + cX
            dD = dD + cY
        else:
            D = D - cX
            dD = dD - cY
        p = not p
    Dinv = np.linalg.inv(D)
    F = Dinv @ F
    dF = Dinv @ (dF - dD @ F)

    for k in range(1, r + 1):
        dF = dF @ F + F @ dF
        F = F @ F

    return F, dF

def softmax(x: np.ndarray, beta: int = 20) -> np.ndarray:
    """
    Computes the softmax of vector x.

    Args:
        x (np.ndarray): Input vector to softmax.

        beta (int): Beta coefficient to scale steepness of softmax.

    Returns:
        np.ndarray: Output vector of softmax.
    """

    shiftx = beta * (x - np.max(x))
    exps = np.exp(shiftx)
    return exps / np.sum(exps)

def dot_product(alpha: RealVector, sigma: RealVector) -> np.ndarray:
    """
    Computes the standard dot product of `alpha` with `sigma`.

    Args:
        alpha (RealVector): The alpha vector.

        sigma (RealVector): The sigma vector.

    Returns:
        np.ndarray: Sum of element-wise multiplication of `alpha`
        and `sigma`.
    """

    return np.array(np.sum([a * s for a, s in zip(alpha, sigma)], 0))

def unitary_log_no_i(U: np.ndarray) -> np.ndarray:
    """
    Solves for H in U = e^{iH}

    Args:
        U (np.ndarray): The unitary to decompose.

    Returns:
        np.ndarray: H in e^{iH} = U.

    Note:
        This assumes the input is unitary but does not check. The output
        is undefined on non-unitary inputs.
    """

    T, Z = sp.linalg.schur(U)
    T = np.diag(T)
    D = T / np.abs(T)
    D = np.diag(np.log(D))
    H0 = -1j * (Z @ D @ Z.conj().T)
    return 0.5 * H0 + 0.5 * H0.conj().T

def pauli_expansion(H: np.ndarray) -> np.ndarray:
    """
    Computes a Pauli expansion of the hermitian matrix H.

    Args:
        H (np.ndarray): The hermitian matrix to expand.

    Returns:
        np.ndarray: The coefficients of a Pauli expansion for H,
        i.e., X dot Sigma = H where Sigma is Pauli matrices of
        same size of H.

    Note:
        This assumes the input is hermitian but does not check. The
        output is undefined on non-hermitian inputs.
    """

    # Change basis of H to Pauli Basis (solve for coefficients -> X)
    n = int(np.log2(len(H)))
    paulis = PauliMatrices(n)
    flatten_paulis = [np.reshape(pauli, 4 ** n) for pauli in paulis]
    flatten_H = np.reshape(H, 4 ** n)
    A = np.stack(flatten_paulis, axis=-1)
    X = np.real(np.matmul(np.linalg.inv(A), flatten_H))
    return np.array(X)


<gh_stars>0
from ._stopping_criterion import StoppingCriterion
from ..accumulate_data import MLQMCData
from ..discrete_distribution import Lattice
from ..true_measure import Gaussian
from ..integrand import MLCallOptions
from ..util import MaxSamplesWarning, ParameterError, MaxLevelsWarning
from numpy import *
from numpy.linalg import lstsq
from scipy.stats import norm
from time import time
import warnings

class CubQMCMLCont(StoppingCriterion):
    """
    Stopping criterion based on continuation multi-level quasi-Monte Carlo.

    >>> mlco = MLCallOptions(Lattice(seed=7))
    >>> sc = CubQMCMLCont(mlco,abs_tol=.05)
    >>> solution,data = sc.integrate()
    >>> solution
    10.394...
    >>> data
    Solution: 10.3948
    MLCallOptions (Integrand Object)
        option          european
        sigma           0.200
        k               100
        r               0.050
        t               1
        b               85
    Lattice (DiscreteDistribution Object)
        d               1
        randomize       1
        order           natural
        seed            561339
        mimics          StdUniform
    Gaussian (TrueMeasure Object)
        mean            0
        covariance      1
        decomp_type     pca
    CubQMCMLCont (StoppingCriterion Object)
        rmse_tol        0.019
        n_init          2^(8)
        n_max           10000000000
        replications    2^(5)
        levels_min      2^(1)
        levels_max      10
        n_tols          10
        tol_mult        1.668
        theta_init      2^(-1)
        theta           2^(-3)
    MLQMCData (AccumulateData Object)
        levels          2^(2)
        dimensions      [1. 2. 4. 8.]
        n_level         [4096.  256.  256.  256.]
        mean_level      [10.053  0.184  0.102  0.055]
        var_level       [8.703e-05 6.794e-05 2.603e-05 8.925e-06]
        bias_estimate   0.035
        n_total         155648
        time_integrate  ...

    References:

        [1] <NAME> and <NAME>. 'Multilevel quasi-Monte Carlo path simulation'.
        pp.165-181 in Advanced Financial Modelling, in Radon Series on Computational and Applied Mathematics,
        de Gruyter, 2009. http://people.maths.ox.ac.uk/~gilesm/files/radon.pdf
    """

    def __init__(self, integrand, abs_tol=.05, alpha=.01, rmse_tol=None, n_init=256., n_max=1e10,
        replications=32., levels_min=2, levels_max=10, n_tols=10, tol_mult=100**(1/9), theta_init=0.5):
        """
        Args:
            integrand (Integrand): integrand with multi-level g method
            abs_tol (float): absolute tolerance
            alpha (float): uncertaintly level.
                If rmse_tol not supplied, then rmse_tol = abs_tol/norm.ppf(1-alpha/2)
            rmse_tol (float): root mean squared error
                If supplied (not None), then absolute tolerance and alpha are ignored
                in favor of the rmse tolerance
            n_max (int): maximum number of samples
            replications (int): number of replications on each level
            levels_min (int): minimum level of refinement >= 2
            levels_max (int): maximum level of refinement >= Lmin
            n_tols (int): number of coarser tolerances to run
            tol_mult (float): coarser tolerance multiplication factor
            theta_init (float) : initial error splitting constant

        """
        self.parameters = ['rmse_tol','n_init','n_max','replications','levels_min',
            'levels_max','n_tols','tol_mult','theta_init','theta']
        # initialization
        if rmse_tol:
            self.target_tol = float(rmse_tol)
        else: # use absolute tolerance
            self.target_tol =  float(abs_tol) / norm.ppf(1-alpha/2)
        self.n_init = float(n_init)
        self.n_max = float(n_max)
        self.replications = float(replications)
        self.levels_min = levels_min
        self.levels_max = levels_max
        self.theta_init = theta_init
        self.theta = theta_init
        self.n_tols = n_tols
        self.tol_mult = tol_mult
        # QMCPy Objs
        self.integrand = integrand
        self.true_measure = self.integrand.true_measure
        self.discrete_distrib = self.integrand.discrete_distrib
        # Verify Compliant Construction
        allowed_levels = ['adaptive-multi']
        allowed_distribs = ["Lattice", "Sobol","Halton"]
        super(CubQMCMLCont,self).__init__(allowed_levels, allowed_distribs)

    def integrate(self):
        # Construct AccumulateData Object to House Integration Data
        self.data = MLQMCData(self, self.integrand, self.true_measure, self.discrete_distrib,
            self.levels_min, self.levels_max, self.n_init, self.replications)
        # Loop over coarser tolerances
        for t in range(self.n_tols):
            self.rmse_tol = self.tol_mult**(self.n_tols-t-1)*self.target_tol # Set new target tolerance
            self._integrate()
        return self.data.solution,self.data

    def _integrate(self):
        """ See abstract method. """
        t_start = time()
        #self.theta = self.theta_init
        self.data.levels = int(self.levels_min+1)

        converged = False
        while not converged:
            # Ensure that we have samples on the finest level
            self.data.update_data()
            self._update_theta()

            while self._varest() > (1-self.theta)*self.rmse_tol**2:
                efficient_level = argmax(self.data.var_cost_ratio_level[:self.data.levels])
                self.data.eval_level[efficient_level] = True

                # Check if over sample budget
                total_next_samples = (self.data.replications*self.data.eval_level*self.data.n_level*2).sum()
                if (self.data.n_total + total_next_samples) > self.n_max:
                    warning_s = """
                    Alread generated %d samples.
                    Trying to generate %d new samples, which would exceed n_max = %d.
                    Stopping integration process.
                    Note that error tolerances may no longer be satisfied""" \
                    % (int(self.data.n_total), int(total_next_samples), int(self.n_max))
                    warnings.warn(warning_s, MaxSamplesWarning)
                    self.data.time_integrate += time() - t_start
                    return

                self.data.update_data()
                self._update_theta()

            # Check for convergence
            converged = self._rmse() < self.rmse_tol
            if not converged:
                if self.data.levels == self.levels_max:
                    warnings.warn(
                        'Failed to achieve weak convergence. levels == levels_max.',
                        MaxLevelsWarning)
                    converged = True
                else:
                    self.data._add_level()

        self.data.time_integrate += time() - t_start

    def set_tolerance(self, abs_tol=None, alpha=.01, rmse_tol=None):
        """
        See abstract method.

        Args:
            integrand (Integrand): integrand with multi-level g method
            abs_tol (float): absolute tolerance. Reset if supplied, ignored if not.
            alpha (float): uncertaintly level.
                If rmse_tol not supplied, then rmse_tol = abs_tol/norm.ppf(1-alpha/2)
            rel_tol (float): relative tolerance. Reset if supplied, ignored if not.
                Takes priority over aboluste tolerance and alpha if supplied.
        """
        if rmse_tol != None:
            self.rmse_tol = float(rmse_tol)
        elif abs_tol != None:
            self.rmse_tol = (float(abs_tol) / norm.ppf(1-alpha/2.))

    def _update_theta(self):
        """Update error splitting parameter"""
        max_levels = len(self.data.n_level)
        A = ones((2,2))
        A[:,0] = range(max_levels-2, max_levels)
        y = ones(2)
        y[0] = log2(abs(self.data.mean_level_reps[max_levels-2].mean()))
        y[1] = log2(abs(self.data.mean_level_reps[max_levels-1].mean()))
        x = lstsq(A, y, rcond=None)[0]
        alpha = maximum(.5,-x[0])
        real_bias = 2**(x[1]+max_levels*x[0]) / (2**alpha - 1)
        self.theta = max(0.01, min(0.125, (real_bias/self.rmse_tol)**2))

    def _rmse(self):
        """Returns an estimate for the root mean square error"""
        return sqrt(self._mse())

    def _mse(self):
        """Returns an estimate for the mean square error"""
        return (1-self.theta)*self._varest() + self.theta*self.data.bias_estimate**2

    def _varest(self):
        """Returns the variance of the estimator"""
        return self.data.var_level[:self.data.levels].sum()


<reponame>daniyaljamal/Personality-prediction-based-on-video-using-CNN<filename>dan.py
########################################################################################
#                                                                                #
# VGG16 implementation in TensorFlow                                                   #
# Details:                                                                             #
# http://www.cs.toronto.edu/~frossard/post/vgg16/                                      #
#                                                                                #
# Model from https://gist.github.com/ksimonyan/211839e770f7b538e2d8#file-readme-md     #
# Weights from Caffe converted using https://github.com/ethereon/caffe-tensorflow      #
########################################################################################

import pickle
import warnings

import numpy as np
import tensorflow as tf
from scipy.io import loadmat

warnings.filterwarnings("ignore")

class DAN:
    def __init__(self, imgs, REG_PENALTY=0, preprocess=None):
        self.imgs = imgs
        if preprocess == "vggface":
            self.mean = [129.1862793, 104.76238251, 93.59396362]
        self.convlayers()
        self.dan_part()
        self.output = tf.nn.sigmoid(self.reg_head, name="output")
        self.cost_reg = REG_PENALTY * tf.reduce_mean(tf.square(self.parameters[-2])) / 2

    def convlayers(self):
        self.parameters = []

        # zero-mean input
        with tf.name_scope("preprocess") as scope:
            mean = tf.constant(
                self.mean, dtype=tf.float32, shape=[1, 1, 1, 3], name="img_mean"
            )
            images = self.imgs - mean

        # conv1_1
        with tf.name_scope("conv1_1") as scope:
            kernel = tf.Variable(
                tf.truncated_normal([3, 3, 3, 64], dtype=tf.float32, stddev=1e-1),
                name="weights",
            )
            conv = tf.nn.conv2d(images, kernel, [1, 1, 1, 1], padding="SAME")
            biases = tf.Variable(
                tf.constant(0.0, shape=[64], dtype=tf.float32),
                trainable=True,
                name="biases",
            )
            out = tf.nn.bias_add(conv, biases)
            self.conv1_1 = tf.nn.relu(out, name=scope)
            self.parameters += [kernel, biases]

        # conv1_2
        with tf.name_scope("conv1_2") as scope:
            kernel = tf.Variable(
                tf.truncated_normal([3, 3, 64, 64], dtype=tf.float32, stddev=1e-1),
                name="weights",
            )
            conv = tf.nn.conv2d(self.conv1_1, kernel, [1, 1, 1, 1], padding="SAME")
            biases = tf.Variable(
                tf.constant(0.0, shape=[64], dtype=tf.float32),
                trainable=True,
                name="biases",
            )
            out = tf.nn.bias_add(conv, biases)
            self.conv1_2 = tf.nn.relu(out, name=scope)
            self.parameters += [kernel, biases]

        # pool1
        self.pool1 = tf.nn.max_pool(
            self.conv1_2,
            ksize=[1, 2, 2, 1],
            strides=[1, 2, 2, 1],
            padding="SAME",
            name="pool1",
        )

        # conv2_1
        with tf.name_scope("conv2_1") as scope:
            kernel = tf.Variable(
                tf.truncated_normal([3, 3, 64, 128], dtype=tf.float32, stddev=1e-1),
                name="weights",
            )
            conv = tf.nn.conv2d(self.pool1, kernel, [1, 1, 1, 1], padding="SAME")
            biases = tf.Variable(
                tf.constant(0.0, shape=[128], dtype=tf.float32),
                trainable=True,
                name="biases",
            )
            out = tf.nn.bias_add(conv, biases)
            self.conv2_1 = tf.nn.relu(out, name=scope)
            self.parameters += [kernel, biases]

        # conv2_2
        with tf.name_scope("conv2_2") as scope:
            kernel = tf.Variable(
                tf.truncated_normal([3, 3, 128, 128], dtype=tf.float32, stddev=1e-1),
                name="weights",
            )
            conv = tf.nn.conv2d(self.conv2_1, kernel, [1, 1, 1, 1], padding="SAME")
            biases = tf.Variable(
                tf.constant(0.0, shape=[128], dtype=tf.float32),
                trainable=True,
                name="biases",
            )
            out = tf.nn.bias_add(conv, biases)
            self.conv2_2 = tf.nn.relu(out, name=scope)
            self.parameters += [kernel, biases]

        # pool2
        self.pool2 = tf.nn.max_pool(
            self.conv2_2,
            ksize=[1, 2, 2, 1],
            strides=[1, 2, 2, 1],
            padding="SAME",
            name="pool2",
        )

        # conv3_1
        with tf.name_scope("conv3_1") as scope:
            kernel = tf.Variable(
                tf.truncated_normal([3, 3, 128, 256], dtype=tf.float32, stddev=1e-1),
                name="weights",
            )
            conv = tf.nn.conv2d(self.pool2, kernel, [1, 1, 1, 1], padding="SAME")
            biases = tf.Variable(
                tf.constant(0.0, shape=[256], dtype=tf.float32),
                trainable=True,
                name="biases",
            )
            out = tf.nn.bias_add(conv, biases)
            self.conv3_1 = tf.nn.relu(out, name=scope)
            self.parameters += [kernel, biases]

        # conv3_2
        with tf.name_scope("conv3_2") as scope:
            kernel = tf.Variable(
                tf.truncated_normal([3, 3, 256, 256], dtype=tf.float32, stddev=1e-1),
                name="weights",
            )
            conv = tf.nn.conv2d(self.conv3_1, kernel, [1, 1, 1, 1], padding="SAME")
            biases = tf.Variable(
                tf.constant(0.0, shape=[256], dtype=tf.float32),
                trainable=True,
                name="biases",
            )
            out = tf.nn.bias_add(conv, biases)
            self.conv3_2 = tf.nn.relu(out, name=scope)
            self.parameters += [kernel, biases]

        # conv3_3
        with tf.name_scope("conv3_3") as scope:
            kernel = tf.Variable(
                tf.truncated_normal([3, 3, 256, 256], dtype=tf.float32, stddev=1e-1),
                name="weights",
            )
            conv = tf.nn.conv2d(self.conv3_2, kernel, [1, 1, 1, 1], padding="SAME")
            biases = tf.Variable(
                tf.constant(0.0, shape=[256], dtype=tf.float32),
                trainable=True,
                name="biases",
            )
            out = tf.nn.bias_add(conv, biases)
            self.conv3_3 = tf.nn.relu(out, name=scope)
            self.parameters += [kernel, biases]

        # pool3
        self.pool3 = tf.nn.max_pool(
            self.conv3_3,
            ksize=[1, 2, 2, 1],
            strides=[1, 2, 2, 1],
            padding="SAME",
            name="pool3",
        )

        # conv4_1
        with tf.name_scope("conv4_1") as scope:
            kernel = tf.Variable(
                tf.truncated_normal([3, 3, 256, 512], dtype=tf.float32, stddev=1e-1),
                name="weights",
            )
            conv = tf.nn.conv2d(self.pool3, kernel, [1, 1, 1, 1], padding="SAME")
            biases = tf.Variable(
                tf.constant(0.0, shape=[512], dtype=tf.float32),
                trainable=True,
                name="biases",
            )
            out = tf.nn.bias_add(conv, biases)
            self.conv4_1 = tf.nn.relu(out, name=scope)
            self.parameters += [kernel, biases]

        # conv4_2
        with tf.name_scope("conv4_2") as scope:
            kernel = tf.Variable(
                tf.truncated_normal([3, 3, 512, 512], dtype=tf.float32, stddev=1e-1),
                name="weights",
            )
            conv = tf.nn.conv2d(self.conv4_1, kernel, [1, 1, 1, 1], padding="SAME")
            biases = tf.Variable(
                tf.constant(0.0, shape=[512], dtype=tf.float32),
                trainable=True,
                name="biases",
            )
            out = tf.nn.bias_add(conv, biases)
            self.conv4_2 = tf.nn.relu(out, name=scope)
            self.parameters += [kernel, biases]

        # conv4_3
        with tf.name_scope("conv4_3") as scope:
            kernel = tf.Variable(
                tf.truncated_normal([3, 3, 512, 512], dtype=tf.float32, stddev=1e-1),
                name="weights",
            )
            conv = tf.nn.conv2d(self.conv4_2, kernel, [1, 1, 1, 1], padding="SAME")
            biases = tf.Variable(
                tf.constant(0.0, shape=[512], dtype=tf.float32),
                trainable=True,
                name="biases",
            )
            out = tf.nn.bias_add(conv, biases)
            self.conv4_3 = tf.nn.relu(out, name=scope)
            self.parameters += [kernel, biases]

        # pool4
        self.pool4 = tf.nn.max_pool(
            self.conv4_3,
            ksize=[1, 2, 2, 1],
            strides=[1, 2, 2, 1],
            padding="SAME",
            name="pool4",
        )

        # conv5_1
        with tf.name_scope("conv5_1") as scope:
            kernel = tf.Variable(
                tf.truncated_normal([3, 3, 512, 512], dtype=tf.float32, stddev=1e-1),
                name="weights",
            )
            conv = tf.nn.conv2d(self.pool4, kernel, [1, 1, 1, 1], padding="SAME")
            biases = tf.Variable(
                tf.constant(0.0, shape=[512], dtype=tf.float32),
                trainable=True,
                name="biases",
            )
            out = tf.nn.bias_add(conv, biases)
            self.conv5_1 = tf.nn.relu(out, name=scope)
            self.parameters += [kernel, biases]

        # conv5_2
        with tf.name_scope("conv5_2") as scope:
            kernel = tf.Variable(
                tf.truncated_normal([3, 3, 512, 512], dtype=tf.float32, stddev=1e-1),
                name="weights",
            )
            conv = tf.nn.conv2d(self.conv5_1, kernel, [1, 1, 1, 1], padding="SAME")
            biases = tf.Variable(
                tf.constant(0.0, shape=[512], dtype=tf.float32),
                trainable=True,
                name="biases",
            )
            out = tf.nn.bias_add(conv, biases)
            self.conv5_2 = tf.nn.relu(out, name=scope)
            self.parameters += [kernel, biases]

        # conv5_3
        with tf.name_scope("conv5_3") as scope:
            kernel = tf.Variable(
                tf.truncated_normal([3, 3, 512, 512], dtype=tf.float32, stddev=1e-1),
                name="weights",
            )
            conv = tf.nn.conv2d(self.conv5_2, kernel, [1, 1, 1, 1], padding="SAME")
            biases = tf.Variable(
                tf.constant(0.0, shape=[512], dtype=tf.float32),
                trainable=True,
                name="biases",
            )
            out = tf.nn.bias_add(conv, biases)
            self.conv5_3 = tf.nn.relu(out, name=scope)
            self.parameters += [kernel, biases]

        # pool5
        self.pool5 = tf.nn.max_pool(
            self.conv5_3,
            ksize=[1, 2, 2, 1],
            strides=[1, 2, 2, 1],
            padding="SAME",
            name="pool5",
        )
        # MaxPool6
        self.maxpool5 = tf.nn.max_pool(
            self.pool5,
            ksize=[1, 7, 7, 1],
            strides=[1, 1, 1, 1],
            padding="SAME",
            name="maxpool5",
        )

        # AvgPool6
        self.avgpool5 = tf.nn.avg_pool(
            self.pool5,
            ksize=[1, 7, 7, 1],
            strides=[1, 1, 1, 1],
            padding="SAME",
            name="avgpool5",
        )

    def dan_part(self):

        # fc1
        with tf.name_scope("reg_head") as scope:
            shape = 2 * int(np.prod(self.maxpool5.get_shape()[1:]))
            fc1w = tf.Variable(
                tf.truncated_normal([shape, 5], dtype=tf.float32, stddev=1e-1),
                name="weights",
            )
            fc1b = tf.Variable(
                tf.constant(1.0, shape=[5], dtype=tf.float32),
                trainable=True,
                name="biases",
            )

            maxpool5_flat = tf.nn.l2_normalize(
                tf.reshape(self.maxpool5, [-1, int(shape / 2)]), 1
            )
            avgpool5_flat = tf.nn.l2_normalize(
                tf.reshape(self.avgpool5, [-1, int(shape / 2)]), 1
            )

            self.concat = tf.concat([maxpool5_flat, avgpool5_flat], 1)
            self.reg_head = tf.nn.bias_add(
                tf.matmul(self.concat, fc1w), fc1b, name=scope
            )
            self.parameters += [fc1w, fc1b]

    def initialize_with_vggface(self, weight_file, sess):
        data = loadmat(weight_file)
        layers = data["layers"][0]
        i = 0
        for layer in layers:
            name = layer[0]["name"][0][0]
            layer_type = layer[0]["type"][0][0]
            if layer_type == "conv" and name[0:2] != "fc":
                kernel, bias = layer[0]["weights"][0][0]
                sess.run(self.parameters[i].assign(kernel))
                sess.run(self.parameters[i + 1].assign(bias.reshape(bias.shape[0])))
                i += 2

    def load_trained_model(self, pickle_file, sess):
        with open(pickle_file, "rb") as pfile:
            param = pickle.load(pfile)
        for i in range(len(param)):
            sess.run(self.parameters[i].assign(param[i]))


'''
punjabi-pca: PCA analysis for Punjabi project (Kochetov, Faytak, Nara).
  WARNING: Largely superceded by ./punjabi-series-pca-lda.py
'''

import argparse
import glob
import matplotlib.pyplot as plt
import numpy as np
import os
import pandas as pd
import re
import sys

from hashlib import sha1
from scipy.ndimage import median_filter
from sklearn.decomposition import PCA
from sklearn.discriminant_analysis import LinearDiscriminantAnalysis as LDA

# read in args
parser = argparse.ArgumentParser()
parser.add_argument("directory", help="Experiment directory containing all subjects")
args = parser.parse_args()

# check for appropriate directory
expdir = args.directory
try:
    assert os.path.exists(args.directory)
except AssertionError:
    # TODO raise exception
    print("\tDirectory provided doesn't exist")
    parser.print_help()
    sys.exit(2)

data_in = os.path.join(expdir,"frames.npy")
data = np.load(data_in)
metadata_in = os.path.join(expdir,"frames_metadata.pickle")
md = pd.read_pickle(metadata_in)

# sanity checks
assert(len(md) == data.shape[0]) # make sure one md row for each frame
assert(md.loc[0, 'sha1'] == sha1(data[0].ravel()).hexdigest()) # checksums
assert(md.loc[len(md)-1,'sha1'] == sha1(data[-1].ravel()).hexdigest())

n_pca = 4
#n_lda = 1
image_shape = data[0].shape # base off of first frame

subject = []
phase = []
trial = []
phone = []

for s in np.unique(md['subject']):
	# subset data by subject ID
	subj_mask = (md['subject'] == s)
	subj_mask = subj_mask.as_matrix()
	model_data = data[subj_mask]
	model_md = md[subj_mask]

	# preallocate array for ultrasound frames for PCA
	model_array = np.empty([model_data.shape[0]] + list(model_data[0].shape)) * np.nan
	model_array = model_array.astype('uint8')

	# fill in the preallocated array, applying median filter (and any other desired transforms)
	for idx,frame in enumerate(model_data):
		filt_frame = median_filter(frame, 5)
		model_array[idx,:,:] = filt_frame # frame

	# run PCA with three PCs
	n_components = int(n_pca)
	pca = PCA(n_components=int(n_components))
	array_reshaped = model_array.reshape([
				model_array.shape[0],
				model_array.shape[1] * model_array.shape[2]
				])
	pca.fit(array_reshaped)
	cumulative_var_exp = sum(pca.explained_variance_ratio_)

	print("Subj.{}: PCA with {} PCs explains {} of variation".format(s,
			n_components,
			round(cumulative_var_exp,4)
			))

	pca_out = pca.transform(array_reshaped)

	# output PC values by acquisition

	# create output table headers
	pc_headers = ["pc"+str(i+1) for i in range(0,n_components)] # n. of PC columns changes acc. to n_components
	meta_headers = list(md.columns.values)
	headers = meta_headers + pc_headers

	# create output table
	headless = np.column_stack((md[meta_headers], pca_out))
	d = np.row_stack((headers, headless))

	out_filename = "{}_pca.csv".format(s)
	out_path = os.path.join(expdir,out_filename)
	np.savetxt(out_path, d, fmt="%s", delimiter ='\t')

	# TODO once relevant, output one table across multiple subjects?

	# output eigentongues
	if n_components < 5:
		n_output_pcs = n_components
	else:
		n_output_pcs = 5

	for n in range(0,n_output_pcs):
		dd = pca.components_[n].reshape(image_shape)
		mag = np.max(dd) - np.min(dd)
		pc_load = (dd-np.min(dd))/mag*255
		plt.title("PC{:} min/max loadings, Subj. {:}".format((n+1),s))
		plt.imshow(pc_load, cmap="Greys_r")
		file_ending = "{:}-pc{:}.pdf".format(s, (n+1))
		savepath = os.path.join(expdir,file_ending) # TODO redefine save path if needed
		plt.savefig(savepath)


<filename>tests/plot/fixtures.py
# fixtures.py
# Copyright (c) 2013-2016 <NAME>
# See LICENSE for details
# pylint: disable=C0111,E0611,W0621

# Standard library imports
from __future__ import print_function
import os
import shutil
import subprocess
# PyPI imports
import numpy
import pytest
import scipy
from scipy.misc import imread
# Putil imports
import putil.plot

###
# Global variables
###
IMGTOL = 1e-3

###
# Fixtures
###
def compare_images(image_file_name1, image_file_name2, no_print=True):
    """ Compare two images by calculating Manhattan and Zero norms """
    # Source: http://stackoverflow.com/questions/189943/
    # how-can-i-quantify-difference-between-two-images
    img1 = imread(image_file_name1).astype(float)
    img2 = imread(image_file_name2).astype(float)
    if img1.size != img2.size:
        m_norm, z_norm = 2*[2*IMGTOL]
    else:
        # Element-wise for Scipy arrays
        diff = img1-img2
        # Manhattan norm
        m_norm = scipy.sum(numpy.abs(diff))
        # Zero norm
        z_norm = scipy.linalg.norm(diff.ravel(), 0)
    result = bool((m_norm < IMGTOL) and (z_norm < IMGTOL))
    if not no_print:
        print(
            'Image 1: {0}, Image 2: {1} -> ({2}, {3}) [{4}]'.format(
                image_file_name1, image_file_name2, m_norm, z_norm, result
            )
        )
    return result

def compare_image_set(tmpdir, images_dict_list, section):
    """ Compare image sets """
    subdir = 'test_images_{0}'.format(section)
    tmpdir.mkdir(subdir)
    global_result = True
    for images_dict in images_dict_list:
        ref_file_name_list = images_dict['ref_fname']
        test_file_name = images_dict['test_fname']
        print('Reference images:')
        for ref_file_name in ref_file_name_list:
            print('   file://{0}'.format(
                    os.path.realpath(ref_file_name)
                )
            )
        print('Actual image:')
        print('   file://{0}'.format(
                os.path.realpath(test_file_name)
            )
        )
        partial_result = []
        for ref_file_name in ref_file_name_list:
            partial_result.append(
                compare_images(ref_file_name, test_file_name)
            )
        result = any(partial_result)
        global_result = global_result and partial_result
        if not result:
            print('Images do not match')
            export_image(test_file_name)
        print('')
    if global_result:
        try:
            tmpdir.remove(subdir)
        except OSError: # pragma: no cover
            pass
    return global_result

@pytest.fixture
def default_source():
    """
    Provides a default source to be used in testing the
    putil.plot.Series class
    """
    return putil.plot.BasicSource(
        indep_var=numpy.array([5, 6, 7, 8]),
        dep_var=numpy.array([0, -10, 5, 4])
    )

@pytest.fixture
def default_series(default_source):
    """
    Provides a default series object to be used in testing the
    putil.plot.Panel class
    """
    return putil.plot.Series(
        data_source=default_source,
        label='test series'
    )

@pytest.fixture
def default_panel(default_series):
    """
    Provides a default panel object to be used in testing the
    putil.plot.Figure class
    """
    return putil.plot.Panel(
        series=default_series,
        primary_axis_label='Primary axis',
        primary_axis_units='A',
        secondary_axis_label='Secondary axis',
        secondary_axis_units='B'
    )

def export_image(fname, method=True):
    tdir = os.path.dirname(
        os.path.dirname(
            os.path.dirname(os.path.abspath(__file__))
        )
    )
    artifact_dir = os.path.join(tdir, 'artifacts')
    if not os.path.exists(artifact_dir):
        os.makedirs(artifact_dir)
    if method:
        src = fname
        dst = os.path.join(artifact_dir, os.path.basename(fname))
        shutil.copyfile(src, dst)
    else:
        if os.environ.get('APPVEYOR', None):
            proc = subprocess.Popen(
                ['appveyor', 'PushArtifact', os.path.realpath(fname)],
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT
            )
            proc.communicate()
        elif os.environ.get('TRAVIS', None):
            # If only a few binary files need to be exported a hex dump works,
            # otherwise the log can grow past 4MB and the process is terminated
            # by Travis
            proc = subprocess.Popen(
                [
                    os.path.join(tdir, 'sbin', 'png-to-console.sh'),
                    os.path.realpath(fname)
                ],
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT
            )
            stdout, _ = proc.communicate()
            print(stdout)


<reponame>whdcumt/exposure_correction

import cv2
import numpy as np
import scipy.sparse as spsparse

		# 1. find all boundaries between two components using adj_mat
		# 2. find the minimum weight between them
		# scratch that, instead need to form and maintain a new adjacency matrix
		# which retains the minimum weight between each component

		# Adjacency matrix is going to shrink by 1 for each component merged
		# at stage 0 the adjacency matrix and components hash are consistent
		# updates need to handle both components hash AND adjacency matrix

		# NEWEST ISSUE:
		# Need to deal with max internal weights in adjacency

		# TODO: solve out ground truth for test image

class ExposureCorrector():
	img = []
	adj_mat = []
	# ^ this guy represents adjacency of COMPONENTS, not individual pixels
	# the weights indicate the minimum weight between neighboring components
	# diagonal = max internal weight

	components = {}
	# self.internal_weights = {}
	w = 0
	h = 0

	k = 300

	def get_adjacency_matrix(self, img):
		self.img = img
		rows = np.shape(img)[0]
		cols = np.shape(img)[1]

		self.w = cols
		self.h = rows

		h_connections = np.ones((cols, 1))
		h_connections[-1] = 0;

		diag_h = np.tile(h_connections, (rows, 1))
		diag_h = diag_h[0:-1]
		print diag_h.shape
		assert(diag_h.shape[0] == rows*cols - 1)

		diag_v = np.ones((cols*(rows-1), 1))
		print diag_v.shape
		assert(diag_v.shape[0] == rows*cols - cols)

		diagonals = [diag_h.T, diag_v.T]

		self.adj_mat = spsparse.diags(diagonals, [1, cols], shape=(rows*cols, rows*cols), format='lil', dtype=np.uint8)

		return self.adj_mat

	def init_components(self):
		# make a hash of the initial components with each pixel in it's own
		# component
		num_pixels = self.w * self.h

		for i in range(num_pixels):
			pixel_row = int(i / self.w);
			pixel_col = i % self.w;
			self.components[i] = [(pixel_row, pixel_col)]

		return self.components

	def init_adj_weights(self):
		# at this point each pixel is in it's own component

		# for each pixel, initial weight = abs(pix[i] - pix[j])
		nonzero = self.adj_mat.nonzero()
		print nonzero[0]
		print nonzero[1]

		for i in range(len(nonzero[0])):
			home_pix = nonzero[0][i]
			neighbor = nonzero[1][i]

			home_i = self.pix_num_to_index(home_pix)
			home_val = self.img[home_i[0], home_i[1]]

			neighbor_i = self.pix_num_to_index(neighbor)
			neighbor_val = self.img[neighbor_i[0], neighbor_i[1]]

			self.adj_mat[home_pix, neighbor] = abs(home_val - neighbor_val)

		# for i in range(len(self.components.keys())):
		# 	self.adj_mat[i, i] = 0

	def segment_img(self):
		num_changes = 1
		while (num_changes != 0):
			keys_to_delete = []
			num_changes = 0

			components = self.components.keys()

			for c in components:
				if (c in self.components):
					# keys = self.update_component(c)
					num_changes += self.update_component(c)

					# for j in keys:
					# 	keys_to_delete.append(j)

			# Delete Keys from merges:
			# for k in keys_to_delete:
			# 	self.components.pop(k, None)

			print "Number of merges: " + str(num_changes)
			print "Total components: " + str(len(self.components.keys()))
			print ""

		return self.components

	def update_component(self, c1):
		keys_to_delete = []
		num_changes = 0
		adj_components = self.adj_mat.getrow(c1).nonzero()
		print adj_components[1]
		for c2 in adj_components[1]:
			# skip if we deleted component already
			if (c2 != c1) and (c2 in self.components):
				if (self.check_boundary(c1, c2, self.k) == False):
					self.merge_components(c1, c2)
					print "merged " + str(c1) + " and " + str(c2)
					num_changes += 1
					# keys_to_delete.append(c2)

		return num_changes

	def check_boundary(self, c1, c2, k):
		boundary = True
		m_int = min([self.adj_mat[c1, c1] + k/len(self.components[c1]), self.adj_mat[c2, c2] + k/len(self.components[c2])])

		if self.adj_mat[c1, c2] > m_int:
			boundary = True
		else:
			boundary = False

		return boundary

	def merge_components(self, c1, c2):
		# copy all elements of c2 into c1
		to_copy = self.components[c2]
		for pixel in to_copy:
			self.components[c1].append(pixel)

		# copy adjacency matrix info
		self.merge_adj_mat(c1, c2)

		# delete component in component hash
		self.components.pop(c2, None)

		print self.adj_mat.toarray()
		return 0

	def merge_adj_mat(self, c1, c2):
		c1_adj = self.adj_mat.getrow(c1)
		c2_adj = self.adj_mat.getrow(c2)

		c2_nonzero = c2_adj.nonzero()

		for component in c2_nonzero[1]:
			if component == c1:
				print "update max"
				# update max internal weight
				c1_adj[0, component] = max(c1_adj[0, component], c2_adj[0, component])
			else:
				# update minimum connection to neighbors
				if (c1_adj[0, component] == 0):
					c1_adj[0, component] = c2_adj[0, component]
				elif (c2_adj[0, component] == 0):
					c1_adj[0, component] = c1_adj[0, component]
				else:
					c1_adj[0, component] = min(c1_adj[0, component], c2_adj[0, component])

		self.adj_mat[c1, :] = c1_adj

		# delete c2's row and col in adj_mat?
		# for now, zero out
		self.adj_mat[c2, :] = 0
		self.adj_mat[:, c2] = 0

		return 0

	def pix_num_to_index(self, pix_num):
		row = int(pix_num / self.w)
		col = pix_num % self.w
		return (row, col)

	def min_x_weight(self, c1, c2):
		# get the minimum weight between two components

		pass


<gh_stars>1-10
"""LQR, iLQR and MPC."""

from deeprl_hw3.controllers import approximate_A, approximate_B
import numpy as np
import scipy.linalg
import time

from ipdb import set_trace as debug

def simulate_dynamics_next(env, x, u):
    """Step simulator to see how state changes.

    Parameters
    ----------
    env: gym.core.Env
      The environment you are try to control. In this homework the 2
      link arm.
    x: np.array
      The state to test. When approximating A you will need to perturb
      this.
    u: np.array
      The command to test. When approximating B you will need to
      perturb this.

    Returns
    -------
    next_x: np.array
    """

    env.state = x.copy()
    x1, _, _, _ = env.step(u)
    return x1

def cost_inter(env, x, u, discrete=False):
    """intermediate cost function

    Parameters
    ----------
    env: gym.core.Env
      The environment you are try to control. In this homework the 2
      link arm.
    x: np.array
      The state to test. When approximating A you will need to perturb
      this.
    u: np.array
      The command to test. When approximating B you will need to
      perturb this.
    discrete: boolean
      if True, multiply the env.dt

    Returns
    -------
    l, l_x, l_xx, l_u, l_uu, l_ux. The first term is the loss, where the remaining terms are derivatives respect to the
    env, corresponding variables, ex: (1) l_x is the first order derivative d l/d x (2) l_xx is the second order derivative
    d^2 l/d x^2
    """
    num_actions = env.action_space.shape[0]
    num_states = env.observation_space.shape[0]
    dt = env.dt if discrete else 1.

    l = np.sum(u**2)*dt
    l_x = np.zeros(num_states)*dt
    l_xx = np.zeros((num_states, num_states))*dt
    l_u = 2 * u*dt
    l_uu = 2 * np.eye(num_actions)*dt
    l_ux = np.zeros((num_actions, num_states))*dt

    return l, l_x, l_xx, l_u, l_uu, l_ux

def cost_final(env, x):
    """cost function of the last step

    Parameters
    ----------
    env: gym.core.Env
      The environment you are try to control. In this homework the 2
      link arm.
    x: np.array
      The state to test. When approximating A you will need to perturb
      this.

    Returns
    -------
    l, l_x, l_xx The first term is the loss, where the remaining terms are derivatives respect to the
    corresponding variables
    """

    num_states = env.observation_space.shape[0]
    target = env.goal.copy()

    # calculate l, l_x, l_xx
    weight = 1e4
    l = weight * np.sum((x-target)**2)
    l_x = 2.*weight*(x-target)
    l_xx = 2.*weight*np.eye(num_states)

    return l, l_x, l_xx

def simulate(env, x0, U):
    tN = U.shape[0]
    num_states = env.observation_space.shape[0]

    X = np.zeros((tN, num_states))
    X[0] = x0
    cost = 0.
    for t in range(tN-1):
        X[t+1] = simulate_dynamics_next(env, X[t], U[t])
        l,_,_,_,_,_ = cost_inter(env, X[t], U[t], discrete=True)
        cost += l

    l_f,_,_ = cost_final(env, X[-1])
    cost += l_f

    return X, cost

def solve(env, x0, U, max_iter, useLM, debug_flag):
    # initialize paramters
    action_dim = env.action_space.shape[0]
    state_dim = env.observation_space.shape[0]

    sim_new_trajectory = True
    lamb = 1.0
    lamb_factor = 10
    lamb_max = 1000
    eps_converge = 0.001
    tN = U.shape[0]

    x0 = env.state.copy()
    for ii in range(int(max_iter)):
        # f, c
        if sim_new_trajectory:

            X, cost = simulate(env, x0, U)
            oldcost = np.copy(cost)

            # linearized dynamics
            f_x = np.zeros((tN, state_dim, state_dim))
            f_u = np.zeros((tN, state_dim, action_dim))

            # quadratized cost function
            l = np.zeros((tN,1))
            l_x = np.zeros((tN, state_dim))
            l_xx = np.zeros((tN, state_dim, state_dim))
            l_u = np.zeros((tN, action_dim))
            l_uu = np.zeros((tN, action_dim, action_dim))
            l_ux = np.zeros((tN, action_dim, state_dim))

            for t in range(tN-1):
                A = approximate_A(env, X[t].copy(), U[t].copy(), dt=env.dt)
                B = approximate_B(env, X[t].copy(), U[t].copy(), dt=env.dt)
                f_x[t] = np.eye(state_dim) + A * env.dt
                f_u[t] = B * env.dt

                (l[t], l_x[t], l_xx[t], l_u[t], l_uu[t], l_ux[t]) = \
                    cost_inter(env, X[t], U[t], discrete=True)

            l[-1], l_x[-1], l_xx[-1] = cost_final(env, X[-1])
            sim_new_trajectory = False

        # f, k update
        V = l[-1].copy() # value function
        V_x = l_x[-1].copy() # dV / dx
        V_xx = l_xx[-1].copy() # d^2 V / dx^2
        k = np.zeros((tN, action_dim)) # feedforward modification
        K = np.zeros((tN, action_dim, state_dim)) # feedback gain

        # backward
        for t in range(tN-2, -1, -1):
            Q_x = l_x[t] + np.dot(f_x[t].T, V_x)
            Q_u = l_u[t] + np.dot(f_u[t].T, V_x)

            Q_xx = l_xx[t] + np.dot(f_x[t].T, np.dot(V_xx, f_x[t]))
            Q_ux = l_ux[t] + np.dot(f_u[t].T, np.dot(V_xx, f_x[t]))
            Q_uu = l_uu[t] + np.dot(f_u[t].T, np.dot(V_xx, f_u[t]))

            # Calculate Q_uu^-1 with regularization term set by
            # Levenberg-Marquardt heuristic (at end of this loop)
            Q_uu_evals, Q_uu_evecs = np.linalg.eig(Q_uu)
            Q_uu_evals[Q_uu_evals < 0] = 0.0
            Q_uu_evals += lamb
            Q_uu_inv = np.dot(Q_uu_evecs,
                np.dot(np.diag(1.0/Q_uu_evals), Q_uu_evecs.T))

            k[t] = -np.dot(Q_uu_inv, Q_u)
            K[t] = -np.dot(Q_uu_inv, Q_ux)

            V_x = Q_x - np.dot(K[t].T, np.dot(Q_uu, k[t]))
            V_xx = Q_xx - np.dot(K[t].T, np.dot(Q_uu, K[t]))

        # forward
        Unew = np.zeros((tN, action_dim))
        xnew = x0.copy()
        for t in range(tN - 1):
            Unew[t] = U[t] + k[t] + np.dot(K[t], xnew - X[t])
            xnew = simulate_dynamics_next(env, xnew, Unew[t])

        # evaluate the new trajectory
        Xnew, costnew = simulate(env, x0, Unew)

        # Levenberg-Marquardt heuristic
        if costnew <= cost:

            if useLM:
                lamb /= lamb_factor

            X = np.copy(Xnew) # update trajectory
            U = np.copy(Unew) # update control signal
            oldcost = np.copy(cost)
            cost = np.copy(costnew)

            sim_new_trajectory = True

            if ii > 0 and ((abs(oldcost-cost)/cost) < eps_converge):
                if debug_flag:
                    print("Converged at iteration = %d; Cost = %.4f;"%(ii,costnew) +
                            " logLambda = %.1f"%np.log(lamb))
                break

        else:
            #  gradient descent
            if useLM:
                lamb *= lamb_factor
            if debug_flag and lamb > lamb_max:
                print("lambda > max_lambda at iteration = %d;"%ii +
                        " Cost = %.4f; logLambda = %.1f"%(cost,
                                                          np.log(lamb)))

    return X, U, cost

U = None
t = -1
def calc_ilqr_input(env, sim_env,
    tN=50, max_iter=1e6, useLM=False, debug_flag=False):
    """Calculate the optimal control input for the given state.

    Parameters
    ----------
    env: gym.core.Env
      This is the true environment you will execute the computed
      commands on. Use this environment to get the Q and R values as
      well as the state.
    sim_env: gym.core.Env
      A copy of the env class. Use this to simulate the dynamics when
      doing finite differences.
    tN: number of control steps you are going to execute
    max_itr: max iterations for optmization

    Returns
    -------
    U: np.array
      The SEQUENCE of commands to execute. The size should be (tN, #parameters)
    """
    global U
    global t

    action_dim = env.action_space.shape[0]

    # update t
    t = np.mod(t+1,tN)
    sim_env.state = env.state.copy()

    if U is None:
        if debug_flag: t0 = time.time()
        U = np.zeros((tN,action_dim))
        x0 = sim_env.state.copy()
        X, U, cost = solve(sim_env, x0, U, max_iter, useLM, debug_flag)
        if debug_flag: print('solve with {}'.format(time.time()-t0))

    u = U[t]

    # reset
    if t == tN-1:
        if debug_flag: print('reset')
        U = None

    return u


import pandas as pd
from scipy.stats import chi2_contingency

class ContingencyChiSquare:
    """
    Helper class for generating contingency table and chi square statistics
    """

    # Defining the function for creating contigency table
    def get_contingency_table(self,column1, column2):
        self.contingeny_table = pd.crosstab(index= column1, columns= column2)
        return self.contingeny_table

    #Definig function for getting chi square test
    def get_chi_square_stats(self,column1,column2):
        self.contingeny_table = self.get_contingency_table(column1,column2)
        self.chi_results = chi2_contingency(self.contingeny_table.values)
        return self.chi_results


'''
Created on Sep 29, 2015

@author: ash
'''

# import libraries
import networkx as nx
import matplotlib.pyplot as plt
import random
import math
import numpy as np
from scipy.interpolate import UnivariateSpline
from scipy.interpolate import splprep, splev
from numpy import linspace
from optparse import OptionParser
import sys
import os
import pysal

shpRoomURL = "/home/ash/Dropbox/OSU/BIM/osuSmithLabShpFile/065-0b.shp"
shpTrajectoryURL = "/home/ash/Dropbox/OSU/BIM/osuSmithLabShpFile/osusmith065-0b.shp"

shpFileURL = "/home/ash/Data/FixedShapeFiles/osusmithlabfloor02/065-02lines.shp"

def showRoomNTrajectories():
    roomd = nx.read_shp(shpRoomURL)
    trajectoryd = nx.read_shp(shpTrajectoryURL)

    room = roomd.to_undirected()
    trajectory = trajectoryd.to_undirected()

    print room.nodes(data=False)
    print trajectory.nodes(data=False)

    nodeList = room.nodes(data=True)
    nNode = len(nodeList)
    pos = []
    for i in xrange(nNode): pos.append(nodeList[i][0])
    roomLayout = dict(zip(room,pos))

    nodeList = trajectory.nodes(data=True)
    nNode = len(nodeList)
    pos = []
    for i in xrange(nNode): pos.append(nodeList[i][0])
    trajectoryLayout = dict(zip(trajectory, pos))

    plt.figure(1)
    nx.draw_networkx_edges(room, pos=roomLayout, edgelist=None, width=2, edge_color='k', style='-', alpha=0.8)
    nx.draw_networkx_edges(trajectory, pos = trajectoryLayout, edgelist=None, width=1, edge_color='b', style='-', alpha=0.2)
    plt.show()
    pass

def showShpFileData(shpFileURL):
    shpData = nx.read_shp(shpFileURL)
    shpData = shpData.to_undirected()

    print shpData.nodes(data=True)
    for edge in shpData.edges(data=True):
        print edge

    nodeList = shpData.nodes(data=True)
    nNode = len(nodeList)

    pos = []
    for i in xrange(nNode):
        pos.append(nodeList[i][0])
        pass

    shpLayout = dict(zip(shpData,pos))

    plt.figure(1)
    nx.draw_networkx_edges(shpData, pos=shpLayout, edgelist=None, width=1, edge_color='b', style='-', alpha=0.5)
    nx.draw_networkx_nodes(shpData, pos=shpLayout, nodelist=None, node_size=2)

    print len(shpData.nodes(data=False))
    print len(shpData.edges(data=False))

    plt.show()

if __name__ == '__main__':
#     showRoomNTrajectories()
    showShpFileData(shpFileURL)
    pass