<filename>fitting.py

import numpy as np
from scipy.linalg import sqrtm
from scipy.interpolate import interp1d
import copy, time
import emcee

def ln_prior(pvals, models):
    """
    Logarithm of the prior (mostly just uniform prior bounds for now).
    """
    # Define priors
    priors = {
        'prob1mbb_Q':   (1., 100.),
        'prob1mbb_U':   (1., 100.),
        'dust_T':       (16., 24.),
        'dust_beta':    (1.4, 1.8),
        'sync_beta':    (-1.6, -0.8),
        'ame_nupeak':   (15., 35.),
        'gdust_beta':   (1.1, 1.8),
        'gdust_dbeta':  (-1.8, 1.8),
        'gdust_Td1':    (5., 30.),
        'gdust_Td2':    (5., 30.),
        'gdust_fI':     (0., 1.),
        'gdust_fQ':     (-2., 2.),
        'gdust_fU':     (-2., 2.),
        'sigma_beta':   (1e-2, 1.),
        'sigma_temp':   (.1, 10.),
    }

    # Make ordered list of parameter names
    param_names = []
    for mod in models:
        param_names += mod.param_names

    # Get a list of amplitude names
    amp_names = []
    for mod in models:
        amp_names += ["%s_%s" % (mod.model, pol) for pol in "IQU"]

    pnames = amp_names + param_names

    # Go through priors and apply them
    ln_prior = 0. # Set default prior value
    for pn in priors.keys():
    #    print 'pnames: ' + str(pnames)
    #    print 'pn: ' + str(pn)
        if pn not in pnames: continue
        pmin, pmax = priors[pn] # Prior bounds
    #    print 'pmin: ' + str(pmin) + ', pmax: ' + str(pmax)
    #    print 'pvals: ' + str(pvals)
        val = pvals[pnames.index(pn)] # Current value of parameter
    #    print 'val: ' + str(val)
        if val < pmin or val > pmax:
            ln_prior = -np.inf
    return ln_prior

def lnprob(pvals, data_spec, models_fit, param_spec, Ninv_sqrt):
    """
    log-probability (likelihood times prior) for a set of parameter values.
    """
    # Retrieve instrument/data model and parameter info
    nu, D_vec, Ninv, beam_mat = data_spec

    # pnames should be amps_names + param_names
    # initial_vals should be amps_vals + param_vals
    pnames, initial_vals, parent_model = param_spec

    # Apply prior
    logpr = ln_prior(pvals, models_fit)
    if not np.isfinite(logpr):
        return -np.inf

    F_fg, F_cmb, F = F_matrix(pvals, nu, models_fit, param_spec)
    H = F_fg.T * Ninv * F_fg

    # GLS solution for component amplitudes
    x_mat = np.linalg.inv(F.T * beam_mat.T * Ninv * beam_mat * F) \
          * F.T * beam_mat.T * Ninv * D_vec # Equation A3

    chi_square = (D_vec - beam_mat * F * x_mat).T * Ninv \
               * (D_vec - beam_mat * F * x_mat) # Equation A4

    # Equation A14
    U, Lambda, VT = np.linalg.svd(Ninv_sqrt*F_fg, full_matrices=False)

    # Equation A16
    N_eff_inv_cmb = F_cmb.T * Ninv_sqrt \
                  * (np.matrix(np.identity(U.shape[0])) - U*U.T) \
                  * Ninv_sqrt * F_cmb

    # Total log posterior
    lnprob = logpr - chi_square - 0.5*np.log(np.linalg.det(H)) \
                              - 0.5*np.log(np.linalg.det(N_eff_inv_cmb))

    # Return log-posterior and GLS amplitudes
    return lnprob, np.array(x_mat.T)[0]

def lnprob_joint(params, data_spec, models_fit, param_spec):
    """
    log-probability (likelihood times prior) for a set of parameter values.
    """
    # Retrieve instrument/data model and parameter info
    nu, D_vec, Ninv, beam_mat = data_spec
    pnames, initial_vals, parent_model = param_spec
    Nmod = len(models_fit)
    Npol = 3

    # Separate amplitude and spectral model parameters
    amps = params[:Nmod*Npol]
    pvals = params[Nmod*Npol:]

    # Apply prior
    logpr = ln_prior(params, models_fit)
    if not np.isfinite(logpr):
        return -np.inf

    # Create new copies of model objects to work with
    #models = [copy.deepcopy(m) for m in models_fit]
    models = models_fit

    # Set new parameter values for the copied model objects, and then get
    # scalings as a function of freq./polarisation
    pstart = 0
    mdata = np.zeros(nu.size * Npol)
    for i in range(len(models)):
        m = models[i]

        # Set new parameter values in the models
        n = m.params().size
        #m.set_params( pvals[pstart:pstart+n] )
        mparams = pvals[pstart:pstart+n]
        pstart += n # Increment for next model

        # Calculate scaling with freq. given new parameter values
        amp = np.outer( amps[3*i:3*(i+1)], np.ones(nu.size) ) # Npol*Nfreq array

        # Apply positivity prior on I amplitudes of all components
        #if m.model == 'ame':
        if np.any(amp[0] < 0.):
            return -np.inf

        # Add to model prediction of data vector
        mdata += (amp * m.scaling(nu, params=mparams)).flatten()

    # Calculate chi-squared with data (assumed beam = 1)
    mdata = np.matrix(mdata).T
    chi_square = (D_vec - mdata).T * Ninv * (D_vec - mdata)

    # Return log-posterior
    #return logpr - 0.5 * chi_square
    return -0.5 * chi_square

def F_matrix(pvals, nu, models_fit, param_spec):
    """
    Foreground spectral dependence operator.
    """
    pnames, initial_vals, parent_model = param_spec

    # Check that the CMB component is the first component in the model list
    if models_fit[0].model != 'cmb':
        raise ValueError("The first model in the models_fit list should be a "
                         "CMB() object.")

    Nband = len(nu) # No. of frequency bands
    Npol = 3 # No. of data components (I, Q, U)
    Ncomp = len(models_fit) # No. of sky components

    F_fg = np.zeros((Npol * Nband, Npol * (Ncomp - 1)))
    F_cmb = np.zeros((Npol * Nband, Npol))
    F = np.zeros((Npol * Nband, Npol * Ncomp))

    # Create new copies of model objects to work with
    #models = [copy.deepcopy(m) for m in models_fit]
    models = models_fit

    # Set new parameter values for the copied model objects, and then get
    # scalings as a function of freq./polarisation
    pstart = 0; k = -1
    for i in range(len(models)):
        m = models[i]

        # Set new parameter values in the models
        n = m.params().size
        #m.set_params( pvals[pstart:pstart+n] )
        mparams = pvals[pstart:pstart+n]
        pstart += n # Increment for next model
        if m.model != 'cmb': k += 1 # Increment for next non-CMB model

        # Calculate scaling with freq. given new parameter values
        scal = m.scaling(nu, params=mparams)

        for j in range(Npol):
            # Fill FG or CMB -matrix with scalings, as appropriate
            if m.model != 'cmb':
                F_fg[j*Nband:(j+1)*Nband, k*Npol + j] = scal[j,:]
            else:
                F_cmb[j*Nband:(j+1)*Nband, j] = scal[j,:]

    # Stack CMB and FG F-matrices together
    F = np.hstack((F_cmb, F_fg))
    return np.matrix(F_fg), np.matrix(F_cmb), np.matrix(F)

def mcmc(data_spec, models_fit, param_spec, nwalkers=50,
         burn=500, steps=1000, sample_file=None):
    """
    Run MCMC to fit model to some simulated data.
    """
    # Retrieve instrument/data model and parameter info
    nu, D_vec, Ninv, beam_mat = data_spec
    pnames, initial_vals, parent_model = param_spec

    # Invert noise covariance matrix
    Ninv_sqrt = np.matrix(sqrtm(Ninv))

    # Get a list of model parameter names (FIXME: Ignores input pnames for now)
    param_names = []
    for mod in models_fit:
        param_names += mod.param_names

    # Get a list of amplitude names
    fg_amp_names = []; cmb_amp_names = []
    for mod in models_fit:
        if mod.model == 'cmb':
            cmb_amp_names += ["%s_%s" % (mod.model, pol) for pol in "IQU"]
        else:
            fg_amp_names += ["%s_%s" % (mod.model, pol) for pol in "IQU"]
    pnames = cmb_amp_names + fg_amp_names + param_names

    # Define starting points
    ndim = len(initial_vals)
    pos = [initial_vals*(1.+1e-3*np.random.randn(ndim)) for i in range(nwalkers)]

    # Run emcee sampler
    sampler = emcee.EnsembleSampler( nwalkers, ndim, lnprob,
                           args=(data_spec, models_fit, param_spec, Ninv_sqrt) )
    sampler.run_mcmc(pos, burn + steps)

    # Recover samples of spectral parameters and amplitudes
    param_samples = sampler.chain[:, burn:, :].reshape((-1, ndim))
    amp_samples = np.swapaxes(np.array(sampler.blobs), 0, 1)
    amp_samples = amp_samples[:, burn:, :].reshape((-1, amp_samples.shape[2]))
    samples = np.concatenate((amp_samples.T, param_samples.T))

    # Save chains to file
    if sample_file is not None:
        np.savetxt(sample_file, samples, fmt="%.6e", header=" ".join(pnames))

    # Summary statistics for fitted parameters
    params_out = np.median(param_samples, axis=0)

    # Return summary statistics and samples
    return params_out, pnames, samples

def joint_mcmc(data_spec, models_fit, param_spec, nwalkers=100,
               burn=500, steps=1000, nthreads=2, sample_file=None):
    """
    Run MCMC to fit model to some simulated data. Fits to all parameters, both
    amplitudes and spectral parameters.
    """
    # Retrieve instrument/data model and parameter info
    nu, D_vec, Ninv, beam_mat = data_spec
    pnames, initial_vals, parent_model = param_spec

    # Get a list of model parameter names (FIXME: Ignores input pnames for now)
    param_names = []
    for mod in models_fit:
        param_names += mod.param_names

    # Get a list of amplitude names
    amp_names = []
    for mod in models_fit:
        amp_names += ["%s_%s" % (mod.model, pol) for pol in "IQU"]
    pnames = amp_names + param_names

    # Define starting points
    ndim = len(initial_vals)
    pos = [initial_vals*(1.+1e-3*np.random.randn(ndim)) for i in range(nwalkers)]

    #print param_spec
    # Run emcee sampler
    sampler = emcee.EnsembleSampler( nwalkers, ndim, lnprob_joint,
                                     args=(data_spec, models_fit, param_spec),
                                     threads=nthreads )
    sampler.run_mcmc(pos, burn + steps)

    # Recover samples of spectral parameters and amplitudes
    samples = sampler.chain[:, burn:, :].reshape((-1, ndim))

    # Recover log(posterior)
    logp = sampler.lnprobability[:,burn:].reshape((-1,))

    # Save chains to file
    if sample_file is not None:
        np.savetxt(sample_file, samples, fmt="%.6e", header=" ".join(pnames))

    # Return summary statistics and samples
    return pnames, samples.T, logp

