text:
"""Contains functions for creating initial and equilibrium density
matrices and evolving them in time."""

from scipy import linalg, constants
import numpy as np
from qutip.nonmarkov.heom import HSolverDL
from qutip import Qobj

from quantum_heom import utilities as util

TEMP_INDEP_MODELS = ['local dephasing lindblad']
TEMP_DEP_MODELS = ['global thermalising lindblad',
                   'local thermalising lindblad',
                   'HEOM']
DYNAMICS_MODELS = TEMP_INDEP_MODELS + TEMP_DEP_MODELS

def initial_density_matrix(dims: int, init_site_pop: list) -> np.ndarray:

    """
    Returns an N x N 2D array corresponding to the density matrix
    of the system at time t=0, where N is the number of sites. Site
    populations are split equally between the sites specified in
    'init_site_pop'.

    Parameters
    ----------
    dims : int
        The dimension (i.e. the number of sites) of the quantum
        system.
    init_site_pop : list of int
        The sites in which to place initial population. For
        example, to place population in a superposition of sites 1
        and 6 (in a 7-site system), the user should pass [1, 6].
        To place twice as much initial population in 3 as in 4 pass
        [3, 3, 4]. Default value is [1], which populates only site 1.
        Example: user passes for a N=3 site system [1, 3].
        Superposition state given by (1 / sqrt(2)) (|1> + |3>).
        Density matrix given by outer product of this state:
        rho (0) = (1/2) (|1><1| + |1><6| + |6><1| + |6><6|)
        Note that there will be coherences present in the inital
        state.

    Returns
    -------
    np.ndarray
        N x N 2D array (where N is the number of sites)
        for the initial density matrix.
    """

    assert isinstance(dims, int), 'dims must be passed as an int'
    assert dims > 1, 'Must specify the dimensions as at least 2.'
    assert isinstance(init_site_pop, list), 'init_site_pop must be a list.'
    for site in init_site_pop:
        if site < 1 or site > dims:
            raise ValueError('Invalid site number.')

    num_states = len(init_site_pop)
    normalisation = 1. / np.sqrt(num_states)
    # Construct the base states
    state = np.zeros(dims)
    for site in init_site_pop:
        state[site - 1] += 1
    state *= normalisation

    return np.outer(state, state)

def equilibrium_state(dynamics_model: str, dims: int, hamiltonian: np.ndarray,
                      temperature: float) -> np.ndarray:

    """
    Returns the equilibirum density matrix of the system. For
    systems described by thermalising models, this is the
    themal equilibirum state, given by:

    .. math::
        \\rho^{(eq)}
            = \\frac{e^{- H / k_B T}}{tr(e^{- H / k_B T})}

    however non-thermalising models ('local dephasing lindblad'),
    this is the maximally mixed state corresponding to a diagonal
    matrix with elements equal to 1/N. This also corresponds to the
    thermal equilibrium state in the infinite temperature limit.

    Parameters
    ----------
    dynamics_model : str
        The model used to describe the system dynamics. Must be
        one of 'local dephasing lindblad', 'local thermalising
        lindblad', 'global thermalising lindblad', or 'HEOM'.
    dims : int
        The dimension (i.e. the number of sites) of the quantum
        system.
    hamiltonian : np.ndarray
        The system Hamiltonian for the open quantum system, with
        dimensions (dims x dims), in units of rad ps^-1. Only
        needs to be passed if dynamics_model is a thermalising
        model.
    temperature : float
        The temperature of the bath, in Kelvin. Need only be
        passed if dynamics_model is a thermalising model.

    Returns
    -------
    np.ndarray
        A 2D square density matrix for the system's equilibrium
        state.
    """

    assert dynamics_model in DYNAMICS_MODELS, (
        'Must choose a dynamics_model from ' + str(DYNAMICS_MODELS))
    if dynamics_model in TEMP_DEP_MODELS:
        assert isinstance(hamiltonian, np.ndarray) and (
            hamiltonian.shape[0] == hamiltonian.shape[1] == dims), (
                'Hamiltonian must be a square np.ndarray with same dimensions'
                ' as specified in "dims", units of rad ps^-1')
        assert isinstance(temperature, float) and temperature > 0., (
            'Must pass temperature as a positive float for thermalising models'
            ' in Kelvin.')

    if dynamics_model in TEMP_INDEP_MODELS:
        # Maximally-mixed state for dephasing model:
        return np.eye(dims, dtype=complex) * 1. / dims

    eigv, eigs = util.eigv(hamiltonian), util.eigs(hamiltonian)
    eq_values = np.zeros_like(eigv)
    for idx, _ in enumerate(eq_values):
        eq_values[idx] = np.exp(- eigv[idx] * 1e12 * constants.hbar
                                / (constants.k * temperature))
    # Divide each exponentiated eigenvalue by the total sum.
    # Forms a vector representation of the diagonalised equilibrium state,
    # in the eigenbasis.
    eq_values = eq_values / np.sum(eq_values)
    # Diagonalize into a sqaure matrix
    eq_values = np.diag(eq_values)
    # Convert from the eigenbasis to site basis.
    return util.basis_change(eq_values, eigs, False)

def evolve_matrix_one_step(dens_mat: np.ndarray, superop: np.ndarray,
                           time_interval: float) -> np.ndarray:

    """
    Evolves a density matrix at time t to time (t + dt), where dt
    is specified by 'time_interval', using an exponential
    propagator, given by:

        \\rho(t + dt) = exp(superop * dt) \\rho(t)

    Typically, 'superop' will be the sum of Hamiltonian and
    Lindbladian superoperators. Assumes quantities passed in
    correct and consistent units; i.e. that the superop is in
    angular frequency units of rad ps^-1, and the time_interval
    is in time units ps.

    Parameters
    ----------
    dens_mat : np.ndarray
        The density matrix to evolve forward in time.
    superop : np.ndarray
        The superoperator that governs the dynamics of the quantum
        system, typically formed from the sum of Hamiltonian and
        Lindbladian superoperators.
    time_interval : float
        The step forward in time to which the density matrix
        will be evolved.

    Returns
    -------
    evolved : np.ndarray
        The input density matrix evolved forward in time by
        time_interval.
    """

    assert isinstance(dens_mat, np.ndarray), 'Input matrix must be a np.ndarray'
    dims = dens_mat.shape[0]
    assert dims == dens_mat.shape[1], 'Input matrix must be square'
    assert isinstance(superop, np.ndarray), 'Superoperator must be a np.ndarray'

    dims = int(np.sqrt(superop.shape[0]))
    # Build the N^2 x N^2 propagator
    propa = linalg.expm(superop * time_interval)
    # Propagate vectorised density matrix
    evolved = np.matmul(propa, dens_mat.flatten('C'))
    # Reshape back to square and return
    return evolved.reshape((dims, dims), order='C')

def time_evo_lindblad(dens_mat: np.ndarray, superop: np.ndarray,
                      timesteps: int, time_interval: float,
                      dynamics_model: str, hamiltonian: np.ndarray,
                      temperature: float) -> np.ndarray:

    """
    Evaluates the time evolution of a starting density matrix over
    multiple time steps for the Lindblad models. Returns an array
    containing the times, density matrices at each timestep, and
    the trace of the density matrix squared and/or the trace
    distance at each step, if specified in 'trace_measure'. Assumes
    quantities passed in correct and consistent units; i.e. that
    the superop is in angular frequency units of rad ps^-1, and the
    time_interval is in time units ps.

    Parameters
    ----------
    dens_mat : np.ndarray
        The initial density matrix to evolve forward in time.
    superop : np.ndarray
        The superoperator that governs the dynamics of the quantum
        system, in units of rad ps^-1. Typically formed from the
        sum of Hamiltonian and Lindbladian superoperators.
    timesteps : int
        The number of timesteps over which to evaluate the density
        matrix.
    time_interval : float
        The step forward in time to which the density matrix
        will be evolved, in fs.
    dynamics_model : str
        The model used to describe the system dynamics. Must be one
        of 'local dephasing lindblad','local thermalising
        lindblad', 'global thermalising lindblad'.
    hamiltonian : np.ndarray
        The system Hamiltonian for the open quantum system, with
        dimensions (dims x dims), in rad ps^-1. Only needs to be
        passed if dynamics_model is a thermalising model.
    temperature : float
        The temperature of the bath, in K. Need only be passed if
        dynamics_model is a thermalising model.

    Returns
    -------
    np.array
        An array where each element corresponds to a timestep in the
        evolution of the density matrix, containing the following
        info, respectively; time, density matrix at time, trace
        squared, trace distance.
    """

    # Check inputs
    assert isinstance(dens_mat, np.ndarray), 'Input matrix must be a np.ndarray'
    dims = dens_mat.shape[0]
    assert dims == dens_mat.shape[1], 'Input matrix must be square'
    assert isinstance(superop, np.ndarray), 'Superoperator must be a np.ndarray'
    assert superop.shape[0] == superop.shape[1], 'Superoperator must be square'
    assert superop.shape[0] == dims**2, (
        'Superoperator dimensions must be the square of the density matrix'
        ' dims.')
    assert isinstance(timesteps, int), 'timesteps must be passed as an int.'
    assert isinstance(time_interval, float), 'time_interval must be a float.'
    assert isinstance(dynamics_model, str), (
        'Must provide the dynamics model to calculate the trace'
        ' distance.')
    if dynamics_model in TEMP_DEP_MODELS:
        assert isinstance(hamiltonian, np.ndarray), (
            'Must provide the system Hamiltonian as a numpy ndarray to'
            ' calculate the trace distance for thermalising models.')
        assert isinstance(temperature, float), (
            'Must provide the temperature as a positive float in order to'
            ' calculate the trace distance for thermalising models.')

    # Convert time fs --> ps to match superoperator units
    time_interval = time_interval * 1e-3
    # Produce time evolution data
    time, evolved = 0., dens_mat
    squared = util.trace_matrix_squared(evolved)
    eq_state = equilibrium_state(dynamics_model, dims, hamiltonian, temperature)
    distance = util.trace_distance(evolved, eq_state)
    evolution = np.empty(timesteps + 1, dtype=np.ndarray)
    evolution[0] = np.array([time, evolved, squared, distance])
    for step in range(1, timesteps + 1):
        time += time_interval
        evolved = evolve_matrix_one_step(evolved, superop, time_interval)
        evolved = util.renormalise_matrix(evolved)
        squared = util.trace_matrix_squared(evolved)
        eq_state = equilibrium_state(dynamics_model, dims,
                                     hamiltonian, temperature)
        distance = util.trace_distance(evolved, eq_state)
        # Add quantities in quantum_HEOM units; i.e. convert time back ps --> fs
        evolution[step] = np.array([time * 1e3, evolved, squared, distance])
    return evolution

