import cvxopt
import cvxopt
from cvxopt import solvers
import random
import matplotlib.pyplot as plt
import torch
from autograd import grad
from autograd import jacobian
import autograd.numpy as np
import autograd.numpy as jnp
import scipy.optimize as optim
from scipy.optimize import minimize, Bounds,LinearConstraint
from scipy.optimize import LinearConstraint,NonlinearConstraint
from scipy.optimize import BFGS

t = 0.02
q_prev = None

device = 'cpu'
model = torch.load('models/model_750_model_epoch_20000.pth', map_location=torch.device('cpu'))  # loaded trained model
q_dim = 6  # q_dim is the dimension of joint space
q_dim_changed = int(0.5 * q_dim)

#value function defnation
weight = []
for key in (model.keys()):
    # print(key)
    weight.append(model[key].cpu().numpy())  # load weight and bias

def leaky_relu(z):
    return np.maximum(0.01 * z, z)

def softplus(z, beta=1):
    return (1 / beta) * np.log(1 + np.exp(z * beta))

def assemble_lower_triangular_matrix(Lo, Ld):
    Lo = Lo.squeeze(0)
    Ld = Ld.squeeze(0)

    assert (2 * Lo.shape[0] == (Ld.shape[0] ** 2 - Ld.shape[0]))
    # pdb.set_trace()
    #     diagonal_matrix = np.diagflat(Ld)
    diagonal_matrix = np.identity(len(Ld)) * np.outer(np.ones(len(Ld)), Ld)

    L = np.tril(np.ones(diagonal_matrix.shape)) - np.eye(q_dim_changed)

    # Set off diagonals

    L = np.array([[0, 0, 0], [1, 0, 0], [0, 0, 0]]) * Lo.reshape(3)[0] + np.array([[0, 0, 0], [0, 0, 0], [1, 0, 0]]) * \
        Lo.reshape(3)[1] + np.array([[0, 0, 0], [0, 0, 0], [0, 1, 0]]) * Lo.reshape(3)[2]
    # print("L now : ", L)
    # Add diagonals
    L = L + diagonal_matrix
    return L

def value(x1):
    global weight, goal
    fc1_w = weight[0]
    fc1_b = weight[1]
    fc2_w = weight[2]
    fc2_b = weight[3]
    fc_Ld_w = weight[4]
    fc_Ld_b = weight[5]
    fc_Lo_w = weight[6]
    fc_Lo_b = weight[7]
    #pdb.set_trace()
    net_input = np.concatenate([np.squeeze(x1), np.squeeze(goal)], axis=0)
    net_input = np.array([net_input])

    z1 = np.dot(net_input, fc1_w.transpose()) + fc1_b
    hidden1 = leaky_relu(z1)
    z2 = np.dot(hidden1, fc2_w.transpose()) + fc2_b
    hidden2 = leaky_relu(z2)
    hidden3 = np.dot(hidden2, fc_Ld_w.transpose()) + fc_Ld_b
    Ld = softplus(hidden3)
    Lo = np.dot(hidden2, fc_Lo_w.transpose()) + fc_Lo_b
    L = assemble_lower_triangular_matrix(Lo, Ld)

    H = L @ L.transpose() + 1e-9 * np.eye(3)
    return H

grad_value = grad(value)
jac_value= jacobian(value)

def fk_franka(q):
    q_1 = q[0]
    q_2 = q[1]
    q_3 = q[2]
    q_4 = q[3]
    q_5 = q[4]
    q_6 = q[5]
    q_7 = q[6]

    x = 0.0825 * jnp.cos(q_1) * jnp.cos(q_2) * jnp.cos(q_3) + 0.384 * jnp.cos(q_1) * jnp.cos(q_4) * jnp.sin(
        q_2) - 0.0825 * jnp.cos(q_1) * jnp.sin(q_2) * jnp.sin(q_4) + 0.316 * jnp.cos(q_1) * jnp.sin(
        q_2) - 0.0825 * jnp.cos(q_4) * (
                jnp.cos(q_1) * jnp.cos(q_2) * jnp.cos(q_3) - jnp.sin(q_1) * jnp.sin(q_3)) + 0.088 * jnp.cos(q_6) * (
                jnp.cos(q_5) * (jnp.cos(q_1) * jnp.sin(q_2) * jnp.sin(q_4) + jnp.cos(q_4) * (
                jnp.cos(q_1) * jnp.cos(q_2) * jnp.cos(q_3) - jnp.sin(q_1) * jnp.sin(q_3))) - jnp.sin(q_5) * (
                        jnp.cos(q_1) * jnp.cos(q_2) * jnp.sin(q_3) + jnp.cos(q_3) * jnp.sin(q_1))) - 0.21 * jnp.cos(
        q_6) * (
                jnp.cos(q_1) * jnp.cos(q_4) * jnp.sin(q_2) - jnp.sin(q_4) * (
                jnp.cos(q_1) * jnp.cos(q_2) * jnp.cos(q_3) - jnp.sin(q_1) * jnp.sin(q_3))) - 0.0825 * jnp.sin(
        q_1) * jnp.sin(q_3) - 0.384 * jnp.sin(q_4) * (
                jnp.cos(q_1) * jnp.cos(q_2) * jnp.cos(q_3) - jnp.sin(q_1) * jnp.sin(q_3)) + 0.21 * jnp.sin(q_6) * (
                jnp.cos(q_5) * (jnp.cos(q_1) * jnp.sin(q_2) * jnp.sin(q_4) + jnp.cos(q_4) * (
                jnp.cos(q_1) * jnp.cos(q_2) * jnp.cos(q_3) - jnp.sin(q_1) * jnp.sin(q_3))) - jnp.sin(q_5) * (
                        jnp.cos(q_1) * jnp.cos(q_2) * jnp.sin(q_3) + jnp.cos(q_3) * jnp.sin(q_1))) + 0.088 * jnp.sin(
        q_6) * (
                jnp.cos(q_1) * jnp.cos(q_4) * jnp.sin(q_2) - jnp.sin(q_4) * (
                jnp.cos(q_1) * jnp.cos(q_2) * jnp.cos(q_3) - jnp.sin(q_1) * jnp.sin(q_3)))
    y = 0.0825 * jnp.cos(q_1) * jnp.sin(q_3) + 0.0825 * jnp.cos(q_2) * jnp.cos(q_3) * jnp.sin(q_1) + 0.384 * jnp.cos(
        q_4) * jnp.sin(q_1) * jnp.sin(q_2) - 0.0825 * jnp.cos(q_4) * (
                jnp.cos(q_1) * jnp.sin(q_3) + jnp.cos(q_2) * jnp.cos(q_3) * jnp.sin(q_1)) + 0.088 * jnp.cos(q_6) * (
                jnp.cos(q_5) * (
                jnp.cos(q_4) * (jnp.cos(q_1) * jnp.sin(q_3) + jnp.cos(q_2) * jnp.cos(q_3) * jnp.sin(q_1)) + jnp.sin(
            q_1) * jnp.sin(q_2) * jnp.sin(q_4)) + jnp.sin(q_5) * (
                        jnp.cos(q_1) * jnp.cos(q_3) - jnp.cos(q_2) * jnp.sin(q_1) * jnp.sin(q_3))) - 0.21 * jnp.cos(
        q_6) * (jnp.cos(q_4) * jnp.sin(q_1) * jnp.sin(q_2) - jnp.sin(q_4) * (
            jnp.cos(q_1) * jnp.sin(q_3) + jnp.cos(q_2) * jnp.cos(q_3) * jnp.sin(q_1))) - 0.0825 * jnp.sin(
        q_1) * jnp.sin(q_2) * jnp.sin(q_4) + 0.316 * jnp.sin(q_1) * jnp.sin(q_2) - 0.384 * jnp.sin(q_4) * (
                jnp.cos(q_1) * jnp.sin(q_3) + jnp.cos(q_2) * jnp.cos(q_3) * jnp.sin(q_1)) + 0.21 * jnp.sin(q_6) * (
                jnp.cos(q_5) * (
                jnp.cos(q_4) * (jnp.cos(q_1) * jnp.sin(q_3) + jnp.cos(q_2) * jnp.cos(q_3) * jnp.sin(q_1)) + jnp.sin(
            q_1) * jnp.sin(q_2) * jnp.sin(q_4)) + jnp.sin(q_5) * (
                        jnp.cos(q_1) * jnp.cos(q_3) - jnp.cos(q_2) * jnp.sin(q_1) * jnp.sin(q_3))) + 0.088 * jnp.sin(
        q_6) * (
                jnp.cos(q_4) * jnp.sin(q_1) * jnp.sin(q_2) - jnp.sin(q_4) * (
                jnp.cos(q_1) * jnp.sin(q_3) + jnp.cos(q_2) * jnp.cos(q_3) * jnp.sin(q_1)))
    z = 0.384 * jnp.cos(q_2) * jnp.cos(q_4) - 0.0825 * jnp.cos(q_2) * jnp.sin(q_4) + 0.316 * jnp.cos(
        q_2) + 0.0825 * jnp.cos(q_3) * jnp.cos(q_4) * jnp.sin(q_2) + 0.384 * jnp.cos(q_3) * jnp.sin(q_2) * jnp.sin(
        q_4) - 0.0825 * jnp.cos(q_3) * jnp.sin(q_2) - 0.21 * jnp.cos(q_6) * (
                jnp.cos(q_2) * jnp.cos(q_4) + jnp.cos(q_3) * jnp.sin(q_2) * jnp.sin(q_4)) + 0.088 * jnp.cos(q_6) * (
                jnp.cos(q_5) * (jnp.cos(q_2) * jnp.sin(q_4) - jnp.cos(q_3) * jnp.cos(q_4) * jnp.sin(q_2)) + jnp.sin(
            q_2) * jnp.sin(q_3) * jnp.sin(q_5)) + 0.088 * jnp.sin(q_6) * (
                jnp.cos(q_2) * jnp.cos(q_4) + jnp.cos(q_3) * jnp.sin(q_2) * jnp.sin(q_4)) + 0.21 * jnp.sin(q_6) * (
                jnp.cos(q_5) * (jnp.cos(q_2) * jnp.sin(q_4) - jnp.cos(q_3) * jnp.cos(q_4) * jnp.sin(q_2)) + jnp.sin(
            q_2) * jnp.sin(q_3) * jnp.sin(q_5)) + 0.33
    cartpos = np.array([x,y,z])
    return cartpos
jac_fk = jacobian(fk_franka)

def traj_cost(trajectory):
    cost = 0
    cart_cost = 0
    for i in range(len(trajectory) - 1):
        cost += np.linalg.norm(np.asarray(trajectory[i]) - np.asarray(trajectory[i + 1]), ord=2)
        # pdb.set_trace()
        current = np.asarray(fk_franka(trajectory[i]))
        next = np.asarray(fk_franka(trajectory[i+1]))
        cart_cost += np.linalg.norm(current - next, ord=2)
    return cost, cart_cost