def noise_model_old(fname="data/CMBpol_extended_noise.dat", scale=1.):
    """
    Load noise model from file and create interpolation function as a fn of
    frequency. This is the noise per pixel, for some arbitrary pixel size.
    """
    # Load from file
    nu, sigma = np.genfromtxt(fname).T

    # Extrapolate at the ends of the frequency range
    if nu[0] > 1.:
        sigma0 = sigma[0] \
               + (sigma[1] - sigma[0]) / (nu[1] - nu[0]) * (1. - nu[0])
        sigman = sigma[-1] \
               + (sigma[-1] - sigma[-2]) / (nu[-1] - nu[-2]) * (1e3 - nu[-1])
        if sigma0 < 0.: sigma0 = sigma[0]
        if sigman < 0.: sigman = sigma[-1]

        # Add to end of range
        nu = np.concatenate(([1.,], nu, [1e3,]))
        sigma = np.concatenate(([sigma0,], sigma, [sigman,]))

    # Rescale by constant overall factor
    sigma *= scale

    # Construct interpolation function
    return interp1d(nu, sigma, kind='linear', bounds_error=False)

def noise_model(fname="data/noise_coreplus_extended.dat", scale=1.):
    """
    Load noise model from file and create interpolation function as a fn of
    frequency. This is the noise per pixel, for some arbitrary pixel size.
    """
    # Load from file
    dat = np.genfromtxt(fname).T
    if dat.shape[0] == 3:
        nu, fwhm, sigma = dat
    elif dat.shape[0] == 2:
        nu, sigma = dat
    else:
        raise ValueError("Unexpected number of columns in noise file.")

    # Rescale by constant overall factor
    sigma *= scale

    # Work in log-space
    sigma = np.log(sigma)

    # Extrapolate at the ends of the frequency range
    if nu[0] > 1.:
        sigma0 = sigma[0] \
               + (sigma[1] - sigma[0]) / (nu[1] - nu[0]) * (1. - nu[0])
        sigman = sigma[-1] \
               + (sigma[-1] - sigma[-2]) / (nu[-1] - nu[-2]) * (1e3 - nu[-1])
        if sigma0 < 0.: sigma0 = sigma[0]
        if sigman < 0.: sigman = sigma[-1]

        # Add to end of range
        nu = np.concatenate(([1.,], nu, [1e3,]))
        sigma = np.concatenate(([sigma0,], sigma, [sigman,]))

    # Construct interpolation function
    _interp = interp1d(nu, sigma, kind='linear', bounds_error=False)
    return lambda freq: np.exp(_interp(freq))

def generate_data(nu, fsigma_T, fsigma_P, components,
                  noise_file="data/core_plus_extended_noise.dat",
                  idx_px = 0):
    """
    Create a mock data vector from a given set of models, including adding a
    noise realization.
    """
    # Loop over components that were included in the data model and calculate
    # the signal at a given frequency (should be in uK_CMB)
    signal = 0
    cmb_signal = 0
    # Disabled for the case of the allsky
    if idx_px == 0:
        pass #print( "(FITTING.PY) Parameters in the input model:" )
    for comp in components:
        if idx_px == 0:
            pass #print comp.param_names

        # Add this component to total signal
        signal += np.atleast_2d(comp.amps()).T * comp.scaling(nu)

        # Store CMB signal separately
        if comp.model == 'cmb':
            cmb_signal = np.atleast_2d(comp.amps()).T * comp.scaling(nu)

    # Construct data vector
    D_vec = np.matrix(signal.flatten()).T

    # Noise rms as a function of frequency
    sigma_interp = noise_model(fname=noise_file, scale=1.)
    sigma_nu = sigma_interp(nu / 1e9)
    fsigma = np.zeros(3*len(nu))
    fsigma[0:len(nu)] = fsigma_T * sigma_nu # Stokes I
    fsigma[len(nu):2*len(nu)] = fsigma_P * sigma_nu # Stokes Q
    fsigma[2*len(nu):] = fsigma_P * sigma_nu # Stokes U

    #noise_mat = np.matrix( np.diagflat(cmb_signal.flatten() * fsigma) )
    #noise_mat = np.matrix( np.diagflat(fsigma) )
    #Ninv = np.linalg.inv(noise_mat)

    # Inverse noise covariance
    noise_mat = np.identity(fsigma.size) * fsigma
    Ninv = np.identity(fsigma.size) / fsigma**2.
    n_vec = (np.matrix(np.random.randn(D_vec.size)) * noise_mat).T

    # Add noise to generated data
    D_vec += n_vec
    return D_vec, Ninv

def model_test(nu, D_vec, Ninv, models_fit, initial_vals=None, burn=500,
               steps=1000, cmb_amp_in=None, sample_file=None):
    """
    Generate simulated data given an input model, and perform MCMC fit using
    another model.
    """
    # Collect together data and noise/instrument model
    Ninv_sqrt = np.matrix(sqrtm(Ninv)) # Invert noise covariance matrix
    beam_mat = np.identity(3*len(nu)) # Beam model
    data_spec = (nu, D_vec, Ninv, beam_mat)

    # Loop over specified component models and set up MCMC parameters for them
    pnames = []; pvals = []; parent_model = []
    for mod in models_fit:
        # Get parameter names, initial parameter values, and component ID
        pn = mod.param_names
        pv = mod.params()

        # Loop through parameters from this component
        for i in range(len(pn)):
            pnames.append( "%s.%s" % (mod.name, pn[i]) )
            pvals.append( pv[i] )
            parent_model.append( mod )

    # Use 'guess' as the initial point for the MCMC if specified
    if initial_vals is None: initial_vals = pvals

    # Collect names, initial values, and parent components for the parameters
    param_spec = (pnames, initial_vals, parent_model)

    # Run MCMC sampler on this model
    t0 = time.time()
    params_out, pnames, samples = mcmc(data_spec, models_fit, param_spec,
                                       burn=burn, steps=steps,
                                       sample_file=sample_file)
    print "MCMC run in %d sec." % (time.time() - t0)

    # Estimate error on recovered CMB amplitudes
    # FIXME: Why estimate error using F_matrix on median!?
    F_fg, F_cmb, F = F_matrix(params_out, nu, models_fit, param_spec)

    H = F_fg.T * Ninv * F_fg

    # Equation A3
    x_mat = np.linalg.inv(F.T * beam_mat.T * Ninv * beam_mat * F) \
          * F.T * beam_mat.T * Ninv * D_vec

    # Equation A14
    U, Lambda, VT = np.linalg.svd(Ninv_sqrt*F_fg, full_matrices=False)

    # Equation A16
    N_eff_inv_cmb = F_cmb.T * Ninv_sqrt \
                  * (np.matrix(np.identity(U.shape[0])) - U*U.T) \
                  * Ninv_sqrt * F_cmb

    N_eff_cmb = np.linalg.inv(N_eff_inv_cmb)
    cmb_noise = np.array([N_eff_cmb[0,0], N_eff_cmb[1,1], N_eff_cmb[2,2]])

    gls_cmb = x_mat[0:3,0]
    cmb_chisq = (np.matrix(cmb_amp_in).T - gls_cmb).T * N_eff_inv_cmb \
              * (np.matrix(cmb_amp_in).T - gls_cmb)

    return gls_cmb, cmb_chisq, cmb_noise


<reponame>kmckiern/scripts
#!/bin/env python

"""
example usage:
        python pore_occupancy.py --trj whereever/cnv.xtc --top 4xdk.pdb --out sf_h2o_p9761_17_14.dat

note: use python version >= 3.4
"""

import argparse
import mdtraj
import numpy as np
import pandas as pd
from scipy.stats import multivariate_normal
import IPython

parser = argparse.ArgumentParser(description='get water and ion occupancy of trek sf')
parser.add_argument('--tf', type=str, help='trajectory file')
parser.add_argument('--gaussian', action='store_true', help='gaussian counting', default=False)
parser.add_argument('--histogram', action='store_true', help='histogram counting', default=False)
parser.add_argument('--start', type=int, help='start bin for sum window')
parser.add_argument('--end', type=int, help='end bin for sum window')
args = parser.parse_args()

# parameters obtained from visual inspection in vmd ... could be better
# help to define the boundary of the selectivity filter
cyl_radius = .55
z_slice = 0.2
cyl_length = 1.0
num_bins = round(cyl_length/z_slice)
full_l = 2.0*cyl_length
nb = 2 * num_bins

res_indxs = [101, 210, 362, 471]

# via:
# http://stackoverflow.com/questions/25720600/generating-3d-gaussian-distribution-in-python
def numerical_gaussian(xmin, xmax, ymin, ymax, mu, sigma):
    # generate numerical grid
    x, y = np.mgrid[xmin:xmax, ymin:ymax]
    xy = np.column_stack([x.flat, y.flat])
    # get gaussian values over grid
    covariance = np.diag(sigma**2)
    gz = multivariate_normal.pdf(xy, mean=mu, cov=covariance)
    gz = g.reshape(x.shape)
    return gz

def relevant_indices(trj):
    # CA of each GLY located centerally to each selectivity filter strand
    # +2 is top, -2 is bottom of strand
    for i in [-3, -2, -1, 1, 2, 3]:
        for j in range(4):
            res_indxs.append(res_indxs[j]+i)
    atom_indxs = []
    for ref in res_indxs:
        atom_indxs += [atom.index for atom in trj.top.atoms if ((atom.residue.resSeq == ref) and (atom.name == 'CA'))]
    ri = np.array(atom_indxs)
    wi = [atom.index for atom in trj.topology.atoms if (atom.residue.is_water and (atom.element.symbol == 'O'))]
    ki = np.array([i.index for i in trj.top.atoms_by_name('K+')])
    return ri, wi, ki