def time_evo_heom(dens_mat: np.ndarray, timesteps: int, time_interval: float,
                  hamiltonian: np.ndarray, coupling_op: np.ndarray,
                  reorg_energy: float, temperature: float, bath_cutoff: int,
                  matsubara_terms: int, cutoff_freq: float,
                  matsubara_coeffs: np.ndarray, matsubara_freqs: np.ndarray
                  ) -> tuple:

    """
    Evaluates the time evolution of a starting density matrix over
    multiple time steps for the HEOM model, interfacing with
    QuTiP's in-built HEOMSolver. Returns a 3-ple (tuple) containing
    the time evolution data, the matsubara coefficients, and the
    matsubara frequencies (respectively) used in the calculation.

    Parameters
    ----------
    dens_mat : np.ndarray
        The initial density matrix to evolve forward in time.
    timesteps : int
        The number of timesteps over which to evaluate the density
        matrix.
    time_interval : float
        The step forward in time to which the density matrix
        will be evolved, in units of ps.
    hamiltonian : np.ndarray
        The system Hamiltonian for the open quantum system, with
        dimensions (dims x dims), in units of rad ps^-1.
    coupling_op : np.ndarray
        The coupling operator for the system-bath interaction.
    reorg_energy : float
        The factor by which the spectral density should be scaled
        by. Should be passed in units of rad ps^-1. Must be a
        non-negative float.
    temperature : float
        The temperature of the bath, in units of rad ps^-1.
    bath_cutoff : int
        The number of bath terms to include in the HEOM evaluation
        of the system dynamics.
    mastsubara_terms : int
        The number of matsubara terms to include in the HEOM
        evaluation of the system dynamics.
    cutoff_freq : float
        The cutoff frequency used in calculating the spectral
        density, in units of rad ps^-1.
    matsubara_coeffs : np.ndarray
        The matsubara coefficients c_k used in calculating the
        spectral density for the HEOM approach. Must be in
        order (largest -> smallest), where the nth coefficient
        corresponds to the nth matsubara term. If None QuTiP's
        HEOMSolver automatically generates them.
    matsubara_freqs: np.ndarray
        The matsubara frequencies v_k used in calculating the
        spectral density for the HEOM approach, in units of rad
        ps^-1.
        Must be in order (smallest -> largest), where the nth
        frequency corresponds to the nth matsubara term. If None;
        QuTiP's HEOMSolver automatically generates them.

    Returns
    -------
    np.array
        An array where each element corresponds to a timestep in the
        evolution of the density matrix, containing the following
        info, respectively; time, density matrix at time, trace
        squared, trace distance.
    """

    assert isinstance(dens_mat, np.ndarray), 'Input matrix must be a np.ndarray'
    dims = dens_mat.shape[0]
    assert dims == dens_mat.shape[1], 'Initial density matrix must be square.'
    assert isinstance(timesteps, int), 'timesteps must be passed as an int.'
    assert isinstance(time_interval, float), 'time_interval must be a float.'
    assert (isinstance(hamiltonian, np.ndarray)
            and hamiltonian.shape[0] == hamiltonian.shape[1]
            and hamiltonian.shape[0] == dims), (
                'Must provide the system Hamiltonian as a square np.ndarray'
                ' with dimensions matching the density matrix.')
    assert (isinstance(coupling_op, np.ndarray)
            and coupling_op.shape[0] == coupling_op.shape[1]
            and coupling_op.shape[0] == dims), (
                'Must provide the coupling operator as a square np.ndarray'
                ' with dimensions matching the density matrix.')
    assert (isinstance(reorg_energy, (int, float)) and reorg_energy >= 0.), (
        'Must provide the coupling strength of the system as a positive float.')
    assert (isinstance(temperature, (int, float)) and temperature > 0.), (
        'Must provide the temperature of the system as a positive float.')
    assert (isinstance(bath_cutoff, int) and bath_cutoff >= 0), (
        'Bath cutoff must be a positive int')
    assert (isinstance(matsubara_terms, int) and matsubara_terms >= 0), (
        'matsubara_terms must be a positive int')
    assert (isinstance(cutoff_freq, (int, float)) and cutoff_freq > 0.), (
        'Must provide the cutoff_freq as a positive float.')
    if matsubara_coeffs is not None:
        check = [isinstance(i, (float, complex)) for i in matsubara_coeffs]
        assert (isinstance(matsubara_coeffs, np.ndarray)
                and check.count(True) == len(check)), (
                    'matsubara_coeffs must be passed as a np.ndarray with all'
                    ' elements as positive floats.')
    if matsubara_freqs is not None:
        check = [isinstance(i, (float, complex)) for i in matsubara_freqs]
        assert (isinstance(matsubara_freqs, np.ndarray)
                and check.count(True) == len(check)), (
                    'matsubara_freqs must be passed as a np.ndarray with all'
                    ' elements as floats.')

    # Build HEOM Solver
    hsolver = HSolverDL(Qobj(hamiltonian),   # rad ps^-1
                        Qobj(coupling_op),
                        reorg_energy,  # rad ps^-1
                        temperature,   # rad ps^-1
                        bath_cutoff,
                        matsubara_terms,
                        cutoff_freq,   # ps^-1
                        planck=1.0,
                        boltzmann=1.0,
                        renorm=True,
                        stats=False)
    # Set the matsubara coeffs and freqs to those passed (if actually passed)
    if matsubara_coeffs is not None:
        hsolver.exp_coeff = matsubara_coeffs
    if matsubara_freqs is not None:
        hsolver.exp_freq = matsubara_freqs
    # Run the simulation over the time interval.
    times = np.array(range(timesteps + 1)) * time_interval  # ps
    result = hsolver.run(Qobj(dens_mat), times)
    # CONVERT BACK TO QUANTUM_HEOM UNITS
    conv_kelvin_to_rad_per_ps = constants.k / (constants.hbar * 1e12)
    temperature = temperature / conv_kelvin_to_rad_per_ps  # Kelvin
    hamiltonian = hamiltonian  # rad ps^-1
    # RETRIEVE THERMAL EQUILIBRIUM STATE
    # equilibrium_state() method requires Hamiltonian in rad ps^-1 and T in K
    eq_state = equilibrium_state('HEOM', dims, hamiltonian, temperature)
    # PROCESS TIME EVOLUTION DATA
    evolution = np.empty(len(result.states), dtype=np.ndarray)
    for i in range(0, len(result.states)):
        dens_matrix = np.array(result.states[i]).T
        dens_matrix = util.renormalise_matrix(dens_matrix)
        evolution[i] = np.array([float(result.times[i]) * 1e3,  # ps --> fs
                                 dens_matrix,
                                 util.trace_matrix_squared(dens_matrix),
                                 util.trace_distance(dens_matrix, eq_state)])
    return evolution, np.array(hsolver.exp_coeff), np.array(hsolver.exp_freq)

def process_evo_data(time_evolution: np.array, elements: [list, None],
                     trace_measure: list):

    """
    Processes a QuantumSystem's time evolution data as produced
    by its time_evolution() method. Returns the time, matrix,
    and trace measure (trace of the matrix squared, and trace
    distance) as separate numpy arrays, ready for use in plotting.

    Parameters
    ----------
    time_evolution : np.array
        As produced by the QuantumSystem's time_evolution() method,
        containing the time, density matrix, and trace measures
        at each timestep in the evolution.
    elements : list
        The elements of the density matrix to extract and return,
        in the format i.e. ['11', '21', ...]. Can also take the
        value None.
    trace_measure : list of str
        The trace measures to extract from the time evolution data.
        Must be a list containing either, both, or neither of
        'squared', 'distance'.

    Returns
    -------
    times : np.array of float
        All the times in the evolution of the system.
    matrix_data : dict of np.array
        Contains {str: np.array} pairs where the str corresponds
        to each of the elements of the density matrix specified
        in elements (i.e. '21'), while the np.array contains all
        the value of that matrix elements at each time step.
        If elements is passed as None, matrix_data is returned as
        None.
    squared : np.array of float
        The value of the trace of the density matrix squared at
        each timestep, if 'squared' was specified in trace_measure.
        If not specified, squared is returned as None.
    distance : np.array of float
        The value of the trace distance of the density matrix at
        each timestep, if 'distance' was specified in
        trace_measure. If not specified, distance is returned as
        None.
    """

    times = np.empty(len(time_evolution), dtype=float)
    matrix_data = ({element: np.empty(len(time_evolution), dtype=complex)
                    for element in elements} if elements else None)
    squared = (np.empty(len(time_evolution), dtype=float)
               if 'squared' in trace_measure else None)
    distance = (np.empty(len(time_evolution), dtype=float)
                if 'distance' in trace_measure else None)
    for idx, (time, rho_t, squ, dist) in enumerate(time_evolution, start=0):
        # Retrieve time
        times[idx] = time  # already in fs
        # Process density matrix data
        if matrix_data is not None:
            for element in elements:
                n, m = int(element[0]) - 1, int(element[1]) - 1
                matrix_data[element][idx] = rho_t[n][m]
        # Process trace measure data
        if squared is not None:
            squared[idx] = squ
        if distance is not None:
            distance[idx] = dist

    return times, matrix_data, squared, distance


text:
# -*- coding: utf-8 -*-

from __future__ import division, print_function

__all__ = []

import os
import pickle
import numpy as np
import matplotlib.pyplot as pl
from scipy.linalg import cho_factor, cho_solve
from itertools import product
import scipy.optimize as op
from multiprocessing import Pool

##### READ IN NORMALIZED APOGEE SPECTRA #####

f_all = pickle.load(open("data/norm_tr_fluxes.p", "r"))
l_all = pickle.load(open("data/tr_label_vals.p", "r"))
ivar_all = pickle.load(open("data/norm_tr_ivars.p", "r"))

# gap pixels are those with flux variance over objects == 0
keep = np.var(f_all, axis=0) > 0
f_all = f_all[:,keep]
ivar_all = ivar_all[:,keep]

nobj = f_all.shape[0]
npix = f_all.shape[1]

# subtract the weighted average (to ignore bad pixels) off flux
# construct a variance array
var_all = np.zeros(ivar_all.shape)
for pix in range(npix):
    f = f_all[:,pix]
    ivar = ivar_all[:,pix]
    var_all[:,pix] = 1./ivar
    f -= np.average(f, weights=ivar)
    f_all[:,pix] = f

def kernel(theta, l_i, l_j, white_noise=True):
    a, tau0, tau1, tau2, s = np.exp(theta)
    K = a**2 * np.exp(-0.5*(((l_i-l_j)[:,:,0])**2/tau0**2 +
                           ((l_i-l_j)[:,:,1])**2/tau1**2 +
                           ((l_i-l_j)[:,:,2])**2/tau2**2))
    if white_noise:
        K[np.diag_indices_from(K)] += s**2
    return K

# Calculate alpha
def compute_l_and_alpha(theta, f, var, l_all):
    # compute the kernel
    K = kernel(theta, l_all[:, None], l_all[None, :])
    # add in the uncertainties in the fluxes to the kernel
    K[np.diag_indices_from(K)] += var
    # factorize the kernel
    L, flag = cho_factor(K, overwrite_a=True)
    return (L, flag), cho_solve((L, flag), f)

# Likelihood function
def ln_likelihood(theta, f, var, l_all):
    (L, flag), alpha = compute_l_and_alpha(theta, f, var, l_all)
    # compute half of the log determinant
    lndet = np.sum(np.log(np.diag(L)))
    return -0.5*np.dot(f, alpha)-lndet

# Evaluate the predictive mean
def predictive_mean(theta, test_labels):
    K_star = kernel(theta, test_labels[:, None],
                    l_all[None, :], white_noise=False)
    alpha = compute_l_and_alpha(theta)[1]
    f_star_bar = np.dot(K_star, alpha)
    return f_star_bar

# Evaluate the likelihood of a data point given model
def test_ln_likelihood_single_pix(label, flux, flux_var, theta, (L,flag), alpha, l_all):
    # Calculate the mean prediction.
    # covariance between test point and n training points
    K_star = kernel(theta, label[None,None,:], l_all[None,:,:],
                    white_noise=False)
    f_star = np.dot(K_star, alpha)

    # Calculate the variance on the prediction.
    label_array = np.array(label[None,None,:])  # Hack so kernel works w scalar.
    pred_var = flux_var + float(kernel(theta, label_array, label_array))
    pred_var -= np.dot(K_star, cho_solve((L,flag), K_star[0]))

    return -0.5 * ((flux - f_star) ** 2 / pred_var + np.log(pred_var))

def test_ln_likelihood(label, kernels, f, var, theta_all, f_all, var_all, l_all):
    #print(label)
    npix = len(kernels)
    ln_likelihoods = np.zeros(npix)
    for pix in range(0,npix):
        print("pix %s" %pix)
        # NO: you should only do this once for each pixel, not every time
        # this function is called!!!
        (L,flag), alpha = kernels[pix]
        ln_likelihoods[pix] = test_ln_likelihood_single_pix(
                label, f[pix], var[pix], theta_all[pix], (L,flag), alpha, l_all)
    print(sum(ln_likelihoods))
    return sum(ln_likelihoods)

def train_single_pix(param):
    print("training pix")
    f, var = param
    # l_all, f, var = param
    #print("training pix %s" %pix)
    nll = lambda theta: -ln_likelihood(theta, f, var, l_all)
    bounds = np.log([(0.001, 0.1), (100, 10000), (0.01, 10), (0.01, 10),
                     (0.0001, 0.1)])
    output = op.minimize(nll, np.log([0.01, 100, 0.1, 0.1, 0.01]),
                         method="L-BFGS-B", bounds=bounds)
    best_hyperparams = output.x
    print("best hyperparams: ")
    print(np.exp(best_hyperparams))
    return best_hyperparams

def infer_labels_single_obj(obj, kernels, hyperparams_all, f_all, var_all, l_all):
    print("test step for object %s" %obj)
    f = f_all[obj,:]
    var = var_all[obj,:]
    # l = l_all[obj,:]
    tll = lambda labels: -test_ln_likelihood(
          labels, kernels, f, var, hyperparams_all, f_all, var_all, l_all)
    p0 = np.array([np.mean(l_all[:,i]) for i in range(l_all.shape[1])])
    #p0 = l_all[obj,:]
    bounds = np.array([(3500,5500), (0,5), (-2.5,0.5)]) # training set dist

    print("fitting for labels")
    output = op.minimize(tll, p0, method="L-BFGS-B", bounds=bounds)
    return output