def constraintfxn(qdot_x_next ):
    global t, q_prev
    q = qdot_x_next[:7]*t  + q_prev
    q_1 = q[0]
    q_2 = q[1]
    q_3 = q[2]
    q_4 = q[3]
    q_5 = q[4]
    q_6 = q[5]
    q_7 = q[6]
    x_next = qdot_x_next[7]
    y_next = qdot_x_next[8]
    z_next = qdot_x_next[9]

    x = 0.0825 * jnp.cos(q_1) * jnp.cos(q_2) * jnp.cos(q_3) + 0.384 * jnp.cos(q_1) * jnp.cos(q_4) * jnp.sin(
        q_2) - 0.0825 * jnp.cos(q_1) * jnp.sin(q_2) * jnp.sin(q_4) + 0.316 * jnp.cos(q_1) * jnp.sin(
        q_2) - 0.0825 * jnp.cos(q_4) * (
                jnp.cos(q_1) * jnp.cos(q_2) * jnp.cos(q_3) - jnp.sin(q_1) * jnp.sin(q_3)) + 0.088 * jnp.cos(q_6) * (
                jnp.cos(q_5) * (jnp.cos(q_1) * jnp.sin(q_2) * jnp.sin(q_4) + jnp.cos(q_4) * (
                jnp.cos(q_1) * jnp.cos(q_2) * jnp.cos(q_3) - jnp.sin(q_1) * jnp.sin(q_3))) - jnp.sin(q_5) * (
                        jnp.cos(q_1) * jnp.cos(q_2) * jnp.sin(q_3) + jnp.cos(q_3) * jnp.sin(q_1))) - 0.21 * jnp.cos(
        q_6) * (
                jnp.cos(q_1) * jnp.cos(q_4) * jnp.sin(q_2) - jnp.sin(q_4) * (
                jnp.cos(q_1) * jnp.cos(q_2) * jnp.cos(q_3) - jnp.sin(q_1) * jnp.sin(q_3))) - 0.0825 * jnp.sin(
        q_1) * jnp.sin(q_3) - 0.384 * jnp.sin(q_4) * (
                jnp.cos(q_1) * jnp.cos(q_2) * jnp.cos(q_3) - jnp.sin(q_1) * jnp.sin(q_3)) + 0.21 * jnp.sin(q_6) * (
                jnp.cos(q_5) * (jnp.cos(q_1) * jnp.sin(q_2) * jnp.sin(q_4) + jnp.cos(q_4) * (
                jnp.cos(q_1) * jnp.cos(q_2) * jnp.cos(q_3) - jnp.sin(q_1) * jnp.sin(q_3))) - jnp.sin(q_5) * (
                        jnp.cos(q_1) * jnp.cos(q_2) * jnp.sin(q_3) + jnp.cos(q_3) * jnp.sin(q_1))) + 0.088 * jnp.sin(
        q_6) * (
                jnp.cos(q_1) * jnp.cos(q_4) * jnp.sin(q_2) - jnp.sin(q_4) * (
                jnp.cos(q_1) * jnp.cos(q_2) * jnp.cos(q_3) - jnp.sin(q_1) * jnp.sin(q_3)))
    y = 0.0825 * jnp.cos(q_1) * jnp.sin(q_3) + 0.0825 * jnp.cos(q_2) * jnp.cos(q_3) * jnp.sin(q_1) + 0.384 * jnp.cos(
        q_4) * jnp.sin(q_1) * jnp.sin(q_2) - 0.0825 * jnp.cos(q_4) * (
                jnp.cos(q_1) * jnp.sin(q_3) + jnp.cos(q_2) * jnp.cos(q_3) * jnp.sin(q_1)) + 0.088 * jnp.cos(q_6) * (
                jnp.cos(q_5) * (
                jnp.cos(q_4) * (jnp.cos(q_1) * jnp.sin(q_3) + jnp.cos(q_2) * jnp.cos(q_3) * jnp.sin(q_1)) + jnp.sin(
            q_1) * jnp.sin(q_2) * jnp.sin(q_4)) + jnp.sin(q_5) * (
                        jnp.cos(q_1) * jnp.cos(q_3) - jnp.cos(q_2) * jnp.sin(q_1) * jnp.sin(q_3))) - 0.21 * jnp.cos(
        q_6) * (jnp.cos(q_4) * jnp.sin(q_1) * jnp.sin(q_2) - jnp.sin(q_4) * (
            jnp.cos(q_1) * jnp.sin(q_3) + jnp.cos(q_2) * jnp.cos(q_3) * jnp.sin(q_1))) - 0.0825 * jnp.sin(
        q_1) * jnp.sin(q_2) * jnp.sin(q_4) + 0.316 * jnp.sin(q_1) * jnp.sin(q_2) - 0.384 * jnp.sin(q_4) * (
                jnp.cos(q_1) * jnp.sin(q_3) + jnp.cos(q_2) * jnp.cos(q_3) * jnp.sin(q_1)) + 0.21 * jnp.sin(q_6) * (
                jnp.cos(q_5) * (
                jnp.cos(q_4) * (jnp.cos(q_1) * jnp.sin(q_3) + jnp.cos(q_2) * jnp.cos(q_3) * jnp.sin(q_1)) + jnp.sin(
            q_1) * jnp.sin(q_2) * jnp.sin(q_4)) + jnp.sin(q_5) * (
                        jnp.cos(q_1) * jnp.cos(q_3) - jnp.cos(q_2) * jnp.sin(q_1) * jnp.sin(q_3))) + 0.088 * jnp.sin(
        q_6) * (
                jnp.cos(q_4) * jnp.sin(q_1) * jnp.sin(q_2) - jnp.sin(q_4) * (
                jnp.cos(q_1) * jnp.sin(q_3) + jnp.cos(q_2) * jnp.cos(q_3) * jnp.sin(q_1)))
    z = 0.384 * jnp.cos(q_2) * jnp.cos(q_4) - 0.0825 * jnp.cos(q_2) * jnp.sin(q_4) + 0.316 * jnp.cos(
        q_2) + 0.0825 * jnp.cos(q_3) * jnp.cos(q_4) * jnp.sin(q_2) + 0.384 * jnp.cos(q_3) * jnp.sin(q_2) * jnp.sin(
        q_4) - 0.0825 * jnp.cos(q_3) * jnp.sin(q_2) - 0.21 * jnp.cos(q_6) * (
                jnp.cos(q_2) * jnp.cos(q_4) + jnp.cos(q_3) * jnp.sin(q_2) * jnp.sin(q_4)) + 0.088 * jnp.cos(q_6) * (
                jnp.cos(q_5) * (jnp.cos(q_2) * jnp.sin(q_4) - jnp.cos(q_3) * jnp.cos(q_4) * jnp.sin(q_2)) + jnp.sin(
            q_2) * jnp.sin(q_3) * jnp.sin(q_5)) + 0.088 * jnp.sin(q_6) * (
                jnp.cos(q_2) * jnp.cos(q_4) + jnp.cos(q_3) * jnp.sin(q_2) * jnp.sin(q_4)) + 0.21 * jnp.sin(q_6) * (
                jnp.cos(q_5) * (jnp.cos(q_2) * jnp.sin(q_4) - jnp.cos(q_3) * jnp.cos(q_4) * jnp.sin(q_2)) + jnp.sin(
            q_2) * jnp.sin(q_3) * jnp.sin(q_5)) + 0.33
    pos_residual = np.array([x - x_next,y-y_next,z - z_next])
    return pos_residual
jac_constraint = jacobian(constraintfxn)

def costfxn(solverVariable,x_pos,goal):
  diff = solverVariable[7:] - goal
  v = value(x_pos)
  cost = np.matmul(diff.transpose(),np.matmul(v, diff ))
  w_des_vel = 0.002
  smoothness_cost = np.sum(solverVariable[0:7]**2,axis = 0)
  #print(f"cost = {cost}, smoothness = {smoothness_cost}")
  return np.add(cost , w_des_vel*smoothness_cost)
jac_cost = jacobian(costfxn)   #jaccost has a shape of (3,)

def trajMetricBased(init_joint,start_cart,end_cartesian):
    global t,q_prev,goal
    # print("Start : ",start_cart)
    # print("Goal : ", end_cartesian)
    goal = np.squeeze(end_cartesian)
    x_des  = end_cartesian
    num_dof = 7
    qdot = np.zeros(num_dof)  ########## initial velocity
    qdotprev = np.zeros(num_dof)
    q_min = np.array([-165, -100, -165, -165, -165, -1.0, -165]) * np.pi / 180
    #q_min = q_min.reshape(7,1)
    q_max = np.array([165, 101, 165, 1.0, 165, 214, 165]) * np.pi / 180
    #q_max = q_max.reshape(7,1)

    qdot_max = np.array([2.1750	,2.1750	,2.1750	,2.1750,	2.6100,	2.6100,	2.6100])
    qdot_min = -1*qdot_max
    qacc_max = np.array([15,	7.5,	10	,12.5	,15	,20,	20])
    qacc_min = -1*qacc_max

    x_pos = np.asarray(fk_franka(init_joint))  ########### initial positions

    q = init_joint
    q_prev = init_joint
      ######### delta t for which the computed velocity is commanded to the robot 3msec
    x_next = fk_franka(qdot*t + q)

    mpc_iter = 200

    q_1 = []
    q_2 = []
    q_3 = []
    q_4 = []
    q_5 = []
    q_6 = []
    q_7 = []
    cost_tracker = np.zeros(mpc_iter)
    trajectory = []
    qdottrajc =[]
    #pdb.set_trace()
    solverVariable = np.hstack((qdot,x_next))
    x_min = -np.ones(3,)*np.Inf #fk_franka(q_min)
    x_max = np.ones(3,)*np.Inf #fk_franka(q_max)
    solver_minbounds= np.hstack((qdot_min , x_min))
    solver_maxbounds = np.hstack((qdot_max , x_max))
    Amat = np.identity(10)
    Bmat = np.identity(10)
    for i in range(0, mpc_iter):
        # print(f"mpc-itr={i}")
        if np.linalg.norm(x_pos - x_des) < 0.01:
            break
            #t = 0.003
            # print("Without value velocity",v_des)

        bnds = Bounds(solver_minbounds,solver_maxbounds)
        nonlinear_constraint = NonlinearConstraint(constraintfxn , np.zeros((3,)), np.zeros((3,)) , jac= jac_constraint, hess=BFGS() , )

        linear_constraint_A = LinearConstraint(Amat*t, np.hstack((q_min-q_prev,x_min)), np.hstack((q_max-q_prev, x_max )))

        linear_constraint_B = LinearConstraint(Bmat ,np.hstack((qacc_min*t + qdotprev,x_min)), np.hstack((qacc_max*t + qdotprev,x_max)))
        defaultopts={ 'maxiter': 100, 'ftol': 1e-06, 'iprint': 1, 'disp': False, 'eps': 1.4901161193847656e-08, 'finite_diff_rel_step': None}
        res = minimize(costfxn ,  solverVariable, args =(x_pos,goal) , method='SLSQP', jac=jac_cost,
                constraints=[nonlinear_constraint,linear_constraint_A, linear_constraint_B], options=defaultopts ,bounds=bnds)  #TODO include linear constraints and smoothness cost
        #changing args for slsq might help ? like tol ??
        solverVariable = np.asarray(res['x']).squeeze()
        cost_tracker[i] = np.linalg.norm(np.hstack((x_pos[0] - x_des[0], x_pos[1] - x_des[1], x_pos[2] - x_des[2])))
        trajectory.append([q[0], q[1], q[2], q[3], q[4], q[5], q[6]])
        q = q + solverVariable[0:7] * t
        x_pos =  solverVariable[7:]
        qdotprev = solverVariable[0:7]
        q_prev = q
        qdottrajc.append(qdotprev)
    JointCost, CartCost = traj_cost(trajectory)
    return JointCost, CartCost, cost_tracker, np.array(trajectory)


import warnings
import numpy as np
from numpy.linalg import LinAlgError
from scipy.optimize._numdiff import approx_derivative
import scipy.stats as stats

from refnx.util import ErrorProp as EP
from refnx._lib import flatten, approx_hess2
from refnx._lib import unique as f_unique
from refnx.dataset import Data1D
from refnx.analysis import (
    is_parameter,
    Parameter,
    possibly_create_parameter,
    is_parameters,
    Parameters,
    Interval,
    PDF,
)

class BaseObjective:
    """Don't necessarily have to use Parameters, could use np.array"""

    def __init__(
        self,
        p,
        logl,
        logp=None,
        fcn_args=(),
        fcn_kwds=None,
        name=None,
        weighted=True,
    ):
        self.name = name
        self.parameters = p
        self.nvary = len(p)
        self._logl = logl
        self._logp = logp
        self.fcn_args = fcn_args
        self.fcn_kwds = {}
        # give the BaseObjective a default value, so that it can be used in a
        # GlobalObjective
        self.weighted = weighted
        if fcn_kwds is not None:
            self.fcn_kwds = fcn_kwds

    def setp(self, pvals):
        """
        Set the parameters from pvals

        Parameters
        ----------
        pvals : np.ndarray
            Array containing the values to be tested.

        """
        self.parameters[:] = pvals

    def nll(self, pvals=None):
        """
        Negative log-likelihood function

        Parameters
        ----------
        pvals : np.ndarray
            Array containing the values to be tested.

        Returns
        -------
        nll : float
            negative log-likelihood

        """
        vals = self.parameters
        if pvals is not None:
            vals = pvals

        return -self.logl(vals)

    def logp(self, pvals=None):
        """
        Log-prior probability function

        Parameters
        ----------
        pvals : np.ndarray
            Array containing the values to be tested.

        Returns
        -------
        logp : float
            log-prior probability

        """
        vals = self.parameters
        if pvals is not None:
            vals = pvals

        if callable(self._logp):
            return self._logp(vals, *self.fcn_args, **self.fcn_kwds)
        return 0

    def logl(self, pvals=None):
        """
        Log-likelihood probability function

        Parameters
        ----------
        pvals : np.ndarray
            Array containing the values to be tested.

        Returns
        -------
        logl : float
            log-likelihood probability.

        """
        vals = self.parameters
        if pvals is not None:
            vals = pvals

        return self._logl(vals, *self.fcn_args, **self.fcn_kwds)

    def logpost(self, pvals=None):
        """
        Log-posterior probability function

        Parameters
        ----------
        pvals : np.ndarray
            Array containing the values to be tested.

        Returns
        -------
        logpost : float
            log-probability.

        Notes
        -----
        The log probability is the sum is the sum of the log-prior and
        log-likelihood probabilities. Does not set the parameter attribute.

        """
        vals = self.parameters
        if pvals is not None:
            vals = pvals

        logpost = self.logp(vals)
        if not np.isfinite(logpost):
            return -np.inf
        logpost += self.logl(vals)
        return logpost

    def nlpost(self, pvals=None):
        """
        Negative log-posterior function

        Parameters
        ----------
        pvals : np.ndarray
            Array containing the values to be tested.

        Returns
        -------
        nlpost : float
            negative log-posterior

        """
        return -self.logpost(pvals)

    def varying_parameters(self):
        """
        Returns
        -------
        varying_parameters : np.ndarray
            The parameters varying in this objective function.

        """
        return self.parameters

    def covar(self, target="nll"):
        """
        Estimates a covariance matrix based on numerical differentiation
        of either the negative log-likelihood or negative log-posterior
        probabilities.

        Parameters
        ----------
        target : str, {"nll", "nlpost"}

        Returns
        -------
        covar : np.ndarray
            The covariance matrix for the fitting system

        Notes
        -----
        Estimation of a covariance matrix can be susceptible to numeric
        instabilities. Critically evaluate the matrix before further use.
        """
        _pvals = np.array(self.varying_parameters())

        if target == "nll":
            fn = self.nll
        elif target == "nlpost":
            fn = self.nlpost

        try:
            # from statsmodels
            # the output from this for the test in test_objective.covar
            # is very similar to numdifftools.Hessian, or a chained version
            # of approx_derivative
            hess = approx_hess2(_pvals, fn)
            covar = np.linalg.inv(hess)
        except LinAlgError:
            sz = np.size(_pvals, 0)
            covar = np.full((sz, sz), np.inf)
        finally:
            self.setp(_pvals)

        return covar