def sf_bc(traj, ri):
    # filter by z
    ri_qz = traj.xyz[:, ri, -1]
    z_cent = np.average(ri_qz, axis=1)
    z_min = z_cent - cyl_length
    z_max = z_cent + cyl_length
    # filter by xy
    ri_qxy = traj.xyz[:, ri, :2]
    xy_cent = np.average(ri_qxy, axis=1)
    return z_min, z_max, xy_cent

def discrete_bins(traj, wi, ki, z_min, z_max, xy_cent):
    nf = len(traj)
    # time series of bin by frame, n_h2o, n_k
    time_series = np.zeros([nb, nf, 2])

    z_lbls = []
    for ndx, atom_var in enumerate([wi, ki]):
        # label above (1), within (0), or below (-1) sf wrt z
        zs = traj.xyz[:, atom_var, -1]
        z_scale = (zs.T - z_min).T
        z_scale[(z_scale >= 0.0) & (z_scale <= full_l)] = 0
        z_scale[z_scale < 0.0] = -1
        z_scale[z_scale > full_l] = 1
        z_lbls.append(z_scale)

        # label if within radius of sf (0), or outside (1)
        xys = traj.xyz[:,atom_var,:2]
        sx_scale = (xys[:,:,0].T - xy_cent[:,0]).T
        sy_scale = (xys[:,:,1].T - xy_cent[:,1]).T
        moduli = sx_scale**2 + sy_scale**2
        moduli[moduli <= cyl_radius] = 0
        moduli[moduli > cyl_radius] = 1

        for frame in range(nf):
            z_good = np.where(z_scale[frame]==0)[0]
            xy_good = np.where(moduli[frame]==0)[0]
            in_sf = np.array(list(set(z_good).intersection(xy_good)))

            final_z = [traj.xyz[frame, atom_var[i], -1] for i in in_sf]

            # histogram the results
            tval, tbin = np.histogram(final_z, bins=nb, range=(z_min[frame], z_max[frame]))
            time_series[:, frame, ndx] = tval

    return time_series, z_lbls

# if state = end = None, full matrix will be summed
def sum_bins(start, end, timeseries):
     totaled = np.sum(timeseries[start:end], axis=0)
     return totaled

# tsave
def write_bins(of, time_series, out_pref):
    if len(time_series.shape) == 3:
        for ndx, bin_vals in enumerate(time_series):
            np.save(out_pref + str(ndx) + '/' + of, bin_vals)
    else:
        np.save(out_pref + '_' + of, time_series)

def main():
    trj_data = open(args.tf).readlines()
    for t_ndx, trjd in enumerate(trj_data):
        # read, load, align
        t, top = trjd.split()
        label = t.strip().split('/')[-1].split('.')[0]
        traj = mdtraj.load(t, top=top)
        pi = traj.top.select('protein')
        traj = traj.superpose(traj, atom_indices=pi)
        x = traj[0]

        ri, wi, ki = relevant_indices(x)
        # get ions and water in the selectivity filter
        if args.gaussian:
            # fit 10 gaussians to SF
            gaussians = fit_gaussians()
            # bin gaussian counts
            bin_timeseries, zs = gaussian_bins(*gaussians)

        if args.histogram:
            # get trj info
            z_min, z_max, xy_cent = sf_bc(traj, ri)
            # histogram data
            bin_timeseries, zs = discrete_bins(traj, wi, ki, z_min, z_max, xy_cent)

        # record bin data
        for ndx, i in enumerate(bin_timeseries):
            if ndx == 0:
                write_bins(label, zs[0], 'atom_resolved_bins/water_')
            if ndx == 1:
                write_bins(label, zs[1], 'atom_resolved_bins/k_')
        write_bins(label, bin_timeseries, 'bins/b')
        # sum over bins for feature vector of SF occupancy
        first = args.start
        last = args.end
        bin_window = sum_bins(first, last, bin_timeseries)
        write_bins(label, bin_window, 'bins/summed/w_' + str(first) + '_' + str(last))

if __name__ == '__main__':
    main()


<filename>modules/spectrograph_optics.py
#!/usr/bin/env python
# vim: set fileencoding=utf-8 ts=4 sts=4 sw=4 et tw=80 :
#
# Routines related to spectrograph optical properties and calculations.
#
# <NAME>
# Created:       2018-12-29
# Last modified: 2019-03-30
#--------------------------------------------------------------------------
#**************************************************************************
#--------------------------------------------------------------------------

## Current version:
__version__ = "0.2.0"

## Python version-agnostic module reloading:
#try:
#    reload                              # Python 2.7
#except NameError:
#    try:
#        from importlib import reload    # Python 3.4+
#    except ImportError:
#        from imp import reload          # Python 3.0 - 3.3

## Modules:
import os
import sys
import time
import numpy as np
#from numpy.lib.recfunctions import append_fields
#import datetime as dt
#from dateutil import parser as dtp
#import scipy.linalg as sla
#import scipy.signal as ssig
#import scipy.ndimage as ndi
import scipy.optimize as opti
#import scipy.interpolate as stp
#import scipy.spatial.distance as ssd
from functools import partial
#from collections import OrderedDict
#import multiprocessing as mp
#np.set_printoptions(suppress=True, linewidth=160)
#import pandas as pd
#import statsmodels.api as sm
#import statsmodels.formula.api as smf
#from statsmodels.regression.quantile_regression import QuantReg
#import theil_sen as ts
#import window_filter as wf
#import itertools as itt

## Glass indexes of refraction:
from glass import Glass
#import glass
#reload(glass)

###--------------------------------------------------------------------------##
### Calculating index of refraction for several glass types:
#class Glass(object):
#
#    def __init__(self, glasstype):
#        self._bcoeffs = {
#            'SiO2':np.array([0.67071081e0, 0.433322857e0, 0.877379057e0]),
#            'LLF1':np.array([1.21640125e0, 1.33664540e-1, 8.83399468e-1]),
#            'PBM2':np.array([1.39446503e0, 1.59230985e-1, 2.45470216e-1]),
#             'LF5':np.array([1.28035628e0, 1.6350597e-1,  8.93930112e-1]), }
#        self._ccoeffs = {
#            'SiO2':np.array([0.00449192312e0, 0.0132812976e0, 95.8899878e0]),
#            'LLF1':np.array([8.57807248e-3,   4.20143003e-2,   1.07593060e+2]),
#            'PBM2':np.array([1.10571872e-2,   5.07194882e-2,   3.14440142e1]),
#             'LF5':np.array([9.29854416e-3,   4.49135769e-2,   1.10493685e2]), }
#        if self._unknown_glass(glasstype):
#            raise
#        self._gtype = glasstype
#        self._coeffs = zip(self._bcoeffs[self._gtype],
#                           self._ccoeffs[self._gtype])
#        return
#
#    def _unknown_glass(self, glasstype):
#        if not glasstype in self._bcoeffs.keys():
#            sys.stderr.write("Unknown glass type: %s\n" % glasstype)
#            return True
#        else:
#            return False
#
#    # Squared index of refraction for specified wavelengths:
#    def refraction_index_squared(self, wlen_um):
#        lam_um_sq = wlen_um**2
#        n_squared = np.ones_like(wlen_um, dtype='float')
#        for bb,cc in self._coeffs:
#            n_squared += (lam_um_sq * bb) / (lam_um_sq - cc)
#        return n_squared
#
#    def refraction_index(self, wlen_um):
#        return np.sqrt(self.refraction_index_squared(wlen_um))
#
#    def glass_dn_dlambda_easy(self, wlen_um, glasstype, epsfrac=1e-5):
#        #if self._unknown_glass(glasstype):
#        #    raise
#
#        wlen_lower = wlen_um * (1.0 - epsfrac)
#        wlen_upper = wlen_um * (1.0 + epsfrac)
#
#        nn_lower = self.refraction_index(wlen_lower) #, glasstype)
#        nn_upper = self.refraction_index(wlen_upper) #, glasstype)
#
#        return (nn_upper - nn_lower) / (wlen_upper - wlen_lower)

##--------------------------------------------------------------------------##
##--------------------------------------------------------------------------##
## Prism deflection:
#def prism_deflection_n(incid_r, apex_r, n):
#    ptemp = np.sqrt(n**2 - np.sin(incid_r)**2) * np.sin(apex_r) \
#                - np.cos(apex_r) * np.sin(incid_r)
#    return incid_r - apex_r + np.arcsin(ptemp)

#def prism_deflection_glass_wl(incid_r, apex_r, gtype, wlen):

#def wiki_prism_deflection_n(i, A, n):
#    return i - A + np.arcsin(n * np.sin(A - np.arcsin(np.sin(i) / n)))

##--------------------------------------------------------------------------##
## Prism object to calculate deflections for the specified material and shape:
class Prism(object):

    def __init__(self, glasstype, apex_deg, n_air=1.0):
        #self._apex_deg = apex_deg
        #self._apex_rad = np.radians(apex_deg)
        self._material = Glass(glasstype)
        self.set_apex_deg(apex_deg)
        self.set_n_air(n_air)    # index of refraction of air in spectrograph
        return

    # -------------------------------------
    # Getters and setters:
    def get_apex_rad(self):
        return self._apex_rad

    def get_apex_deg(self):
        return self._apex_deg

    def set_apex_deg(self, apex_deg):
        self._apex_deg = apex_deg
        self._apex_rad = np.radians(apex_deg)
        return

    def set_n_air(self, index_refr):
        self._n_air = index_refr
        return

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

    @staticmethod
    def _wiki_deflection_rad_nr(i, A, nr):
        """Deflection angle using formula from Wikipedia (which seems
        to be identical to others used here. Inputs are:
        i  -- incidence angle (RADIANS)
        A  -- prism apex angle (RADIANS)
        nr -- index of refraction RATIO at desired wavelength(s):
                    n_glass / n_air
        """
        return i - A + np.arcsin(nr * np.sin(A - np.arcsin(np.sin(i) / nr)))

    def deflection_rad_nr(self, incidence_r, nr):
        """Deflection angle as a function of incidence angle and
        index of refraction RATIO (n_glass / n_air)."""
        return self._wiki_deflection_rad_nr(incidence_r, self._apex_rad, nr)

    def deflection_rad_nr2(self, incidence_r, nr2):
        """Calculate deflection angle from incidence and SQUARED index of
        refraction RATIO (glass / air). Units of RADIANS used throughout."""
        ptemp = np.sqrt(nr2 - np.sin(incidence_r)**2) * np.sin(self._apex_rad) \
                - np.cos(self._apex_rad) * np.sin(incidence_r)
        return incidence_r - self._apex_rad + np.arcsin(ptemp)

    def deflection_rad_wl(self, incidence_r, wavelength_um):
        """Calculate deflection angle given incidence angle and wavelength
        in microns. Units of RADIANS used throughout."""
        n2_glass = self._material.refraction_index_squared(wavelength_um)
        n2_air   = self._n_air**2
        return self.deflection_rad_nr2(incidence_r, n2_glass/n2_air)
        #ptemp = np.sqrt(n2 - np.sin(incidence_r)**2) * np.sin(self._apex_rad) \
        #        - np.cos(self._apex_rad) * np.sin(incidence_r)
        #return incidence_r - self._apex_rad + np.arcsin(ptemp)

    def deflection_deg_wl(self, incidence_d, wavelength_um):
        """Calculate deflection angle given incidence angle and wavelength
        in microns. Units of RADIANS used throughout."""
        incidence_r = np.radians(incidence_deg)
        return np.degrees(self.deflection_rad_wl(incidence_r, wavelength_um))