##### TRAINING STEP #####

hyperfn = "best_hyperparams.p"
if not os.path.exists(hyperfn):
    print("fitting for hyperparams")
    # Optimize for each pixel independently
    # start 4 workers
    pool = Pool(processes=4)

    # Dan's fix:
    # t_s_p = lambda args: train_single_pix(l_all, *args)
    best_hyperparams_all = pool.map(train_single_pix, zip(f_all.T, var_all.T))
    best_hyperparams_all = np.array(best_hyperparams_all)

    # Brani's fix:
    # pack parameters
    # params = [(f, var, l_all) for f, var in zip(f_all.T, var_all.T)]
    # best_hyperparams_all = pool.imap(train_single_pix, params, chunksize=10)

    # kill workers
    # pool.terminate()

    # t_s_p = lambda f, var: train_single_pix(f, var, l_all)
    # best_hyperparams_all = p.map(t_s_p, f_all.T, var_all.T)
    pickle.dump(best_hyperparams_all, open(hyperfn, "wb"), -1)

else:
    print("loading")
    best_hyperparams_all = pickle.load(open(hyperfn, "r"))

##### TEST STEP #####

# Now optimize over the test label space

# kernel calculation for each pixel
npix = best_hyperparams_all.shape[0]
kernels = []

# maybe just one pixel?
pix = 1383
(L,flag), alpha = compute_l_and_alpha(best_hyperparams_all[pix,:], f_all[:,pix],
        var_all[:,pix], l_all)
kernels = [(L,flag), alpha]
obj = 10
labels = infer_labels_single_obj(obj, kernels, best_hyperparams_all,
        f_all, ivar_all, l_all)

Ls = np.zeros((7212,544,544))
flags = np.zeros(7212)
alphas = np.zeros((7212,544))

for pix in range(0,10):
    print("pix %s" %pix)
    (L,flag), alpha = compute_l_and_alpha(
            best_hyperparams_all[pix,:], f_all[:,pix], var_all[:,pix], l_all)
    Ls[pix,:,:] = L
    flags[pix] = flag
    alphas[pix,:] = alpha

# obj = 0
# infer_labels_single_obj(obj, kernels, best_hyperparams_all, f_all, ivar_all, l_all)

# axis = 2
# test_labels[:,axis] = np.linspace(np.min(l_all[:,axis]),
#                                   np.max(l_all[:,axis]),
#                                   len(test_labels[:,axis]))
# mu = predictive_mean(best_hyperparams, test_labels)
# pl.scatter(l_all[:,axis], f, c="k")
# pl.plot(test_labels[:,axis], mu, "g", lw=2)
# pl.show()

# (L,flag), alpha = compute_l_and_alpha(best_hyperparams)
# pl.figure()
# for obj in range(5,9):
#     lls = [test_ln_likelihood(test_labels[i,:], f[obj], var[obj],
#            best_hyperparams, (L,flag), alpha) for i in range(0,len(test_labels))]
#     pl.plot(test_labels[:,axis], lls, label="%s"%np.round(l_all[obj,axis],2))
#     pl.legend()
#     pl.gca().axvline(l_all[obj,axis])
# pl.show()

ext: py

max_stars_repo_name: annayqho/gp-cannon

max_issues_repo_name: annayqho/gp-cannon

max_forks_repo_name: annayqho/gp-cannon

reason: import numpy,import scipy,from scipy


text:
"""
This is the script for constructing pressure-potential phase diagrams (pressure-dependent Pourbaix diagrams
at given pH) of binary metal-hydrogen systems. The input files are pressure-depedent Gibbs energy of the most
stable structure for each considered composition, named such as 'hmin_Pd','hmin_PdH','hmin_Pd3H4', etc.. The
common file for each system is hmin_H for pure H. The format of each input file is like:
#P (eV/A^3)   Hmin (eV/atom)
P1            G1
P2            G2
...
Pn            Gn
where the first column is the pressure, and the second is the correponding Gibbs energy. In cases where the
finite-temperature effects are ignored, you can use enthalpy to approximate Gibbs energy. But for pure H, the
finite-temperature effects usually cannot be ignored.
The code will create folder "pb" storing results, containg a "pbmap_pH" file for the Pourbaix diagram at given
pH. The code will also plot and prompt the diagram.
"""

from phasediagram import Pourbaix_H
import numpy as np
from scipy import interpolate
import os

# Define conditions.
kB=1.38065E-23;T=300;P0=1e5
plim=[1e-4,200] #GPa
#plim=[1e-4,200]
nnp=400
#HER=-0.09 #Pd
#HER=-0.611323393 #Y
#HER=-0.419984919 #La
HER=-0.1257775 #LaH3
#HER=-0.217096776 #Mg
#HER=-0.397163924 #Ca
#HER=0
pH=0 # set a pH value
nu=400;nph=300;ulim=[-1,0.2];phlim=[-2,14]
scale='log' # log, linear
if scale=='linear':
	dp=(plim[1]-plim[0])/nnp
	plim=np.array(plim)*1e9*1e-30/(1.6*1e-19)
	dp=dp*1e9*1e-30/(1.6*1e-19)
	p=[plim[0]+i*dp for i in range(nnp+1)]

if scale=='log':
	plim=np.array(plim)*1e9*1e-30/(1.6*1e-19)
	p=[plim[0]*(plim[1]/plim[0])**(i/nnp) for i in range(nnp+1)]

def mkdir(path):
    folder = os.path.exists(path)
    if not folder:
        os.makedirs(path)
    else:
        print('---  There is this folder!  ---')

"""
Define studied compositions.
comp: compositions to be considered except the pure H, consistent with the input file names, e.g.,
comp=['Ca','CaH2','CaH6'] corresponds to input files 'hmin_Ca','hmin_CaH2','hmin_CaH6' in addition to
'hmin_H'.
x: hydrogen content of the compositions defined in comp.
nat: number of atoms of the compositions defined in comp.
"""

"""Pd
comp=['Pd','PdH','Pd3H4','PdH6','PdH7','PdH8','PdH10','PdH12']
x=[0,0.5,4/7,6/7,7/8,8/9,10/11,12/13]
nat=[1,2,7,7,8,9,11,13]
"""
"""Y
comp=['Y','YH2','YH3','YH4','YH6','YH7','YH9','YH10']
x=[0,2/3,0.75,4/5,6/7,7/8,9/10,10/11]
nat=[1,3,4,5,7,8,10,11]
"""
"""Mg
comp=['Mg','MgH2','MgH4','MgH6','MgH12','MgH16']
x=[0,2/3,4/5,6/7,12/13,16/17]
nat=[1,3,5,7,13,17]
"""
"""La
comp=['La','LaH2','LaH3','LaH4','LaH5','LaH6','LaH8','LaH10','LaH11','LaH16']
x=[0,2/3,0.75,4/5,5/6,6/7,8/9,10/11,11/12,16/17]
nat=[1,3,4,5,6,7,9,11,12,17]
"""
#"""La-remove LaH11
comp=['La','LaH2','LaH3','LaH4','LaH5','LaH6','LaH8','LaH10','LaH16']
x=[0,2/3,0.75,4/5,5/6,6/7,8/9,10/11,16/17]
nat=[1,3,4,5,6,7,9,11,17]
#"""
"""Ca
comp=['Ca','CaH2','CaH6']
x=[0,2/3,6/7]
nat=[1,3,7]
"""

# read pressure-Gibbs energy data
p_each=[[] for i in comp]
h=[[] for i in comp]
h_ip=['' for i in comp]

for i in range(len(comp)):
	data=np.loadtxt('hmin_'+comp[i],usecols=(0,1))
	p_each[i]=data[:,0]
	h[i]=data[:,1]
	h_ip[i]=interpolate.interp1d(p_each[i], h[i],kind='cubic')

hset=[[] for i in p]
xset=[[] for i in p]

for i in range(len(p)):
	for j in range(len(comp)):
		try:
			hset[i].append(float(h_ip[j](p[i])))
			xset[i].append(x[j])
		except:
			print(comp[j])
			pass

data_H2=np.loadtxt('hmin_H',usecols=(0,1))
p_H2=data_H2[:,0]
G_H2=data_H2[:,1]
G300_H2=interpolate.interp1d(p_H2, G_H2)
h_H=[0 for i in p]
for j in range(len(p)):
	h_H[j]=float(G300_H2(p[j]))

#calculate Pourbaix diagrams
mkdir('pb')
d=['' for j in range(len(p))]
names=['' for j in range(len(p))]
text=['' for j in range(len(p))]
# generate a Pourbaix diagram for each pressure
for j in range(len(p)):
    try:
        print(p[j])
        y=hset[j]
        xx=xset[j]
        h_Pd=hset[j][0]
        y=[(y[k]-(1-xx[k])*h_Pd-xx[k]*h_H[j])*nat[k] for k in range(len(y))]
        d[j], names[j]=Pourbaix_H(comp,y,nu=nu,nph=nph,ulim=ulim,phlim=phlim)
        np.savetxt('pb/pbmap_'+str(p[j]/(1e-30)*(1.6*1e-19)/1e9)+'GPa',d[j],fmt='%5s',header=' P='+str(p[j])+'eV/A^3 '+str(names[j]))
    except:
        print(j)

# Merge the calculated phase equilibria at different pressures.
import glob

c=glob.glob('pb/pbmap*GPa')
s=[]
for i in range(len(c)):
	with open(c[i],'r') as f:
		cc=f.readline().split('eV/A^3')
		phases=cc[-1].strip('\n').replace('[','').replace(']','').replace('\'','').replace(' ','').split(',')
		p=float(cc[0].strip('#').replace('P=','').replace(' ',''))
	a0=np.loadtxt(c[i])[:,int(nph*(pH-phlim[0])/(phlim[1]-phlim[0]))] # Slice the Pourbaix diagram at given pH
	a=[phases[int(j)] for j in a0]
	s.append((a,p)) # Pourbaix diagrams at given pH paired with pressures

s=sorted(s, key=lambda x: x[1])
map=[i[0] for i in s]
phases=[]
for i in range(len(map)):
	phases=phases+map[i] # All the phases at given pH at all the pressures

phases_reduced = []
null=[phases_reduced.append(i) for i in phases if not i in phases_reduced]
from pymatgen.core.composition import Composition
from pymatgen.core.periodic_table import get_el_sp, Element, Specie, DummySpecie
comps_reduced = [Composition(i) for i in phases_reduced]
comps_H=[i.get_atomic_fraction(Element("H")) for i in comps_reduced]
comps_phases_reduced=list(zip(comps_H,phases_reduced))
comps_phases_reduced.sort()
phases_reduced=[i[1] for i in comps_phases_reduced] # Unduplicated phases ordered by their hydrogen content

# Map the Pourbaix diagram at given pH to numbers corresponding to the order of phases by their hydrogen content.
map_number=[[] for i in map]
for i in range(len(map_number)):
	map_number[i]=[phases_reduced.index(j) for j in map[i]]

map_number=np.array(map_number)
map_number=map_number.transpose()
# save the Pourbaix diagram at given pH mapped to numbers corresponding to the order of phases listed in the header.
np.savetxt('pb/pbmap_pH'+str(pH),map_number,fmt='%5s',header=' pH='+str(pH)+' '+str(phases_reduced))

# plot the diagram.
import matplotlib
import matplotlib.pyplot as plt
from matplotlib import ticker,cm
from matplotlib.colors import ListedColormap, LinearSegmentedColormap

plt.rcParams["font.weight"] = "bold"
plt.rcParams["axes.labelweight"] = "bold"

map_number=np.loadtxt('pb/pbmap_pH'+str(pH))
cmp=cm.get_cmap('rainbow')
# Linear or log scale in pressure:
if scale=='linear':
	plt.imshow(map_number,extent=[plim[0]/(1e-30)*(1.6*1e-19)/1e9,plim[1]/(1e-30)*(1.6*1e-19)/1e9,ulim[0],ulim[1]],origin='lower',aspect='auto',cmap=cmp)
	plt.hlines(HER, plim[0]/(1e-30)*(1.6*1e-19)/1e9,plim[1]/(1e-30)*(1.6*1e-19)/1e9, colors = "r", linestyles = "dashed")
	plt.xlabel("P (GPa)",fontsize=15)