class Objective(BaseObjective):
    """
    Objective function for using with curvefitters such as
    `refnx.analysis.curvefitter.CurveFitter`.

    Parameters
    ----------
    model : refnx.analysis.Model
        the generative model function. One can also provide an object that
        inherits `refnx.analysis.Model`.
    data : refnx.dataset.Data1D
        data to be analysed.
    lnsigma : float or refnx.analysis.Parameter, optional
        Used if the  experimental uncertainty (`data.y_err`) underestimated by
        a constant fractional amount. The experimental uncertainty is modified
        as:

        `s_n**2 = y_err**2 + exp(lnsigma * 2) * model**2`

        See `Objective.logl` for more details.
    use_weights : bool
        use experimental uncertainty in calculation of residuals and
        logl, if available. If this is set to False, then you should also
        set `self.lnsigma.vary = False`, it will have no effect on the fit.
    transform : callable, optional
        the model, data and data uncertainty are transformed by this
        function before calculating the likelihood/residuals. Has the
        signature `transform(data.x, y, y_err=None)`, returning the tuple
        (`transformed_y, transformed_y_err`).
    logp_extra : callable, optional
        user specifiable log-probability term. This contribution is in
        addition to the log-prior term of the `model` parameters, and
        `model.logp`, as well as the log-likelihood of the `data`. Has
        signature:
        `logp_extra(model, data)`. The `model` will already possess
        updated parameters. Beware of including the same log-probability
        terms more than once.
    auxiliary_params : {sequence, Parameters}, optional
        Extra Parameter objects that are involved with curvefitting, but
        aren't directly included as part of the `model`. See notes for more
        details.
    name : str
        Name for the objective.

    Notes
    -----
    For parallelisation `logp_extra` needs to be picklable.

    `auxiliary_params` are included in calculating the `Objective.logp`
    term, are present in `Objective.varying_parameters()`, and are modified by
    Curvefitter during an analysis. Their main purpose is to aid in making
    constraints in models.
    """

    def __init__(
        self,
        model,
        data,
        lnsigma=None,
        use_weights=True,
        transform=None,
        logp_extra=None,
        auxiliary_params=(),
        name=None,
    ):
        self.model = model
        # should be a Data1D instance
        if isinstance(data, Data1D):
            self.data = data
        else:
            self.data = Data1D(data=data)

        self.lnsigma = lnsigma
        if lnsigma is not None:
            self.lnsigma = possibly_create_parameter(lnsigma, "lnsigma")

        if isinstance(auxiliary_params, Parameters):
            self.auxiliary_params = auxiliary_params
        else:
            self.auxiliary_params = Parameters(auxiliary_params)

        self._use_weights = use_weights
        self.transform = transform
        self.logp_extra = logp_extra
        self.name = name
        if name is None:
            self.name = id(self)

    def __str__(self):
        s = ["{:_>80}".format("")]
        s.append("Objective - {0}".format(self.name))

        # dataset name
        if self.data.name is None:
            s.append("Dataset = {0}".format(self.data))
        else:
            s.append("Dataset = {0}".format(self.data.name))

        s.append("datapoints = {0}".format(self.npoints))
        s.append("chi2 = {0}".format(self.chisqr()))
        s.append("Weighted = {0}".format(self.weighted))
        s.append("Transform = {0}".format(self.transform))
        s.append(str(self.parameters))

        return "\n".join(s)

    def __repr__(self):
        return (
            "Objective({model!r}, {data!r},"
            " lnsigma={lnsigma!r},"
            " use_weights={_use_weights},"
            " transform={transform!r},"
            " logp_extra={logp_extra!r},"
            " name={name!r})".format(**self.__dict__)
        )

    @property
    def weighted(self):
        """
        **bool** Does the data have weights (`data.y_err`), and is the
        objective using them?

        """
        return self.data.weighted and self._use_weights

    @weighted.setter
    def weighted(self, use_weights):
        self._use_weights = bool(use_weights)

    @property
    def npoints(self):
        """
        **int** the number of points in the dataset.

        """
        return self.data.y.size

    def varying_parameters(self):
        """
        Returns
        -------
        varying_parameters : refnx.analysis.Parameters
            The varying Parameter objects allowed to vary during the fit.

        """
        # create and return a Parameters object because it has the
        # __array__ method, which allows one to quickly get numerical values.
        p = Parameters()
        p.data = list(f_unique(p for p in flatten(self.parameters) if p.vary))
        return p

    def _data_transform(self, model=None):
        """
        Returns
        -------
        y, y_err, model: tuple of np.ndarray
            The y data, its uncertainties, and the model, all put through the
             transform.
        """
        x = self.data.x
        y = self.data.y

        y_err = 1.0
        if self.weighted:
            y_err = self.data.y_err

        if self.transform is None:
            return y, y_err, model
        else:
            if model is not None:
                model, _ = self.transform(x, model)

            y, y_err = self.transform(x, y, y_err)
            if self.weighted:
                return y, y_err, model
            else:
                return y, 1, model

    def generative(self, pvals=None):
        """
        Calculate the generative (dependent variable) function associated with
        the model.

        Parameters
        ----------
        pvals : array-like or refnx.analysis.Parameters
            values for the varying or entire set of parameters

        Returns
        -------
        model : np.ndarray

        """
        self.setp(pvals)
        return self.model(self.data.x, x_err=self.data.x_err)

    def residuals(self, pvals=None):
        """
        Calculates the residuals for a given fitting system.

        Parameters
        ----------
        pvals : array-like or refnx.analysis.Parameters
            values for the varying or entire set of parameters

        Returns
        -------
        residuals : np.ndarray
            Residuals, `(data.y - model) / y_err`.

        """
        self.setp(pvals)

        model = self.model(self.data.x, x_err=self.data.x_err)
        # TODO add in varying parameter residuals? (z-scores...)

        y, y_err, model = self._data_transform(model)

        if self.lnsigma is not None:
            s_n = np.sqrt(
                y_err * y_err + np.exp(2 * float(self.lnsigma)) * model * model
            )
        else:
            s_n = y_err

        return (y - model) / s_n

    def chisqr(self, pvals=None):
        """
        Calculates the chi-squared value for a given fitting system.

        Parameters
        ----------
        pvals : array-like or refnx.analysis.Parameters
            values for the varying or entire set of parameters

        Returns
        -------
        chisqr : np.ndarray
            Chi-squared value, `np.sum(residuals**2)`.

        """
        # TODO reduced chisqr? include z-scores for parameters? DOF?
        self.setp(pvals)
        res = self.residuals(None)
        return np.sum(res * res)

    @property
    def parameters(self):
        """
        :class:`refnx.analysis.Parameters`, all the Parameters contained in the
        fitting system.

        """
        if is_parameter(self.lnsigma):
            return self.lnsigma | self.auxiliary_params | self.model.parameters
        elif len(self.auxiliary_params):
            return self.auxiliary_params | self.model.parameters
        else:
            return self.model.parameters

    def setp(self, pvals):
        """
        Set the parameters from pvals.

        Parameters
        ----------
        pvals : array-like or refnx.analysis.Parameters
            values for the varying or entire set of parameters

        """
        if pvals is None:
            return

        # set here rather than delegating to a Parameters
        # object, because it may not necessarily be a
        # Parameters object
        _varying_parameters = self.varying_parameters()
        if len(pvals) == len(_varying_parameters):
            for idx, param in enumerate(_varying_parameters):
                param.value = pvals[idx]
            return

        # values supplied are enough to specify all parameter values
        # even those that are repeated
        flattened_parameters = list(flatten(self.parameters))
        if len(pvals) == len(flattened_parameters):
            for idx, param in enumerate(flattened_parameters):
                param.value = pvals[idx]
            return

        raise ValueError(
            f"Incorrect number of values supplied ({len(pvals)})"
            f", supply either the full number of parameters"
            f" ({len(flattened_parameters)}, or only the varying"
            f" parameters ({len(_varying_parameters)})."
        )

    def prior_transform(self, u):
        """
        Calculate the prior transform of the system.

        Transforms uniform random variates in the unit hypercube,
        `u ~ uniform[0.0, 1.0)`, to the parameter space of interest, according
        to the priors on the varying parameters.

        Parameters
        ----------
        u : array-like
            Size of the varying parameters

        Returns
        -------
        pvals : array-like
            Scaled parameter values

        Notes
        -----
        If a parameter has bounds, `x ~ Unif[-10, 10)` then the scaling from
        `u` to `x` is done as follows:

        .. code-block:: python

            x = 2. * u - 1.  # scale and shift to [-1., 1.)
            x *= 10.  # scale to [-10., 10.)

        """
        var_pars = self.varying_parameters()
        pvals = np.empty(len(var_pars), dtype=np.float64)

        for i, var_par in enumerate(var_pars):
            pvals[i] = var_par.bounds.invcdf(u[i])

        return pvals

    def logp(self, pvals=None):
        """
        Calculate the log-prior of the system

        Parameters
        ----------
        pvals : array-like or refnx.analysis.Parameters
            values for the varying or entire set of parameters

        Returns
        -------
        logp : float
            log-prior probability

        Notes
        -----
        The log-prior is calculated as:

        .. code-block:: python

            logp = np.sum(param.logp() for param in
                             self.varying_parameters())

        """
        self.setp(pvals)

        logp = np.sum(
            [
                param.logp()
                for param in f_unique(
                    p for p in flatten(self.parameters) if p.vary
                )
            ]
        )

        if not np.isfinite(logp):
            return -np.inf

        return logp

    def logl(self, pvals=None):
        """
        Calculate the log-likelhood of the system

        The major component of the log-likelhood probability is from the data.
        Extra potential terms are added on from the Model, `self.model.logp`,
        and the user specifiable `logp_extra` function.

        Parameters
        ----------
        pvals : array-like or refnx.analysis.Parameters
            values for the varying or entire set of parameters

        Returns
        -------
        logl : float
            log-likelihood probability

        Notes
        -----
        The log-likelihood is calculated as:

        .. code-block:: python

            logl = -0.5 * np.sum(((y - model) / s_n)**2
                                 + np.log(2 * pi * s_n**2))
            logp += self.model.logp()
            logp += self.logp_extra(self.model, self.data)

        where

        .. code-block:: python

            s_n**2 = y_err**2 + exp(2 * lnsigma) * model**2

        """
        self.setp(pvals)

        model = self.model(self.data.x, x_err=self.data.x_err)

        logl = 0.0

        y, y_err, model = self._data_transform(model)

        if self.lnsigma is not None:
            var_y = (
                y_err * y_err + np.exp(2 * float(self.lnsigma)) * model * model
            )
        else:
            var_y = y_err ** 2

        # TODO do something sensible if data isn't weighted
        if self.weighted:
            logl += np.log(2 * np.pi * var_y)

        logl += (y - model) ** 2 / var_y

        # nans play havoc
        if np.isnan(logl).any():
            raise RuntimeError("Objective.logl encountered a NaN.")

        # add on extra 'potential' terms from the model.
        extra_potential = self.model.logp()

        if self.logp_extra is not None:
            extra_potential += self.logp_extra(self.model, self.data)

        return -0.5 * np.sum(logl) + extra_potential

    def nll(self, pvals=None):
        """
        Negative log-likelihood function

        Parameters
        ----------
        pvals : array-like or refnx.analysis.Parameters
            values for the varying or entire set of parameters

        Returns
        -------
        nll : float
            negative log-likelihood

        """
        self.setp(pvals)
        return -self.logl()

    def logpost(self, pvals=None):
        """
        Calculate the log-probability of the curvefitting system

        Parameters
        ----------
        pvals : array-like or refnx.analysis.Parameters
            values for the varying or entire set of parameters

        Returns
        -------
        logpost : float
            log-probability

        Notes
        -----
        The overall log-probability is the sum of the log-prior and
        log-likelihood. The log-likelihood is not calculated if the log-prior
        is impossible (`logp == -np.inf`).

        """
        self.setp(pvals)
        logpost = self.logp()

        # only calculate the probability if the parameters have finite
        # log-prior
        if not np.isfinite(logpost):
            return -np.inf

        logpost += self.logl()
        return logpost

    def covar(self, target="residuals"):
        """
        Estimates the covariance matrix of the Objective.

        Parameters
        ----------
        target : {"residuals", "nll", "nlpost"}
            Specifies what approach should be used to estimate covariance.

        Returns
        -------
        covar : np.ndarray
            Covariance matrix

        Notes
        -----
        For most purposes the Jacobian of the `'residuals'` should be used to
        calculate the covariance matrix, estimated as J.T x J.
        If an Objective cannot calculate residuals then the covariance matrix
        can be estimated by inverting a Hessian matrix created from either the
        `'nll'` or `'nlpost'` methods.
        The default `'residuals'` approach falls back to `'nll'` if a problem
        is experienced.
        The default `'residuals'` setting is preferred as the other settings
        can sometimes experience instabilities during Hessian estimation with
        numerical differentiation.
        """
        if target == "residuals":
            try:
                covar = self._covar_from_residuals()
            except Exception:
                # fallback to "nll"
                target = "nll"

        if target in ["nll", "nlpost"]:
            covar = super().covar(target)

        pvar = np.diagonal(covar).copy()
        psingular = np.where(pvar == 0)[0]

        if len(psingular) > 0:
            var_params = self.varying_parameters()
            singular_params = [var_params[ps] for ps in psingular]

            raise LinAlgError(
                "The following Parameters have no effect on"
                " Objective.residuals, please consider fixing"
                " them.\n" + repr(singular_params)
            )

        return covar

    def _covar_from_residuals(self):
        _pvals = np.array(self.varying_parameters())

        used_residuals_scaler = False

        def fn_scaler(vals):
            return np.squeeze(self.residuals(_pvals * vals))

        try:
            # we should be able to calculate a Jacobian for a parameter whose
            # value is zero. However, the scaling approach won't work.
            # This will force Jacobian calculation by unscaled parameters
            if np.any(_pvals == 0):
                raise FloatingPointError()

            with np.errstate(invalid="raise"):
                jac = approx_derivative(fn_scaler, np.ones_like(_pvals))
            used_residuals_scaler = True
        except FloatingPointError:
            jac = approx_derivative(self.residuals, _pvals)
        finally:
            # using approx_derivative changes the state of the objective
            # parameters have to make sure they're set at the end
            self.setp(_pvals)

        # need to create this because GlobalObjective may not have
        # access to all the datapoints being fitted.
        n_datapoints = np.size(jac, 0)

        # covar = J.T x J

        # from scipy.optimize.minpack.py
        # eliminates singular parameters
        _, s, VT = np.linalg.svd(jac, full_matrices=False)
        threshold = np.finfo(float).eps * max(jac.shape) * s[0]
        s = s[s > threshold]
        VT = VT[: s.size]
        covar = np.dot(VT.T / s ** 2, VT)

        if used_residuals_scaler:
            # unwind the scaling.
            covar = covar * np.atleast_2d(_pvals) * np.atleast_2d(_pvals).T

        scale = 1.0
        # scale by reduced chi2 if experimental uncertainties weren't used.
        if not (self.weighted):
            scale = self.chisqr() / (
                n_datapoints - len(self.varying_parameters())
            )

        return covar * scale

    def pgen(self, ngen=1000, nburn=0, nthin=1):
        """
        Yield random parameter vectors from the MCMC samples. The objective
        state is not altered.

        Parameters
        ----------
        ngen : int, optional
            the number of samples to yield. The actual number of samples
            yielded is `min(ngen, chain.size)`
        nburn : int, optional
            discard this many steps from the start of the chain
        nthin : int, optional
            only accept every `nthin` samples from the chain

        Yields
        ------
        pvec : np.ndarray
            A randomly chosen parameter vector

        """
        yield from self.parameters.pgen(ngen=ngen, nburn=nburn, nthin=nthin)

    def plot(self, pvals=None, samples=0, parameter=None, fig=None):
        """
        Plot the data/model.

        Requires matplotlib be installed.

        Parameters
        ----------
        pvals : np.ndarray, optional
            Numeric values for the Parameter's that are varying
        samples: number
            If the objective has been sampled, how many samples you wish to
            plot on the graph.
        parameter: refnx.analysis.Parameter
            Creates an interactive plot for the Parameter in Jupyter. Requires
            ipywidgets be installed. Use with %matplotlib notebook/qt.
        fig: Figure instance, optional
            If `fig` is not supplied then a new figure is created. Otherwise
            the graph is created on the current axes on the supplied figure.

        Returns
        -------
        fig, ax : :class:`matplotlib.Figure`, :class:`matplotlib.Axes`
            `matplotlib` figure and axes objects.

        """
        self.setp(pvals)

        if fig is None:
            import matplotlib.pyplot as plt

            fig = plt.figure()
            ax = fig.add_subplot(111)
        else:
            ax = fig.gca()

        y, y_err, model = self._data_transform(model=self.generative())

        # add the data (in a transformed fashion)
        if self.weighted:
            ax.errorbar(
                self.data.x,
                y,
                y_err,
                color="blue",
                label=self.data.name,
                marker="o",
                ms=3,
                lw=0,
                elinewidth=2,
            )
        else:
            ax.scatter(self.data.x, y, color="blue", s=3, label=self.data.name)

        if samples > 0:
            saved_params = np.array(self.parameters)
            # Get a number of chains, chosen randomly, set the objective,
            # and plot the model.
            for pvec in self.pgen(ngen=samples):
                y, y_err, model = self._data_transform(
                    model=self.generative(pvec)
                )

                ax.plot(self.data.x, model, color="k", alpha=0.01)

            # put back saved_params
            self.setp(saved_params)

        # add the fit
        generative_plot = ax.plot(self.data.x, model, color="red", zorder=20)

        if parameter is None:
            return fig, ax

        # create an interactive plot in a Jupyter notebook.
        def f(val):
            if parameter is not None:
                parameter.value = float(val)
            y, y_err, model = self._data_transform(model=self.generative())
            generative_plot[0].set_data(self.data.x, model)
            fig.canvas.draw()

        import ipywidgets

        return fig, ax, ipywidgets.interact(f, val=float(parameter))

    def corner(self, **kwds):
        """
        Corner plot of the chains belonging to the Parameters.
        Requires the `corner` and `matplotlib` packages.

        Parameters
        ----------
        kwds: dict
            passed directly to the `corner.corner` function

        Returns
        -------
        fig : :class:`matplotlib.Figure` object.
        """
        import corner

        var_pars = self.varying_parameters()
        chain = np.array([par.chain for par in var_pars])
        labels = [par.name for par in var_pars]
        chain = chain.reshape(len(chain), -1).T
        kwds["labels"] = labels
        kwds["quantiles"] = [0.16, 0.5, 0.84]
        return corner.corner(chain, **kwds)