##--------------------------------------------------------------------------##
## Notes on notation, relations, identities, etc.:

# theta_B  --> blaze angle
# lambda_B --> blaze wavelength
# theta    --> facet illumination angle
# alpha    --> angle of incidence
# beta     --> angle of diffraction
# gamma    --> angle of incidence relative to facet normal SPECIFICALLY
#                   in the plane that is parallel to grooves

# Related by:
# alpha = theta_B + theta
# beta  = theta_B - theta

# Note that:
# * theta is SMALL in Littrow configuration, alpha ~= beta

## Grating methods implemented:
class GratingTools(object):

    #groove_spacing_units = {'mm':1.0, 'um':1e3, 'nm':1e6}

    def __init__(self, gratio=None, lines_per_mm=None):
        # Placeholders:
        self._blaze_angle_rad = 0.0
        self._g_spacing_units = {'mm':1.0, 'um':1e3, 'nm':1e6}
        self._groove_spacing = {x:0.0 for x in self._g_spacing_units.keys()}

        # Handle useful inputs:
        if gratio:
            self._blaze_angle_rad = self._calc_blaze_angle(gratio)
        if lines_per_mm:
            for kk,mult in self._g_spacing_units.items():
                self._groove_spacing[kk] = mult / lines_per_mm

        return

    # --------------------------------------
    # Blaze angle and wavelength:

    @staticmethod
    def _calc_blaze_angle(ratio):
        """For R4 grating, use ratio = 4.  Blaze angle returns in RADIANS."""
        return np.arctan(ratio)

    @staticmethod
    def _calc_blaze_wavelength(spec_orders, groove_spacing,
            theta_B_rad, facet_rad, gamma_rad):
        """
        Compute blaze wavelength for an array of order numbers.
        NOTE: wavelength computed in same units as groove spacing.
        """
        ord_vec = np.atleast_1d(spec_orders)
        if not np.all(ord_vec == np.int_(ord_vec)):
            sys.stderr.write("Error: non-integer spec_orders provided!\n")
            return np.nan * spec_orders

        ang_tmp = np.sin(theta_B_rad) * np.cos(facet_rad) * np.cos(gamma_rad)
        return 2.0 * groove_spacing * ang_tmp / np.float_(spec_orders)

    def get_blaze_wavelengths(self, spec_orders, facet_rad=0.0, gamma_rad=0.0,
            units='nm'):
        return self._calc_blaze_wavelength(spec_orders,
                self._groove_spacing[units], self._blaze_angle_rad,
                facet_rad=facet_rad, gamma_rad=gamma_rad)

    @staticmethod
    def _calc_line_tilt_ctr(blaze_rad, gamma_rad):
        """Calculate line tilt at order center (at blaze wavelength).
        Formulae from Barnes (2004) thesis, page 11:
        tan(φ) = (sin(α) + sin(β)) / cos(β) * (sin(γ) / cos(γ))
               = λ * dβ/dλ * tan(γ)
        tan(φ) = 2 * tan(θ_B) * tan(γ)      # at blaze wavelength λ_B
        """
        return np.arctan(2.0 * np.tan(blaze_rad) * np.tan(gamma_rad))

    @staticmethod
    def _calc_line_tilt_any(alpha_rad, beta_rad, gamma_rad):
        """Calculate line tilt at order center (at blaze wavelength).
        Formulae from Barnes (2004) thesis, page 11:
        tan(φ) = (sin(α) + sin(β)) / cos(β) * (sin(γ) / cos(γ))
               = λ * dβ/dλ * tan(γ)
        tan(φ) = 2 * tan(θ_B) * tan(γ)      # at blaze wavelength λ_B
        """
        numer = (np.sin(alpha_rad) + np.sin(beta_rad)) * np.sin(gamma_rad)
        denom = np.cos(beta_rad) * np.cos(gamma_rad)
        return np.arctan(numer / denom)

    @staticmethod
    def _calc_order_FSR(wlen_cen, spec_orders):
        return wlen_cen / spec_orders

    @staticmethod
    def _calc_order_width(wlen_cen, groove_spacing,
            theta_B_rad, facet_rad, gamma_rad):
        """
        Compute angular span of the order (free spectral range) with central
        wavelength wlen_cen.

        NOTES:
        * groove_spacing and wlen_cen need same units for correct result!
        * diff_angle_rad should correspond to central wavelength (wlen_cen)
        * wlen_cen supports numpy arrays
        """
        diff_angle_rad = theta_B_rad - facet_rad
        cos_diff_gamma = np.cos(diff_angle_rad) * np.cos(gamma_rad)
        return wlen_cen / (groove_spacing * cos_diff_gamma)

    def get_order_params(self, spec_orders, facet_rad=0.0, gamma_rad=0.0,
            units='nm'):
        """
        Computes central (blaze) wavelength and angular extent of specified
        orders using chosen geometry/angles.

        Outputs:
        central_wlen -- corresponding central wavelengths in requested units
        order_FSR    -- order free spectral ranges (units of central_wlen)
        order_angwid -- order angular width (RADIANS)
        """
        use_spacing = self._groove_spacing[units]
        central_wlen = self._calc_blaze_wavelength(spec_orders, use_spacing,
                self._blaze_angle_rad, facet_rad, gamma_rad)
        order_FSR = self._calc_order_FSR(central_wlen, spec_orders)
        order_angwid = self._calc_order_width(central_wlen, use_spacing,
                self._blaze_angle_rad, facet_rad, gamma_rad)
        return (central_wlen, order_FSR, order_angwid)

##--------------------------------------------------------------------------##
##--------------------------------------------------------------------------##
## Grating class (to be overloaded for double pass:

##--------------------------------------------------------------------------##
##--------------------------------------------------------------------------##
## More sophisticated grating+prism class (closer to real design):
class DoublePassPrismGratingPrism(object):

    def __init__(self, config={}):
        # PARAMETERS LIST:
        # 0. coordinate system:
        #       * x-axis points "right" towards input fiber (grating at -X).
        #       * y-axis points "ahead" from camera to grating (optical axis)
        #       * z-axis points "up" towards sky from bench surface
        # 1. prism:
        #       * glass type
        #       * apex angle (degrees)
        #
        # 2. grating:
        #       * blaze ratio (i.e., R4)
        #       * groove spacing in lines/mm
        #
        # 3. layout:
        #       * angle of initial beam w.r.t. optical axis (assume 0?)
        #       * rotation of prism apex w.r.t. optical axis. Let 0 degrees
        #           represent apex perpendicular to optical axis.
        #       * direction to grating w.r.t. optical axis. This is really
        #           the prism deflection angle that "points" to the grating
        #       * alpha or sine(alpha), tilt of grating w.r.t. optical bench.
        #           --> could also be expressed as facet angle
        #       * rotation of grating w.r.t. optical bench normal. This is
        #           effectively the nominal gamma angle of the ray pointing
        #           directly at the grating. EQUIVALENTLY can use orientation
        #           of grating on optical bench (same units as prism base)

        # For prism apex 'perpendicular' to optical axis, incidence ~ apex/2
        self.prism_glass    = config.get('prism_glass', "PBM2")
        self.apex_angle_deg = config.get('apex_angle_deg', 55.0)

        # How far prism is "turned" w.r.t. main optical axis (def from CAD):
        self.prism_turn_deg = config.get('prism_turn_deg', 23.507)

        # Turn of incident beam (from fiber/collimator)
        # w.r.t. main optical axis (positive angles towards grating):
        # MAY VARY WITH WAVELENGTH???
        self.input_turn_deg = config.get('input_turn_deg', 2.0)

        # Allow non-unity index of refraction for air:
        self.air_index_refr = config.get('air_index_refr', 1.0)

        # Front and rear prism face orientations w.r.t. z-axis:
        self.prism_face1_deg = self.prism_turn_deg + 0.5 * self.apex_angle_deg
        self.prism_face2_deg = self.prism_turn_deg - 0.5 * self.apex_angle_deg

        # Angle of incidence of input beam w.r.t. prism face1 normal:
        self.prism_front_incid_deg = self.prism_face1_deg - self.input_turn_deg

        # Create prism object with given settings:
        self.pr_obj = Prism(self.prism_glass, self.apex_angle_deg)
        self.pr_obj.set_n_air(self.air_index_refr)

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

        # Grating ruling in lines / millimeter:
        self.grating_ruling_lmm = config.get('grating_ruling_lmm', 41.59)
        self.grating_spacing_um = 1e3 / self.grating_ruling_lmm

        # Grating face angle w.r.t. optical axis (turn about bench normal):
        self.grating_turn_deg = config.get('grating_turn_deg', 44.827)

        # Grating face tilt w.r.t. optical bench:
        self.grating_tilt_deg = config.get('grating_tilt_deg', 13.786)

        # Derived quantities:
        self.input_turn_rad   = np.radians(self.input_turn_deg)
        self.grating_turn_rad = np.radians(self.grating_turn_deg)
        self.grating_tilt_rad = np.radians(self.grating_tilt_deg)

        self.alpha_angle_rad  = np.radians(90.0 - self.grating_tilt_deg)
        self.blaze_angle_rad  = config.get('blaze_angle_rad', np.arctan(4.))
        self.facet_angle_rad  = self.alpha_angle_rad - self.blaze_angle_rad

        self.fixed_geometry   = 2.0 * self.grating_spacing_um \
                * np.cos(self.facet_angle_rad) * np.sin(self.blaze_angle_rad)

        ## deal with angle changes between prism/grating for second prism pass:
        #pass1_out_heading = self.input_turn_deg + np.degrees(deflect_r)
        #gamma_eff_grating = pass1_out_heading - self.grating_turn_rad
        #pass1_out_incid_d = pass1_out_heading - self.prism_face2_deg    # emergent
        #pass2_new_incid_d = pass1_out_incid_d - 2.0 * gamma_eff_grating

        # Lastly include some distances:
        self.coll_focallen_mm = config.get('coll_focallen_mm', 380.0)
        #self.coll_focallen_mm = 375.0
        #self.prism_grating_mm = 284.0       # approximately
        #self.prism_grating_mm = 275.0       # approximately
        #self.prism_grating_mm = 250.0       # approximately
        self.prism_grating_mm = config.get('prism_grating_mm', 100.0)
        self.lens_compression = config.get('lens_compression',   2.0)
        return

    # ---------------------------------------------------------
    # Joint prism+grating central wavelength and deflection:
    def _prism_pass1_deflection(self, wavelength_um):
        """
        Calculate the deflection angle and new direction of travel for the
        specified wavelength after first pass through prism (using configured
        prism and layout). Assumes all rays incident on prism are parallel.
        Output angles are in RADIANS.
        """
        incid_1_r = np.radians(self.prism_front_incid_deg)
        deflect_r = self.pr_obj.deflection_rad_wl(incid_1_r, wavelength_um)
        return deflect_r

    def _deflection_gamma(self, wavelength_um):
        """
        Calculate the effective gamma angle at the grating for the specified
        wavelength using configured prism and layout. Assumes all rays
        incident on prism are parallel. Output angle in RADIANS.
        """
        deflect_r = self._prism_pass1_deflection(wavelength_um)
        #incid_1_r = np.radians(self.prism_front_incid_deg)
        #deflect_r = self.pr_obj.deflection_rad_wl(incid_1_r, wavelength_um)
        #heading_r = self.input_turn_deg + deflect_r     # outbound direction
        return self.input_turn_rad + deflect_r - self.grating_turn_rad

    def _lamcen_residual(self, wavelength_um, order=0):
        ls = wavelength_um * order
        rs = self.fixed_geometry * np.cos(self._deflection_gamma(wavelength_um))
        return ls - rs

    def _iter_calc_lamcen(self, order):
        kw = {'order':order}
        runme = partial(self._lamcen_residual, **kw)
        return opti.bisect(runme, 0.0, 10.0)

    def calc_central_wlen_um(self, spec_order_list):
        return np.array([self._iter_calc_lamcen(x) for x in spec_order_list])

    def gamma_from_wlen_um(self, wlen_list_um):
        return np.array([self._deflection_gamma(x) for x in wlen_list_um])

    def fancy_deflections(self, spec_order_list):
        order_ctr_wlen_um = self.calc_central_wlen_um(spec_order_list)
        order_ctr_gamma_r = self.gamma_from_wlen_um(order_ctr_wlen_um)
        return order_ctr_wlen_um, order_ctr_gamma_r

    def two_pass_deflection(self, wavelength_um):
        """
        incid_1_r   --> angle of incidence of light onto prism (first pass)
        deflect_1_r --> deflection angle produced by first pass
        heading_1_r --> direction of deflected ray in benchtop units
        gamma_eff_r --> effective gamma angle at grating for this wavelength
        excid_1_r   --> angle of deflected ray to exit face normal

        incid_2_r   --> angle of incidence on prism face 2 (second pass)
        deflect_2_r --> deflection angle produced by second pass
        """
        incid_1_r   = np.radians(self.prism_front_incid_deg)
        deflect_1_r = self.pr_obj.deflection_rad_wl(incid_1_r, wavelength_um)
        heading_1_r = self.input_turn_rad + deflect_1_r
        gamma_eff_r = heading_1_r - self.grating_turn_rad
        excid_1_r   = heading_1_r - np.radians(self.prism_face2_deg)

        # Non-zero gamma causes change in incidence for second pass:
        incid_2_r = excid_1_r - 2.0 * gamma_eff_r
        refl_beam_r = heading_1_r + np.pi - 2 * gamma_eff_r
        deflect_2_r = self.pr_obj.deflection_rad_wl(incid_2_r, wavelength_um)
        heading_2_r = refl_beam_r - deflect_2_r

        # Total Y-shift produced by a combination of prism<-->grating offset
        # and deflection over the collimator focal length:
        pg_yshift =  2.0 * self.prism_grating_mm * np.tan(gamma_eff_r)
        pc_yshift = -1.0 * self.coll_focallen_mm * np.tan(heading_2_r - np.pi)
        wlen_nm = 1e3 * wavelength_um
        gamma_deg = np.degrees(gamma_eff_r)
        sys.stderr.write("λ=%6.1f nm, γ= %6.3f: PG,PC shifts: %6.2f, %6.2f | i2=%5.2f\n"
                % (wlen_nm, gamma_deg, pg_yshift, pc_yshift, np.degrees(incid_2_r)))
        return heading_2_r, pg_yshift, pc_yshift
        #return deflect_1_r, refl_beam_r, deflect_2_r, heading_2_r

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

######################################################################
# CHANGELOG (spectrograph_optics.py):
#---------------------------------------------------------------------
#
#  2018-12-29:
#     -- Increased __version__ to 0.1.0.
#     -- First created spectrograph_optics.py.
#


<filename>ex4/ex4.py<gh_stars>10-100
import scipy.io as sio
import numpy as np
from scipy.optimize import minimize

from common_functions import add_zero_feature, sigmoid, matrix_args_array_only

def sigmoid_gradient(z):
    return np.multiply(sigmoid(z), 1-sigmoid(z))

@matrix_args_array_only
def cf_nn(nn_params, input_layer_size, hidden_layer_size, num_labels, X, Y, lambda_coef):

    Theta1 = nn_params[0, :hidden_layer_size * (input_layer_size + 1)].reshape((hidden_layer_size, (input_layer_size + 1)))
    Theta2 = nn_params[0, hidden_layer_size * (input_layer_size + 1):].reshape((num_labels, (hidden_layer_size + 1)))

    m = Y.shape[1]
    Y = Y.A

    A_1 = X
    Z_2 = Theta1*A_1.T
    A_2 = sigmoid(Z_2)
    A_2 = add_zero_feature(A_2, axis=0)
    Z_3 = Theta2*A_2
    A_3 = sigmoid(Z_3)
    H = A_3.A

    J = np.sum(-Y*np.log(H) - (1-Y)*np.log(1-H))/m

    reg_J = 0.0
    reg_J += np.sum(np.power(Theta1, 2)[:, 1:])
    reg_J += np.sum(np.power(Theta2, 2)[:, 1:])

    J += reg_J*(float(lambda_coef)/(2*m))

    return J

@matrix_args_array_only
def gf_nn(nn_params, input_layer_size, hidden_layer_size, num_labels, X, Y, lambda_coef):

    Theta1 = nn_params[0, :hidden_layer_size * (input_layer_size + 1)].reshape((hidden_layer_size, (input_layer_size + 1)))
    Theta2 = nn_params[0, hidden_layer_size * (input_layer_size + 1):].reshape((num_labels, (hidden_layer_size + 1)))

    m = Y.shape[1]

    A_1 = X
    Z_2 = Theta1*A_1.T
    A_2 = sigmoid(Z_2)
    A_2 = add_zero_feature(A_2, axis=0)
    Z_3 = Theta2*A_2
    A_3 = sigmoid(Z_3)

    DELTA_3 = A_3 - Y
    DELTA_2 = np.multiply((Theta2.T*DELTA_3)[1:, :], sigmoid_gradient(Z_2))
    Theta1_grad = (DELTA_2 * A_1)/m
    Theta2_grad = (DELTA_3 * A_2.T)/m

    lambda_coef = float(lambda_coef)
    Theta1_grad[:, 1:] += (lambda_coef/m)*Theta1[:, 1:]
    Theta2_grad[:, 1:] += (lambda_coef/m)*Theta2[:, 1:]

    return np.concatenate((Theta1_grad.A1, Theta2_grad.A1))

def rand_initialize_weights(L_in, L_out):
    epsilon_init = 0.12
    return np.random.rand(L_out, 1 + L_in) * 2 * epsilon_init - epsilon_init

if __name__ == '__main__':
    data = sio.loadmat('ex4data1.mat')
    y = data['y']
    X = data['X']

    X = add_zero_feature(X)

    data = sio.loadmat('ex4weights.mat')
    Theta1 = data['Theta1']
    Theta2 = data['Theta2']

    nn_params = np.concatenate((Theta1.ravel(), Theta2.ravel()))

    input_layer_size  = 400
    hidden_layer_size = 25
    num_labels = 10
    m = len(y)

    Y = (np.arange(num_labels)[:, np.newaxis] == (y.T-1)).astype(float)

    for lambda_coef in (0, 1):
        print 'Cost function = {}, lambda = {}'.format(cf_nn(nn_params, input_layer_size, hidden_layer_size, num_labels, X, Y, lambda_coef), lambda_coef)

    initial_Theta1 = rand_initialize_weights(input_layer_size, hidden_layer_size)
    initial_Theta2 = rand_initialize_weights(hidden_layer_size, num_labels)

    initial_nn_params = np.concatenate((initial_Theta1.ravel(), initial_Theta2.ravel()))

    res = minimize(cf_nn, initial_nn_params, method='L-BFGS-B', jac=gf_nn, options={'disp': True, 'maxiter':100},
                        args=(input_layer_size, hidden_layer_size, num_labels, X, Y, lambda_coef)).x

    Theta1 = res[:hidden_layer_size * (input_layer_size + 1)].reshape((hidden_layer_size, (input_layer_size + 1)))
    Theta2 = res[hidden_layer_size * (input_layer_size + 1):].reshape((num_labels, (hidden_layer_size + 1)))

    h1 = sigmoid(np.dot(X, Theta1.T))
    h2 = sigmoid(np.dot(add_zero_feature(h1), Theta2.T))
    y_pred = np.argmax(h2, axis=1)+1

    print 'Training Set Accuracy: {}'.format(np.mean(y_pred == y.ravel(), ) * 100)