if scale=='log':
	plt.imshow(map_number,extent=[np.log10(plim[0]/(1e-30)*(1.6*1e-19)/1e9),np.log10(plim[1]/(1e-30)*(1.6*1e-19)/1e9),ulim[0],ulim[1]],origin='lower',aspect='auto',cmap=cmp)
	plt.hlines(HER, np.log10(plim[0]/(1e-30)*(1.6*1e-19)/1e9),np.log10(plim[1]/(1e-30)*(1.6*1e-19)/1e9),
		colors = "k", linestyles = "dashed")
	plt.xlabel("logP (GPa)",fontsize=15)

plt.ylabel('U_RHE (V)',fontsize=15)
plt.tick_params(axis='both',which='major',labelsize=13,width=2)
plt.rcParams["axes.labelweight"] = "bold"
plt.xticks(color='w')
plt.yticks(color='w')
plt.show()

reason: import numpy,from scipy


text:
from statsmodels.tools.data import _is_using_pandas

def _maybe_get_pandas_wrapper(X, trim=None):
    """
    If using pandas returns a function to wrap the results, e.g., wrapper(X)
    trim is an integer for the symmetric truncation of the series in some
    filters.
    otherwise returns None
    """
    if _is_using_pandas(X, None):
        index = X.index
        if trim is not None:
            index = X.index[trim:-trim]
        if hasattr(X, "columns"):
            return lambda x : X.__class__(x, index=index, columns=X.columns)
        else:
            return lambda x : X.__class__(x, index=index, name=X.name)
    else:
        return

max_issues_repo_name: AnaMP/statsmodels

max_forks_repo_name: AnaMP/statsmodels

reason: from statsmodels


text:
"""Code for evaluating a dynamics model"""

import numpy as np
import tensorflow as tf

import env_info
from envs import NumpyReward
import log
import reporter
from utils import rate_limit

class DynamicsMetrics:
    """
    The dynamics metrics evaluate how accurate a planning controller was
    in its predictions.

    An evaluation of the predictive power of the dynamics tests
    how close the dynamics are to what the open-loop planner
    thought it'd be doing. At each step i, the open-loop planner that the
    controller uses had an H-step plan of actions a_i^0, a_i^1, ..., a_i^H.
    It gets to execute a_i = a_i^0, but not the rest with MPC.
    However, the open-loop planner took a_i^0 because of how it perceived
    the future would look like based on the dynamics model after taking
    actions a_i^j.

    For i in [T], using the true environment M, we can evaluate the states
    that *would have occurred* had the planner executed its actions. Thus,
    starting at t_i^0 = s_i, we get the "true planned states"
    t_i^{j+1} = M(t_i^j, a_i^j). Then compute the corresponding h-step
    prediction MSE between t_i^j and s_i^j.

    Finally, we can also measure, to some extent, the nature of the error
    between t_i^j and s_i^j by evaluating the H-step reward according
    to the planner (i.e., from trajectory s_i^0 to s_i^H) versus
    what would have really happened (from t_i^0 to t_i^H). This
    gives a measure of the planner over-optimism or pessimism.

    The metrics object accepts as arguments the maximum horizon H,
    number environments for ground-truth evaluation, and the discount factor.
    """

    def __init__(self, planning_horizon, nenvs, discount):
        self._venv = env_info.make_venv(nenvs)
        self._venv.reset()
        self._env = env_info.make_env()
        self._np_reward = NumpyReward(self._env)
        self._num_envs = nenvs
        self._horizon = planning_horizon
        self._discount = discount

    @staticmethod
    def _mse_h(true_obs_nhs, pred_obs_nhs):
        # n = batch size, h = horizon, s = state dim
        return np.square(true_obs_nhs - pred_obs_nhs).mean(axis=2).mean(axis=0)

    def evaluate(self, obs, planned_acs, planned_obs, prefix=''):
        """
        Report H-step standardized dynamics accuracy.
        """
        if planned_acs.size == 0:
            return

        acs_nha = planned_acs
        assert acs_nha.shape[1] == self._horizon, (
            acs_nha.shape, self._horizon)
        obs_ns = obs
        obs_nhs, mask = self._eval_open_loop(obs_ns, acs_nha)
        if np.sum(mask) == 0:
            log.debug('all open loop evaluations terminated early -- '
                      'skipping open-loop dynamics eval')
            return
        acs_nha = acs_nha[mask]
        obs_nhs = obs_nhs[mask]
        obs_ns = obs_ns[mask]
        planned_obs_nhs = planned_obs[mask]

        mse_h = self._mse_h(obs_nhs, planned_obs_nhs)
        prefix = prefix + ('/' if prefix else '') + 'dynamics/'

        self._print_mse_prefix(prefix, mse_h)

        obs_n1s = obs_ns[:, np.newaxis, :]
        prev_obs_nhs = np.concatenate([obs_n1s, obs_nhs[:, :-1]], axis=1)
        prev_planned_obs_nhs = np.concatenate(
            [obs_n1s, planned_obs_nhs[:, :-1]], axis=1)
        rew_n = self._rewards_from_transitions(prev_obs_nhs, acs_nha, obs_nhs)
        planned_rew_n = self._rewards_from_transitions(
            prev_planned_obs_nhs, acs_nha, planned_obs_nhs)
        rew_bias = (rew_n - planned_rew_n).mean()
        rew_mse = np.square(rew_n - planned_rew_n).mean()
        reporter.add_summary('reward bias', rew_bias)
        reporter.add_summary('reward mse', rew_mse)

    def _rewards_from_transitions(self, prev_states_nhs, acs_nha, states_nhs):
        # define N = n * h
        prev_states_Ns = self._merge_axes(prev_states_nhs)
        acs_Ns = self._merge_axes(acs_nha)
        states_Ns = self._merge_axes(states_nhs)
        rew_N = self._np_reward.np_reward(prev_states_Ns, acs_Ns, states_Ns)
        rew_nh = self._split_axes(rew_N)
        rew_n = np.zeros(rew_nh.shape[0])
        for i in reversed(range(rew_nh.shape[1])):
            rew_n = rew_nh[:, i] + rew_n * self._discount
        return rew_n

    def _merge_axes(self, arr):
        assert arr.ndim == 3, arr.ndim
        a, b, c = arr.shape
        assert b == self._horizon, (b, self._horizon)
        return arr.reshape(a * b, c)

    def _split_axes(self, arr):
        assert arr.ndim == 1, arr.ndim
        a = arr.shape[0]
        assert a % self._horizon == 0, (a, self._horizon)
        b, c = a // self._horizon, self._horizon
        return arr.reshape(b, c)

    def _print_mse_prefix(self, prefix, mse_h):
        fmt = len(str(self._horizon))
        fmt = '{:' + str(fmt) + 'd}'
        fmt_str = prefix + fmt + '-step mse'
        assert len(mse_h) == self._horizon, (len(mse_h), self._horizon)
        for h, mse in enumerate(mse_h, 1):
            print_str = fmt_str.format(h)
            hide = h not in [1, self._horizon]
            reporter.add_summary(print_str, mse, hide)

    def _eval_open_loop(self, states_ns, acs_nha):
        states_nhs, mask_n = rate_limit(
            self._num_envs, self._eval_open_loop_limited,
            states_ns, acs_nha)
        ndone = np.sum(~mask_n)
        if ndone > 0:
            log.debug('WARNING: {} early terminations during open-loop'
                      ' eval', ndone)
        return states_nhs, mask_n

    def _eval_open_loop_limited(self, states_ns, acs_nha):
        self._venv.set_state_from_ob(states_ns)
        acs_hna = np.swapaxes(acs_nha, 0, 1)
        states_hns, _, done_hn = self._venv.multi_step(acs_hna)
        done_n = done_hn.sum(axis=0)
        states_nhs = np.swapaxes(states_hns, 0, 1)
        return states_nhs, ~done_n

    def close(self):
        """Close the internal simulation environment"""
        self._venv.close()
        self._env.close()

class TFStateUnroller:
    """
    This class constructs a graph for unrolling a horizon-length
    rollout according the given policy function and dynamics model.
    """

    def __init__(self, horizon, tf_action, tf_predict_dynamics,
                 scope='unroll_dynamics'):
        self._unroll_states_ph_ns = tf.placeholder(
            tf.float32, shape=[None, env_info.ob_dim()])
        n = tf.shape(self._unroll_states_ph_ns)[0]
        with tf.variable_scope(scope):
            actions_hna = tf.get_variable(
                'dynamics_actions', initializer=tf.zeros(
                    [horizon, n, env_info.ac_dim()]),
                dtype=tf.float32, trainable=False, collections=[],
                validate_shape=False)
            states_hns = tf.get_variable(
                'dynamics_states', initializer=tf.zeros(
                    [horizon, n, env_info.ob_dim()]),
                dtype=tf.float32, trainable=False, collections=[],
                validate_shape=False)
        self._unroll_loop = _tf_unroll(
            horizon, self._unroll_states_ph_ns,
            tf_action, tf_predict_dynamics,
            actions_hna, states_hns)
        self._initializer = [
            actions_hna.initializer, states_hns.initializer]

    def eval_dynamics(self, data, nsamples):
        """
        Unroll the actor and dynamics by the given horizon length.
        return the initial observations, planned actions, and
        planned resulting observations.

        First axis always corresponds to the nsamples batching
        dimension. Second axis is the horizon.
        """
        ixs = np.random.randint(data.size, size=(nsamples,))

        obs = data.obs[ixs]
        tf.get_default_session().run(
            self._initializer, feed_dict={
                self._unroll_states_ph_ns: obs})

        acs_hna, obs_hns = tf.get_default_session().run(
            self._unroll_loop, feed_dict={
                self._unroll_states_ph_ns: obs})

        planned_acs = np.swapaxes(acs_hna, 0, 1)
        planned_obs = np.swapaxes(obs_hns, 0, 1)
        return obs, planned_acs, planned_obs

def _tf_unroll(h, initial_states_ns, act, dynamics, acs_hna, states_hns):
    def _body(t, state_ns):
        action_na = act(state_ns)
        save_action_op = tf.scatter_update(acs_hna, t, action_na)
        next_state_ns = dynamics(state_ns, action_na)
        save_state_op = tf.scatter_update(
            states_hns, t, next_state_ns)
        with tf.control_dependencies([save_action_op, save_state_op]):
            return [t + 1, next_state_ns]

    loop_vars = [0, initial_states_ns]
    loop, _ = tf.while_loop(lambda t, _: t < h, _body,
                            loop_vars, back_prop=False)

    with tf.control_dependencies([loop]):
        return acs_hna.read_value(), states_hns.read_value()

reason: import numpy


text:
from numpy import array
from pybimstab.slope import NaturalSlope
terrainCoords = array(
    [[0, 10, 18, 28], [16.571, 16.571,  4.571,  4.571]])
slope = NaturalSlope(terrainCoords)
fig = slope.plot()

max_stars_repo_name: eamontoyaa/pybimstab

reason: from numpy


text:
# -*- coding: utf-8 -*-

import contextlib
import datetime

import hypothesis.extra.numpy as hyp_np
import hypothesis.strategies as hyp_st
import numpy as np
import pandas as pd

from kartothek.core.uuid import gen_uuid_object

try:
    from freezegun.api import real_date as date
except ImportError:
    from datetime import date

try:
    # python 3:
    from unittest import mock
except ImportError:
    import mock

TIME_TO_FREEZE = datetime.datetime(2000, 1, 1, 1, 1, 1, 1)
TIME_TO_FREEZE_ISO = "2000-01-01T01:01:01.000001"
TIME_TO_FREEZE_ISO_QUOTED = "2000-01-01T01%3A01%3A01.000001"

def get_dataframe_alltypes():
    """
    Return a pandas DataFrame of length one with a column for each commonly used data types
    """
    # fmt: off
    not_nested = get_dataframe_not_nested()
    nested_types = pd.DataFrame(
        {
            "array_int8": pd.Series([np.array([1], dtype=np.int8)], dtype=object),
            "array_int16": pd.Series([np.array([1], dtype=np.int16)], dtype=object),
            "array_int32": pd.Series([np.array([1], dtype=np.int32)], dtype=object),
            "array_int64": pd.Series([np.array([1], dtype=np.int64)], dtype=object),
            "array_uint8": pd.Series([np.array([1], dtype=np.uint8)], dtype=object),
            "array_uint16": pd.Series([np.array([1], dtype=np.uint16)], dtype=object),
            "array_uint32": pd.Series([np.array([1], dtype=np.uint32)], dtype=object),
            "array_uint64": pd.Series([np.array([1], dtype=np.uint64)], dtype=object),
            "array_float32": pd.Series([np.array([1], dtype=np.float32)], dtype=object),
            "array_float64": pd.Series([np.array([1], dtype=np.float64)], dtype=object),
            "array_unicode": pd.Series([np.array(["Ö"], dtype=object)], dtype=object),
        }
    )

    return pd.concat([not_nested, nested_types], axis=1).reset_index(drop=True).sort_index(axis=1)
    # fmt: on