class GlobalObjective(Objective):
    """
    Global Objective function for simultaneous fitting with
    `refnx.analysis.CurveFitter`

    Parameters
    ----------
    objectives : list
        list of :class:`refnx.analysis.Objective` objects

    """

    def __init__(self, objectives):
        self.objectives = objectives
        weighted = [objective.weighted for objective in objectives]

        self._weighted = np.array(weighted, dtype=bool)

        if len(np.unique(self._weighted)) > 1:
            raise ValueError(
                "All the objectives must be either weighted or"
                " unweighted, you cannot have a mixture."
            )

    def __str__(self):
        s = ["{:_>80}".format("\n")]
        s.append("--Global Objective--")
        for obj in self.objectives:
            s.append(str(obj))
            s.append("\n")
        return "\n".join(s)

    def __repr__(self):
        return "GlobalObjective({0})".format(repr(self.objectives))

    @property
    def weighted(self):
        """
        **bool** do all the datasets have y_err, and are all the objectives
        wanting to use weights?

        """
        return self._weighted.all()

    @property
    def npoints(self):
        """
        **int** number of data points in all the objectives.

        """
        npoints = 0
        for objective in self.objectives:
            npoints += objective.npoints
        return npoints

    def residuals(self, pvals=None):
        """
        Concatenated residuals for each of the
        :meth:`refnx.analysis.Objective.residuals`.

        Parameters
        ----------
        pvals : array-like or refnx.analysis.Parameters
            values for the varying or entire set of parameters

        Returns
        -------
        residuals : np.ndarray
            Concatenated :meth:`refnx.analysis.Objective.residuals`

        """
        self.setp(pvals)

        residuals = []
        for objective in self.objectives:
            residual = objective.residuals()
            residuals.append(residual)

        return np.concatenate(residuals)

    @property
    def parameters(self):
        """
        :class:`refnx.analysis.Parameters` associated with all the objectives.

        """
        # TODO this is probably going to be slow.
        # cache and update strategy?
        p = Parameters(name="global fitting parameters")

        for objective in self.objectives:
            p.append(objective.parameters)

        return p

    def logp(self, pvals=None):
        """
        Calculate the log-prior of the system

        Parameters
        ----------
        pvals : array-like or refnx.analysis.Parameters, optional
            values for the varying or entire set of parameters

        Returns
        -------
        logp : float
            log-prior probability

        """
        self.setp(pvals)

        logp = 0.0
        for objective in self.objectives:
            logp += objective.logp()
            # shortcut if one of the priors is impossible
            if not np.isfinite(logp):
                return -np.inf

        return logp

    def logl(self, pvals=None):
        """
        Calculate the log-likelhood of the system

        Parameters
        ----------
        pvals : array-like or refnx.analysis.Parameters
            values for the varying or entire set of parameters

        Returns
        -------
        logl : float
            log-likelihood probability

        """
        self.setp(pvals)
        logl = 0.0

        for objective in self.objectives:
            logl += objective.logl()

        return logl

    def plot(self, pvals=None, samples=0, parameter=None, fig=None):
        """
        Plot the data/model for all the objectives in the GlobalObjective.

        Matplotlib must be installed to use this method.

        Parameters
        ----------
        pvals : np.ndarray, optional
            Numeric values for the Parameter's that are varying
        samples: number, optional
            If the objective has been sampled, how many samples you wish to
            plot on the graph.
        parameter: refnx.analysis.Parameter, optional
            Creates an interactive plot for the Parameter in Jupyter. Requires
            ipywidgets be installed. Use with %matplotlib notebook/qt.
        fig: Figure instance, optional
            If `fig` is not supplied then a new figure is created. Otherwise
            the graph is created on the current axes on the supplied figure.

        Returns
        -------
        fig, ax : :class:`matplotlib.Figure`, :class:`matplotlib.Axes`
            `matplotlib` figure and axes objects.

        """

        self.setp(pvals)

        if fig is None:
            import matplotlib.pyplot as plt

            fig = plt.figure()
            ax = fig.add_subplot(111)
        else:
            ax = fig.gca()

        generative_plots = []

        if samples > 0:
            saved_params = np.array(self.parameters)

            # Get a number of chains, chosen randomly, set the objectives,
            # and plot the model.
            for pvec in self.pgen(ngen=samples):
                self.setp(pvec)

                for objective in self.objectives:
                    y, y_err, model = objective._data_transform(
                        model=objective.generative()
                    )

                    ax.plot(objective.data.x, model, color="k", alpha=0.01)

            # put back saved_params
            self.setp(saved_params)

        for objective in self.objectives:
            # add the data (in a transformed fashion)
            y, y_err, model = objective._data_transform(
                model=objective.generative()
            )

            if objective.weighted:
                ax.errorbar(
                    objective.data.x,
                    y,
                    y_err,
                    label=objective.data.name,
                    ms=3,
                    lw=0,
                    elinewidth=2,
                    marker="o",
                )
            else:
                ax.scatter(objective.data.x, y, label=objective.data.name)

            # add the fit
            generative_plots.append(
                ax.plot(objective.data.x, model, color="r", lw=1.5, zorder=20)[
                    0
                ]
            )

        if parameter is None:
            return fig, ax

        # create an interactive plot in a Jupyter notebook.
        def f(val):
            if parameter is not None:
                parameter.value = float(val)
            for i, objective in enumerate(self.objectives):
                y, y_err, model = objective._data_transform(
                    model=objective.generative()
                )

                generative_plots[i].set_data(objective.data.x, model)
            fig.canvas.draw()

        import ipywidgets

        return fig, ax, ipywidgets.interact(f, val=float(parameter))

        return fig, ax