<filename>audio/audio_record.py
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu May 20 21:00:52 2021

@author: omar
"""

import sounddevice as sd
from scipy.io.wavfile import write
from audio.audio_converter import convert_audio
from audio.audio_recognition import speech_rec

def record(audio_needed = False):

    fs = 44100  # Sample rate
    seconds = 6  # Duration of recording

    myrecording = sd.rec(int(seconds * fs), samplerate=fs, channels=2)
    print("Recording...")
    sd.wait()  # Wait until recording is finished
    print("End of Record Session.")

    write("recorded.wav", fs, myrecording)  # Save as WAV file

    convert_audio()

    result = speech_rec(audio_needed)

    return result


"""
Code to do fractal geometry analysis:
    surface area of simple curves
    volume of simple curves
    find all line segment intersections
"""

import numpy as np
from scipy.spatial import ConvexHull
#from shapely.geometry import LineString
#from shapely.geometry.polygon import LinearRing

def poly_point_pairs(curve):
    pairs = [(x_n, x_np1) for x_n, x_np1 in zip(curve[:-1], curve[1:])]
    pairs.append((curve[-1], curve[0]))
    return pairs

def pppair_2_det(pppair):
    pt_1, pt_2 = pppair
    x1, y1 = pt_1
    x2, y2 = pt_2
    return x1*y2 - x2*y1

def pppair_2_len(pppair):
    pt_1, pt_2 = pppair
    return np.linalg.norm(pt_2 - pt_1)

def simple_surface_area(curve):
    pppairs = poly_point_pairs(curve)
    lens = [pppair_2_len(pppair) for pppair in pppairs]
    return np.sum(lens)

def simple_volume(curve):
    pppairs = poly_point_pairs(curve)
    dets = [pppair_2_det(pppair) for pppair in pppairs]
    return 0.5 * np.sum(dets)

def curve_2_tuples(curve):
    return [(x,y) for x,y in curve]

def is_simple(curve):
    is_simple = False
    tuples = curve_2_tuples(curve)
    lr = LinearRing(tuples)
    try:
        test = lr.area
        is_simple = True
    finally:
        return is_simple

def do_intersect_shp(p_n, p_np1, p_m, p_mp1):
    ln = LineString([tuple(p_n), tuple(p_np1)])
    lm = LineString([tuple(p_m), tuple(p_mp1)])
    return ln.intersects(lm)

def get_intersect(p_n, p_np1, p_m, p_mp1):
    xm1, ym1 = p_m
    xm2, ym2 = p_mp1
    xn1, yn1 = p_n
    xn2, yn2 = p_np1
    slope_n = (yn2 - yn1)/(xn2 - xn1)
    slope_m = (ym2 - ym1)/(xm2 - xm1)
    x_star = (ym1 - yn1 + slope_n * xn1 - slope_m * xm1)/(slope_n - slope_m)
    y_star = slope_n * (x_star - xn1) + yn1
    return [x_star, y_star]

def does_intersect(p_n, p_np1, p_m, p_mp1):
    xm1, ym1 = p_m
    xm2, ym2 = p_mp1
    xn1, yn1 = p_n
    xn2, yn2 = p_np1
    slope_n = (yn2 - yn1)/(xn2 - xn1)
    slope_m = (ym2 - ym1)/(xm2 - xm1)
    x_star = (ym1 - yn1 + slope_n * xn1 - slope_m * xm1)/(slope_n - slope_m)
    y_star = slope_n * (x_star - xn1) + yn1
    cond_x = min(xn1,xn2) <= x_star <= max(xn1,xn2)
    cond_y = min(yn1,yn2) <= y_star <= max(yn1,yn2)
    return cond_x and cond_y

def cclockwise_order(p1,p2,p3):
    slope_12 = (p3[1] - p1[1]) * (p2[0] - p1[0])
    slope_13 = (p2[1] - p1[1]) * (p3[0] - p1[0])
    # those are not errors, i think
    return slope_12 > slope_13

def do_intersect_ccw(p_n, p_np1, p_m, p_mp1):
    cond_n = cclockwise_order(p_n,p_m,p_mp1) != cclockwise_order(p_np1,p_m,p_mp1)
    cond_m = cclockwise_order(p_n,p_np1,p_m) != cclockwise_order(p_n,p_np1,p_mp1)
    return cond_n and cond_m

def intersects(curve):
    '''
    returns list of intersections, if at=None then there was a div by 0 issue
    '''
    pppairs = poly_point_pairs(curve)
    buffer = 2
    intersect_list = []
    for n, (p_n, p_np1) in enumerate(pppairs[:-buffer]):
        pppairs_ahead = pppairs[n + buffer:]
        for k, (p_m, p_mp1) in enumerate(pppairs_ahead):
            m = n + k + buffer
            if np.linalg.norm(p_n-p_m)<2.5 and k>5 and does_intersect(p_n, p_np1, p_m, p_mp1):
                try:
                    point = get_intersect(p_n, p_np1, p_m, p_mp1)
                except:
                    point = None
                desc = {'n':n, 'm':m, 'at':point}
                intersect_list.append(desc)
    return intersect_list

def complex_to_simples(curve, intersect_list):
    '''
    https://web.archive.org/web/20100805164131/http://www.cis.southalabama.edu/~hain/general/Theses/Subramaniam_thesis.pdf

    '''
    simples = []
    smallest = sorted(intersect_list, key=lambda d: d['m']-d['n'])
    base = curve
    for intersect in intersect_list:
        pass

def convex_hull(curve, len_line_seg=None):
    # need to configure to keep maximum hull line segment length small
    # enough to include non-overlapped points

    if len_line_seg is None:
        deltas = curve[:-1] - curve[1:]
        dists = np.linalg.norm(deltas, axis=1)
        dmin, dmax, dmean = np.amin(dists), np.amax(dists), np.mean(dists)

        double_deltas = curve[:-2] - curve[2:]
        double_dists = np.linalg.norm(double_deltas, axis=1)
        ddmin, ddmax, ddmean = np.amin(double_dists), np.amax(double_dists), np.mean(double_dists)

        #expectations:

        # dmean:
        # dmean = dmin = dmax; give or take some rounding error

        # ddmax:
        # ddmax = 2*dmean (2 line segments, no change in heading)

        # ddmin:
        # ddmin = 2 line segments, one bend at angle alpha
        # ddmin = len of 3rd leg of isoceles triangle where angle is alpha and leg len is dmean
        # ddmin = dmean * sin(alpha) / sin((pi - alpha)/2)
        # ddmin = dmean * 2 sin(alpha / 2)

        # ddmean:
        # dmean * 2sin(alpha/2) <= ddmean <= 2 dmean
        # let the ratio of 1's to 0's be r_c
        # ddmean = 2 dmean (r_c + (1-r_c) sin(alpha/2))

    hull = ConvexHull(curve)
    return hull


<reponame>metocean/tcrm
"""
:mod:`interp3d` -- interpolate to a set of points in 3-dimensional space
========================================================================

.. module:: interp3d
    :synopsis: Interpolate data to a set of points in 3-d space.

.. moduleauthor:: <NAME> <<EMAIL>>

Use :func:`scipy.ndimage.interpolation.map-coordinates` to interpolate
data in three dimensions.