def get_dataframe_not_nested():
    return pd.DataFrame(
        {
            "bool": pd.Series([1], dtype=np.bool),
            "int8": pd.Series([1], dtype=np.int8),
            "int16": pd.Series([1], dtype=np.int16),
            "int32": pd.Series([1], dtype=np.int32),
            "int64": pd.Series([1], dtype=np.int64),
            "uint8": pd.Series([1], dtype=np.uint8),
            "uint16": pd.Series([1], dtype=np.uint16),
            "uint32": pd.Series([1], dtype=np.uint32),
            "uint64": pd.Series([1], dtype=np.uint64),
            "float32": pd.Series([1.0], dtype=np.float32),
            "float64": pd.Series([1.0], dtype=np.float64),
            "date": pd.Series([date(2018, 1, 1)], dtype=object),
            "datetime64": pd.Series(["2018-01-01"], dtype="datetime64[ns]"),
            "unicode": pd.Series(["Ö"], dtype=np.unicode),
            "null": pd.Series([None], dtype=object),
            # Adding a byte type with value as byte sequence which can not be encoded as UTF8
            "byte": pd.Series([gen_uuid_object().bytes], dtype=object),
        }
    ).sort_index(axis=1)

def get_scalar_dtype_strategy(exclude=None):
    """
    A `hypothesis` strategy yielding
    """
    possible_strategies = {
        "datetime": hyp_np.datetime64_dtypes(max_period="ms", min_period="ns"),
        "uint": hyp_np.unsigned_integer_dtypes(),
        "int": hyp_np.integer_dtypes(),
        "float": hyp_np.floating_dtypes(),
        "byte": hyp_np.byte_string_dtypes(),
        "unicode": hyp_np.unicode_string_dtypes(),
    }
    if exclude is None:
        exclude = {}
    elif not isinstance(exclude, list):
        exclude = [exclude]
    for ex in exclude:
        if ex in possible_strategies:
            del possible_strategies[ex]
        else:
            raise ValueError(
                "Strategy {} unknown. Possible values are {}".format(
                    ex, possible_strategies.keys()
                )
            )
    return hyp_st.one_of(*list(possible_strategies.values()))

def get_numpy_array_strategy(
    shape=10, exclude_dtypes=None, unique=False, sort=False, allow_nan=True
):
    # the text example generation has quite some overhead when called the first time.
    # we don't want this in our test sample generation since the HealthCheck of hypothesis
    # might be triggered.
    hyp_st.text().example()

    dtype_strategy = get_scalar_dtype_strategy(exclude_dtypes)
    array_strategy = hyp_np.arrays(dtype=dtype_strategy, shape=shape, unique=unique)

    if exclude_dtypes is None or "date" not in exclude_dtypes:
        date_strat = hyp_st.lists(
            hyp_st.dates(
                min_value=datetime.date(1970, 1, 1), max_value=datetime.date(2100, 1, 1)
            ),
            min_size=shape,
            max_size=shape,
            unique=unique,
        )
        date_strat = date_strat.map(np.array)
        one_of_strategies = [array_strategy] + [date_strat]
        array_strategy = hyp_st.one_of(one_of_strategies)

    def _restrict_datetime_ranges(arr):
        if np.issubdtype(arr.dtype, np.datetime64):
            return all(
                (arr < np.datetime64("2200-01-01"))
                & (arr > np.datetime64("1970-01-01"))
            )
        return True

    if exclude_dtypes is None or "datetime" not in exclude_dtypes:
        array_strategy = array_strategy.filter(_restrict_datetime_ranges)
    if not allow_nan:

        def _check_for_nan(arr):
            if np.issubdtype(arr.dtype, np.floating):
                return not any(np.isnan(arr))
            return True

        array_strategy = array_strategy.filter(_check_for_nan)
    if unique and allow_nan:

        def _maximum_single_nan(arr):
            if np.issubdtype(arr.dtype, np.floating):
                return sum(np.isnan(arr)) <= 1
            return True

        array_strategy = array_strategy.filter(_maximum_single_nan)

    if sort:
        array_strategy = array_strategy.map(np.sort)
    return array_strategy

@contextlib.contextmanager
def cm_frozen_time(time_to_freeze):
    """
    Context manager to monkeypatch kartothek.core._time.* to return
    a fixed datetime value `time_to_freeze`.
    """
    with mock.patch("kartothek.core._time.datetime_now") as mock_now, mock.patch(
        "kartothek.core._time.datetime_utcnow"
    ) as mock_utcnow:
        mock_now.return_value = time_to_freeze
        mock_utcnow.return_value = time_to_freeze
        yield

max_stars_repo_name: tim-werner/kartothek

max_issues_repo_name: tim-werner/kartothek

max_forks_repo_name: tim-werner/kartothek


text:
import datetime as dt
import logging
import os
import pickle
import re
from collections import OrderedDict
from typing import Dict, List

import mlflow
import numpy as np
import pandas as pd
from metpy import calc
from metpy.units import units
from scipy.constants import convert_temperature
from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.compose import ColumnTransformer
from sklearn.feature_selection import SelectKBest, mutual_info_regression
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import PowerTransformer
from yellowbrick.target import FeatureCorrelation

from wind_power_forecasting.nodes import utils

def _get_wind_speed(x: pd.DataFrame) -> float:
    """Function to get wind speed from wind velocity components U and V.

    Args:
        x: Features data frame containing components U and V as columns.

    Returns:
        Wind speed for each pair <U,V>.

    """
    return float(
        calc.wind_speed(
            x.U * units.meter / units.second, x.V * units.meter / units.second
        ).magnitude
    )

def _get_wind_dir(x: pd.DataFrame) -> float:
    """Function to get wind direction from wind velocity components U and V.

    Args:
        x: Features data frame containing components U and V as columns.

    Returns:
        Wind direction for each pair <U,V>.

    """

    return float(
        calc.wind_direction(
            x.U * units.meter / units.second,
            x.V * units.meter / units.second,
            convention="from",
        ).magnitude
    )

def _encode_cyclic(data: pd.DataFrame, col: str, max_val: int) -> None:
    """Encondes cyclic features using sinusoidal functions.

    Args:
        data: the data frame containing the cyclic features.
        col: the feature name to be enconded.
        max_val: the maximum value in each case (for instance: 24 for hour, 12 for month.)

    Returns:
        None, it performs the enconding inplace.

    """
    data[col + "_sin"] = np.sin(2 * np.pi * data[col] / max_val)
    data[col + "_cos"] = np.cos(2 * np.pi * data[col] / max_val)

def _temperature_in_celsius(df, temp_col):
    """Converts temperature units from Kelvin to Celsius.

    Args:
        df: the data frame containing a temperature feature.
        temp_col: the temperature feature name.

    Returns:
        The initial dataframe with the temperature feature in Celsius.

    """
    df[col] = pd.Series(
        convert_temperature(np.array(df[temp_col]), "Kelvin", "Celsius")
    )
    return df

def _add_interaction_terms(df, var_list):
    """Adds the possible interaction terms from the list elements.

    Args:
        var_list: list of features.

    Returns:
        The data frame with the interaction terms as new columns.

    """
    print("Function to be implemented!")

class NewFeaturesAdder(BaseEstimator, TransformerMixin):
    """Scikit-learn custom transformer that allows to add new features
    derived from the original ones.
    """

    def __init__(
        self,
        add_time_feat=False,
        add_cycl_feat=False,
        add_inv_T=False,
        add_interactions=False,
    ):
        self.add_time_feat = add_time_feat
        self.add_cycl_feat = add_cycl_feat
        self.add_inv_T = add_inv_T
        self.add_interactions = add_interactions

    def fit(self, documents, y=None):
        return self

    def transform(self, x_dataset) -> pd.DataFrame:

        # Velocity derived features
        x_dataset["wspeed"] = x_dataset.apply(_get_wind_speed, axis=1)
        x_dataset["wdir"] = x_dataset.apply(_get_wind_dir, axis=1)

        if self.add_time_feat:
            # Time derived features
            x_dataset["hour"] = x_dataset["Time"].dt.hour
            x_dataset["month"] = x_dataset["Time"].dt.month

        if self.add_cycl_feat:
            if self.add_time_feat == False:
                _encode_cyclic(x_dataset, "wdir", 360)
            else:
                # Hour
                _encode_cyclic(x_dataset, "hour", 24)

                # Month
                _encode_cyclic(x_dataset, "month", 12)

                # Wind direction
                _encode_cyclic(x_dataset, "wdir", 360)

        if self.add_inv_T:
            x_dataset["inv_T"] = 1 / x_dataset["T"]

        if self.add_interactions:
            if not self.add_cycl_feat:
                x_dataset["wspeed_wdir"] = x_dataset["wspeed"] * x_dataset["wdir"]
                x_dataset["wspeed_T"] = x_dataset["wspeed"] * x_dataset["T"]
                x_dataset["wspeed_wdir_T"] = (
                    x_dataset["wspeed"] * x_dataset["wdir"] * x_dataset["T"]
                )

                if self.add_inv_T:
                    x_dataset["wspeed_invT"] = x_dataset["wspeed"] * x_dataset["inv_T"]
                    x_dataset["wspeed_wdir_invT"] = (
                        x_dataset["wspeed"] * x_dataset["wdir"] * x_dataset["inv_T"]
                    )
            else:
                x_dataset["wspeed_T"] = x_dataset["wspeed"] * x_dataset["T"]
                x_dataset["wspeed_wdirsin"] = (
                    x_dataset["wspeed"] * x_dataset["wdir_sin"]
                )
                x_dataset["wspeed_wdircos"] = (
                    x_dataset["wspeed"] * x_dataset["wdir_cos"]
                )
                x_dataset["wspeed_wdirsin_T"] = (
                    x_dataset["wspeed"] * x_dataset["wdir_sin"] * x_dataset["T"]
                )
                x_dataset["wspeed_wdircos_T"] = (
                    x_dataset["wspeed"] * x_dataset["wdir_cos"] * x_dataset["T"]
                )
                if self.add_inv_T:
                    x_dataset["wspeed_invT"] = x_dataset["wspeed"] * x_dataset["inv_T"]
                    x_dataset["wspeed_wdirsin_invT"] = (
                        x_dataset["wspeed"] * x_dataset["wdir_sin"] * x_dataset["inv_T"]
                    )
                    x_dataset["wspeed_wdircos_invT"] = (
                        x_dataset["wspeed"] * x_dataset["wdir_cos"] * x_dataset["inv_T"]
                    )

        return x_dataset

def _walklevel(some_dir, level):
    """
    It works just like os.walk, but you can pass it a level parameter
    that indicates how deep the recursion will go.
    If depth is -1 (or less than 0), the full depth is walked.
    """
    some_dir = some_dir.rstrip(os.path.sep)
    assert os.path.isdir(some_dir)
    num_sep = some_dir.count(os.path.sep)

    for root, dirs, files in os.walk(some_dir):
        yield root, dirs, files
        num_sep_this = root.count(os.path.sep)

        if num_sep + level <= num_sep_this:
            del dirs[:]

def _load_csv_files(loc: str, level=0, header=0) -> Dict:
    """It loads all csv files in a location and returns a list of data frames
    created from those files.

    Args:
        loc: directory where the csv files are located.
        level: how deep the recursion will go in the directory. By default is 0.
    """
    df_dict = {}
    for dirname, _, filenames in _walklevel(loc, level):
        for filename in filenames:
            df_dict[filename] = pd.read_csv(dirname + filename, sep=",", header=header)

    return df_dict