class Transform:
    r"""
    Mathematical transforms of numeric data.

    Parameters
    ----------
    form : None or str
        One of:

            - 'lin'
                No transform is made
            - 'logY'
                log10 transform
            - 'YX4'
                YX**4 transform
            - 'YX2'
                YX**2 transform
            - None
                No transform is made

    Notes
    -----
    You ask for a transform to be carried out by calling the Transform object
    directly.

    >>> x = np.linspace(0.01, 0.1, 11)
    >>> y = np.linspace(100, 1000, 11)
    >>> y_err = np.sqrt(y)
    >>> t = Transform('logY')
    >>> ty, te = t(x, y, y_err)
    >>> ty
    array([2.        , 2.2787536 , 2.44715803, 2.56820172, 2.66275783,
           2.74036269, 2.80617997, 2.86332286, 2.91381385, 2.95904139,
           3.        ])

    """

    def __init__(self, form):
        types = [None, "lin", "logY", "YX4", "YX2"]
        self.form = None

        if form in types:
            self.form = form
        else:
            raise ValueError(
                "The form parameter must be one of [None, 'lin',"
                " 'logY', 'YX4', 'YX2']"
            )

    def __repr__(self):
        return "Transform({0})".format(repr(self.form))

    def __call__(self, x, y, y_err=None):
        """
        Calculate the transformed data

        Parameters
        ----------
        x : array-like
            x-values
        y : array-like
            y-values
        y_err : array-like
            Uncertainties in `y` (standard deviation)

        Returns
        -------
        yt, et : tuple
            The transformed data

        Examples
        --------
        >>> x = np.linspace(0.01, 0.1, 11)
        >>> y = np.linspace(100, 1000, 11)
        >>> y_err = np.sqrt(y)
        >>> t = Transform('logY')
        >>> ty, te = t(x, y, y_err)
        >>> ty
        array([2.        , 2.2787536 , 2.44715803, 2.56820172, 2.66275783,
               2.74036269, 2.80617997, 2.86332286, 2.91381385, 2.95904139,
               3.        ])

        """
        return self.__transform(x, y, y_err=y_err)

    def __transform(self, x, y, y_err=None):
        r"""
        Transform the data passed in

        Parameters
        ----------
        x : array-like

        y : array-like

        y_err : array-like

        Returns
        -------
        yt, et : tuple
            The transformed data
        """

        if y_err is None:
            etemp = np.ones_like(y)
        else:
            etemp = y_err

        if self.form in ["lin", None]:
            yt = np.copy(y)
            et = np.copy(etemp)
        elif self.form == "logY":
            yt, et = EP.EPlog10(y, etemp)
            if not np.isfinite(yt).all():
                warnings.warn(
                    "Some of the transformed data was non-finite."
                    " Please check your datasets for points with zero or"
                    " negative values.",
                    RuntimeWarning,
                )
        elif self.form == "YX4":
            yt = y * np.power(x, 4)
            et = etemp * np.power(x, 4)
        elif self.form == "YX2":
            yt = y * np.power(x, 2)
            et = etemp * np.power(x, 2)
        if y_err is None:
            return yt, None
        else:
            return yt, et

def pymc3_model(objective):
    """
    Creates a pymc3 model from an Objective.

    Requires theano and pymc3 be installed. This is an experimental feature.

    Parameters
    ----------
    objective: refnx.analysis.Objective

    Returns
    -------
    model: pymc3.Model

    Notes
    -----
    The varying parameters are renamed 'p0', 'p1', etc, as it's vital in pymc3
    that all parameters have their own unique name.

    """
    import pymc3 as pm
    import theano.tensor as tt
    from refnx._lib._pymc3 import _LogLikeWithGrad

    basic_model = pm.Model()

    pars = objective.varying_parameters()
    wrapped_pars = []
    with basic_model:
        # Priors for unknown model parameters
        for i, par in enumerate(pars):
            name = "p%d" % i
            p = _to_pymc3_distribution(name, par)
            wrapped_pars.append(p)

        # Expected value of outcome
        try:
            # Likelihood (sampling distribution) of observations
            pm.Normal(
                "y_obs",
                mu=objective.generative,
                sigma=objective.data.y_err,
                observed=objective.data.y,
            )
        except Exception:
            # Falling back, theano autodiff won't work on function object
            theta = tt.as_tensor_variable(wrapped_pars)
            logl = _LogLikeWithGrad(objective.logl)
            pm.Potential("log-likelihood", logl(theta))

    return basic_model

def _to_pymc3_distribution(name, par):
    """
    Create a pymc3 continuous distribution from a Bounds object.

    Parameters
    ----------
    name : str
        Name of parameter
    par : refnx.analysis.Parameter
        The parameter to wrap

    Returns
    -------
    d : pymc3.Distribution
        The pymc3 distribution

    """
    import pymc3 as pm
    import theano.tensor as T
    from theano.compile.ops import as_op

    dist = par.bounds
    # interval and both lb, ub are finite
    if isinstance(dist, Interval) and np.isfinite([dist.lb, dist.ub]).all():
        return pm.Uniform(name, dist.lb, dist.ub)
    # no bounds
    elif (
        isinstance(dist, Interval)
        and np.isneginf(dist.lb)
        and np.isinf(dist.lb)
    ):
        return pm.Flat(name)
    # half open uniform
    elif isinstance(dist, Interval) and not np.isfinite(dist.lb):
        return dist.ub - pm.HalfFlat(name)
    # half open uniform
    elif isinstance(dist, Interval) and not np.isfinite(dist.ub):
        return dist.lb + pm.HalfFlat(name)

    # it's a PDF
    if isinstance(dist, PDF):
        dist_gen = getattr(dist.rv, "dist", None)

        if isinstance(dist.rv, stats.rv_continuous):
            dist_gen = dist.rv

        if isinstance(dist_gen, type(stats.uniform)):
            if hasattr(dist.rv, "args"):
                p = pm.Uniform(
                    name, dist.rv.args[0], dist.rv.args[1] + dist.rv.args[0]
                )
            else:
                p = pm.Uniform(name, 0, 1)
            return p

        # norm from scipy.stats
        if isinstance(dist_gen, type(stats.norm)):
            if hasattr(dist.rv, "args"):
                p = pm.Normal(name, mu=dist.rv.args[0], sigma=dist.rv.args[1])
            else:
                p = pm.Normal(name, mu=0, sigma=1)
            return p

    # not open, uniform, or normal, so fall back to DensityDist.
    d = as_op(itypes=[T.dscalar], otypes=[T.dscalar])(dist.logp)
    r = as_op(itypes=[T.dscalar], otypes=[T.dscalar])(dist.rvs)
    p = pm.DensityDist(name, d, random=r)

    return p


<filename>DynaMETE_Rfunctions_FlexibleFunctions.py
'''
This file defines all of the necessary functions for DynaMETE, including the transition functions and
the structure function R. This function does NOT include sums over n, since it is designed to be a
more flexible version incorporating different transition functions. This will be very slow for large N or E.
It also defines the METE constraint for beta, which is needed, and a function to obtain mete_lambdas.

To change the functional form of the transition functions, you need only change f, h, and/or q, and the
corresponding function dfdt, dhdt, and/or dqdt.

This version specifically replaces d0/E_c n/e^(-1/3) with d0/E_c n^2/e^(-1/3) to test adding a new degree of
freedom in n dependence.
'''

# Import
import numpy as np
import pandas as pd
from scipy.optimize import fsolve
from scipy import integrate

# METE functions
def beta_constraint(b,s):
    '''This is the beta constraint in METE with give state variables. Use this as a function call to get beta.
    Inputs s as state variables, call S, N, or E
    Also inputs beta
    outputs beta constraint to minimize'''
    return b*np.log(1/(1-np.exp(-b)))-s['S']/s['N']

def get_beta(s,b0=0.0001):
    '''This returns beta from METE. Inputs s as state variables.'''
    return fsolve(beta_constraint,b0,args=s)[0]

def mete_lambdas(s,b0=0.0001):
    '''This returns the METE lambdas for a given set of state variables.
    Inputs s as state variables, call S, N, or E
    Optional input of an initial beta, if we know it's going to be somewhere other than small positive.
    outputs array of lambdas'''
    beta = get_beta(s,b0)
    l2 = s['S']/(s['E']-s['N'])
    ls = np.array([beta-l2,l2,0,0,0])
    return ls

# Transition functions
# The idea here is to make everything easy to change by changing only these functions.
# For f
def fb0(s,p):
    return p['b0']
def fd0(s,p):
    return -p['d0']*s['E']/p['Ec']
def fm0(s,p):
    return p['m0']/s['N']
def f(n,e,s,p):
    '''Transition function for dN/dt. n and e are microscopic variables.
    s are state variables, call S, N, or E
    p are parameters, call b0, d0, m0, w0, w10, Ec, or mu '''
    return (fb0(s,p)+fd0(s,p)*n)*n/e**(1/3)+fm0(s,p)*n

# For h
def hw0(s,p):
    return p['w0']
def hd0(s,p):
    return -p['d0']*s['E']/p['Ec']
def hw10(s,p):
    b0i=0.0001
    beta = get_beta(s,b0i)
    return -p['w10']/np.log(1/beta)**(2/3)
def hm0(s,p):
    return p['m0']/s['N']
def h(n,e,s,p):
    '''Transition function for dE/dt. n and e are microscopic variables.
    s are state variables, call S, N, or E
    p are parameters, call b0, d0, m0, w0, w1, Ec, or mu '''
    return (hw0(s,p)+hd0(s,p)*n)*n*e**(2/3)+hw10(s,p)*n*e+hm0(s,p)*n

# For q
def qc(s,p):
    return p['m0']*np.exp(-p['mu']*s['S']-np.euler_gamma)
def qd0(s,p):
    return -s['S']*p['d0']*s['E']/p['Ec']
def q(n,e,s,p):
    '''Transition function for dS/dt. n and e are microscopic variables.
    s are state variables, call S, N, or E
    p are parameters, call b0, d0, m0, w0, w10, Ec, or mu
    For now this doesn't implement speciation models, ie. s1=s2=0'''
    # Set up kronecker delta in an easy way. Round n to nearest int, if it's 1 then include term
    # I actually need this to be vectorized, so let's do it slightly differently.
    # First check if n is scalar
    if np.isscalar(n):
        kn1 = int(np.rint(n))==1
    else:
        kn1 = np.zeros(len(n))
        kn_arg = np.where(np.rint(n)==1)
        # I included the rounded int here because really this kronecker delta can be defined in continuous space
        # In that case there should also be a correction factor though, but let's ignore that.
        # The good news here is that below we only pass in arange, which by default passes in integers
        # So we should be ok as long as we are only using arange for passing in ranges of n
        # That is because arange rounds the variable it takes in so it can take steps of length 1
        kn1[kn_arg] = 1
    return qc(s,p) + qd0(s,p)*kn1/e**(1/3)

# Also need derivatives for lambda dynamics. Note that these have to be manually editted for alternate f,h,q
def dfdt(n,e,s,p,ds):
    return fd0(s,p)/s['E']*ds['dE']*n**2/e**(1/3) - fm0(s,p)*ds['dN']/s['N']*n

def dhdt(n,e,s,p,ds):
    return hd0(s,p)/s['E']*ds['dE']*n**2*e**(2/3) - hm0(s,p)*ds['dN']/s['N']*n

def dqdt(n,e,s,p,ds):
    # See q for how the kronecker delta works.
    if np.isscalar(n):
        kn1 = int(np.rint(n))==1
    else:
        kn1 = np.zeros(len(n))
        kn_arg = np.where(np.rint(n)==1)
        kn1[kn_arg] = 1
    return -qc(s,p)*ds['dS']*p['mu'] + qd0(s,p)*(ds['dS']/s['S']+ds['dE']/s['E'])*kn1/e**(1/3)

# R itself
def R(n,e,l,s,p):
    '''Unnormalized struture function for DynaMETE.
    n,e are microscopic variables.
    l are lambdas
    s are state variables, call S, N, or E
    p are parameters, call b0, d0, m0, w0, w10, Ec, or mu '''
    return np.exp(-l[0]*n-l[1]*n*e-l[2]*f(n,e,s,p)-l[3]*h(n,e,s,p)-l[4]*q(n,e,s,p))