"""

import numpy as np
from scipy.ndimage.interpolation import map_coordinates

def interp3d(input_array, coords,
             scale=[360., 180., 365.],
             offset=[0., -90., 0.],
             prefilter=True):
    """
    Wrapper to :func:`scipy.ndimage.interpolation.map_coordinates`, which
    converts coordinates of points to indices that correspond to the
    array.  We assume that one is working with lon, lat, day data
    (i.e. initially designed to work with daily long term mean sea
    level pressure)

    :param input_array: A 3-d array of data at regular intervals,
                        representing the data to be evaluated.
    :param coords: A 3-by-n array of coordinates at which the data
                   in ``input_array`` will be interpolated to.
    :param list scale: A (list of) scale factor(s) that reduces the
                       ``coords`` values to the range of indices in
                       ``input_array``.
    :param list offset: A (list of) offset factor(s) that is
                        subtracted from the ``coords`` values before
                        adjusting the ``scale`` (above).
    :param boolean prefilter: If ``True`` (default), then apply a spline
                              filter before interpolation (necessary for
                              spline interpolation of order > 1).  If
                              ``False``, it is assumed that the input is
                              already filtered. Default is ``True``.

    :returns: 1-d array of values corresponding to the interpolated
              values at the points given in ``coords``.

    Example::

      >>> vals = interp3d(data, coords, scale=[360., 180., 365.],
      ...                 offset=[0., -90., 0.])

    """

    if input_array.ndim != 3:
        raise ValueError('Input array has incorrect shape')
    if coords.shape[0] != 3:
        raise ValueError('Coordinates of points must be 3-d')

    dims = input_array.shape
    indices = [d*(c - o) / s for d, c, o, s in
               zip(dims, coords, offset, scale)]

    values = map_coordinates(input_array, indices, mode='wrap',
                             prefilter=prefilter)
    dtype = input_array.dtype
    return np.array(values, dtype)

def _interp(data, coords, scale=[360., 180.], offset=[0., -90.]):
    """
    Wrapper to scipy.ndimage.interpolation.map_coordinates, which converts
    coordinates of points to indices that correspond to the array.
    We assume that one is working with lon, latdata (i.e. initially
    designed to work with daily long term mean sea level pressure)

    Input:
    input_array - a 2-d array of data at regular intervals, representing the
                  data to be evaluated
    coords - a 2xn array of coordinates at which the data in input_array
             will be interpolated
    scale - a scale factor that reduces the coords values to the range of
            indices in input_array
    offset - an offset factor that is subtracted from the coords values
             before adjusting the scale (above)

    Output:
    1-d array of values corresponding to the interpolated values
    at the points given in 'coords'

    Example: vals = interp2d( data, coords, scale=[360., 180.],
                              offset=[0.,-90.] )
    """

    if data.ndim != np.asarray(coords).ndim:
        raise ValueError('Input array and coordinates do not have matching dimensions')

    dims = np.array(data.shape)

    indices = [d*(c - o) / s for d, c, o, s in
               zip(dims, coords, offset, scale)]

    values = map_coordinates(data, indices, mode='wrap')

    return values


<gh_stars>0
from scipy.stats import kurtosis, skew
from torch.utils.data import Dataset
from sklearn.datasets import make_spd_matrix
from sklearn.covariance import empirical_covariance
from sklearn.metrics import mean_squared_error
from torch.utils.data import DataLoader
import numpy as np
import itertools
import torch.nn.functional as F
#from .set_transformer.models import SmallSetTransformer, SmallDeepSamples
import torch.nn as nn
import torch
import copy
import matplotlib.pyplot as plt
import os
from scipy.special import logsumexp
import math
from sklearn.preprocessing import StandardScaler
from statsmodels.distributions.empirical_distribution import ECDF

from utils import str_to_bool_arg, QuantileScaler

from deep_samples.model import DeepSamples

from src.dataset import FullLargeDataset

class BasicDeepSet(nn.Module):
    def __init__(self, n_inputs=2, n_outputs=1, n_enc_layers=4, n_hidden_units=64, n_dec_layers=1,
                 multiplication=True,ln=False, bn=False, activation=nn.ReLU, instance_norm=False, sample_norm=False, n_samples=1000, **kwargs):
        """ Note: sample_norm = True first tranposes the data so that the sample_dim is last to reuse existing norm implementations """
        super().__init__()
        if sample_norm and any([bn, ln, instance_norm]):
            raise ValueError("Cannot have sample_norm and other norms")
        enc_layers = []
        for i in range(n_enc_layers):
            if i == 0:
                if sample_norm:
                    enc_layers.append(nn.ConvTranspose1d(n_inputs, n_hidden_units, 1))
                else:
                    enc_layers.append(nn.Linear(in_features=n_inputs, out_features=n_hidden_units))
            else:
                if sample_norm:
                    enc_layers.append(nn.ConvTranspose1d(n_hidden_units, n_hidden_units, 1))
                else:
                    enc_layers.append(nn.Linear(in_features=n_hidden_units, out_features=n_hidden_units))
            if ln:
                enc_layers.append(nn.LayerNorm(n_hidden_units))
            if bn:
                enc_layers.append(nn.BatchNorm1d(n_samples))
            if instance_norm:
                enc_layers.append(nn.InstanceNorm1d(n_samples))
            if sample_norm:
                if i == 0:
                    enc_layers.append(nn.InstanceNorm1d(n_hidden_units, affine=True))
            enc_layers.append(activation())
        # remove last relu
        enc_layers = enc_layers[:-1]
        self.enc = nn.Sequential(*enc_layers)
        dec_layers = []
        # for i in range(n_dec_layers - 1):
        #     dec_layers.append(nn.Linear(in_features=n_hidden_units, out_features=n_hidden_units))
        #     dec_layers.append(activation())
        # dec_layers.append(nn.Linear(in_features=n_hidden_units, out_features=n_outputs))
        for i in range(n_dec_layers):
            if i == n_dec_layers - 1:
                dec_layers.append(nn.Linear(in_features=n_hidden_units, out_features=n_outputs))
            else:
                dec_layers.append(nn.Linear(in_features=n_hidden_units, out_features=n_hidden_units))
                if ln:
                    dec_layers.append(nn.LayerNorm(n_hidden_units))
                dec_layers.append(activation())
        self.dec = nn.Sequential(*dec_layers)
        self.multiplication=multiplication
        self.sample_norm = sample_norm

    def forward(self, x):
        if len(x.shape) == 4 and x.shape[1] > 1:
            encoded = []
            for j in range(x.shape[1]):
                a = x[:, j, :, :].squeeze(1)
                if self.sample_norm:
                    encoded.append(torch.transpose(self.enc(torch.transpose(a, 1, 2)), 1, 2))
                else:
                    encoded.append(self.enc(a))
            x = torch.cat(encoded, 1)
        else:
            x = x.squeeze(1)
            if self.sample_norm:
                out = torch.transpose(self.enc(torch.transpose(x, 1, 2)), 1, 2)
            else:
                out = self.enc(x)
            #x = torch.mul(x, out)
        return out

class BasicDeepSetMean(BasicDeepSet):
    def forward(self, x, length=None):
        x = super().forward(x)
        if self.sample_norm:
            x = super().forward(torch.transpose(x, 1, 2))
        else:
            x = super().forward(x)
        if self.multiplication:
            x = torch.mul(x, x)
        if self.sample_norm:
            x = torch.transpose(x, 1, 2)
        x = x.mean(dim=-2)
        x = self.dec(x)
        return x

class BasicDeepSetMeanRC(BasicDeepSet):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        dec_layers = []
        for i in range(kwargs['n_dec_layers']):
            if i == kwargs['n_dec_layers'] - 1:
                dec_layers.append(nn.Linear(in_features=kwargs['n_hidden_units'] + kwargs['n_inputs'], out_features=kwargs['n_outputs']))
            else:
                dec_layers.append(nn.Linear(in_features=kwargs['n_hidden_units'] + kwargs['n_inputs'], out_features=kwargs['n_hidden_units'] + kwargs['n_inputs']))
                if kwargs['ln']:
                    dec_layers.append(nn.LayerNorm(kwargs['n_hidden_units'] + kwargs['n_inputs']))
                dec_layers.append(kwargs['activation']())
        self.dec = nn.Sequential(*dec_layers)

    def forward(self, x, length=None):
#         x = super().forward(x)
        means = torch.mean(x, axis=1)
        # print(means.shape)
        x -= means.unsqueeze(1)
        if self.sample_norm:
            x = self.enc(torch.transpose(x, 1, 2))
        else:
            x = self.enc(x)
        if self.multiplication:
            x = torch.mul(x, x)
        if self.sample_norm:
            x = torch.transpose(x, 1, 2)
        x = x.mean(dim=-2)
        x = torch.cat([x, means], axis=1)  # [b, hidden + features_per_sample]
        # print('x', x.shape)
        x = self.dec(x)
        return x

class BasicDeepSetSum(BasicDeepSet):
    def forward(self, x, length=None):
#         x = super().forward(x)
        if self.sample_norm:
            x = self.enc(torch.transpose(x, 1, 2))
        else:
            x = self.enc(x)
        if self.multiplication:
            x = torch.mul(x, x)
        if self.sample_norm:
            x = torch.transpose(x, 1, 2)
        x = x.sum(dim=-2)
        x = self.dec(x)
        return x

def train_nn(model, name, optimizer, scheduler, train_generator, test_generator, classification=False,
             n_epochs=10, outputs=[], use_wandb=False, plot_gradients=False, seed=0):

    np.random.seed(seed)
    torch.manual_seed(seed)
    if use_wandb:
        import wandb
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print(device)
    print(classification)
    model = model.to(device)

    if use_wandb and plot_gradients:
        wandb.watch(model, log='all')
    # by default, reduction = mean when multiple outputs
    #criterion = nn.MSELoss()
    if classification:
        criterion = nn.BCELoss()
    else:
        criterion = nn.MSELoss(reduction="none")
    step = 0
    best_loss_ts = None
    best_loss_tr = None
    losses_tr = []
    losses_ts = []
    dtype = torch.cuda.FloatTensor if torch.cuda.is_available() else torch.FloatTensor
    for epoch in range(n_epochs):
        print(epoch)
        train_aux = []
        for x, y, lengths in train_generator:
            # print(x.shape)
            x, y, lengths = x.type(dtype).to(device), y.type(dtype).to(device), lengths.to(device)
            preds = model(x, lengths)
            preds = preds.reshape(x.shape[0], len(outputs))
            assert preds.shape == y.shape, "{} {}".format(preds.shape, y.shape)
            loss_elements = criterion(preds, y)
            loss = loss_elements.mean()
            if np.isnan(loss.detach().cpu().numpy()):
                raise ValueError("Train loss is nan: ", loss)
            train_aux.append(loss.detach().cpu().numpy())
            # TODO: maybe we don't want to log at every step
            if use_wandb:
                wandb.log({f"{name} train loss per step": loss}, step=step)
            if len(outputs) > 1:
                outputs_loss = loss_elements.mean(dim=0)
                # print(outputs)
                # print(outputs_loss)
                assert len(outputs) == len(outputs_loss)
                per_output_loss = outputs_loss
                if use_wandb:
                    for i in range(len(outputs)):
                        wandb.log({outputs[i]: per_output_loss[i]}, step=step)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            preds = model(x, lengths)
            step += 1
            if step % 20 == 0:
                # losses_tr.append(per_output_loss.detach().cpu().numpy())

                aux = []
                accuracy = []
                for x, y, lengths in test_generator:
                    x, y, lengths = x.type(dtype).to(device), y.type(dtype).to(device), lengths.to(device)
                    loss_elements = criterion(model(x, lengths), y)
                    loss = loss_elements.mean()
                    if np.isnan(loss.detach().cpu().numpy()):
                        raise ValueError("Test loss is nan: ", loss)
                    if classification:
                        accuracy.append(accuracy_score(model(x, lengths).detach().cpu().numpy(),
                                                       y.detach().cpu().numpy().astype(np.int8)))
                    aux.append(loss.detach().cpu().numpy())
                test_loss = np.nanmean(aux)
                if use_wandb:
                    wandb.log({f"{name} test loss per step": test_loss}, step=step)
                if len(outputs) > 1:
                    outputs_loss = loss_elements.mean(dim=0)
                    assert len(outputs) == len(outputs_loss)
                    per_output_loss = outputs_loss
                    if use_wandb:
                        for i in range(len(outputs)):
                            wandb.log({outputs[i]: per_output_loss[i]}, step=step)
                train_loss = train_aux[-1]
#                 train_loss = np.nanmean(train_aux)
#                 print(train_aux)
                train_aux = []
                losses_tr.append(train_loss)
#                 print(train_loss)
                if not np.isnan(train_loss) and not best_loss_tr or (train_loss < best_loss_tr):
                    if use_wandb:
                        wandb.run.summary["best_tr_loss"] = train_loss
                    best_loss_tr = train_loss
                scheduler.step()
                if classification:
                    print('Train loss: '+str(train_loss)+", test loss: "+str(test_loss)
                        +'test accuracy: ' + np.nanmean(accuracy))
                else:
                    print('Train loss: '+str(train_loss)+", test loss: "+str(test_loss))
                # losses_ts.append(per_output_loss.detach().cpu().numpy())
                losses_ts.append(test_loss)
                if not np.isnan(train_loss) and not best_loss_ts or (test_loss < best_loss_ts):
                    if use_wandb:
                        wandb.run.summary["best_loss"] = test_loss
                    best_loss_ts = test_loss
            #print(list(model.parameters())[4])
    return model, best_loss_tr, best_loss_ts, losses_tr, losses_ts

if __name__ == "__main__":
    import argparse
    import wandb

    parser = argparse.ArgumentParser(description='Results summary')
    parser.add_argument('-lr', '--learning_rate', default=.01, type=float)
    parser.add_argument('-bs', '--batch_size', default=64, type=int)
    parser.add_argument('-e', '--enc_layers', default=2, type=int)
    parser.add_argument('-d', '--dec_layers', default=1, type=int)
    parser.add_argument('-ol', '--output_layers', default=1, type=int)
    parser.add_argument('-u', '--hidden_units', default=64, type=int)
    parser.add_argument('-s', '--step_size', default=60, type=int)
    parser.add_argument('-g', '--gamma', default=.9, type=float)
    parser.add_argument('-f', '--features', default=2, type=int)
    parser.add_argument('-n', '--sample_size', default=1000, type=int)
    parser.add_argument('-eo', '--enc_layers_outer', default=2, type=int)
    parser.add_argument('-do', '--dec_layers_outer', default=1, type=int)
    parser.add_argument('-ei', '--enc_layers_inner', default=2, type=int)
    parser.add_argument('-di', '--dec_layers_inner', default=1, type=int)
    parser.add_argument('-uo', '--hidden_units_outer', default=64, type=int)
    parser.add_argument('-ui', '--hidden_units_inner', default=64, type=int)
    parser.add_argument('--normalization', default="true", type=str)
    parser.add_argument('--connect_decoder', default="true", type=str)

    parser.add_argument('--layer_norm', default='false', type=str)
    parser.add_argument('--batch_norm', default='false', type=str)
    parser.add_argument('--instance_norm', default='false', type=str)
    parser.add_argument('--sample_norm', default='false', type=str)
    parser.add_argument('--mean_center', default='false', type=str)
    parser.add_argument('--quantile_scaling', default='false', type=str)
    parser.add_argument('--seed_weights', default=0, type=int)
    parser.add_argument('-a', '--activation', default='relu', help='relu|elu', type=str)
    parser.add_argument('-m', '--model', default='deepsets', type=str, help='deepsets|settransformer|deepsamples')
    parser.add_argument('--path', default='distribution_plots/', type=str)
    parser.add_argument('--wandb_test', action='store_true')
    parser.add_argument('--cpu', action='store_true')
    parser.add_argument('--epochs', default=100, type=int)
    parser.add_argument('--train_size', default=10000, type=int)
    args = parser.parse_args()

    layer_norm = str_to_bool_arg(args.layer_norm, 'layer_norm')
    batch_norm = str_to_bool_arg(args.batch_norm, 'batch_norm')
    instance_norm = str_to_bool_arg(args.instance_norm, 'instance_norm')
    sample_norm = str_to_bool_arg(args.sample_norm, 'sample_norm')
    mean_center = str_to_bool_arg(args.mean_center, 'mean_center')
    ensemble_network = str_to_bool_arg(args.ensemble_network, 'ensemble_network')
    quantile_scaling = str_to_bool_arg(args.quantile_scaling, 'quantile_scaling')
    args.output_name = args.output_name.split()
    normalization = str_to_bool_arg(args.normalization, 'normalization')
    connect_decoder = str_to_bool_arg(args.connect_decoder, 'connect_decoder')

    wandb.init(project='hematocrit')
    data_config = {
        'inputs': ["rbc", "retics", "plt", "basos", "perox"],
        'outputs': ["Hematocrit"],
        'id_file': "/misc/vlgscratch5/RanganathGroup/lily/blood_dist/data_large/id_files/Hematocrit_rbc,retics,plt,basos,perox_January-1-2021.txt",
        'num_samples': 1000,
        'num_subsamples': 1000,
        'permute_subsamples': False,
        'normalizer': "all",
        'imputation': "zero"
    }
    Dataset = FullLargeDataset
    train = Dataset(test=False, **data_config)
    test = Dataset(test=True, **data_config)
    num_workers = 32
    n_dists = 5
    n_final_outputs = 1
    output_names = ['hematocrit']

    train_generator = DataLoader(train,
                                    batch_size=args.batch_size,
                                    shuffle=True,
                                    num_workers=num_workers,
                                    pin_memory=False,
                                    drop_last=True)
    test_generator = DataLoader(test,
                            batch_size=args.batch_size,
                            shuffle=True,
                            num_workers=num_workers,
                            pin_memory=False,
                            drop_last=True)

    device = 'cpu' if args.cpu else torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    if args.activation == 'relu':
        activation = nn.ReLU
    elif args.activation == 'elu':
        activation = nn.ELU

    n_outputs = 1
    if args.model == 'deepsample':
        model = DeepSamples(n_inputs=n_inputs, n_outputs=n_final_outputs, n_enc_layers_outer=args.enc_layers_outer, n_hidden_units_outer=args.hidden_units_outer,
        n_enc_layers_inner=args.enc_layers_inner, n_hidden_units_inner=args.hidden_units_inner, n_dec_layers_inner=args.dec_layers_inner, n_dists=n_dists,
        activation=activation, normalization=normalization,
        connect_decoder=connect_decoder).to(device)
        n_inputs = args.features
    else:
        if args.model == 'settransformer':
            model_unit = SmallSetTransformer
            n_inputs = args.features
        elif args.model == 'deepsets-sum':
            model_unit = BasicDeepSetSum
            n_inputs = args.features
        elif args.model == 'deepsets-rc':
            model_unit = BasicDeepSetMeanRC
            n_inputs = args.features
        else:
            model_unit = BasicDeepSetMean
            n_inputs = args.features
        model = model_unit(n_inputs=n_inputs, n_outputs=n_outputs, n_enc_layers=args.enc_layers, n_hidden_units=args.hidden_units, n_dec_layers=args.dec_layers, ln=layer_norm, bn=batch_norm, activation=activation, instance_norm=instance_norm, n_samples=args.sample_size, sample_norm=sample_norm).to(device)

    print(model)
    optimizer = torch.optim.Adam(model.parameters(),lr=args.learning_rate)
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=args.step_size, gamma=args.gamma, last_epoch=-1)

    output_names = ['hematocrit']
    print(output_names)
    model, train_score, test_score, losses_tr, losses_ts = train_nn(model, 'tentative', optimizer, scheduler,
                                            train_generator, test_generator, n_epochs=args.epochs,
                                            outputs=output_names, use_wandb=True, plot_gradients=False, seed=args.seed_weights)


#!/usr/bin/env python3
from __future__ import print_function
from __future__ import division
import rospy
import rosbag
import numpy as np
import matplotlib.pyplot as plt
import matplotlib; matplotlib.use('TkAgg')
from waypoint_generation_library import WaypointGen
from scipy import linalg
from nav_msgs.msg import Odometry
from geometry_msgs.msg import Quaternion
from sensor_msgs.msg import Imu
from tf.transformations import euler_from_quaternion, quaternion_from_euler
from mpl_toolkits import mplot3d

WaypointGeneration = WaypointGen()
waypoints, desVel, desAcc, timeVec = WaypointGeneration.waypoint_calculation()
desiredPos = WaypointGeneration.desiredPos

# plot the waypoints
figPos = plt.figure()
axPos = plt.axes(projection = '3d')
axPos.plot3D(desiredPos[:,0], desiredPos[:,1], desiredPos[:,2], 'ro')

pnt3d = axPos.scatter(waypoints[:,0], waypoints[:,1], waypoints[:,2], c = timeVec)
cbar = plt.colorbar(pnt3d)
cbar.set_label("Time [sec]")
# label the axes and give title
axPos.set_xlabel('X-Axis [m]')
axPos.set_ylabel('Y-Axis [m]')
axPos.set_zlabel('Z-Axis [m]')
axPos.set_title('Minimum Jerk Position Waypoints')

# plot the desired kinematics
figOtherKinematics = plt.figure()
figOtherKinematics.suptitle('Desired Kinematics in Inertial Frame')
# desired position waypoints
axPos = plt.subplot(311)
axPos.plot(timeVec, waypoints[:,0], '-r', label = '$x_b$')
axPos.plot(timeVec, waypoints[:,1], '-k', label = '$y_b$')
axPos.plot(timeVec, waypoints[:,2], '-b', label = '$z_b$')
# add the yaw legend
axPos.plot(np.nan, '-g', label = 'yaw')
axPos.legend(loc = 0)
plt.grid()
plt.xlabel('Time [sec]')
plt.ylabel('Position [m]')
# plt.title('Desired Position in Inertial Frame')
# desired yaw
axYaw = axPos.twinx()
axYaw.plot(timeVec, waypoints[:,3], '-g')
axYaw.set_ylabel('Yaw [rad]')

# desired velocity waypoints
axVel = plt.subplot(312)
axVel.plot(timeVec, desVel[:,0], '-r', label = '$v_{x,b}$')
axVel.plot(timeVec, desVel[:,1], '-k', label = '$v_{y,b}$')
axVel.plot(timeVec, desVel[:,2], '-b', label = '$v_{z,b}$')
# add the yaw legend
axVel.plot(np.nan, '-g', label = '$yaw_{rate}$')
axVel.legend(loc = 0)
plt.grid()
plt.xlabel('Time [sec]')
plt.ylabel('Velocity [m/s]')
# plt.title('Desired Velocity in Inertial Frame')
# desired yaw
axYawRate = axVel.twinx()
axYawRate.plot(timeVec, desVel[:,3], '-g')
axYawRate.set_ylabel('Yaw [rad/s]')

# desired acceleration waypoints
axAcc = plt.subplot(313)
axAcc.plot(timeVec, desAcc[:,0], '-r', label = '$a_{x,b}$')
axAcc.plot(timeVec, desAcc[:,1], '-k', label = '$a_{y,b}$')
axAcc.plot(timeVec, desAcc[:,2], '-b', label = '$a_{z,b}$')
# add the yaw legend
axAcc.plot(np.nan, '-g', label = '$yaw_{acc}$')
axAcc.legend(loc = 0)
plt.grid()
plt.xlabel('Time [sec]')
plt.ylabel('Acceleration [$m/s^2$]')
# plt.title('Desired Acceleration in Inertial Frame')
# desired yaw
axYawRate = axAcc.twinx()
axYawRate.plot(timeVec, desAcc[:,3], '-g')
axYawRate.set_ylabel('Yaw [$rad/s^2$]')
plt.show()