def feature_engineering(
    wf: str,
    data_source: str,
    add_time_feat: bool,
    add_cycl_feat: bool,
    add_inv_T: bool,
    add_interactions: bool,
) -> (np.ndarray, np.ndarray, List):
    """It performs feature engineering pipeline with these steps:
        - Add new features derived from original ones.
        - Delete innecesary features.
        - Apply a power transform + standarization.

    Args:
        data_source: folder where the primary data is located.
        add_time_feat: (True/False) to add or not features derived from Time.
        add_cycle_feat: (True/False) whether to encode cyclic features or not.
        add_inv_T: (True/False) to add or not 1/T derived feature from T.
        add_interactions: (True/False) to add or not interaction terms.

    Returns:
        X_train and X_test pickle objects prepared for modeling.
        A list with the dropped unnecessary columns.

    """
    # Import data sets depending on provided Wind Farm (wf)
    folder = data_source + wf + "/"
    X_train = _load_csv_files(folder).get("X_train.csv")
    X_test = _load_csv_files(folder).get("X_test.csv")
    X_train["Time"] = pd.to_datetime(X_train["Time"], format="%d/%m/%Y %H:%M")
    X_test["Time"] = pd.to_datetime(X_test["Time"], format="%d/%m/%Y %H:%M")

    feat_adder = NewFeaturesAdder(
        add_time_feat=add_time_feat,
        add_cycl_feat=add_cycl_feat,
        add_inv_T=add_inv_T,
        add_interactions=add_interactions,
    )

    ###  Save intermediate expanded sets to be used in EDA
    X_train_expanded = feat_adder.fit_transform(X_train)
    X_test_expanded = feat_adder.transform(X_test)

    os.makedirs("./data/04_feature/" + wf, exist_ok=True)

    with open(
        "./data/04_feature/{}/X_train_expanded.pickle".format(wf), "wb"
    ) as handle:
        pickle.dump(X_train_expanded, handle, protocol=pickle.HIGHEST_PROTOCOL)

    with open("./data/04_feature/{}/X_test_expanded.pickle".format(wf), "wb") as handle:
        pickle.dump(X_test_expanded, handle, protocol=pickle.HIGHEST_PROTOCOL)
    ##########

    drop_lst = []
    if feat_adder.get_params().get("add_cycl_feat"):
        if feat_adder.get_params().get("add_inv_T"):
            drop_lst = ["ID", "Time", "U", "V", "wdir", "hour", "month", "T"]
        else:
            drop_lst = ["ID", "Time", "U", "V", "wdir", "hour", "month"]
    else:
        if feat_adder.get_params().get("add_inv_T"):
            drop_lst = ["ID", "Time", "U", "V", "T"]
        else:
            drop_lst = ["ID", "Time", "U", "V"]

    pre_process = ColumnTransformer(
        remainder="passthrough", transformers=[("drop_columns", "drop", drop_lst)]
    )

    # full pipeline to get preperaded data for ML
    feat_eng_pipeline = Pipeline(
        steps=[("feature_adder", feat_adder), ("pre_processing", pre_process),]
    )

    X_train_pped = feat_eng_pipeline.fit_transform(X_train)
    X_test_pped = feat_eng_pipeline.transform(X_test)

    # Get the names of the features
    feature_names = X_train.drop(drop_lst, axis=1).columns

    return X_train_pped, X_test_pped, feature_names

def save_prepared_data(
    folder: str, X_train, X_test, feature_names: List, WF: str
) -> None:
    """Saves the prepared data after feature engineering.

    Args:
        folder: the folder where the pickle objects will be saved.
        WF: Wind Farm identification.

    Returns:
        None.
    """
    os.makedirs(folder + WF, exist_ok=True)

    with open(folder + "{}/X_train_pped.pickle".format(WF), "wb") as handle:
        pickle.dump(X_train, handle, protocol=pickle.HIGHEST_PROTOCOL)

    with open(folder + "{}/X_test_pped.pickle".format(WF), "wb") as handle:
        pickle.dump(X_test, handle, protocol=pickle.HIGHEST_PROTOCOL)

    with open(folder + "{}/feature_names.pickle".format(WF), "wb") as handle:
        pickle.dump(feature_names, handle, protocol=pickle.HIGHEST_PROTOCOL)

def show_feature_importance(
    wf: str,
    data_src: str,
    data_dst: str,
    X: np.ndarray,
    feature_names: List,
    k_best: str,
):
    """It finds the k most influecers features on target by performing
    mutual information regression.

    Args:
        wf: Wind Farm identification.
        data_src: location of data files.
        data_dst: plots location
        X: features data set.
        feature_names: the name of the features in X.
        k_best: the number of features to return.

    Returns:
        A descending ordered list of the k best features.
        If k='all' the list can be considered as a feature importance list.

    """
    # Load target data (y_train)
    y = _load_csv_files(data_src + wf + "/", header=None).get("y_train.csv")
    y = y.to_numpy().reshape(-1)

    if k_best == "all":
        k = X.shape[1]
    else:
        k = int(k_best)

    # Perform MI regression to get the k-best features.
    selec_k_best = SelectKBest(mutual_info_regression, k=k)
    selec_k_best.fit(X, y)

    # Get the names of the k-best features
    fnames = feature_names
    mask = selec_k_best.get_support()
    scores = selec_k_best.scores_
    selected_feat = {}

    for bool, feature, score in zip(mask, fnames, scores):
        if bool:
            selected_feat[feature] = score

    sorted_sel_feat = {
        k: v
        for k, v in reversed(sorted(selected_feat.items(), key=lambda item: item[1]))
    }

    i = 1
    print("Feature importance in descending order: ")
    for k, v in sorted_sel_feat.items():
        print("{0}. {1}: {2:.2f}".format(i, k, v))
        i += 1

    logger = logging.getLogger(__name__)
    logger.info("Feature importance in descending order: {}".format(sorted_sel_feat))

    # Create a list of the feature names
    features = feature_names

    # Instantiate the visualizer
    visualizer = FeatureCorrelation(
        method="mutual_info-regression",
        labels=features,
        size=(1080, 720),
        title="Importancia de variables",
    )

    visualizer.fit(X, y, random_state=0)
    visualizer.show(outpath=data_dst + "figures/" + wf + "/feature_importance.png",)
    visualizer.show(clear_figure=True)

    mlflow.log_artifacts(data_dst + "figures/" + wf)

max_stars_repo_name: vchaparro/wind-power-forecasting


text:
import numpy as np
import onnx
import onnx.helper
import pytest

from dnnv.nn.operations import *
from dnnv.nn.operations.base import *

def test_str():
    assert str(Operation()) == "Operation"
    assert str(Input(None, None)) == "Input"
    assert str(Add(None, None)) == "Add"
    assert str(Cast(None, None)) == "Cast"
    assert str(Gemm(np.empty((1, 5)), np.empty((5, 2)))) == "Gemm"

def test_from_onnx_Input():
    onnx_node = onnx.helper.make_tensor_value_info(
        "test_input", onnx.TensorProto.FLOAT, ["N", 3, 4, 4]
    )
    op_from_onnx = Operation.from_onnx(onnx_node)
    input_op_from_init = Input(np.array([-1, 3, 4, 4]), np.dtype(np.float32))
    assert type(op_from_onnx) is Input
    assert np.all(op_from_onnx.shape == input_op_from_init.shape)
    assert op_from_onnx.dtype == input_op_from_init.dtype

def test_from_onnx_Add():
    input_op = Input(np.array([-1, 5]), np.dtype(np.float32))

    add_node = onnx.helper.make_node(
        "Add", inputs=["input", "const"], outputs=["add"], name="add"
    )
    op_from_onnx = Operation.from_onnx(
        add_node, input_op, np.ones((1, 5), dtype=np.float32)
    )

    assert type(op_from_onnx) is Add
    assert op_from_onnx.a is input_op
    assert np.all(op_from_onnx.b == np.ones((1, 5)))

def test_from_onnx_Mul_constants(caplog):
    # TODO : is this a necessary feature?
    add_node = onnx.helper.make_node(
        "Mul", inputs=["const1", "const2"], outputs=["mul"], name="mul"
    )
    op_from_onnx = Operation.from_onnx(add_node, 2, 5)

    assert type(op_from_onnx) is Mul
    assert op_from_onnx.a == 2
    assert op_from_onnx.b == 5

    assert "Operation on constant inputs returned non-constant." in caplog.text

def test_from_onnx_Elu():
    input_op = Input(np.array([-1, 5]), np.dtype(np.float32))

    add_node = onnx.helper.make_node(
        "Elu", inputs=["input"], outputs=["elu"], name="elu", alpha=0.5
    )
    op_from_onnx = Operation.from_onnx(add_node, input_op)

    assert type(op_from_onnx) is Elu
    assert op_from_onnx.x is input_op
    assert op_from_onnx.alpha == 0.5

def test_from_onnx_Elu():
    input_op = Input(np.array([-1, 5]), np.dtype(np.float32))

    add_node = onnx.helper.make_node(
        "Fake", inputs=["input"], outputs=["fake"], name="fake"
    )
    with pytest.raises(ValueError) as excinfo:
        op_from_onnx = Operation.from_onnx(add_node, input_op)
    assert str(excinfo.value) == "Unimplemented operation type: Fake"

def test_inputs():
    input_op_0 = Input((1, 4), np.dtype(np.float32))
    add_op = Add(input_op_0, np.ones((1, 4), dtype=np.float32))
    add_op_inputs = add_op.inputs
    assert len(add_op_inputs) == 1
    assert input_op_0 in add_op_inputs

    input_op_1 = Input((1, 3, 2, 2), np.dtype(np.float32))
    flatten_op = Flatten(input_op_1)
    concat_op = Concat([add_op, flatten_op], axis=0)
    concat_op_inputs = concat_op.inputs
    assert len(concat_op_inputs) == 2
    assert add_op in concat_op_inputs
    assert flatten_op in concat_op_inputs

    reshape_op = Reshape(input_op_1, [1, 12])
    reshape_op_inputs = reshape_op.inputs
    assert len(reshape_op_inputs) == 1
    assert input_op_1 in reshape_op_inputs

def test_match_none():
    assert list(Operation.match([])) == []
    assert list(Input.match([])) == []
    assert list(Add.match([])) == []

def test_match_true():
    assert list(Operation.match([Operation()])) == [[]]
    assert list(Input.match([Input(None, None)])) == [[]]

    input_op_0 = Input(None, None)
    input_op_1 = Input(None, None)
    add_op = Add(input_op_0, input_op_1)
    assert list(Add.match([add_op])) == [[input_op_0, input_op_1]]

def test_match_false():
    assert list(Input.match([Operation()])) == []
    assert list(Add.match([Input(None, None)])) == []

    input_op_0 = Input(None, None)
    input_op_1 = Input(None, None)
    add_op = Add(input_op_0, input_op_1)
    assert list(Add.match([add_op, input_op_0])) == []

max_stars_repo_name: samysweb/dnnv

max_issues_repo_name: samysweb/dnnv

max_forks_repo_name: samysweb/dnnv


text:
# -*- coding: UTF-8 -*-
#
# Copyright (C) 2010-2011 Yung-Yu Chen <yyc@solvcon.net>.
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.

"""
Anisotropic velocity-stress equations solver using the linear CESE method with
CUDA support.
"""

from solvcon.gendata import TypeNameRegistry
from solvcon.anchor import Anchor
from solvcon.hook import BlockHook
from solvcon.kerpak.cuse import CUDA_RAISE_ON_FAIL, CuseBC
from solvcon.kerpak.lincuse import (LincuseSolver, LincuseCase,
    PlaneWaveSolution)

###############################################################################
# Solver.
###############################################################################

class VslinSolver(LincuseSolver):
    """
    Basic elastic solver.

    @ivar mtrldict: map from names to material objects.
    @itype mtrldict: dict
    @ivar mtrllist: list of all material objects.
    @itype mtrllist: list
    """
    from solvcon.dependency import getcdll
    __clib_vslin_c = {
        2: getcdll('vslin2d_c', raise_on_fail=False),
        3: getcdll('vslin3d_c', raise_on_fail=False),
    }
    __clib_vslin_cu = {
        2: getcdll('vslin2d_cu', raise_on_fail=CUDA_RAISE_ON_FAIL),
        3: getcdll('vslin3d_cu', raise_on_fail=CUDA_RAISE_ON_FAIL),
    }
    del getcdll
    @property
    def _clib_vslin_c(self):
        return self.__clib_vslin_c[self.ndim]
    @property
    def _clib_vslin_cu(self):
        return self.__clib_vslin_cu[self.ndim]
    @property
    def _gdlen_(self):
        return 9 * 9 * self.ndim
    def __init__(self, *args, **kw):
        super(VslinSolver, self).__init__(*args, **kw)
        self.mtrldict = kw.pop('mtrldict', {})
        self.mtrllist = None
    def make_grpda(self):
        self.mtrllist = self._build_mtrllist(self.grpnames, self.mtrldict)
        for igrp in range(len(self.grpnames)):
            mtrl = self.mtrllist[igrp]
            jaco = self.grpda[igrp].reshape(self.neq, self.neq, self.ndim)
            mjacos = mtrl.get_jacos()
            for idm in range(self.ndim):
                jaco[:,:,idm] = mjacos[idm,:,:]
    @staticmethod
    def _build_mtrllist(grpnames, mtrldict):
        """
        Build the material list out of the mapping dict.

        @type grpnames: list
        @param mtrldict: the map from names to material objects.
        @type mtrldict: dict
        @return: the list of material object.
        @rtype: Material
        """
        mtrllist = list()
        default_mtuple = mtrldict.get(None, None)
        for grpname in grpnames:
            try:
                mtrl = mtrldict.get(grpname, default_mtuple)
            except KeyError as e:
                args = e.args[:]
                args.append('no material named %s in mtrldict'%grpname)
                e.args = args
                raise
            mtrllist.append(mtrl)
        return mtrllist