# For calculating a single mean with specific powers of n and e
def mean_pow(npow,epow,l,s,p,z=1):
    '''
    This function returns the mean of n^npow*e^epow over the R function.
    It is NOT normalized, but it does take in z as an optional argument to normalize.
    This function uses quad integral over log e for each n then sums over n.
    Note that npow=epow=0 corresponds to Z, so by default these are not normalized.
    l are lambdas
    s are state variables, call S, N, or E
    p are parameters, call b0, d0, m0, w0, w10, Ec, or mu
    '''
    nrange = np.arange(s['N'])+1
    eint = integrate.quad_vec(lambda loge: np.exp(loge*(1+epow))*R(nrange,np.exp(loge),l,s,p),0,np.log(s['E']))[0]
    return np.sum(nrange**npow*eint)/z

# For calculating a covariance with specific powers of n and e for each function
def cov_pow(npow,epow,l,s,p,z):
    '''
    This function returns the covariance of two functions with the form n^npow*e^epow over the R function.
    You have to pass in the normalization so that things are faster than calculating normalization each time.
    npow and epow should both be 2d arrays with the functions.
    For example, if you want COV(n^2,ne), pass npow=[2,1], epow=[0,1]
    This function uses quad integral over log e for each n then sums over n.
    z is the normalization
    l are lambdas
    s are state variables, call S, N, or E
    p are parameters, call b0, d0, m0, w0, w10, Ec, or mu
    '''
    nrange = np.arange(s['N'])+1
    # Get integral over both functions
    ffeint = integrate.quad_vec(lambda loge: np.exp(loge*(1+np.sum(epow)))*R(nrange,np.exp(loge),l,s,p),0,np.log(s['E']))[0]
    ff = np.sum(nrange**np.sum(npow)*ffeint)/z
    # Get integral over each function
    f1f2 = 1
    for nn,ee in zip(npow,epow):
        feint = integrate.quad_vec(lambda loge: np.exp(loge*(1+ee))*R(nrange,np.exp(loge),l,s,p),0,np.log(s['E']))[0]
        f1f2 *= np.sum(nrange**nn*feint)/z
    return ff-f1f2

# For calculating a single mean over an arbitrary function
# Use mean_pow for non-functions
def mean(func,l,s,p,*args,z=1):
    '''
    This function returns the mean of an arbitrary function over the R function.
    It is NOT normalized, but it does take in z as an optional argument to normalize.
    Because I put *args first, you have to use z=z0 if you want to put in a normalization.
    The arbitrary function must take arguments of the form (n,e,s,p) for this to work.
    This is the form of the f,h, and q functions above.
    You can pass additional arguments as required for the function (ie. pass ds for df/dt)
    To pass in n or n*e, use lambda n,e,s,p: n or lambda n,e,s,p: n*e, or similar.
    Alternatively, use mean_pow
    This function uses quad integral over log e for each n then sums over n.
    l are lambdas
    s are state variables, call S, N, or E
    p are parameters, call b0, d0, m0, w0, w10, Ec, or mu
    z is the normalization
    '''
    nrange = np.arange(s['N'])+1
    # Below is to make this easier for lambda functions, but it isn't worth it. Just require s and p passed,
    # and let other things be passed as args if needed.
    # Check if we need args by looking at function passed in
#    funcargs = func.__code__.co_argcount
#    if funcargs >= 4:
#        args = s,p,args
    eint = integrate.quad_vec(lambda loge: np.exp(loge)*R(nrange,np.exp(loge),l,s,p)*func(nrange,np.exp(loge),s,p,*args),0,np.log(s['E']))[0]
    return np.sum(eint)/z

# For calculating a covariance
# Note if you want to do this with non-functions, use cov_pow
def cov(func1,func2,l,s,p,z,*args):
    '''
    This function returns the covariance of two arbitrary functions over the R function.
    You have to pass in the normalization so that things are faster than calculating normalization each time.
    The arbitrary functions must take arguments of the form (n,e,s,p) for this to work.
    This is the form of the f,h, and q functions above.
    You can pass additional arguments as required for the function (ie. pass ds for df/dt)
    To pass in n or n*e, use lambda n,e,s,p: n or lambda n,e,s,p: n*e, or similar.
    This function uses quad integral over log e for each n then sums over n.
    l are lambdas
    s are state variables, call S, N, or E
    p are parameters, call b0, d0, m0, w0, w10, Ec, or mu
    z is the normalization
    '''
    nrange = np.arange(s['N'])+1
    # Get integral over both functions
    ffeint = integrate.quad_vec(lambda loge: np.exp(loge)*R(nrange,np.exp(loge),l,s,p)*func1(nrange,np.exp(loge),s,p,*args)*func2(nrange,np.exp(loge),s,p,*args),0,np.log(s['E']))[0]
    ff = np.sum(ffeint)/z
    # Get integral over each function
    f1f2 = 1
    for func in [func1,func2]:
        feint = integrate.quad_vec(lambda loge: np.exp(loge)*R(nrange,np.exp(loge),l,s,p)*func(nrange,np.exp(loge),s,p,*args),0,np.log(s['E']))[0]
        f1f2 *= np.sum(feint)/z
    return ff-f1f2

def get_dXdt(l,s,p):
    '''
    Returns the time derivatives of the state variables. This makes it easier than calling mean three times
    every time I want to see the derivatives.
    Inputs lambdas, state variables, and parameters.
    Outputs a pandas series of ds.
    '''
    # Create storage
    ds = pd.Series(np.zeros(3),index=['dS','dN','dE'])
    # To normalize
    z = mean_pow(0,0,l,s,p)
    ds['dS'] = mean(q,l,s,p,z=z)
    ds['dN'] = s['S']*mean(f,l,s,p,z=z)
    ds['dE'] = s['S']*mean(h,l,s,p,z=z)
    return ds


<filename>data_analysis/Statistical_Analysis.py
import cPickle
import numpy as np
import pandas as pd
from statsmodels.stats import multitest
from scipy.ndimage import gaussian_filter1d
from scipy.stats import mannwhitneyu, binom_test, wilcoxon

EXP_NAME = "XENO_3"
PICKLE_DIR = "/home/sam/Projects/research_code/evosoro/data_analysis/results"

inverted = 0
with open('{0}/XENO_3_Robustness_DataFrame_{1}.pickle'.format(PICKLE_DIR, inverted), 'rb') as handle:
    df = cPickle.load(handle)
X, Y, Z = df["TraceX"], df["TraceY"], df["TraceZ"]
time = df["Time"]
X, Y, time = X[time < 61], Y[time < 61], time[time < 61]  # only use one minute in silico
X, Y = (X[:len(X)*25/32+1]-0.05*4)*2.5, (Y[:len(X)*25/32+1]-0.05*4)*2.5

inverted = 1
with open('{0}/XENO_3_Robustness_DataFrame_{1}.pickle'.format(PICKLE_DIR, inverted), 'rb') as handle:
    df_inv = cPickle.load(handle)
X_inv, Y_inv, Z_inv = df_inv["TraceX"], df_inv["TraceY"], df_inv["TraceZ"]
time_inv = df_inv["Time"]
X_inv, Y_inv, time_inv = X_inv[time_inv < 61], Y_inv[time_inv < 61], time_inv[time_inv < 61]
X_inv, Y_inv = (X_inv[:len(X_inv)*25/32+1]-0.05*4)*2.5, (Y_inv[:len(Y_inv)*25/32+1]-0.05*4)*2.5

# vivo_data = pd.read_csv("/home/sam/Desktop/vivo_data.csv")
vivo_data = pd.read_csv("../data_analysis/vivo_data.csv")  # additional cardiacBot

vivo_data["X"] -= vivo_data["Start_X"]
vivo_data["Y"] -= vivo_data["Start_Y"]

X_vivo_inv, Y_vivo_inv = vivo_data["X"][vivo_data["Upright"] == 0], vivo_data["Y"][vivo_data["Upright"] == 0]
vivo_time_inv = vivo_data["Time"][vivo_data["Upright"] == 0]
vivo_run_inv = vivo_data["Run"][vivo_data["Upright"] == 0]
vivo_robot_inv = vivo_data["Robot"][vivo_data["Upright"] == 0]

X_vivo, Y_vivo = vivo_data["X"][vivo_data["Upright"] == 1], vivo_data["Y"][vivo_data["Upright"] == 1]
vivo_time = vivo_data["Time"][vivo_data["Upright"] == 1]
vivo_run = vivo_data["Run"][vivo_data["Upright"] == 1]
vivo_robot = vivo_data["Robot"][vivo_data["Upright"] == 1]

vivo_final_dist = []
robot = []
run = []

for n in range(1, len(vivo_time)):
    if vivo_time.values[n] < vivo_time.values[n-1]:
        vivo_final_dist += [(X_vivo.values[n-1]**2 + Y_vivo.values[n-1]**2)**0.5]
        robot += [vivo_robot.values[n-1]]
        run += [vivo_run.values[n-1]]

    if n == len(vivo_time)-1:
        vivo_final_dist += [(X_vivo.values[n-1]**2 + Y_vivo.values[n-1]**2)**0.5]
        robot += [vivo_robot.values[n]]
        run += [vivo_run.values[n]]

vivo_final_dist_inv = []
robot_inv = []
run_inv = []

for n in range(1, len(vivo_time_inv)):
    if vivo_time_inv.values[n] < vivo_time_inv.values[n-1]:
        vivo_final_dist_inv += [(X_vivo_inv.values[n-1]**2 + Y_vivo_inv.values[n-1]**2)**0.5]
        robot_inv += [vivo_robot.values[n-1]]
        run_inv += [vivo_run.values[n-1]]

    if n == len(vivo_time)-1:
        vivo_final_dist_inv += [(X_vivo_inv.values[n-1]**2 + Y_vivo_inv.values[n-1]**2)**0.5]
        robot_inv += [vivo_robot.values[n]]
        run_inv += [vivo_run.values[n]]

diff = []
for bot, r, d in zip(robot, run, vivo_final_dist):
    for bot_inv, r_inv, d_inv in zip(robot_inv, run_inv, vivo_final_dist_inv):
        if bot == bot_inv:
            if r == r_inv or (r_inv == 1 and bot < 3):
                diff += [d-d_inv]

print len(np.unique(diff))
print len(diff), len(run)
p1 = wilcoxon(diff)[1]  # *3
print "vivo ", p1

silico_final_dist = []

for n in range(1, len(X)):
    if time.values[n] < time.values[n-1]:
        silico_final_dist += [(X.values[n-1]**2 + Y.values[n-1]**2)**0.5]

    if n == len(X)-1:
        silico_final_dist += [(X.values[n-1]**2 + Y.values[n-1]**2)**0.5]

silico_final_dist_inv = []

for n in range(1, len(X_inv)):
    if time_inv.values[n] < time_inv.values[n-1]:
        silico_final_dist_inv += [(X_inv.values[n-1]**2 + Y_inv.values[n-1]**2)**0.5]

    if n == len(X_inv)-1:
        silico_final_dist_inv += [(X_inv.values[n-1]**2 + Y_inv.values[n-1]**2)**0.5]

# print len(silico_final_dist), len(silico_final_dist_inv)

# print "silco 5x4 ", wilcoxon(silico_final_dist, silico_final_dist_inv[:len(silico_final_dist)*4/5])[1]*3

p2 = wilcoxon(silico_final_dist, silico_final_dist_inv)[1]  # *3
print "silco ", p2

# p = []
# for _ in range(5000):
#     p += [wilcoxon(silico_final_dist, np.random.choice(silico_final_dist_inv, len(silico_final_dist)))[1]]
# print "boot slico ", np.mean(p)*3

# print "vivo ", wilcoxon(vivo_final_dist, vivo_final_dist_inv)[1]*3
# print "vivo ", wilcoxon(np.random.choice(vivo_final_dist, len(vivo_final_dist_inv)), vivo_final_dist_inv)[1]*3

# p = []
# for _ in range(5000):
#     p += [wilcoxon(vivo_final_dist, np.random.choice(vivo_final_dist_inv, len(vivo_final_dist)))[1]]
# print "boot vivo ", np.mean(p)*3

# print "binom: ", binom_test(4, 5, 0.25)*3  # original 5
p3 = binom_test(5, 6, 0.25)  # *3  # with 6th
print "binom: ", p3  # with 6th
print "fdr: ", multitest.fdrcorrection([p1, p2, p3], alpha=0.0001, method='indep', is_sorted=False)

# np.random.choice(vivo_final_dist, )


import myo
import sys
from threading import Thread
import time
import numpy as np
import os
import matplotlib
matplotlib.use('GTKAgg')
import scipy

from math import ceil, floor

e = []
emg_correctmean = []
emg_filtered = []
emg_rectified = []
emg_envelope = []
normalized = []
import config
count =0
lastLine = ""
import scipy as sp
from scipy import signal
low_pass = 4
sfreq = 1000
#delete training data from last run
from sklearn.preprocessing import MinMaxScaler

start_time = time.time()
#exit process
def buildData(name):
	m = myo.Myo(sys.argv[1] if len(sys.argv) >= 2 else None)
	_dir = os.path.join('/home/venkatesh/Desktop', 'sub1')

	if not os.path.exists(_dir):
	    	os.makedirs(_dir)

	f = open("tra/sub1/" + name + ".txt", "a")

	#Callback for EMG data from Myo (8 words)
	def proc_emg(emg, moving, times=[]):
		global e,emg_correctmean,emg_filtered,emg_rectified,low_pass,sfreq,emg_envelope
		e = emg
		#emg_correctmean = e - np.mean(e)
		emg_correctmean = scipy.signal.detrend(e)
		high = 20/(1000/2)
		low = 450/(1000/2)
		b, a = sp.signal.butter(4, [high,low], btype='bandpass')
		emg_filtered = sp.signal.filtfilt(b, a, e, method = 'gust')
		emg_rectified = abs(emg_filtered)
		l = float(low_pass / (sfreq/2))
		b2, a2 = sp.signal.butter(4, l, btype='lowpass')
		emg_envelope =sp.signal.filtfilt(b2, a2, emg_rectified,method = 'gust')
		emg_envelope = emg_envelope * 100

		#print(emg_envelope)

	#Callback for other motion data, including accelerometer and gycroscope
	def proc_imu(quat, acc, gyro, times=[]):
		global q,a,g,b,t,c,count
		q = quat
		a = acc
		g = gyro
		if count < config.samples:
			if len(emg_envelope) > 0:
				c = list(emg_envelope)
				print(str(c) + "\n")
				f.write(str(c) + "\n")
				#plt.plot(emg_rectified)
				count = count + 1

	m.connect()
	m.add_emg_handler(proc_emg)
	m.add_imu_handler(proc_imu)

	for x in range(0,config.samples):
		m.run()
	m.disconnect()

for i in range(0,config.numGestures):
	thread = Thread(target = buildData, args = ("s" + str(i+9), ))
	thread.deamon = True
	thread.start()
	thread.join()

	print ("Next attempt")
	input("Press Enter to continue...")
	count = 0


import numpy as np
from scipy.interpolate import interp1d

def main_sun_fun(time,stellar_mass,tsat_XUV,beta_XUV,fsat):

    if stellar_mass == 1.0:
        stellar_data = np.loadtxt('Baraffe3.txt',skiprows=31) # for reproducing sun exactly
    else:
        print ("This version of code is only set up for solar mass stars")
        return [time*0,time*0,time*0,time*0]

    stellar_array=[]
    for i in range(0,len(stellar_data[:,0])):
        if stellar_data[i,0] == stellar_mass:
            stellar_array.append(stellar_data[i,:])

    stellar_array=np.array(stellar_array)

    min_time = np.min(stellar_array[:,1])
    max_time = np.max(stellar_array[:,1])

    if (min_time>np.min(time) ) or (max_time<np.max(time)):
        print ("Problem: exceeding time range for stellar data")

    time_array = stellar_array[:,1]

    Total_Lum = (10**stellar_array[:,4])

    ratio_out = [] # For XUV ratio
    for i in range(0,len(time_array)):
        if time_array[i]<tsat_XUV:
            ratio= fsat
        else:
            ratio = fsat*(time_array[i]/tsat_XUV)**beta_XUV
        ratio_out.append(ratio)

    XUV_Lum = ratio_out*Total_Lum

    Total_fun = interp1d(time_array,Total_Lum)
    XUV_fun = interp1d(time_array,XUV_Lum)
    Relative_total_Lum = Total_fun(time)
    Relative_XUV_lum = XUV_fun(time)
    Absolute_total_Lum = Relative_total_Lum*3.828e26
    Absolute_XUV_Lum = Relative_XUV_lum*3.828e26

    return [Relative_total_Lum,Relative_XUV_lum,Absolute_total_Lum,Absolute_XUV_Lum ]


from functools import partial
import copy

import numpy as np
from scipy import optimize

DEBUG = False

class Elliptic:
    """Generate an orthogonal curvilinear grid by solving the inverse form of
    two laplace equations.
    Source: Basic Structured Grid Generation, <NAME> and <NAME>
    """

    def __init__(self, boundary):
        """Summary

        Args:
            boundary (dictionary): contains x,y-coordinates of all 4 boundaries
        """
        self.xw = boundary['west'][0]
        self.yw = boundary['west'][1]
        self.xe = boundary['east'][0]
        self.ye = boundary['east'][1]
        self.xn = boundary['north'][0]
        self.yn = boundary['north'][1]
        self.xs = boundary['south'][0]
        self.ys = boundary['south'][1]

    def solver(self, solver_type, iterations):

        # knowns
        f_o = copy.copy(self.f[:, 0])
        u_o = copy.copy(self.u[:, 0])

        # extra parameters for 'function'
        # are later wrapped with partial
        etamax = copy.copy(self.etamax)

        #
        def F(unknowns, F_args=[etamax]):
            """Summary

            Args:
                unknowns (np.array): x, y

            Returns:
                TYPE: Description
            """

            # unknowns
            if DEBUG:
                print('unknowns', unknowns)

            # (x, y) = unknowns
            x = unknowns[0: etamax]
            y = unknowns[etamax:2 * etamax]

            eq1 = np.zeros_like(x)
            eq2 = np.zeros_like(x)

            # boundary conditions
            x[0] = 0.0
            y[0] = 0.0

            deta = 1.0
            dgsi = 1.0

            # array slicing for index j
            # [1:] means index j
            # [:-1] means index j-1

            dxdgsi = (x[1:, :] - x[:-1, :]) / (2 * dgsi)
            dydgsi = (y[1:, :] - y[:-1, :]) / (2 * dgsi)
            dxdeta = (x[1:, :] - x[:-1, :]) / (2 * deta)
            dydeta = (y[1:, :] - y[:-1, :]) / (2 * deta)

            # components of the covariant metric tensor
            g11 = dxdgsi**2 + dydgsi**2
            g22 = dxdeta**2 + dydeta**2
            g12 = dxdgsi * dxdeta + dydgsi * dydeta

            a = g22 / dgsi**2
            b = 2 * g22 / dgsi**2 + 2 * g11 / deta**2
            c = a
            d = g11 / deta**2 * (x[:, :-1] + x[:, 1:]) - \
                2 * g12 * (x[:-1, :-1] + x[1:, 1:] - x[1:, :-1] - x[:-1, 1:]) \
                / (4 * dgsi * deta)
            e = g11 / deta**2 * (y[:, :-1] + y[:, 1:]) - \
                2 * g12 * (y[:-1, :-1] + y[1:, 1:] - y[1:, :-1] - y[:-1, 1:]) \
                / (4 * dgsi * deta)

            # Winslow x
            eq1[1:] = -a * x[1:, :] + b * x[:, :] - c * x[:-1, :] - d

            # Winslow y
            eq2[1:] = -a * y[1:, :] + b * y[:, :] - c * y[:-1, :] - e

            # boundary conditions make up another 2 equations
            # put them on the 0-th element of all 2 equations
            eq1[0] = x[0]
            eq2[0] = y[0]

            return np.array([eq1, eq2]).ravel()

        # initial guess
        guess = np.array([f_o, u_o]).ravel()

        F_partial = partial(F,
                            F_args=[etamax])

        solution = optimize.fsolve(F_partial, guess,
                                   full_output=True, xtol=1e-06)

        solver_message = solution[3]
        print('  Solver: {}'.format(solver_message))

        return solution, solver_message

    def shift_profiles(self):

        self.x[:, 0] = copy.copy(self.solution[0][0 *
                                 self.etamax:1 * self.etamax])
        self.y[:, 0] = copy.copy(self.solution[0][1 *
                                 self.etamax:2 * self.etamax])

    def main(self, solver_type='fsolve', iterations=10):

        # initial velocity profile
        self.boundary_conditions()

        for self.nx in range(1, self.gsimax):
            self.solution, self.solver_message = \
                self.solver(solver_type, iterations)
            self.shift_profiles()


## @ingroup Methods-Performance
# propeller_range_endurance_speeds.py
#
# Created: Dec 2020, <NAME>
# Modified:

# ----------------------------------------------------------------------
#  Imports
# ----------------------------------------------------------------------

from SUAVE.Core import Data
import numpy as np
import scipy as sp

import SUAVE

# ----------------------------------------------------------------------
#  Propeller Range and Endurance Speeds
# ----------------------------------------------------------------------

def propeller_range_endurance_speeds(analyses,altitude,CL_max,up_bnd,delta_isa):
        """ Computes L/D max and CL^3/2 / CD max at a given altitude. This runs a mini mission wrapped by an
        optimizer to find the L/D max. up_bnd is the fastest airspeed that the optimizer can try. The output is a
        dictionary containing the maximum values as well as the airspeeds.

        Assumptions:
        No propulsion source is given

        Source:
        N/A

        Inputs:
        analyses.atmosphere        [-]
        analyses.aerodynamics      [-]
        altitude                   [m]
        CL_max                     [float]
        up_bnd                     [m/s]
        delta_isa                  [deg C]

        Outputs:
        results.CL32.air_speed     [m/s]
        results.CL32.value         [-]
        results.L_D_max.air_speed  [m/s]
        results.L_D_max.value      [-]

        Properties Used:
        N/A
        """

        # setup a mission that runs a single point segment without propulsion
        def mini_mission():

                # ------------------------------------------------------------------
                #   Initialize the Mission
                # ------------------------------------------------------------------
                mission = SUAVE.Analyses.Mission.Sequential_Segments()
                mission.tag = 'the_mission'

                # ------------------------------------------------------------------
                #  Single Point Segment 1: constant Speed, constant altitude
                # ------------------------------------------------------------------
                segment = SUAVE.Analyses.Mission.Segments.Single_Point.Set_Speed_Set_Altitude_No_Propulsion()
                segment.tag = "single_point"
                segment.analyses.extend(analyses)
                segment.altitude              = altitude
                segment.temperature_deviation = delta_isa

                # add to misison
                mission.append_segment(segment)

                return mission

        # This is what's called by the optimizer for CL**3/2 /CD Max
        def single_point_3_halves(X):

                # Update the mission
                mission.segments.single_point.air_speed = X

                # Run the Mission
                point_results = mission.evaluate()

                CL = point_results.segments.single_point.conditions.aerodynamics.lift_coefficient
                CD = point_results.segments.single_point.conditions.aerodynamics.drag_coefficient

                three_halves = -(CL**(3/2))/CD # Negative because optimizers want to make things small

                return three_halves

        # This is what's called by the optimizer for L/D Max
        def single_point_LDmax(X):

                # Modify the mission for the next iteration
                mission.segments.single_point.air_speed = X

                # Run the Mission
                point_results = mission.evaluate()

                CL = point_results.segments.single_point.conditions.aerodynamics.lift_coefficient
                CD = point_results.segments.single_point.conditions.aerodynamics.drag_coefficient

                L_D = -CL/CD # Negative because optimizers want to make things small

                return L_D

        # ------------------------------------------------------------------
        #   Run the optimizer to solve
        # ------------------------------------------------------------------

        # Setup the a mini mission
        mission = mini_mission()

        # Takeoff mass:
        mass = analyses.aerodynamics.geometry.mass_properties.takeoff

        # Calculate the stall speed
        Vs = stall_speed(analyses,mass,CL_max,altitude,delta_isa)[0][0]

        # The final results to save
        results = Data()

        # Wrap an optimizer around both functions to solve for CL**3/2 /CD max
        outputs_32 = sp.optimize.minimize_scalar(single_point_3_halves,bounds=(Vs,up_bnd),method='bounded')

        # Pack the results
        results.CL32 = Data()
        results.CL32.air_speed = outputs_32.x
        results.CL32.value     = -outputs_32.fun[0][0]

        # Wrap an optimizer around both functions to solve for L/D Max
        outputs_ld = sp.optimize.minimize_scalar(single_point_LDmax,bounds=(Vs,up_bnd),method='bounded')

        # Pack the results
        results.L_D_max = Data()
        results.L_D_max.air_speed = outputs_ld.x
        results.L_D_max.value     = -outputs_ld.fun[0][0]

        return results

def stall_speed(analyses,mass,CL_max,altitude,delta_isa):

        # Unpack
        atmo  = analyses.atmosphere
        S     = analyses.aerodynamics.geometry.reference_area

        # Calculations
        atmo_values       = atmo.compute_values(altitude,delta_isa)
        rho               = atmo_values.density
        sea_level_gravity = atmo.planet.sea_level_gravity

        W = mass*sea_level_gravity

        V = np.sqrt(2*W/(rho*S*CL_max))

        return V


<filename>Python/pervasiveheartmonitor/pervasiveheartmonitor.py<gh_stars>10-100
from statistics import mean
def is_numeric(txt):
    try:
        float(txt)
        return True
    except ValueError:
        return False