###############################################################################
# Case.
###############################################################################

class VslinCase(LincuseCase):
    """
    Case for anisotropic elastic solids.
    """
    defdict = {
        'execution.neq': 9,
        'solver.solvertype': VslinSolver,
        'solver.mtrldict': dict,
    }
    def make_solver_keywords(self):
        kw = super(VslinCase, self).make_solver_keywords()
        # setup material mapper.
        kw['mtrldict'] = self.solver.mtrldict
        return kw

###############################################################################
# Boundary conditions.
###############################################################################

class VslinBC(CuseBC):
    """
    Basic BC class for elastic problems.
    """
    from solvcon.dependency import getcdll
    __clib_vslinb_c = {
        2: getcdll('vslinb2d_c', raise_on_fail=False),
        3: getcdll('vslinb3d_c', raise_on_fail=False),
    }
    __clib_vslinb_cu = {
        2: getcdll('vslinb2d_cu', raise_on_fail=CUDA_RAISE_ON_FAIL),
        3: getcdll('vslinb3d_cu', raise_on_fail=CUDA_RAISE_ON_FAIL),
    }
    del getcdll
    @property
    def _clib_vslinb_c(self):
        return self.__clib_vslinb_c[self.svr.ndim]
    @property
    def _clib_vslinb_cu(self):
        return self.__clib_vslinb_cu[self.svr.ndim]

class VslinTraction(VslinBC):
    vnames = [
        'bfcsys', 'tau1', 'tau2', 'tau3', 'freq', 'phase',
    ]
    vdefaults = {
        'bfcsys': 0.0,
        'tau1': 0.0, 'tau2': 0.0, 'tau3': 0.0, 'freq': 0.0, 'phase': 0.0,
    }
    _ghostgeom_ = 'compress'
    def soln(self):
        from ctypes import byref, c_int
        self._clib_boundcond.bound_traction_soln(
            byref(self.svr.exd),
            c_int(self.facn.shape[0]),
            self.facn.ctypes._as_parameter_,
            c_int(self.value.shape[1]),
            self.value.ctypes._as_parameter_,
        )
    def dsoln(self):
        from ctypes import byref, c_int
        self._clib_boundcond.bound_traction_dsoln(
            byref(self.svr.exd),
            c_int(self.facn.shape[0]),
            self.facn.ctypes._as_parameter_,
        )

class VslinTractionFree(VslinBC):
    _ghostgeom_ = 'mirror'
    def soln(self):
        from ctypes import byref, c_int
        self._clib_boundcond.bound_traction_free_soln(
            byref(self.svr.exd),
            c_int(self.facn.shape[0]),
            self.facn.ctypes._as_parameter_,
        )
    def dsoln(self):
        from ctypes import byref, c_int
        self._clib_boundcond.bound_traction_free_dsoln(
            byref(self.svr.exd),
            c_int(self.facn.shape[0]),
            self.facn.ctypes._as_parameter_,
        )

class VslinTractionFree2(VslinBC):
    _ghostgeom_ = 'mirror'
    def soln(self):
        from ctypes import byref, c_int
        self._clib_boundcond.bound_traction_free2_soln(
            byref(self.svr.exd),
            c_int(self.facn.shape[0]),
            self.facn.ctypes._as_parameter_,
        )
    def dsoln(self):
        from ctypes import byref, c_int
        self._clib_boundcond.bound_traction_free2_dsoln(
            byref(self.svr.exd),
            c_int(self.facn.shape[0]),
            self.facn.ctypes._as_parameter_,
        )

################################################################################
# Plane wave solution.
################################################################################

class VslinPWSolution(PlaneWaveSolution):
    def _calc_eigen(self, **kw):
        from numpy import sqrt
        from numpy.linalg import eig
        wvec = kw['wvec']
        mtrl = kw['mtrl']
        idx = kw['idx']
        nml = wvec/sqrt((wvec**2).sum())
        jacos = mtrl.get_jacos()
        jaco = jacos[0] * nml[0]
        for idm in range(1, len(nml)):
            jaco += jacos[idm] * nml[idm]
        evl, evc = eig(jaco)
        srt = evl.argsort()
        evl = evl[srt[idx]].real
        evc = evc[:,srt[idx]].real
        evc *= evc[0]/abs(evc[0]+1.e-200)
        return evl, evc

################################################################################
# Anchor.
################################################################################

class VslinOAnchor(Anchor):
    """
    Calculate total energy, i.e., the summation of kinetic energy and strain
    energy.
    """
    def _calculate_physics(self):
        from ctypes import byref
        from numpy import empty
        from numpy.linalg import inv
        svr = self.svr
        # input arrays.
        rhos = empty(svr.ngroup, dtype=svr.fpdtype)
        comps = empty((svr.ngroup, 6, 6), dtype=svr.fpdtype)
        for igp in range(svr.ngroup):
            mtrl = svr.mtrllist[igp]
            rhos[igp] = mtrl.rho
            comps[igp,:,:] = inv(mtrl.stiff).T
        # output arrays.
        svr._clib_vslin_c.calc_energy(
            byref(svr.exd),
            rhos.ctypes._as_parameter_,
            comps.ctypes._as_parameter_,
            svr.der['energy'].ctypes._as_parameter_,
        )
    def provide(self):
        from numpy import empty
        svr = self.svr
        svr.der['energy'] = empty(svr.ngstcell+svr.ncell, dtype=svr.fpdtype)
        self._calculate_physics()
    def postfull(self):
        self._calculate_physics()

################################################################################
# Material definition.
################################################################################

mltregy = TypeNameRegistry()  # registry singleton.
class MaterialMeta(type):
    """
    Meta class for material class.
    """
    def __new__(cls, name, bases, namespace):
        newcls = super(MaterialMeta, cls).__new__(cls, name, bases, namespace)
        # register.
        mltregy.register(newcls)
        return newcls

class Material(object):
    """
    Material properties.  The constitutive relation needs not be symmetric.

    @cvar _zeropoints_: list of tuples for indices where the content should be
        zero.
    @ctype _zeropoints_: list
    @ivar rho: density
    @ivar al: alpha angle.
    @ivar be: beta angle.
    @ivar ga: gamma angle.
    @ivar origstiff: stiffness matrix in the crystal coordinate.
    @ivar stiff: stiffness matrix in the transformed global coordinate.
    """

    __metaclass__ = MaterialMeta

    _zeropoints_ = []

    from numpy import array
    K = array([ [
            [1, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 1],
            [0, 0, 0, 0, 1, 0],
        ], [
            [0, 0, 0, 0, 0, 1],
            [0, 1, 0, 0, 0, 0],
            [0, 0, 0, 1, 0, 0],
        ], [
            [0, 0, 0, 0, 1, 0],
            [0, 0, 0, 1, 0, 0],
            [0, 0, 1, 0, 0, 0],
        ], ], dtype='float64')
    del array

    def __init__(self, *args, **kw):
        from numpy import empty, dot
        self.rho = kw.pop('rho')
        self.al = kw.pop('al')
        self.be = kw.pop('be')
        self.ga = kw.pop('ga')
        # set stiffness matrix.
        origstiff = empty((6,6), dtype='float64')
        origstiff.fill(0.0)
        for key in kw.keys():   # becaues I pop out the key.
            if len(key) == 4 and key[:2] == 'co':
                try:
                    i = int(key[2])-1
                    j = int(key[3])-1
                except:
                    continue
                assert i < origstiff.shape[0]
                assert j < origstiff.shape[1]
                val = kw.pop(key)
                origstiff[i,j] = val
        self.origstiff = origstiff
        # check for zeros.
        self._check_origstiffzero(self.origstiff)
        # compute the stiffness matrix in the transformed global coordinate
        # system.
        bondmat = self.get_bondmat()
        self.stiff = dot(bondmat, dot(self.origstiff, bondmat.T))
        super(Material, self).__init__(*args, **kw)

    def __getattr__(self, key):
        if len(key) == 4 and key[:2] == 'co':
            i = int(key[2])
            j = int(key[3])
            if 1 <= i <= 6 and 1 <= j <= 6:
                return self.origstiff[i-1,j-1]
        elif len(key) == 3 and key[0] == 'c':
            i = int(key[1])
            j = int(key[2])
            if 1 <= i <= 6 and 1 <= j <= 6:
                return self.stiff[i-1,j-1]
        else:
            raise AttributeError

    def __str__(self):
        from math import pi
        return '[%s: al=%.2f be=%.2f ga=%.2f (deg)]' % (self.__class__.__name__,
            self.al/(pi/180), self.be/(pi/180), self.ga/(pi/180))

    @classmethod
    def _check_origstiffzero(cls, origstiff):
        """
        Check for zero in original stiffness matrix.

        @note: no assumed symmetry.
        """
        for i, j in cls._zeropoints_:
            assert origstiff[i,j] == 0.0

    def get_rotmat(self):
        """
        Coordinate transformation matrix for three successive rotations through
        the Euler angles.

        @return: the transformation matrix.
        @rtype: numpy.ndarray
        """
        from numpy import array, cos, sin, dot
        al = self.al; be = self.be; ga = self.ga
        almat = array([
            [cos(al), sin(al), 0],
            [-sin(al), cos(al), 0],
            [0, 0, 1],
        ], dtype='float64')
        bemat = array([
            [1, 0, 0],
            [0, cos(be), sin(be)],
            [0, -sin(be), cos(be)],
        ], dtype='float64')
        gamat = array([
            [cos(ga), sin(ga), 0],
            [-sin(ga), cos(ga), 0],
            [0, 0, 1],
        ], dtype='float64')
        return dot(gamat, dot(bemat, almat))

    def get_bondmat(self):
        """
        The Bond's matrix M as a shorthand of coordinate transformation for the
        6-component stress vector.

        @return: the Bond's matrix.
        @rtype: numpy.ndarray
        """
        from numpy import empty
        rotmat = self.get_rotmat()
        bond = empty((6,6), dtype='float64')
        # upper left.
        bond[:3,:3] = rotmat[:,:]**2
        # upper right.
        bond[0,3] = 2*rotmat[0,1]*rotmat[0,2]
        bond[0,4] = 2*rotmat[0,2]*rotmat[0,0]
        bond[0,5] = 2*rotmat[0,0]*rotmat[0,1]
        bond[1,3] = 2*rotmat[1,1]*rotmat[1,2]
        bond[1,4] = 2*rotmat[1,2]*rotmat[1,0]
        bond[1,5] = 2*rotmat[1,0]*rotmat[1,1]
        bond[2,3] = 2*rotmat[2,1]*rotmat[2,2]
        bond[2,4] = 2*rotmat[2,2]*rotmat[2,0]
        bond[2,5] = 2*rotmat[2,0]*rotmat[2,1]
        # lower left.
        bond[3,0] = rotmat[1,0]*rotmat[2,0]
        bond[3,1] = rotmat[1,1]*rotmat[2,1]
        bond[3,2] = rotmat[1,2]*rotmat[2,2]
        bond[4,0] = rotmat[2,0]*rotmat[0,0]
        bond[4,1] = rotmat[2,1]*rotmat[0,1]
        bond[4,2] = rotmat[2,2]*rotmat[0,2]
        bond[5,0] = rotmat[0,0]*rotmat[1,0]
        bond[5,1] = rotmat[0,1]*rotmat[1,1]
        bond[5,2] = rotmat[0,2]*rotmat[1,2]
        # lower right.
        bond[3,3] = rotmat[1,1]*rotmat[2,2] + rotmat[1,2]*rotmat[2,1]
        bond[3,4] = rotmat[1,0]*rotmat[2,2] + rotmat[1,2]*rotmat[2,0]
        bond[3,5] = rotmat[1,1]*rotmat[2,0] + rotmat[1,0]*rotmat[2,1]
        bond[4,3] = rotmat[0,1]*rotmat[2,2] + rotmat[0,2]*rotmat[2,1]
        bond[4,4] = rotmat[0,0]*rotmat[2,2] + rotmat[0,2]*rotmat[2,0]
        bond[4,5] = rotmat[0,1]*rotmat[2,0] + rotmat[0,0]*rotmat[2,1]
        bond[5,3] = rotmat[0,1]*rotmat[1,2] + rotmat[0,2]*rotmat[1,1]
        bond[5,4] = rotmat[0,0]*rotmat[1,2] + rotmat[0,2]*rotmat[1,0]
        bond[5,5] = rotmat[0,1]*rotmat[1,0] + rotmat[0,0]*rotmat[1,1]
        return bond

    def get_jacos(self):
        """
        Obtain the Jacobian matrices for the solid.

        @param K: the K matrix.
        @type K: numpy.ndarray
        @return: the Jacobian matrices
        @rtype: numpy.ndarray
        """
        from numpy import zeros, dot
        rho = self.rho
        sf = self.stiff
        jacos = zeros((3,9,9), dtype='float64')
        for idm in range(3):
            K = self.K[idm]
            jaco = jacos[idm]
            jaco[:3,3:] = K/(-rho)  # the upper right submatrix.
            jaco[3:,:3] = -dot(sf, K.T) # the lower left submatrix.
        return jacos