while(True):
    try:
        line = input().split()
        print(mean([float(x) for x in line if is_numeric(x)]), " ".join([x for x in line if x.isalpha()]))
    except EOFError:
        break


#!/usr/bin/env python3

"""
This module provides the basic lattice model as used in e.g. cluster expansion, lattice dynamics
"""

import numpy as np
import logging
import scipy.sparse
from collections import Counter

from .cluster import Cluster
from .atomic_model import AtomicModel
from .util.mathtool import relativePosition
from .util.tool import pad_right, matrix2text
from _c_util import get_structure_ordering
logger = logging.getLogger(__name__)

class BasicLatticeModel(AtomicModel):
    """
        A generic model with variables on the lattice:
    """
    def __init__(self, prim, raw_clusters, irange=None, **kwargs):
        """

        :param prim:
        :param raw_clusters:
        :return:
        """
        self.clusters = raw_clusters
        self.prim = prim
        self.maxorder=kwargs['maxorder']
        self.proper_range=kwargs['prange']
        self.cluster_filter = kwargs['filter']
        self.orbits = None
        self.imp_range = irange

    @classmethod
    def from_file(cls, prim, clusf, **kwargs):
        raw = Cluster.read_cluster(prim, clusf)
        return cls(prim, raw, **kwargs)

    def save_clusters(self, fname):
        self.save_clusters_to_file(self.clusters, fname)

    def save_clusters_to_file(self, clusters, fname):
        scale = self.prim.lattice._scale
        elements= self.prim.elements
        sublat = self.prim.orbit_of_l
        out_s = []
        out_ijkl = []
        for cl in clusters:
            out_s.append("1\n%.8f\n%s\n%s\n" % (cl.diameter/scale, " ".join([str(i) for i in cl.orders] +
                                ([str(i[0]) for i in cl.id] if hasattr(cl, 'id') else [])), matrix2text(pad_right(
                np.array(cl.frac_coords, dtype=object), (cl.order, 5)))))
            out_ijkl.append(", ".join(["%d %d %d %d %s%d"%tuple(i.tolist()+[elements[i[-1]],sublat[i[-1]]]) for i in cl._ijkls_np]))
            out_ijkl[-1]+= "  frac: "+ matrix2text(cl.frac_coords, ' ')
            out_ijkl[-1]+= "  cart: "+ matrix2text(cl.coords, ' ')
        open(fname, "w").write("\n".join(out_s))
        open(fname+'_ijkl', "w").write("\n".join(out_ijkl))

    @classmethod
    def generate_clusters(cls, prim, **kwargs):
        """
        Generating PROPER clusters.
        Iteratively generation starting from empty and point clusters
        :param prim: Primitive cell
        :return: LDModel with certain clusters
        """
        sites = kwargs['sites'] if 'sites' in kwargs.keys() else list(range(prim.num_sites))
        proper_range = kwargs['prange']
        maxord = max([ord for ord in proper_range.keys() if proper_range[ord]>0])
        raw= []
        clus_per_ord = []
        for ord in range(maxord+1):
            clus_per_ord = cls.generate_n(prim, ord, proper_range[ord], clus_per_ord, sites)
            raw.extend(clus_per_ord)
        return cls(prim, raw, **kwargs)

    def generate_improper(self):
        """

        :return: Improper clusters
        """
        improp = []
        assert self.imp_range is not None, ValueError('improper range not specified')
        for ord, cut in self.imp_range.items():
            if ord <= 1:
                continue
            for clus in self.clusters:
                if (clus.diameter <= cut) and (clus.order< ord) and clus.order >0:
                    # print("extending", ord, cut, clus)
                    improp.extend(Cluster.improper_from_proper(clus, ord, self.prim))
        print("  generated %d improper clusters" %(len(improp)), self.tally_clusters(improp, 2, self.maxorder))
        self.clusters.extend(improp)

    def tally(self, minord=0):
        return self.tally_clusters(self.clusters, minord, self.maxorder+1)

    @staticmethod
    def tally_clusters(clusters, minord=0, maxord=0):
        return dict(Counter([c.order for c in clusters]))

    @staticmethod
    def generate_n(prim, n, cutoff, clus0, sites):
        """
        Input:
        prim cell
        order n
        cutoff for diameter of clusters
        clus0 clusters at order n-1
        :param sites: allowed sites in the primitive cell. Default all sites
        """
        pts = [[0, 0, 0, l] for l in sites]
        if n == 0:
            uc = [Cluster([], prim)]
        elif n == 1:
            uc = [Cluster.from_ijkl([pt], prim) for pt in pts]
            uc = Cluster.remove_equivalent_clusters(uc)
            #assert len(uc) == prim.n_orb
        elif n >= 2:
            print("  generating order %d clusters..."% (n))
            uc = []
            for clus in clus0:
                if clus.diameter > cutoff:
                    continue
                ijkl0 = clus.ijkls
                # each cluster is now sorted by the orbit (sub-lattice) type
                max_id = prim.orbit_of_l[ijkl0[-1][-1]]
                sumPts = prim.find_nb_cluster(np.array(clus.ijkls), cutoff)
                clus_new = []
                for pt_n in sumPts:
                    pt = pt_n.tolist()
                    if not pt[-1] in sites:
                        continue
                    if prim.orbit_of_l[pt[-1]] < max_id:
                        continue
                    if pt in ijkl0:
                        continue
                    clusSum = clus.append_site(pt)
                    if clusSum.diameter > cutoff:
                        continue
                    clus_new.append(clusSum)
                clus_new = Cluster.remove_equivalent_clusters(clus_new)
                uc.extend(clus_new)
            uc = Cluster.remove_equivalent_clusters(uc)
        return uc

    def __str__(self):
        if self.orbits is None:
            outs= ["  No.  ord uniq  diam"]
            for i, c in enumerate(self.clusters):
                outs.append("%4d %4d %4d %9.5f" % (i, c.order, c.order_uniq, c.diameter))
            return "\n".join(outs)
        else:
            outs= ["  No.  ord  uniq mult  diam"]
            for i, orb in enumerate(self.orbits):
                c = orb.cluster
                outs.append("%5d %4d %4d %4d %9.5f" % (i, c.order, c.order_uniq, len(orb.clusters), c.diameter))
            return "\n".join(outs)

    def cleanup(self):
        """
        1. filter through clusters
        2. move first atoms of each cluster to primitive cell
        3. sort clusters
        :return:
        """
        self.clusters= list(filter(self.cluster_filter, self.clusters))
        for i in range(len(self.clusters)):
            self.clusters[i].move_to_prim()
        self.clusters.sort()

    def l_point_cls(self):
        return [orb.cluster.ijkls[0][-1] for orb in self.orbits if orb.cluster.order==1]

    def symmetrize(self):
        self.Cmat = scipy.sparse.identity(self.ncorr_full)

    @property
    def ncorr(self): return self.Cmat.shape[0]

    def prepare_index_full(self):
    # implementations should set the number of correlations for each cluster.
    # by default set to 1 per cluster
        for orb in self.orbits:
            if orb.ncorr_full<0:
                print("WARNING: ncorr_full NOT set for cluster and set to 1:",orb.cluster)
                orb.ncorr_full = 1
        self.orb_idx_full=np.cumsum([0]+[orb.ncorr_full for orb in self.orbits])
        self.ncorr_full= self.orb_idx_full[-1]

    def prepare_index(self):
        npara_ord = [0]*(self.maxorder+1)
        for orb in self.orbits:
            npara_ord[orb.cluster.order]+= orb.ncorr_full
        self.ord_idx=np.cumsum([0]+npara_ord)
        self.orb_idx=self.orb_idx_full

    def get_orbit_isotropy(self):
        """
        1. Find orbit of each cluster, modulus lattice translation
        2. Find isotropy group of each representative (first) cluster of an orbit
        :return: list of [orbit, isotropy]
            orbit is a list of Clusters
            isotropy is list of [ig, relativePos]
                ig = symmop index
                relativePos indexing function under this symmop
        """
        if self.orbits is None:
            self.orbits = Cluster.clusters_to_orbit_isotropy(self.prim, self.clusters)
        return self.orbits

    def translate_to_supercell(self, sc, orb_idx=None, clus_idx=None):
        """

        :param sc: SupercellStructure
        :return: all clusters in the supercell. Each cluster is give as
                [list of the indices of the vertices, i.d. of orbit, i.d. of symop]
        orb_idx: default (None) to process all; explicit e.g. [0, 3, 5] to process selected few
        clus_idx: default (None) to process all; explicit e.g. [0] to process selected few
        """
        # natom = self.prim.num_sites

        # def match_ordering(l1, l2):
        #     """
        #     ordering such that l1[ordering] == l2
        #     """
        #     n= len(l1)
        #     ordering = list(range(n))

        orb_range = orb_idx if orb_idx is not None else range(len(self.orbits))
        allclus = []
        for iorb in orb_range:
            orbit = self.orbits[iorb]
            clus_range = clus_idx if clus_idx is not None else range(len(orbit.clusters))
            for ic in clus_range:
                clus = orbit.clusters[ic]
                ig = orbit.clusters_ig[ic]
                allclus.extend([cx, iorb, ig, ic] for cx in self.translate_cluster_to_supercell(sc, clus))
        return allclus

    @staticmethod
    def translate_cluster_to_supercell(sc, clus):
        """

        :param sc: SupercellStructure
        :return: all clusters in the supercell
        """
        if clus.order == 0:
            return [[]]*sc.n_cell
        coords= clus.frac_coords
        use_compiled_code = True
        if use_compiled_code:
            from f_util import f_util
            allclus= f_util.tr_cl_sc(sc._ijkl.T, sc.sc_mat.T, sc.inv_sc_mat.T, sc.sc_ref.T, sc.prim.frac_coords.T, clus._ijkls_np.T).T-1
        else:
            allclus = [get_structure_ordering((coords+ijk[None,:]).dot(sc.inv_sc_mat), sc.frac_coords,0) for ijk in sc.ijk_ref]
        return allclus

    def identify_cluster(self, clus_in):
        """
        Find if and which the input cluster matches the model clusters
        :param clus_try:
        :return: [matched cluster, iorb, ig, clus_sorted]
        """
        for iorb, orb in enumerate(self.orbits):
            # Quick return if not match obviously *)
            if orb.cluster._must_differ(clus_in):
                continue
            # logger.debug("  proceed with %d"%(iorb))
            for ic, clus_match in enumerate(orb.clusters):
                [flag, foundOrb] = clus_match.equivalent_by_lattranslation(clus_in.ijkls)
                if flag:
                    ig = orb.clusters_ig[ic]
                    # logger.debug("debug>  sum "+repr(clus_in)+" == orbit %d cluster %d sym %d matching="%(iorb, ic, ig) +str(foundOrb))
                    # logger.debug(" calc FC Trans mem=")#, MemoryInUse()/1E6)
                    return [True, iorb, ic, ig, relativePosition(clus_match.ijkls, foundOrb, dim=len(clus_in.orders))]
        return [False, -1, -1, None, None]

    def load_solution(self, sol_f):
        return np.loadtxt(sol_f, ndmin=2)

    def process_name_ord(self, name_ord, all_ord):
        if len(name_ord) == 0:
            name_ord = [['All', all_ord]]
        else:
            for i in name_ord:
                try:
                    i[1]= list(map(int, i[1]))
                except:
                    continue
                if i[0][-6:].lower()=='except':
                    i[1] = [j for j in all_ord if j not in i[1]]
                else:
                    i[1] = [j for j in i[1] if j in all_ord]
        return name_ord

    def get_submodels(self, name_ord, knownsol=None, **kwargs):
        """
        :param name_ord: list of [name, fct order], e.g. pair 0,2
                         or [name, filename] where filename is a matrix
        :return: list of [name, matrix] defining the different fittings
        """
        def down_select(o, sol0):
            if isinstance(o, list):
                return sum([list(range(self.ord_idx[i], self.ord_idx[i+1])) for i in o if i<=self.maxorder], [])
            elif o=='except-known':
                return np.where(sol0==0)[0]
            else:
                return np.loadtxt(o).reshape((-1))

        all_ord=list(set([o.cluster.order for o in self.orbits]))
        name_ord = self.process_name_ord(name_ord, all_ord)
        print("No. of parameters", {o: self.ord_idx[o+1]-self.ord_idx[o] for o in range(self.maxorder+1)})
        sol0 = np.zeros(self.ncorr)
        if knownsol:
             print("  Reading previous solution from %s"%(knownsol))
             input_sol = self.load_solution(knownsol).reshape(-1)
             sol0[:min(sol0.size, input_sol.size)] = input_sol[:min(sol0.size, input_sol.size)]
        return [[nm,
                scipy.sparse.identity(self.ncorr).tocsr()[:,
                  down_select(o, sol0)], sol0]
                for (nm, o) in name_ord]