################################################################################
# Symmetry.
################################################################################

class Triclinic(Material):
    """
    The stiffness matrix has to be symmetric.
    """
    _zeropoints_ = []
    def __init__(self, *args, **kw):
        for key in kw.keys():   # becaues I modify the key.
            if len(key) == 4 and key[:2] == 'co':
                try:
                    i = int(key[2])
                    j = int(key[3])
                except:
                    continue
                symkey = 'co%d%d' % (j, i)
                if i != j:
                    assert symkey not in kw
                kw[symkey] = kw[key]
        super(Triclinic, self).__init__(*args, **kw)
    @classmethod
    def _check_origstiffzero(cls, origstiff):
        for i, j in cls._zeropoints_:
            assert origstiff[i,j] == origstiff[j,i] == 0.0

class Monoclinic(Triclinic):
    _zeropoints_ = [
        (0,3), (0,5),
        (1,3), (1,5),
        (2,3), (2,5),
        (3,4), (4,5),
    ]

class Orthorhombic(Triclinic):
    _zeropoints_ = [
        (0,3), (0,4), (0,5),
        (1,3), (1,4), (1,5),
        (2,3), (2,4), (2,5),
        (3,4), (3,5), (4,5),
    ]

class Tetragonal(Triclinic):
    _zeropoints_ = [
        (0,3), (0,4),
        (1,3), (1,4),
        (2,3), (2,4), (2,5),
        (3,4), (3,5), (4,5),
    ]
    def __init__(self, *args, **kw):
        kw['co22'] = kw['co11']
        kw['co23'] = kw['co13']
        kw['co26'] = -kw.get('co16', 0.0)
        kw['co55'] = kw['co44']
        super(Tetragonal, self).__init__(*args, **kw)

class Trigonal(Triclinic):
    _zeropoints_ = [
        (0,5), (1,5),
        (2,3), (2,4), (2,5),
        (3,4),
    ]
    def __init__(self, *args, **kw):
        kw['co15'] = -kw.get('co25', 0.0)
        kw['co22'] = kw['co11']
        kw['co23'] = kw['co13']
        kw['co24'] = -kw.get('co14', 0.0)
        kw['co46'] = kw.get('co25', 0.0)
        kw['co55'] = kw['co44']
        kw['co56'] = kw.get('co14', 0.0)
        kw['co66'] = (kw['co11'] - kw['co12'])/2
        super(Trigonal, self).__init__(*args, **kw)

class Hexagonal(Trigonal):
    _zeropoints_ = [
        (0,3), (0,4), (0,5),
        (1,3), (1,4), (1,5),
        (2,3), (2,4), (2,5),
        (3,4), (3,5), (4,5),
    ]

class Cubic(Triclinic):
    _zeropoints_ = [
        (0,3), (0,4), (0,5),
        (1,3), (1,4), (1,5),
        (2,3), (2,4), (2,5),
        (3,4), (3,5), (4,5),
    ]
    def __init__(self, *args, **kw):
        kw['co13'] = kw['co12']
        kw['co22'] = kw['co11']
        kw['co23'] = kw['co12']
        kw['co33'] = kw['co11']
        kw['co55'] = kw['co44']
        kw['co66'] = kw['co44']
        super(Cubic, self).__init__(*args, **kw)

class Isotropic(Triclinic):
    _zeropoints_ = [
        (0,3), (0,4), (0,5),
        (1,3), (1,4), (1,5),
        (2,3), (2,4), (2,5),
        (3,4), (3,5), (4,5),
    ]
    def __init__(self, *args, **kw):
        kw['co12'] = kw['co11']-2*kw['co44']
        kw['co13'] = kw['co11']-2*kw['co44']
        kw['co22'] = kw['co11']
        kw['co23'] = kw['co11']-2*kw['co44']
        kw['co33'] = kw['co11']
        kw['co55'] = kw['co44']
        kw['co66'] = kw['co44']
        super(Isotropic, self).__init__(*args, **kw)

################################################################################
# Material properties.
################################################################################

class GaAs(Cubic):
    def __init__(self, *args, **kw):
        kw.setdefault('rho', 5307.0)
        kw.setdefault('co11', 11.88e10)
        kw.setdefault('co12', 5.38e10)
        kw.setdefault('co44', 5.94e10)
        super(GaAs, self).__init__(*args, **kw)

class ZnO(Hexagonal):
    def __init__(self, *args, **kw):
        kw.setdefault('rho', 5680.0)
        kw.setdefault('co11', 20.97e10)
        kw.setdefault('co12', 12.11e10)
        kw.setdefault('co13', 10.51e10)
        kw.setdefault('co33', 21.09e10)
        kw.setdefault('co44', 4.247e10)
        super(ZnO, self).__init__(*args, **kw)

class CdS(Hexagonal):
    def __init__(self, *args, **kw):
        kw.setdefault('rho', 4820.0)
        kw.setdefault('co11', 9.07e10)
        kw.setdefault('co12', 5.81e10)
        kw.setdefault('co13', 5.1e10)
        kw.setdefault('co33', 9.38e10)
        kw.setdefault('co44', 1.504e10)
        super(CdS, self).__init__(*args, **kw)

class Zinc(Hexagonal):
    def __init__(self, *args, **kw):
        kw.setdefault('rho', 7.1*1.e-3/(1.e-2**3))
        kw.setdefault('co11', 14.3e11*1.e-5/(1.e-2**2))
        kw.setdefault('co12', 1.7e11*1.e-5/(1.e-2**2))
        kw.setdefault('co13', 3.3e11*1.e-5/(1.e-2**2))
        kw.setdefault('co33', 5.0e11*1.e-5/(1.e-2**2))
        kw.setdefault('co44', 4.0e11*1.e-5/(1.e-2**2))
        super(Zinc, self).__init__(*args, **kw)

class Beryl(Hexagonal):
    def __init__(self, *args, **kw):
        kw.setdefault('rho', 2.7*1.e-3/(1.e-2**3))
        kw.setdefault('co11', 26.94e11*1.e-5/(1.e-2**2))
        kw.setdefault('co12', 9.61e11*1.e-5/(1.e-2**2))
        kw.setdefault('co13', 6.61e11*1.e-5/(1.e-2**2))
        kw.setdefault('co33', 23.63e11*1.e-5/(1.e-2**2))
        kw.setdefault('co44', 6.53e11*1.e-5/(1.e-2**2))
        super(Beryl, self).__init__(*args, **kw)

class Albite(Triclinic):
    def __init__(self, *args, **kw):
        #kw.setdefault('rho', )
        kw.setdefault('co11', 69.9e9)
        kw.setdefault('co22', 183.5e9)
        kw.setdefault('co33', 179.5e9)
        kw.setdefault('co44', 24.9e9)
        kw.setdefault('co55', 26.8e9)
        kw.setdefault('co66', 33.5e9)
        kw.setdefault('co12', 34.0e9)
        kw.setdefault('co13', 30.8e9)
        kw.setdefault('co14', 5.1e9)
        kw.setdefault('co15', -2.4e9)
        kw.setdefault('co16', -0.9e9)
        kw.setdefault('co23', 5.5e9)
        kw.setdefault('co24', -3.9e9)
        kw.setdefault('co25', -7.7e9)
        kw.setdefault('co26', -5.8e9)
        kw.setdefault('co34', -8.7e9)
        kw.setdefault('co35', 7.1e9)
        kw.setdefault('co36', -9.8e9)
        kw.setdefault('co45', -2.4e9)
        kw.setdefault('co46', -7.2e9)
        kw.setdefault('co56', 0.5e9)
        super(Albite, self).__init__(*args, **kw)

class Acmite(Monoclinic):
    def __init__(self, *args, **kw):
        kw.setdefault('rho', 3.5e3)
        kw.setdefault('co11', 185.8e9)
        kw.setdefault('co22', 181.3e9)
        kw.setdefault('co33', 234.4e9)
        kw.setdefault('co44', 62.9e9)
        kw.setdefault('co55', 51.0e9)
        kw.setdefault('co66', 47.4e9)
        kw.setdefault('co12', 68.5e9)
        kw.setdefault('co13', 70.7e9)
        kw.setdefault('co15', 9.8e9)
        kw.setdefault('co23', 62.9e9)
        kw.setdefault('co25', 9.4e9)
        kw.setdefault('co35', 21.4e9)
        kw.setdefault('co46', 7.7e9)
        super(Acmite, self).__init__(*args, **kw)

class AlphaUranium(Orthorhombic):
    def __init__(self, *args, **kw):
        #kw.setdefault('rho', )
        kw.setdefault('rho', 8.2e3) # a false value.
        kw.setdefault('co11', 215.e9)
        kw.setdefault('co22', 199.e9)
        kw.setdefault('co33', 267.e9)
        kw.setdefault('co44', 124.e9)
        kw.setdefault('co55', 73.e9)
        kw.setdefault('co66', 74.e9)
        kw.setdefault('co12', 46.e9)
        kw.setdefault('co13', 22.e9)
        kw.setdefault('co23', 107.e9)
        super(AlphaUranium, self).__init__(*args, **kw)

class BariumTitanate(Tetragonal):
    def __init__(self, *args, **kw):
        kw.setdefault('rho', 6.2e3)
        kw.setdefault('co11', 275.0e9)
        kw.setdefault('co33', 165.0e9)
        kw.setdefault('co44', 54.3e9)
        kw.setdefault('co66', 113.0e9)
        kw.setdefault('co12', 179.0e9)
        kw.setdefault('co13', 151.0e9)
        super(BariumTitanate, self).__init__(*args, **kw)

class AlphaQuartz(Trigonal):
    def __init__(self, *args, **kw):
        kw.setdefault('rho', 2.651e3)
        kw.setdefault('co11', 87.6e9)
        kw.setdefault('co33', 106.8e9)
        kw.setdefault('co44', 57.2e9)
        kw.setdefault('co12', 6.1e9)
        kw.setdefault('co13', 13.3e9)
        kw.setdefault('co14', 17.3e9)
        super(AlphaQuartz, self).__init__(*args, **kw)

class RickerSample(Isotropic):
    def __init__(self, *args, **kw):
        kw.setdefault('rho', 2200.e0)
        kw.setdefault('co11', 3200.e0**2*2200.e0)
        kw.setdefault('co44', 1847.5e0**2*2200.e0)
        super(RickerSample, self).__init__(*args, **kw)
class RickerSampleLight(Isotropic):
    def __init__(self, *args, **kw):
        scale = 1.e-3
        kw.setdefault('rho', 2200.e0*scale)
        kw.setdefault('co11', 3200.e0**2*2200.e0*scale)
        kw.setdefault('co44', 1847.5e0**2*2200.e0*scale)
        super(RickerSampleLight, self).__init__(*args, **kw)