<filename>students_final_projects/group-f/utilities/basic/coordinate.py
'''
Utility functions for positions and velocities.

@author: <NAME> <<EMAIL>>
'''

# system ----
from __future__ import absolute_import, division, print_function  # python 2 compatability
import numpy as np
from scipy import spatial
# local ----
from . import constant

#===================================================================================================
# coordinate transformation
#===================================================================================================
def get_positions_in_coordinate_system(
    position_vectors, system_from='cartesian', system_to='cylindrical'):
    '''
    Convert input 3-D position vectors from (cartesian, cylindrical, spherical) to
    (cartesian, cylindrical, spherical):
        cartesian : x, y, z
        cylindrical : R (along major axes, absolute/unsigned), Z (along minor axis, signed),
            angle phi [0, 2 * pi)
        spherical : r (absolute/unsigned), angle theta [0, pi), angle phi [0, 2 * pi)

    Parameters
    ----------
    position_vectors : array (object number x 3) : position[s]/distance[s] wrt a center

    Returns
    -------
    positions_new : array (object number x 3) : position[s]/distance[s] in new coordiante system
    '''
    assert system_from in ('cartesian', 'cylindrical', 'spherical')
    assert system_to in ('cartesian', 'cylindrical', 'spherical')

    if system_from == system_to:
        return position_vectors

    position_vectors = np.asarray(position_vectors)
    if np.ndim(position_vectors) == 1:
        position_vectors = np.asarray([position_vectors])

    assert np.shape(position_vectors)[1] == 3

    positions_new = np.zeros(position_vectors.shape, dtype=position_vectors.dtype)

    if system_from == 'cartesian':
        if system_to == 'cylindrical':
            # R = sqrt(x^2 + y^2)
            positions_new[:, 0] = np.sqrt(np.sum(position_vectors[:, [0, 1]] ** 2, 1))
            # Z = z
            positions_new[:, 1] = position_vectors[:, 2]
            # phi = arctan(y / x)
            positions_new[:, 2] = np.arctan2(position_vectors[:, 1], position_vectors[:, 0])
            positions_new[:, 2][positions_new[:, 2] < 0] += 2 * np.pi  # convert to [0, 2 * pi)
        elif system_to == 'spherical':
            # r = sqrt(x^2 + y^2 + z^2)
            positions_new[:, 0] = np.sqrt(np.sum(position_vectors ** 2, 1))
            # theta = arccos(z / r)
            positions_new[:, 1] = np.arccos(position_vectors[:, 2] / positions_new[:, 0])
            # phi = arctan(y / x)
            positions_new[:, 2] = np.arctan2(position_vectors[:, 1], position_vectors[:, 0])
            positions_new[:, 2][positions_new[:, 2] < 0] += 2 * np.pi  # convert to [0, 2 * pi)

    elif system_from == 'cylindrical':
        if system_to == 'cartesian':
            # x = R * cos(phi)
            positions_new[:, 0] = position_vectors[:, 0] * np.cos(position_vectors[:, 2])
            # y = R * sin(phi)
            positions_new[:, 1] = position_vectors[:, 0] * np.sin(position_vectors[:, 2])
            # z = Z
            positions_new[:, 2] = position_vectors[:, 1]
        elif system_to == 'spherical':
            # r = sqrt(R^2 + Z^2)
            positions_new[:, 0] = np.sqrt(position_vectors[:, 0] ** 2 + position_vectors[:, 1] ** 2)
            # theta = arctan(R / Z)
            positions_new[:, 1] = np.arctan2(position_vectors[:, 0], position_vectors[:, 1])
            # phi = phi
            positions_new[:, 2] = position_vectors[:, 2]

    elif system_from == 'spherical':
        if system_to == 'cartesian':
            # x = r * sin(theta) * cos(phi)
            positions_new[:, 0] = (position_vectors[:, 0] * np.sin(position_vectors[:, 1]) *
                                   np.cos(position_vectors[:, 2]))
            # y = r * sin(theta) * sin(phi)
            positions_new[:, 1] = (position_vectors[:, 0] * np.sin(position_vectors[:, 1]) *
                                   np.sin(position_vectors[:, 2]))
            # z = r * cos(theta)
            positions_new[:, 2] = position_vectors[:, 0] * np.cos(position_vectors[:, 1])
        elif system_to == 'cylindrical':
            # R = r * sin(theta)
            positions_new[:, 0] = position_vectors[:, 0] * np.sin(position_vectors[:, 1])
            # Z = r * cos(theta)
            positions_new[:, 1] = position_vectors[:, 0] * np.cos(position_vectors[:, 1])
            # phi = phi
            positions_new[:, 2] = position_vectors[:, 2]

    # if only one position vector, return as 1-D array
    if len(positions_new) == 1:
        positions_new = positions_new[0]

    return positions_new

def get_velocities_in_coordinate_system(
    velocity_vectors, position_vectors, system_from='cartesian', system_to='cylindrical'):
    '''
    Convert input 3-D velocity vectors from (cartesian, cylindrical, spherical) to
    (cartesian, cylindrical, spherical).
        cartesian : velocity along x, y, z
        cylindrical : velocity along R (major axes), Z (minor axis), angle phi
        spherical : velocity along r, angle theta, angle phi

    Parameters
    ----------
    velocity_vectors : array (object number x 3) : velocity[s] wrt a center
    position_vectors : array (object number x 3) : position[s]/distance[s] wrt a center

    Returns
    -------
    velocity_vectors_new : array (object number x 3) : velocity[s] in new coordiante system
    '''
    assert system_from in ('cartesian', 'cylindrical', 'spherical')
    assert system_to in ('cartesian', 'cylindrical', 'spherical')

    if system_from == system_to:
        return velocity_vectors

    velocity_vectors = np.asarray(velocity_vectors)
    if np.ndim(velocity_vectors) == 1:
        velocity_vectors = np.asarray([velocity_vectors])

    position_vectors = np.asarray(position_vectors)
    if np.ndim(position_vectors) == 1:
        position_vectors = np.asarray([position_vectors])

    assert np.shape(velocity_vectors)[1] == 3 and np.shape(position_vectors)[1] == 3

    velocities_new = np.zeros(velocity_vectors.shape, dtype=velocity_vectors.dtype)

    if system_from == 'cartesian':
        # convert position vectors
        # R = {x,y}
        R = position_vectors[:, [0, 1]]
        R_norm = np.zeros(R.shape, position_vectors.dtype)
        # R_total = sqrt(x^2 + y^2)
        R_total = np.sqrt(np.sum(R ** 2, 1))
        masks = np.where(R_total > 0)[0]
        # need to do this way
        R_norm[masks] = np.transpose(R[masks].transpose() / R_total[masks])

        if system_to == 'cylindrical':
            # v_R = dot(v_{x,y}, R_norm)
            velocities_new[:, 0] = np.sum(velocity_vectors[:, [0, 1]] * R_norm, 1)
            # v_Z = v_z
            velocities_new[:, 1] = velocity_vectors[:, 2]
            # v_phi = cross(R_norm, v_{x,y})
            velocities_new[:, 2] = np.cross(R_norm, velocity_vectors[:, [0, 1]])
        elif system_to == 'spherical':
            # convert position vectors
            position_vectors_norm = np.zeros(position_vectors.shape, position_vectors.dtype)
            position_vectors_total = np.sqrt(np.sum(position_vectors ** 2, 1))
            masks = np.where(position_vectors_total > 0)[0]
            # need to do this way
            position_vectors_norm[masks] = np.transpose(
                position_vectors[masks].transpose() / position_vectors_total[masks])

            # v_r = dot(v, r)
            velocities_new[:, 0] = np.sum(velocity_vectors * position_vectors_norm, 1)
            # v_theta
            a = np.transpose([R_norm[:, 0] * position_vectors_norm[:, 2],
                              R_norm[:, 1] * position_vectors_norm[:, 2],
                              -R_total / position_vectors_total])
            velocities_new[:, 1] = np.sum(velocity_vectors * a, 1)
            # v_phi = cross(R_norm, v_{x,y})
            velocities_new[:, 2] = np.cross(R_norm, velocity_vectors[:, [0, 1]])

    elif system_from == 'cylindrical':
        raise ValueError('not yet support conversion from {} to {}'.format(system_from, system_to))

    elif system_from == 'spherical':
        raise ValueError('not yet support conversion from {} to {}'.format(system_from, system_to))

    # if only one velocity vector, return as 1-D array
    if len(velocities_new) == 1:
        velocities_new = velocities_new[0]

    return velocities_new

#===================================================================================================
# rotation of position or velocity
#===================================================================================================
def get_coordinates_rotated(coordinate_vectors, rotation_tensor=None, rotation_angles=None):
    '''
    Get 3-D coordinate[s] (distance or velocity vector[s]) that are rotated by input rotation
    vectors or input rotation angles.
    If rotation_tensor, need to input vectors that are orthogonal.
    If rotation_angles, rotate by rotation_angles[0] about x-axis, then by rotation_angles[1] about
    y-axis, then by rotation_angles[2] about z-axis.

    Parameters
    ----------
    coordinate_vectors : array : coordinate[s] (distance[s] or velocity[s]) wrt a center of rotation
        (object number x dimension number)
    rotation_tensor : array : *orthogonal* rotation vectors (such as max, med, min eigen-vectors)
    rotation_angles : array : rotation angles about x-axis, y-axis, z-axis [radians]

    Returns
    -------
    coordinate[s] (distance[s] or velocity[s]) in rotated basis :
        array (object number x dimension number)
    '''
    if rotation_tensor is not None:
        # sanity check - ensure input rotation vectors are orthogonal
        tolerance = 1e-6
        if (np.abs(np.dot(rotation_tensor[0], rotation_tensor[1])) > tolerance or
                np.abs(np.dot(rotation_tensor[0], rotation_tensor[2])) > tolerance or
                np.abs(np.dot(rotation_tensor[1], rotation_tensor[2])) > tolerance):
            raise ValueError('input rotation_tensor is not orthogonal')

    elif rotation_angles is not None:
        m11 = np.cos(rotation_angles[1]) * np.cos(rotation_angles[2])
        m12 = (np.cos(rotation_angles[0]) * np.sin(rotation_angles[2]) +
               np.sin(rotation_angles[0]) * np.sin(rotation_angles[1]) * np.cos(rotation_angles[2]))
        m13 = (np.sin(rotation_angles[0]) * np.sin(rotation_angles[2]) -
               np.cos(rotation_angles[0]) * np.sin(rotation_angles[1]) * np.cos(rotation_angles[2]))
        m21 = -np.cos(rotation_angles[1]) * np.sin(rotation_angles[2])
        m22 = (np.cos(rotation_angles[0]) * np.cos(rotation_angles[2]) -
               np.sin(rotation_angles[0]) * np.sin(rotation_angles[1]) * np.sin(rotation_angles[2]))
        m23 = (np.sin(rotation_angles[0]) * np.cos(rotation_angles[2]) +
               np.cos(rotation_angles[0]) * np.sin(rotation_angles[1]) * np.sin(rotation_angles[2]))
        m31 = np.sin(rotation_angles[1])
        m32 = -np.sin(rotation_angles[0]) * np.cos(rotation_angles[1])
        m33 = np.cos(rotation_angles[0]) * np.cos(rotation_angles[1])

        rotation_tensor = np.array([
            [m11, m12, m13],
            [m21, m22, m23],
            [m31, m32, m33]],
            dtype=coordinate_vectors.dtype
        )

    else:
        raise ValueError('need to input either rotation angles or rotation vectors')

    # have to do this way
    coordinate_vectors_rotated = np.asarray(
        np.dot(coordinate_vectors, rotation_tensor.transpose()), dtype=coordinate_vectors.dtype)

    return coordinate_vectors_rotated

def get_principal_axes(position_vectors, weights=None, print_results=True):
    '''
    Compute principal axes of input position_vectors (which should be wrt a center),
    defined via the moment of inertia tensor.
    Get reverse-sorted eigen-vectors, eigen-values, and axis ratios of these principal axes.

    Parameters
    ----------
    position_vectors : array (object number x dimension number) :
        position[s]/distance[s] wrt a center
    weights : array : weight for each position (usually mass) - if None, assume all have same weight
    print_results : bool : whether to print axis ratios

    Returns
    -------
    eigen_vectors : array : max, med, min eigen-vectors
    eigen_values : array : max, med, min eigen-values
    axis_ratios : array : ratios of principal axes
    '''
    if weights is None or not len(weights):
        weights = 1
    else:
        weights = weights / np.median(weights)

    if position_vectors.shape[1] == 3:
        # 3-D
        xx = np.sum(weights * position_vectors[:, 0] ** 2)
        yy = np.sum(weights * position_vectors[:, 1] ** 2)
        zz = np.sum(weights * position_vectors[:, 2] ** 2)
        xy = yx = np.sum(weights * position_vectors[:, 0] * position_vectors[:, 1])
        xz = zx = np.sum(weights * position_vectors[:, 0] * position_vectors[:, 2])
        yz = zy = np.sum(weights * position_vectors[:, 1] * position_vectors[:, 2])

        moi_tensor = [[xx, xy, xz],
                      [yx, yy, yz],
                      [zx, zy, zz]]

    elif position_vectors.shape[1] == 2:
        # 2-D
        xx = np.sum(weights * position_vectors[:, 0] ** 2)
        yy = np.sum(weights * position_vectors[:, 1] ** 2)
        xy = yx = np.sum(weights * position_vectors[:, 0] * position_vectors[:, 1])

        moi_tensor = [[xx, xy],
                      [yx, yy]]

    eigen_values, eigen_vectors = np.linalg.eig(moi_tensor)

    # order eigen-vectors by eigen-values, from largest to smallest
    eigen_indices_sorted = np.argsort(eigen_values)[::-1]
    eigen_values = eigen_values[eigen_indices_sorted]
    eigen_values /= eigen_values.max()  # renormalize to 1
    # make eigen_vectors[0] corresponds to vector of eigen_values[0]
    eigen_vectors = eigen_vectors.transpose()[eigen_indices_sorted]

    if position_vectors.shape[1] == 3:
        axis_ratios = np.sqrt(
            [eigen_values[2] / eigen_values[0],
             eigen_values[2] / eigen_values[1],
             eigen_values[1] / eigen_values[0]]
        )

        if print_results:
            print('* principal axes:  min/maj = {:.3f}, min/med = {:.3f}, med/maj = {:.3f}'.format(
                  axis_ratios[0], axis_ratios[1], axis_ratios[2]))

    elif position_vectors.shape[1] == 2:
        axis_ratios = eigen_values[1] / eigen_values[0]

        if print_results:
            print('* principal axes:  min/maj = {:.3f}'.format(axis_ratios))

    return eigen_vectors, eigen_values, axis_ratios

#===================================================================================================
# position distances
#===================================================================================================
def get_positions_periodic(positions, periodic_length=None):
    '''
    Get position in range [0, periodic_length).

    Parameters
    ----------
    positions : float or array
    periodic_length : float : periodicity length (if none, return array as is)
    '''
    if periodic_length is None:
        return positions

    if np.isscalar(positions):
        if positions >= periodic_length:
            positions -= periodic_length
        elif positions < 0:
            positions += periodic_length
    else:
        positions[positions >= periodic_length] -= periodic_length
        positions[positions < 0] += periodic_length

    return positions

def get_position_differences(position_difs, periodic_length=None):
    '''
    Get distance / separation vector, in range [-periodic_length/2, periodic_length/2).

    Parameters
    ----------
    position_difs : array : position difference[s]
    periodic_length : float : periodicity length (if none, return array as is)
    '''
    if not periodic_length:
        return position_difs
    else:
        if np.isscalar(periodic_length) and periodic_length <= 1:
            print('! got unusual periodic_length = {}'.format(periodic_length))

    if np.isscalar(position_difs):
        if position_difs >= 0.5 * periodic_length:
            position_difs -= periodic_length
        elif position_difs < -0.5 * periodic_length:
            position_difs += periodic_length
    else:
        position_difs[position_difs >= 0.5 * periodic_length] -= periodic_length
        position_difs[position_difs < -0.5 * periodic_length] += periodic_length

    return position_difs

def get_distances(
    positions_1=None, positions_2=None, periodic_length=None, scalefactor=None,
    total_distance=False):
    '''
    Get vector or total/scalar distance[s] between input position vectors.
    If input scale-factors, will convert distance from comoving to physical.

    Parameters
    ----------
    positions_1 : array : position[s]
    positions_2 : array : position[s]
    periodic_length : float : periodic length (if none, not use periodic)
    scalefactor : float or array : expansion scale-factor (to convert comoving to physical)
    total : bool : whether to compute total/scalar (instead of vector) distance

    Returns
    -------
    distances : array (object number x dimension number, or object number) :
        vector or total/scalar distance[s]
    '''
    if not isinstance(positions_1, np.ndarray):
        positions_1 = np.array(positions_1)
    if not isinstance(positions_2, np.ndarray):
        positions_2 = np.array(positions_2)

    if len(positions_1.shape) == 1 and len(positions_2.shape) == 1:
        shape_pos = 0
    else:
        shape_pos = 1

    distances = get_position_differences(positions_1 - positions_2, periodic_length)

    if total_distance:
        distances = np.sqrt(np.sum(distances ** 2, shape_pos))

    if scalefactor is not None:
        if scalefactor > 1 or scalefactor <= 0:
            print('! got unusual scalefactor = {}'.format(scalefactor))
        distances *= scalefactor

    return distances

def get_distances_angular(positions_1=None, positions_2=None, sphere_angle=360):
    '''
    Get angular separation[s] between input positions, valid for small separations.

    Parameters
    ----------
    positions_1, positions_2 : arrays : positions in [RA, dec]
    sphere_angle : float : angular size of sphere 360 [degrees], 2 * pi [radians]

    Returns
    -------
    angular distances : array (object number x angular dimension number)
    '''
    if sphere_angle == 360:
        angle_scale = constant.radian_per_degree
    elif sphere_angle == 2 * np.pi:
        angle_scale = 1
    else:
        raise ValueError('angle of sphere = {} does not make sense'.format(sphere_angle))

    if np.ndim(positions_1) == 1 and positions_1.size == 2:
        ras_1, decs_1 = positions_1[0], positions_1[1]
    else:
        ras_1, decs_1 = positions_1[:, 0], positions_1[:, 1]

    if np.ndim(positions_2) == 1 and positions_2.size == 2:
        ras_2, decs_2 = positions_2[0], positions_2[1]
    else:
        ras_2, decs_2 = positions_2[:, 0], positions_2[:, 1]

    return np.sqrt((get_position_differences(ras_1 - ras_2, sphere_angle) *
                    np.cos(angle_scale * 0.5 * (decs_1 + decs_2))) ** 2 + (decs_1 - decs_2) ** 2)

#===================================================================================================
# velocity conversion
#===================================================================================================
def get_velocity_differences(
    velocity_vectors_1=None, velocity_vectors_2=None,
    position_vectors_1=None, position_vectors_2=None, periodic_length=None,
    scalefactor=None, hubble_time=None,
    total_velocity=False):
    '''
    Get relative velocity[s] [km / s] between input velocity vectors.
    If input positions as well, add Hubble flow to velocities.

    Parameters
    ----------
    velocity_vectors_1 : array : velocity[s] (object number x dimension number) [km / s]
    velocity_vectors_2 : array : velocity[s] (object number x dimension number) [km / s]
    position_vectors_1 : array : position[s] associated with velocity_vector_1
        (object number x dimension number) [kpc comoving]
    position_vectors_2 : array : position[s] associated with velocity_vector_2
        (object number x dimension number) [kpc comoving]
    periodic_length : float : periodicity length [kpc comoving]
    scalefactor : float : expansion scale-factor
    hubble_time : float : 1 / H(z) [Gyr]
    total_velocity : bool : whether to compute total/scalar (instead of vector) velocity

    Returns
    -------
    velocity_difs  : array (object number x dimension number, or object number) :
        velocity differences [km / s]
    '''
    if np.ndim(velocity_vectors_1) == 1 and np.ndim(velocity_vectors_1) == 1:
        dimension_shape = 0
    else:
        dimension_shape = 1

    velocity_difs = velocity_vectors_1 - velocity_vectors_2  # [km / s]

    if position_vectors_1 is not None and position_vectors_2 is not None:
        # add hubble flow: dr/dt = a * dx/dt + da/dt * x = a(t) * dx/dt + r * H(t)
        # [kpc / Gyr]
        vels_hubble = (scalefactor / hubble_time *
                       get_distances(position_vectors_1, position_vectors_2, periodic_length))
        vels_hubble *= constant.km_per_kpc / constant.sec_per_Gyr  # [km / s]
        velocity_difs += vels_hubble

    if total_velocity:
        velocity_difs = np.sqrt(np.sum(velocity_difs ** 2, dimension_shape))

    return velocity_difs

#===================================================================================================
# center of mass: position and velocity
#===================================================================================================
def get_center_position_zoom(
    positions, weights=None, periodic_length=None, position_number_min=32, center_position=None,
    distance_max=np.Inf):
    '''
    Get position of center of mass, using iterative zoom-in.

    Parameters
    ----------
    positions : array (particle number x dimension number) : position[s]
    weights : array : weight for each position (usually mass) - if None, assume all have same weight
    periodic_length : float : periodic box length
    position_number_min : int : minimum number of positions within distance to keep zooming in
    center_position : array : initial center position to use
    distance_max : float : maximum distance to consider initially

    Returns
    -------
    center_position : array : position vector of center of mass
    '''
    distance_bins = np.array([
        np.Inf, 1000, 700, 500, 300, 200, 150, 100,
        70, 50, 30, 20, 15, 10,
        7, 5, 3, 2, 1.5, 1,
        0.7, 0.5, 0.3, 0.2, 0.15, 0.1,
        0.07, 0.05, 0.03, 0.02, 0.015, 0.01,
        0.007, 0.005, 0.003, 0.002, 0.0015, 0.001,
    ])
    distance_bins = distance_bins[distance_bins <= distance_max]

    if weights is not None:
        assert positions.shape[0] == weights.size
        # normalizing weights by median seems to improve numerical stability
        weights = np.asarray(weights) / np.median(weights)

    if center_position is None or not len(center_position):
        center_position = np.zeros(positions.shape[1], positions.dtype)
    else:
        center_position = np.array(center_position, positions.dtype)

    if positions.shape[0] > 2147483647:
        idtype = np.int64
    else:
        idtype = np.int32
    part_indices = np.arange(positions.shape[0], dtype=idtype)

    for dist_i, dist_max in enumerate(distance_bins):
        ## direct method ----------
        distance2s = get_position_differences(
            positions[part_indices] - center_position, periodic_length) ** 2
        distance2s = np.sum(distance2s, 1)

        # get particles within distance max
        masks = (distance2s < dist_max ** 2)
        part_indices_dist = part_indices[masks]

        # store particles slightly beyond distance max for next interation
        masks = (distance2s < (1.5 * dist_max) ** 2)
        part_indices = part_indices[masks]

        """
        ## kd-tree method ----------
        if dist_i == 0:
            # does not handle periodic boundaries, but should be ok for zoom-in
            KDTree = spatial.cKDTree(positions)
            particle_number_max = positions.shape[0]

        distances, indices = KDTree.query(
            center_position, particle_number_max, distance_upper_bound=dist_max)

        masks = (distances < dist_max)
        part_indices_dist = indices[masks]
        particle_number_max = part_indices_dist.size
        """

        # check whether reached minimum total number of particles within distance
        # but force at least one loop over distance bins to get *a* center
        if part_indices_dist.size <= position_number_min and dist_i > 0:
            return center_position

        if weights is None:
            weights_use = weights
        else:
            weights_use = weights[part_indices_dist]

        # ensure that np.average uses 64-bit internally for accuracy, but returns as input dtype
        center_position = np.average(
            positions[part_indices_dist].astype(np.float64), 0, weights_use).astype(positions.dtype)

    return center_position

def get_center_velocity(
    velocities, weights=None, positions=None, center_position=None, distance_max=20,
    periodic_length=None):
    '''
    Get velocity of center of mass.
    If no input masses, assume all masses are the same.

    Parameters
    ----------
    velocities : array (particle number x 3) : velocity[s]
    weights : array : weight for each position (usually mass) - if None, assume all have same weight
    positions : array : positions (particle number x dimension number), if want to select by this
    center_position : array : center position, if want to select by this
    distance_max : float : maximum position difference from center to use particles
    periodic_length : float : periodic box length

    Returns
    -------
    center_velocity : array : velocity vector of center of mass
    '''
    masks = np.full(velocities.shape[0], True, dtype=bool)

    # ensure that use only finite values
    for dimen_i in range(velocities.shape[1]):
        masks *= np.isfinite(velocities[:, dimen_i])

    if positions is not None and center_position is not None and len(center_position):
        assert velocities.shape == positions.shape
        distance2s = np.sum(
            get_position_differences(positions - center_position, periodic_length) ** 2, 1)
        masks *= (distance2s < distance_max ** 2)

    if weights is not None:
        assert velocities.shape[0] == weights.size
        # normalizing weights by median seems to improve numerical stability
        weights = weights[masks] / np.median(weights[masks])

    if not masks.any():
        print('! could not compute center velocity')
        print('  no particles within distance_max = {} kpc (nearest = {} kpc)'.format(
            distance_max, np.sqrt(distance2s.min())))
        return np.r_[np.nan, np.nan, np.nan]

    # ensure that np.average uses 64-bit internally for accuracy, but returns as input dtype
    return np.average(velocities[masks].astype(np.float64), 0, weights).astype(velocities.dtype)

#===================================================================================================
# volume of region
#===================================================================================================
def get_volume_of_convex_hull(positions):
    '''
    Compute volume of convex hull that encloses input positions.

    Parameters
    ----------
    positions : array : positions (object number x dimension number)

    Returns
    -------
    volume within convex hull around positions : float
    '''

    def get_tetrahedron_volume(a, b, c, d):
        return np.abs(np.einsum('ij,ij->i', a - d, np.cross(b - d, c - d))) / 6

    ConvexHull = spatial.ConvexHull(positions)
    DelaunayTes = spatial.Delaunay(positions[ConvexHull.vertices])
    tets = DelaunayTes.points[DelaunayTes.simplices]

    return np.sum(get_tetrahedron_volume(tets[:, 0], tets[:, 1], tets[:, 2], tets[:, 3]))

#===================================================================================================
# coordinates in redshift space
#===================================================================================================
def convert_velocity_redshift(value_kind, values, solve_exact=True):
    '''
    Get velocity/redshift along the line of sight from redshift/velocity [km / s].
    Independent of cosmology.

    Parameters
    ----------
    value_kind: str : 'redshift', 'velocity'
    values : float or array : redshift/velocity value[s]
    solve_exact : bool : whether to use exact solution or faster approximation

    Returns
    -------
    velocities [km / s] or redshifts : array (object number)
    '''
    if value_kind == 'redshift':
        # input redshift, get velocity
        if solve_exact:
            return (((1 + values) ** 2 - 1) / ((1 + values) ** 2 + 1) * constant.speed_light *
                    constant.kilo_per_centi)
        else:
            return constant.speed_light * values * constant.kilo_per_centi

    elif value_kind == 'velocity':
        # input velocity, get redshift
        if solve_exact:
            return ((1 + values * constant.centi_per_kilo / constant.speed_light) /
                    (1 - values * constant.centi_per_kilo / constant.speed_light)) ** 0.5 - 1
        else:
            return values * constant.centi_per_kilo / constant.speed_light

def get_position_difs_from_redshift_difs(redshift_difs, hubble_times):
    '''
    Get position difference [kpc comoving] from redshift difference (redshift-space distortion).
    *** distance is *approximate*, valid in non-relativistic limit.

    Parameters
    ----------
    redshift_difs : float or array : redshift difference[s]
    hubble_times : float or array : hubble time[s] = 1 / H [Gyr]

    Returns
    -------
    position differences [kpc comoving] : array (object number)
    '''
    return convert_velocity_redshift('redshift', redshift_difs, solve_exact=False) * hubble_times

def get_position_difs_from_velocity_difs(velocity_difs, hubble_times, redshifts):
    '''
    Get position difference [kpc comoving] from velocity difference (redshifts-space distortion).

    Parameters
    ----------
    velocity_difs : floar or array : peculiar velocity[s] [km / s]
    hubble_time : float or array : hubble time[s] = 1 / H [Gyr]
    redshifts : float or array

    Returns
    -------
    position differences [kpc comoving]: array (object number)
    '''
    return ((velocity_difs / (1 + redshifts) * constant.kpc_per_km / constant.Gyr_per_sec) *
            hubble_times)

def get_positions_in_redshift_space(
    positions, velocities, hubble_times, redshifts, periodic_length=None):
    '''
    Get position[s] [kpc comoving] in redshifts space, convolving real position[s] with
    redshifts-space distortion.

    Parameters
    ----------
    positions : float or array : actual position[s] [kpc comoving]
    velocities : float or array : peculiar velocity[s] [km / s]
    hubble_times : float or array : hubble time[s] = 1 / H [Gyr]
    periodic_length : float : periodicity length

    Returns
    -------
    positions [kpc comoving] : array (object number)
    '''
    return positions(
        positions + get_position_difs_from_velocity_difs(velocities, hubble_times, redshifts),
        periodic_length)


<reponame>snumrl/skate
import numpy as np
import math
from scipy.optimize import minimize

class IKsolver(object):
    def __init__(self, skel1, skel2, left_foot_traj, right_foot_traj, t, pend):
        self.left_foot_traj = left_foot_traj
        self.right_foot_traj = right_foot_traj
        self.skel1 = skel1
        self.skel2 = skel2
        self.gtime = t
        self.pendulum = pend

    def update_target(self, ):
        # print("self.gtime:", self.gtime)
        # self.target_foot = np.array([skel1.q["j_cart_x"], -0.92, skel1.q["j_cart_z"]])

        # ground_height = -1.1
        ground_height = -0.92
        self.target_left_foot = np.array(
            [self.skel1.q["j_cart_x"], ground_height, self.left_foot_traj[1][self.gtime] - 0.05 - 0.3])
        self.target_right_foot = np.array(
            [self.skel1.q["j_cart_x"], ground_height, self.right_foot_traj[1][self.gtime] + 0.05 + 0.3])
        # l = self.skeletons[1].body('h_pole').local_com()[1]
        l = self.pendulum.length
        # self.target_pelvis = np.array([0, 0, 0])
        self.target_pelvis = np.array(
            [self.skel1.q["j_cart_x"] - l * math.sin(self.pendulum.theta), -0.92 + l * math.cos(self.pendulum.theta), 0])

    def set_params(self, x):
        q = self.skel2.q
        q = x
        # q[6:] = x
        self.skel2.set_positions(q)

    def f(self, x):

        def rotationMatrixToAxisAngles(R):
            temp_r = np.array([R[2,1]-R[1,2], R[0,2]-R[2,0], R[1,0]-R[0,1]])
            angle = math.acos((R[0, 0]+R[1, 1]+R[2, 2] - 1)/2.)
            temp_r_norm = np.linalg.norm(temp_r)
            if temp_r_norm < 0.000001:
                return np.zeros(3)

            return temp_r/temp_r_norm * angle

        def rotationMatrixToEulerAngles(R):

            # assert (isRotationMatrix(R))

            sy = math.sqrt(R[0, 0] * R[0, 0] + R[1, 0] * R[1, 0])

            singular = sy < 1e-6

            if not singular:
                x = math.atan2(R[2, 1], R[2, 2])
                y = math.atan2(-R[2, 0], sy)
                z = math.atan2(R[1, 0], R[0, 0])
            else:
                x = math.atan2(-R[1, 2], R[1, 1])
                y = math.atan2(-R[2, 0], sy)
                z = 0

            return np.array([x, y, z])

        self.set_params(x)
        pelvis_state = self.skel2.body("h_pelvis").to_world([0.0, 0.0, 0.0])
        # pelvis_state = self.skel2.body("h_abdomen").to_world([0.0, 0.0, 0.0])
        left_foot_state = self.skel2.body("h_blade_left").to_world([0.0, 0.0, 0.0])
        left_foot_ori = self.skel2.body("h_blade_left").world_transform()
        right_foot_state = self.skel2.body("h_blade_right").to_world([0.0, 0.0, 0.0])
        right_foot_ori = self.skel2.body("h_blade_right").world_transform()

        abdomen_ori = self.skel2.body("h_abdomen").world_transform()

        # print("foot_ori: ", left_foot_ori[:3, :3], type(left_foot_ori[:3, :3]))

        # todo : calculate the x_axis according to the spline

        # x_axis = np.array([0., 0., self.left_der[0][self.gtime]])
        x_axis = np.array([1.0, 0., 1.0])
        x_axis = x_axis / np.linalg.norm(x_axis)  # normalize
        y_axis = np.array([0., 1., 0.])
        z_axis = np.cross(x_axis, y_axis)
        z_axis = z_axis / np.linalg.norm(z_axis)  # normalize
        R_des = np.stack((x_axis, y_axis, z_axis), axis=-1)
        # print("x_axis: \n")
        # print(x_axis)
        # print("y_axis: \n")
        # print(y_axis)
        # print("z_axis: \n")
        # print(z_axis)
        # print("R: \n")
        # print(R_des)

        # left_axis_angle =rotationMatrixToAxisAngles(left_foot_ori[:3, :3])
        # right_axis_angle = rotationMatrixToAxisAngles(right_foot_ori[:3, :3])

        left_res_rot = np.transpose(R_des).dot(left_foot_ori[:3, :3])
        right_res_rot = np.transpose(R_des).dot(right_foot_ori[:3, :3])
        # print("vec res: ", left_res_rot)
        left_axis_angle = rotationMatrixToAxisAngles(left_res_rot)
        right_axis_angle = rotationMatrixToAxisAngles(right_res_rot)

        abdomen_axis_angle = rotationMatrixToAxisAngles(abdomen_ori[:3, :3])

        # print("axis_angle", axis_angle)

        # return 0.5 *  np.linalg.norm(left_foot_state - self.target_foot) ** 2
        return 0.5 * (np.linalg.norm(pelvis_state - self.target_pelvis) ** 2 + np.linalg.norm(
            left_foot_state - self.target_left_foot) ** 2 + np.linalg.norm(
            right_foot_state - self.target_right_foot) ** 2 + np.linalg.norm(
            left_axis_angle - np.array([0., 0., 0.])) ** 2 + np.linalg.norm(
            right_axis_angle - np.array([0., 0., 0.])) ** 2 + np.linalg.norm(
            abdomen_axis_angle - np.array([0., 0., 0.])) ** 2)

    def g(self, x):
        self.set_params(x)

        pelvis_state = self.skel2.body("h_pelvis").to_world([0.0, 0.0, 0.0])
        # pelvis_state = self.skel2.body("h_abdomen").to_world([0.0, 0.0, 0.0])
        left_foot_state = self.skel2.body("h_blade_left").to_world([0.0, 0.0, 0.0])
        right_foot_state = self.skel2.body("h_blade_right").to_world([0.0, 0.0, 0.0])

        J_pelvis = self.skel2.body("h_pelvis").linear_jacobian()
        J_left_foot = self.skel2.body("h_blade_left").linear_jacobian()
        J_right_foot = self.skel2.body("h_blade_right").linear_jacobian()

        J_temp = np.vstack((J_pelvis, J_left_foot))
        J = np.vstack((J_temp, J_right_foot))
        AA = np.append(pelvis_state - self.target_pelvis, left_foot_state - self.target_left_foot)
        AAA = np.append(AA, right_foot_state - self.target_right_foot)

        # print("pelvis", pelvis_state - self.target_pelvis)
        # print("J", J_pelvis)
        #
        # print("AA", AA)
        # print("J", J)

        # g = AA.dot(J)
        g = AAA.dot(J)
        # g = AA.dot(J)[6:]

        return g

    def solve(self, ):
        q_backup = self.skel2.q
        res = minimize(self.f, x0= self.skel2.q, jac=self.g, method="SLSQP")
        # res = minimize(self.f, x0=self.skeletons[2].q[6:], jac=self.g, method="SLSQP")
        # print(res)
        self.skel2.set_positions(q_backup)

        return res['x']


<reponame>RobbinBouwmeester/oxPTMScanner_cmd<filename>oxPTM_scanner.py
"""
Copyright 2017 <NAME>

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

"""
This software can be used to flag MS2 spectra based on belonging to one
of two classes of spectra. The two classes this code has been tested on
oxPTM VS native spectra. The code can be used to train a new model that
discriminates between two classes. After a model has been trained
the code can be used to make predictions for unseen spectra.

Library versions:

Python 2.7.13
xgboost.__version__ = '0.6'
sklearn.__version__ = '0.19.0'
scipy.__version__ = '0.19.1'
numpy.__version__ = '1.13.3'
pandas.__version__ = '0.20.3'

This project was made possible by MASSTRPLAN. MASSTRPLAN received funding
from the Marie Sklodowska-Curie EU Framework for Research and Innovation
Horizon 2020, under Grant Agreement No. 675132.
"""

# TODO make webversion

__author__ = "<NAME>"
__copyright__ = "Copyright 2017"
__credits__ = ["<NAME>","<NAME>","Prof. <NAME>"]
__license__ = "Apache License, Version 2.0"
__version__ = "1.0"
__maintainer__ = "<NAME>"
__email__ = "<EMAIL>"
__status__ = "Beta version; ready for in-field application"

#Native library
import argparse
import pickle
import copy
from operator import itemgetter
from itertools import combinations
import os

#Pandas
import pandas as pd

#Numpy
import numpy as np
from numpy.random import ranf
import numpy as np
import numpy.random as np_random

#Matplotlib
import matplotlib.pyplot as plt

#SciPy
import scipy.stats as st
import scipy

#ML
from sklearn.model_selection import RandomizedSearchCV
from sklearn.metrics import roc_curve
from sklearn.model_selection import cross_val_predict
from sklearn.metrics import auc
from sklearn.model_selection import StratifiedKFold
import xgboost as xgb

gui_object = None

class Dummy(): pass

def parse_msp(msp_entry,tic_normalization=True,min_perc=False,windowed_mode=False,top=10,window_size=100):
    """
    Parse an MSP entry and return the identifier, peaks and intensity. Normalization
    can be enabled and filtering on the intensity of peaks.

    Parameters
    ----------
    msp_entry : list
        list with the lines of the MSP entry
    tic_normalization : bool
        Return tic normalized intensities
    min_perc : bool
        Flag to use a minimal percentage intensity to filter peaks
    windowed_mode : bool
        Flag to use windowed mode to return the top intensity peaks
    top : int
        The top intensity peaks to filter (in windowed mode it will return the top peaks within the window)
    window_size : int
        The size of the window in windowed mode

    Returns
    -------
    str
        identifier of the msp entry
    list
        the m/z values in the msp entry
    list
        the intensity values in the msp entry
    """

    identifier = ""
    mz_list = []
    intensity_list = []
    if tic_normalization: tot_tic = 0.0

    #Iterate over the lines in the MSP entry and record the identifiers, m/z and intensities
    for line in msp_entry:
        line = line.rstrip()
        if line == "": continue
        if line.startswith("Name: "):
            identifier = line.lstrip("Name: ")
            continue
        if ":" in line: continue

        splitline = line.split("\t")

        mz_list.append(float(splitline[0]))
        intensity_list.append(float(splitline[1]))
        if tic_normalization: tot_tic += intensity_list[-1]

    #In the case of tic normalization iterate over the values and divide by total intensity
    if tic_normalization:
        for index,intens in enumerate(intensity_list):
            intensity_list[index] = intens/tot_tic

    #Filter based on the top intensities
    gr_mz_list,gr_intensity_list = get_top_spec(mz_list,
                                                intensity_list,
                                                min_perc=min_perc,
                                                windowed_mode=windowed_mode,
                                                top=top,
                                                window_size=window_size)

    return(identifier,gr_mz_list,gr_intensity_list)

def parse_mgf(mgf_entry,tic_normalization=True,min_perc=False,top=10,window_size=100,windowed_mode=False):
    """
    Parse an MGF entry and return the identifier, peaks and intensity. Normalization
    can be enabled and filtering on the intensity of peaks.

    Parameters
    ----------
    mgf_entry : list
        list with the lines of the MGF entry
    tic_normalization : bool
        Return tic normalized intensities
    min_perc : bool
        Flag to use a minimal percentage intensity to filter peaks
    windowed_mode : bool
        Flag to use windowed mode to return the top intensity peaks
    top : int
        The top intensity peaks to filter (in windowed mode it will return the top peaks within the window)
    window_size : int
        The size of the window in windowed mode

    Returns
    -------
    str
        identifier of the MGF entry
    list
        the m/z values in the MGF entry
    list
        the intensity values in the MGF entry
    """

    identifier = ""
    mz_list = [0.0]
    intensity_list = [1.0]
    if tic_normalization: tot_tic = 0.0

    #Iterate over the lines in the MGF entry
    for line in mgf_entry:
        line = line.rstrip()
        if line == "": continue
        if line.startswith("TITLE="):
            identifier = line.lstrip("TITLE=")
            continue
        if "=" in line: continue

        if "\t" in line: splitline = line.split("\t")
        else: splitline = line.split(" ")

        mz_list.append(float(splitline[0]))
        intensity_list.append(float(splitline[1]))
        if tic_normalization: tot_tic += intensity_list[-1]

    #In the case of tic normalization iterate over the values and divide by total intensity
    if tic_normalization:
        for index,intens in enumerate(intensity_list):
            intensity_list[index] = intens/tot_tic

    #Filter based on the top intensities
    gr_mz_list,gr_intensity_list = get_top_spec(mz_list,
                                                intensity_list,
                                                min_perc=min_perc,
                                                windowed_mode=windowed_mode,
                                                top=top,
                                                window_size=window_size)

    return(identifier,gr_mz_list,gr_intensity_list)

def get_top_spec(mz_list,intensity_list,min_perc=False,windowed_mode=False,top=10,window_size=100,add_dummy_peak=True):
    """
    Filter in multiple ways on the intensity of peaks.

    Parameters
    ----------
    mz_list : list
        The m/z values of a spectrum in a list; equal length to the intensity list
    intensity_list : list
        The intensity values of a spectrum in a list; equal length to the m/z list
    min_perc : bool
        Flag to use a minimal percentage intensity to filter peaks
    windowed_mode : bool
        Flag to use windowed mode to return the top intensity peaks
    top : int
        The top intensity peaks to filter (in windowed mode it will return the top peaks within the window)
    window_size : int
        The size of the window in windowed mode
    add_dummy_peak : bool
        Flag to add a dummy peak at 0.0 m/z

    Returns
    -------
    list
        the filtered m/z values from the spectrum
    list
        the filtered intensity values from the spectrum
    """
    gr_intensity_list = []
    gr_mz_list = []

    #In the case of minimal percentage... calculate perc intensity and filter
    if min_perc:
        for i,mz in zip(intensity_list,mz_list):
            if i > min_perc:
                gr_intensity_list.append(i)
                gr_mz_list.append(mz)

    #In the case of windowed mode... iterate over the possible windows and intensity values; take the top per window
    if windowed_mode:
        start_index = 0
        for w in range(window_size,int(max(mz_list)),window_size):
            temp_mz = []
            temp_intens = []
            temp_start_index = 0

            #Iterate over all m/z values and see if they fall within the window
            for mz,intens in zip(mz_list[start_index:],intensity_list[start_index:]):
                if mz > w and mz <= w+window_size:
                    temp_start_index += 1
                    temp_mz.append(mz)
                    temp_intens.append(intens)
                if mz > w+window_size:
                    break
            #Next window ignore all these lower values
            start_index = copy.deepcopy(temp_start_index)

            #Use all if there are less peaks than the top number of peaks it should select
            if len(temp_mz) <= top:
                gr_mz_list.extend(temp_mz)
                gr_intensity_list.extend(temp_intens)
                continue

            #Get the indexes of the top peaks
            idxs = np.sort(np.argpartition(np.array(temp_intens), -top)[-top:])
            gr_mz_list.extend([temp_mz[idx] for idx in idxs])
            gr_intensity_list.extend([temp_intens[idx] for idx in idxs])

    #If not windowed or min perc use a simple top peaks
    if not windowed_mode and not min_perc:
        if len(intensity_list) > top:
            #Get the indexes of the top peaks
            idxs = np.sort(np.argpartition(np.array(intensity_list), -top)[-top:])
            gr_mz_list = [mz_list[idx] for idx in idxs]
            gr_intensity_list = [intensity_list[idx] for idx in idxs]
        else:
            #If there are less peaks than top peaks; return all
            gr_mz_list = mz_list
            gr_intensity_list = intensity_list

    #If needed add a dummy peak; this is important later since I want to take into account immonium ions and small fragments
    if add_dummy_peak:
        gr_mz_list.insert(0,0.0)
        gr_intensity_list.insert(0,1.0)

    return(gr_mz_list,gr_intensity_list)

def get_feats(mz_list,intensity_list,feat_matrix,instance_index,feats,max_dist=275,allowed_c=[],pseudocount=0.000000001):
    """
    Retrieve features and write them to a matrix.

    Parameters
    ----------
    mz_list : list
        The m/z values of a spectrum in a list; equal length to the intensity list
    intensity_list : list
        The intensity values of a spectrum in a list; equal length to the m/z list
    feat_matrix : lil or csr matrix
        Sparse matrix that should be used to fill in the features from the m/z and intensity lists
    instance_index : int
        Row number in the matrix where the features should be filled in (indexing starts at 0)
    feats : list
        The bins used for features; should be sorted!; features are assigned if value is in between values of this list
    max_dist : int
        Maximum distance between peaks
    allowed_c : list
        Allowed bins (selected features); not used any more

    Returns
    -------
    matrix
        A sparse matrix (lil or csr) is returned with filled in features
    """
    # UNCOMMENT var below if standard library combinations is used
    #allowed_c = set(allowed_c)

    spectrum = zip(mz_list,intensity_list)
    dists_mz = []
    dists_mz_intens = []
    prev_analyzed = set()

    #Make deepcopy since we are going to change the spectra!
    spec_one = list(copy.deepcopy(spectrum))
    spec_two = list(copy.deepcopy(spectrum))

    #Iterate over the peaks and measure the distance in m/z between all combinations
    for peak_one in spec_one:
        if len(spec_two) == 1: continue
        spec_two = spec_two[1:]
        for peak_two in spec_two:
            dist_mz = abs(peak_one[0]-peak_two[0])
            if dist_mz > max_dist: break
            dists_mz.append(dist_mz)
            dists_mz_intens.append(peak_one[1]+peak_two[1])

    # UNCOMMENT code below if standard library combinations is used
    #for c in combinations(spectrum,2):
    #   dist_mz = abs(c[0][0]-c[1][0])
    #   if c[0][0] in prev_analyzed: continue
    #   if dist_mz > max_dist:
    #       prev_analyzed.add(c[0][0])
    #       continue
    #   if len(allowed_c) != 0:
    #       if dist_mz not in allowed_c: continue
    #   dists_mz.append(dist_mz)
    #   dists_mz_intens.append(c[0][1]+c[1][1])

    #Digitize the delta m/z; assign bins for all delta m/z s
    index_bins = np.digitize(dists_mz,feats)

    #Iterate over assigned bins and sum the intensity for possible existing values
    for index,intens in zip(index_bins,dists_mz_intens):
        feat_matrix[instance_index,index-1] += intens

    feat_matrix[instance_index,feat_matrix.shape[1]-1] += pseudocount
    return(feat_matrix)

def read_msp(infile_name,feat_lim_file="",
             sum_feats=False,selected_features=[],
             max_dist=275,step_size=0.005,feat_bins=[],
             top_peaks=50,windowed_mode=False):
    """
    Read an MSP file and put the features into a matrix.

    Parameters
    ----------
    infile_name : list
        The infile MSP file.
    feat_lim_file : list
        Old variable with the name of a file that contains the features.
    sum_feats : bool
        Old variable used to sum features of the two classes.
    selected_features : list
        Old variable for selected features; use feat_bins.
    max_dist : int
        Maximum distance between peaks
    step_size : float
        Size between the m/z values for bins.
    feat_bins : list
        Bins to use for features.
    top_peaks : int
        Number of peaks to select based on the intensity
    windowed_mode : bool
        Flag to used windowed mode for selecting the highest intensity peaks

    Returns
    -------
    matrix
        A sparse matrix (csr) is returned with filled in features
    list
        Used features for binning
    list
        Identifiers of all entries in the MSP file
    int
        Number of analyzed MSP entries
    """
    print(" Reading the MSP file ... ")

    infile = open(infile_name)

    if len(feat_lim_file) > 0:
        selected_features = [float(f.strip()) for f in open(feat_lim_file).readlines()]

    counter = 0
    temp_entry = []
    instance_names = []
    num_instances = num_instances_msp(infile_name)

    if len(feat_bins) == 0: feat_bins = np.arange(0,max_dist+step_size,step_size)

    #Initialize the feature matrix, must be lil since scr is slow when mutating values!
    feat_matrix = scipy.sparse.lil_matrix((num_instances, len(feat_bins)),dtype=np.float32)

    #Iterate over the file and filter out single entries
    for line in infile:
        if line.startswith("Name: "):
            if len(temp_entry) == 0:
                temp_entry.append(line.strip())
                continue
            #For this entry get identifier,m/z,intensities
            identifier,mz_list,intensity_list = parse_msp(temp_entry,top=top_peaks,windowed_mode=windowed_mode)
            instance_names.append(identifier)
            #Fill in the feature matrix
            feat_matrix = get_feats(mz_list,intensity_list,feat_matrix,counter,feat_bins,allowed_c=selected_features,max_dist=max_dist)

            #Make sure the current line is still used for the next entry
            temp_entry = [line]
            perc = round((float(counter)/num_instances)*1000,1)
            if perc.is_integer():
                try:
                    gui_object.progress_bar.setProperty("value", perc/10.0)
                except:
                    pass

            counter += 1

        temp_entry.append(line.strip())

    #If everything is empty; return
    if len(temp_entry) == 0:
        temp_entry.append(line.strip())
        return(feat_matrix.asformat("csr"),feat_bins,instance_names,counter)

    #Analyse the last record; since we do not know when the spectra ends
    identifier,mz_list,intensity_list = parse_msp(temp_entry,top=top_peaks,windowed_mode=windowed_mode)
    instance_names.append(identifier)
    feat_matrix = get_feats(mz_list,intensity_list,feat_matrix,counter,feat_bins,allowed_c=selected_features)

    print(" Total number of entries read: %s "  % (counter))
    counter += 1

    return(feat_matrix.asformat("csr"),feat_bins,instance_names,counter)

def num_instances_msp(infile_name):
    """
    Count the number of entries in the MSP file.

    Parameters
    ----------
    infile_name : list
        The infile MSP file.

    Returns
    -------
    int
        Number of analyzed MSP entries
    """
    infile = open(infile_name)
    num_instances = 0
    for line in infile:
        if line.startswith("Name: "):
            num_instances += 1
    return(num_instances)

def read_mgf(infile_name,feat_lim_file="",
             sum_feats=False,selected_features=[],
             max_dist=275,step_size=0.005,feat_bins=[],
             top_peaks=50,windowed_mode=False):
    """
    Read an MGF file and put the features into a matrix.

    Parameters
    ----------
    infile_name : list
        The infile MGF file.
    feat_lim_file : list
        Old variable with the name of a file that contains the features.
    sum_feats : bool
        Old variable used to sum features of the two classes.
    selected_features : list
        Old variable for selected features; use feat_bins.
    max_dist : int
        Maximum distance between peaks
    step_size : float
        Size between the m/z values for bins.
    feat_bins : list
        Bins to use for features.
    top_peaks : int
        Number of peaks to select based on the intensity
    windowed_mode : bool
        Flag to used windowed mode for selecting the highest intensity peaks

    Returns
    -------
    matrix
        A sparse matrix (csr) is returned with filled in features
    list
        Used features for binning
    list
        Identifiers of all entries in the MGF file
    int
        Number of analyzed MGF entries
    """

    infile = open(infile_name)

    if len(feat_lim_file) > 0:
        selected_features = [float(f.strip()) for f in open("selected_features.txt").readlines()]

    counter = 0
    temp_entry = []
    instance_names = []
    num_instances = num_instances_mgf(infile_name)
    #print(num_instances)

    if len(feat_bins) == 0: feat_bins = np.arange(0,max_dist+step_size,step_size)

    #Initialize the feature matrix, must be lil since scr is slow when mutating values!
    feat_matrix = scipy.sparse.lil_matrix((num_instances, len(feat_bins)),dtype=np.float32)

    #Iterate over the file and filter out single entries
    for line in infile:
        if line.startswith("END IONS"):
            #For this entry get identifier,m/z,intensities
            identifier,mz_list,intensity_list = parse_mgf(temp_entry,top=top_peaks,windowed_mode=windowed_mode)
            instance_names.append(identifier)
            #Fill in the feature matrix
            feat_matrix = get_feats(mz_list,intensity_list,feat_matrix,counter,feat_bins,allowed_c=selected_features,max_dist=max_dist)
            counter += 1
            temp_entry = []
            continue
        if line.startswith("BEGIN IONS"):
            continue
        temp_entry.append(line)

    return(feat_matrix.asformat("csr"),feat_bins,instance_names,counter)

def num_instances_mgf(infile_name):
    """
    Count the number of entries in the MGF file.

    Parameters
    ----------
    infile_name : list
        The infile MGF file.

    Returns
    -------
    int
        Number of analyzed MGF entries
    """
    infile = open(infile_name)
    num_instances = 0
    for line in infile:
        if line.startswith("BEGIN IONS"):
            num_instances += 1
    return(num_instances)

def train_xgb(X,y,n_jobs=32,n_iter_search=20):
    """
    Train an XGBoost model with hyper parameter optimization.

    Parameters
    ----------
    X : matrix
        Matrix with all the features, every instance should be coupled to the y-value
    y : vector
        Vector with the class, every value should be coupled to an x-vector with features

    Returns
    -------
    object
        Trained XGBoost model
    object
        Cross-validation results
    """

    xgb_handle = xgb.XGBClassifier()

    one_to_left = st.beta(10, 1)
    from_zero_positive = st.expon(0, 50)

    #Define distributions to sample from for hyper parameter optimization
    param_dist = {
        "n_estimators": st.randint(3, 80),
        "max_depth": st.randint(3, 40),
        "learning_rate": st.uniform(0.05, 0.4),
        "colsample_bytree": one_to_left,
        "subsample": one_to_left,
        "gamma": st.uniform(0, 10),
        "reg_alpha": from_zero_positive,
        "min_child_weight": from_zero_positive
    }

    n_iter_search = n_iter_search
    random_search = RandomizedSearchCV(xgb_handle, param_distributions=param_dist,
                                       n_iter=n_iter_search,verbose=10,scoring="roc_auc",
                                       n_jobs=n_jobs,cv=StratifiedKFold(n_splits=10, shuffle=True,random_state=42))

    #print(X)
    #print(X.shape)
    #print(y)
    random_search_res_xgb = random_search.fit(X, y)

    #Get the best model that was retrained on all data
    xgb_model = random_search_res_xgb.best_estimator_

    return(xgb_model,random_search_res_xgb)

def train_xgb_lim(X,y,params_dist,out_dir="res/"):
    """
    Train an XGBoost model with set hyper parameters.

    Parameters
    ----------
    X : matrix
        Matrix with all the features, every instance should be coupled to the y-value
    y : vector
        Vector with the class, every value should be coupled to an x-vector with features
    params_dist : dict
        The hyperparameters to use
    out_dir : str
        String value that points to a directory used for output

    Returns
    -------
    list
        The cross-validated predictions.
    """
    #There is a need to unpack the hyperparameter dictionary with "**"
    xgb_handle = xgb.XGBClassifier(**params_dist)
    #Using predict_proba since ROC-curve
    test_preds = cross_val_predict(xgb_handle,X,y,method="predict_proba")
    plot_roc(X,y,test_preds[:,1],fname=out_dir+"roc.png")
    return(test_preds)

def plot_feat_imp(feats_index,feat_names,X,y,top_imp=10,out_dir="res/"):
    """
    Plot the most important features in a boxplot and seperate on class (y)

    Parameters
    ----------
    feats_index : list
        Indexes of the features coupled to the X matrix
    feat_names : list
        Names of the features coupled to specific indices
    X : matrix
        Matrix with all the features, every instance should be coupled to the y-value
    y : vector
        Vector with the class, every value should be coupled to an x-vector with features
    top_imp : int
        Plot this number of top features in a boxplot
    out_dir : str
        String value that points to a directory used for output

    Returns
    -------

    """
    for fi in feats_index[0:top_imp]:
        #Need for a dense matrix when plotting
        #print(fi)
        #print(X.todense()[y==1,:][:,fi])
        #print(X.todense()[y==0,:][:,fi])
        plt.boxplot([list(X.todense()[y==1,:][:,fi]),list(X.todense()[y==0,:][:,fi])])
        plt.title(feat_names[fi])
        plt.savefig(out_dir+"%s_feat_groups.png" % (feat_names[fi]), bbox_inches='tight')
        plt.close()

def plot_train_distr(xgb_model,X,y,out_dir="res/"):
    """
    Plot probability distributions for the input matrix.

    Parameters
    ----------
    xgb_model : object
        Trained XGBoost model
    X : matrix
        Matrix with all the features, every instance should be coupled to the y-value
    y : vector
        Vector with the class, every value should be coupled to an x-vector with features
    out_dir : str
        String value that points to a directory used for output

    Returns
    -------

    """

    #Get the predicted probabilities for both classes (store them seperately)
    probs_oxid = xgb_model.predict_proba(X[y==1])[:,1]
    probs_native = xgb_model.predict_proba(X[y==0])[:,1]

    #Plot density distribution for probailities
    pd.Series(probs_oxid).plot(kind="density")
    pd.Series(probs_native).plot(kind="density")
    axes = plt.gca()
    axes.set_xlim([0.0,1.0])
    plt.savefig(out_dir+"density_groups.png", bbox_inches='tight')
    plt.close()

    #Plot density distribution for probailities; zoom in more so the y-axis is readable
    pd.Series(probs_oxid).plot(kind="density")
    pd.Series(probs_native).plot(kind="density")
    axes = plt.gca()
    axes.set_xlim([0.0,1.0])
    axes.set_ylim([0.0,1.0])
    plt.savefig(out_dir+'density_groups_zoomed.png', bbox_inches='tight')
    plt.close()

    #Plot probability distributions in histogram
    plt.hist(probs_native,bins=100)
    plt.hist(probs_oxid,bins=100)
    plt.savefig(out_dir+'hist_groups.png', bbox_inches='tight')
    plt.close()

    #Plot probability distributions in histogram; zoom in more so the y-axis is readable
    plt.hist(probs_native,bins=100)
    plt.hist(probs_oxid,bins=100)
    axes = plt.gca()
    axes.set_ylim([0.0,1000.0])
    plt.savefig(out_dir+'hist_groups_zoomed.png', bbox_inches='tight')
    plt.close()

def xgboost_to_wb(xgboost,params=None,outfile="model.pickle"):
    """
    Pickle a trained XGBoost model.

    Parameters
    ----------
    xgboost : object
        Trained XGBoost model
    outfile : str
        Location of the pickle

    Returns
    -------

    """
    if params != None: xgboost.params_train = params
    pickle.dump(xgboost, open(outfile,"wb"))

def plot_roc(X,y,test_preds,fname="res/roc.png"):
    """
    Plot an ROC-curve and write to a file

    Parameters
    ----------
    X : matrix
        Matrix with all the features, every instance should be coupled to the y-value
    y : vector
        Vector with the class, every value should be coupled to an x-vector with features
    test_preds : vector
        Predicted probabilities for classes

    Returns
    -------

    """
    #Retrieve multiple fpr and tpr values for different thresholds
    fpr, tpr, thresholds = roc_curve(y,test_preds)
    plt.plot(fpr, tpr)
    plt.title(auc(fpr, tpr))
    plt.savefig(fname, bbox_inches='tight')
    plt.close()

def get_min_diff(zero_f="NIST/human_hcd_synthetic_oxidized.msp",
                 one_f="NIST/human_hcd_synthetic_native.msp",
                 outfile="res_small/selected_features_diff.txt",
                 top_mean = 1000,
                 top_peaks = 50,
                 max_distance = 275,
                 distance_bins = 0.005,
                 windowed_mode=False):
    """
    Function that is used to get the most important bins in terms of mean intensity between the two classes.

    Parameters
    ----------
    zero_f : str
        Filename (and dir) that contains the spectra for class 0
    one_f : str
        Filename (and dir) that contains the spectra for class 1
    top_mean : int
        Return this amount of features (bins)
    top_peaks : int
        The top intensity peaks to filter (in windowed mode it will return the top peaks within the window)
    max_distance : int
        Maximum distance in m/z to use for features
    windowed_mode : bool
        Flag to use windowed mode to return the top intensity peaks
    distance_bins : float
        Distance in m/z between the bins

    Returns
    -------
    list
        The most important bins that are on average different between the classes
    """

    #Check the file extension and parse to get features for class zero
    if zero_f.endswith(".mgf"): feats_zero_sum,feat_bins,instance_names,count_zero = read_mgf(zero_f,
                                                                                sum_feats=True,
                                                                                max_dist=max_distance,
                                                                                step_size=distance_bins,
                                                                                top_peaks=top_peaks)
    elif zero_f.endswith(".msp"): feats_zero_sum,feat_bins,instance_names,count_zero = read_msp(zero_f,
                                                                                sum_feats=True,
                                                                                max_dist=max_distance,
                                                                                step_size=distance_bins,
                                                                                top_peaks=top_peaks)
    else: return(False)

    #Check the file extension and parse to get features for class one
    if one_f.endswith(".mgf"): feats_one_sum,feat_bins,instance_names,count_one = read_mgf(one_f,
                                                                                sum_feats=True,
                                                                                max_dist=max_distance,
                                                                                step_size=distance_bins,
                                                                                top_peaks=top_peaks)
    elif one_f.endswith(".msp"): feats_one_sum,feat_bins,instance_names,count_one = read_msp(one_f,
                                                                                sum_feats=True,
                                                                                max_dist=max_distance,
                                                                                step_size=distance_bins,
                                                                                top_peaks=top_peaks)
    else: return(False)

    #Get average difference between the bins of both groups
    diffs = [abs(m1-m2) for m1,m2 in zip(feats_zero_sum.mean(axis=0).tolist()[0],feats_one_sum.mean(axis=0).tolist()[0])]

    #Get the indexes of the biggest differences in bins
    indexes_diff = sorted(list(enumerate(diffs)),key=itemgetter(1),reverse=True)
    selected_features_diff = [feat_bins[ind] for ind,val in indexes_diff[0:top_mean]]
    selected_features_diff.sort()

    #For the important bins we need the next number to create a closed bin; calculated vals; extend to bins
    diff_bins = [sfd+distance_bins for sfd in selected_features_diff]
    diff_bins.extend(selected_features_diff)
    diff_bins.sort()

    #Remove duplicate values
    diff_bins = list(set(diff_bins))
    diff_bins.sort()

    #Write feats to a file
    outfile_feats = open(outfile,"w")
    outfile_feats.write("\n".join(map(str,diff_bins)))
    outfile_feats.close()

    return(diff_bins)

def train_initial_classifier(zero_f="NIST/human_hcd_synthetic_oxidized.msp",
                             one_f="NIST/human_hcd_synthetic_native.msp",
                             selected_features_diff=[],
                             top_mean = 1000,
                             top_peaks = 100,
                             max_distance = 275,
                             distance_bins = 0.005,
                             windowed_mode = False,
                             out_dir="res/",
                             out_pickle="model.pickle",
                             n_jobs=32,
                             n_iter_search=20):
    """
    Function that is used to train an XGBoost model to discrimate between the MS2 spectra from two files.

    Parameters
    ----------
    zero_f : str
        Filename (and dir) that contains the spectra for class 0
    one_f : str
        Filename (and dir) that contains the spectra for class 1
    selected_features_diff : list
        List with m/z bin values that should be used as features
    top_mean : int
        Return this amount of features (bins)
    top_peaks : int
        The top intensity peaks to filter (in windowed mode it will return the top peaks within the window)
    max_distance : int
        Maximum distance in m/z to use for features
    windowed_mode : bool
        Flag to use windowed mode to return the top intensity peaks
    distance_bins : float
        Distance in m/z between the bins
    out_dir : str
        Directory to write the results to

    Returns
    -------
    dict
        Used parameters in the XGBoost model
    list
        Most important features according to the F-score in the XGBoost model
    """
    #Check the file extension and parse to get features for class zero
    if zero_f.endswith(".mgf"): feats_zero,feat_bins,instance_names,count_zero = read_mgf(zero_f,sum_feats=False,
                                                                                feat_bins=selected_features_diff,
                                                                                max_dist=max_distance,
                                                                                step_size=distance_bins,
                                                                                top_peaks=top_peaks)
    elif zero_f.endswith(".msp"): feats_zero,feat_bins,instance_names,count_zero = read_msp(zero_f,
                                                                                sum_feats=False,
                                                                                feat_bins=selected_features_diff,
                                                                                max_dist=max_distance,
                                                                                step_size=distance_bins,
                                                                                top_peaks=top_peaks)
    else: return(False) # TODO display error!

    #Check the file extension and parse to get features for class one
    if one_f.endswith(".mgf"): feats_one,feat_bins,instance_names,count_one = read_mgf(one_f,
                                                                                sum_feats=False,
                                                                                feat_bins=selected_features_diff,
                                                                                max_dist=max_distance,
                                                                                step_size=distance_bins,
                                                                                top_peaks=top_peaks)
    elif one_f.endswith(".msp"): feats_one,feat_bins,instance_names,count_one = read_msp(one_f,
                                                                                sum_feats=False,
                                                                                feat_bins=selected_features_diff,
                                                                                max_dist=max_distance,
                                                                                step_size=distance_bins,
                                                                                top_peaks=top_peaks)
    else: return(False) # TODO display error!

    #Prepare labels equal to length class zero and one
    y = [0]*(count_zero)
    y.extend([1]*(count_one))

    y = np.array(y)

    #Stack the feature matrices of both classes
    X = scipy.sparse.vstack((feats_zero,feats_one))

    #Train optimizing the hyperparameters
    xgb_model,random_search_res_xgb = train_xgb(X,y,n_jobs=n_jobs,n_iter_search=n_iter_search)
    print(random_search_res_xgb.best_params_)
    print(random_search_res_xgb.best_score_)

    #Train use selected hyperparameters
    train_xgb_lim(X,y,random_search_res_xgb.best_params_,out_dir=out_dir)
    plot_train_distr(xgb_model,X,y,out_dir=out_dir)

    #Make a dummy object that contains the params used for training
    dummy = Dummy()
    dummy.selected_features_diff = selected_features_diff
    dummy.top_mean = top_mean
    dummy.top_peaks = top_peaks
    dummy.max_distance = max_distance
    dummy.distance_bins = distance_bins
    dummy.windowed_mode = windowed_mode

    #Flush to pickle
    xgboost_to_wb(random_search_res_xgb,params=dummy,outfile=os.path.join(out_dir,out_pickle))

    random_search_res_xgb = pickle.load(open(os.path.join(out_dir,out_pickle),"rb"))

    #Plot some of the feature importances and probs
    fscores = xgb_model.booster().get_fscore()
    fscores_list = sorted(list(fscores.items()),key=itemgetter(1),reverse=True)
    selected_features_indexes = list(map(int,[f.replace("f","") for f,n in fscores_list]))
    selected_features_xgboost  = [selected_features_diff[sfp] for sfp in selected_features_indexes]
    plot_feat_imp(selected_features_indexes,selected_features_diff,X,y,out_dir=out_dir)

    return(random_search_res_xgb.best_params_,selected_features_xgboost)

def apply_model(infile_pred,
                infile_model,
                infile_features,
                threshold_prob=0.5,
                windowed_peak_picking=False,
                out_dir="res/",
                top_peaks=50,
                max_distance=275,
                distance_bins=0.005):
    """
    Make predictions to discriminate between classes of spectra using pretrained models.

    Parameters
    ----------
    infile_pred : str
        Filename (and dir) that contain the spectra we need to make predictions for
    infile_model : str
        Filename (and dir) that points to the trained XGBoost model
    infile_features : str
        Filename (and dir) that contains the m/z bins (features)
    top_peaks : int
        The top intensity peaks to filter (in windowed mode it will return the top peaks within the window)
    max_distance : int
        Maximum distance in m/z to use for features
    windowed_peak_picking : bool
        Flag to use windowed mode to return the top intensity peaks
    distance_bins : float
        Distance in m/z between the bins
    threshold_prob : float
        Threshold to determine an instance belongs to class zero or one
    out_dir : str
        Directory to write the results to

    Returns
    -------

    """
    print(" Making predictions... ")

    if isinstance(infile_features, str):
        #Read the to be used m/z bins (features)
        features = [f.strip() for f in open(infile_features).readlines()]
    else:
        features = infile_features

    #Check the file extension and parse to get features
    if infile_pred.endswith(".mgf"):
        feats,feat_bins,instance_names,count_inst = read_mgf(infile_pred,
                                                             sum_feats=False,
                                                             feat_bins=features,
                                                             max_dist=max_distance,
                                                             step_size=distance_bins,
                                                             top_peaks=top_peaks)
    elif infile_pred.endswith(".msp"):
        feats,feat_bins,instance_names,count_inst = read_msp(infile_pred,
                                                             sum_feats=False,
                                                             feat_bins=features,
                                                             max_dist=max_distance,
                                                             step_size=distance_bins,
                                                             top_peaks=top_peaks)
    else: return(False)

    #Load the XGBoost model
    random_search_res_xgb = pickle.load(open(infile_model,"rb"))

    print(" Formatting predictions ... ")
    #Format the preds
    preds = pd.DataFrame(random_search_res_xgb.predict_proba(feats),index=instance_names,columns=["Prob_class_0","Prob_class_1"])

    print(" Plotting predictions ... ")
    #Plot the probabilities
    pd.Series(preds["Prob_class_1"]).plot(kind="density")
    axes = plt.gca()
    axes.set_xlim([0.0,1.0])
    axes.set_ylim([0.0,1.0])
    plt.savefig(os.path.join(out_dir,"density_groups_zoomed.png"), bbox_inches="tight")
    plt.close()

    #Threshold probabilities
    #print(list(preds.index[preds["Prob_class_1"] > threshold_prob]))

    print(" Writing predictions ... ")
    preds.to_csv(os.path.join(out_dir,"predictions.csv"))

def parse_argument():
    """
    Read arguments from the command line

    Parameters
    ----------

    Returns
    -------

    """
    parser = argparse.ArgumentParser()

    parser.add_argument("--top_peaks", type=int, dest="top_peaks", default=100,
                        help="Number of peaks to extract and consider for combinations in a spectrum")

    parser.add_argument("--top_mean", type=int, dest="top_mean", default=2000,
                        help="The top bins in different mean between group A and B to learn on")

    parser.add_argument("--max_distance", type=int, dest="max_distance", default=300,
                        help="The maximum difference between peaks (maximum bin value)")

    parser.add_argument("--n_jobs", type=int, dest="n_jobs", default=32,
                        help="The maximum number of threads to spawn")

    parser.add_argument("--n_iter_search", type=int, dest="n_iter_search", default=20,
                        help="Number of random samplings from the hyperparameter sets")

    parser.add_argument("--distance_bins", type=float, dest="distance_bins", default=0.1,
                        help="Distance in m/z of the bins")

    parser.add_argument("--file_a", type=str, dest="file_a",default="data/human_hcd_synthetic_native_sample.msp", #"NIST/human_hcd_synthetic_native.msp"
                        help="The mgf or msp of group A")

    parser.add_argument("--file_b", type=str, dest="file_b", default="data/human_hcd_synthetic_oxidized_sample.msp", #"NIST/human_hcd_synthetic_oxidized.msp"
                        help="The mgf or msp of group B")

    parser.add_argument("--file_pred", type=str, dest="file_pred", default="NIST/human_hcd_synthetic_native.msp",
                        help="The mgf or msp to make predictions for")

    parser.add_argument("--out_dir", type=str, dest="out_dir", default="res/",
                        help="Directory where the results are written. WILL OVERWITE EXISTING FILES!")

    parser.add_argument("--out_pickle", type=str, dest="out_pickle", default="model.pickle",
                        help="File where the model is pickled.")

    parser.add_argument("--make_pred", action="store_true",
                        help="Flag that can be included to indicate predictions are desired instead of training a model")

    parser.add_argument("--windowed_peak_picking", action="store_true",
                        help="Flag that can be included to use windowed peak picking per 100 m/z")

    parser.add_argument("--model", type=str, dest="model", default="res/model.pickle",
                        help="File that refers to a model that is used for predictions")

    parser.add_argument("--feats", type=str, dest="feats", default="res/selected_features_diff.txt",
                        help="File that refers to the features that are used in the model")

    parser.add_argument("--version", action="version", version="%(prog)s 1.0")

    results = parser.parse_args()

    return(results)

def main(arguments=None,gui_object_local=None):
    global gui_object
    gui_object = gui_object_local

    if arguments == None:
        #Get command line arguments
        argu = parse_argument()
    else:
        argu = arguments

    #Train a new model
    if not argu.make_pred:
        #print(os.getcwd())
        #print(os.path.join(os.getcwd(),"selected_features.txt"))
        selected_features_diff = get_min_diff(zero_f=argu.file_a,
                                              one_f=argu.file_b,
                                              outfile=os.path.join(os.getcwd(),"selected_features.txt"),
                                              top_peaks=argu.top_peaks,
                                              top_mean=argu.top_mean,
                                              max_distance=argu.max_distance,
                                              distance_bins=argu.distance_bins,
                                              windowed_mode=argu.windowed_peak_picking)
        # UNCOMMENT line below and comment above call to bypass initial feature selection
        #selected_features_diff = [float(f.strip()) for f in open("res/selected_features.txt").readlines()]

        random_search_params,selected_features_xgb = train_initial_classifier(zero_f=argu.file_a,
                                                                              one_f=argu.file_b,
                                                                              selected_features_diff=selected_features_diff,
                                                                              top_peaks=argu.top_peaks,
                                                                              top_mean=argu.top_mean,
                                                                              max_distance=argu.max_distance,
                                                                              distance_bins=argu.distance_bins,
                                                                              windowed_mode=argu.windowed_peak_picking,
                                                                              out_pickle=argu.out_pickle,
                                                                              n_jobs=argu.n_jobs,
                                                                              n_iter_search=argu.n_iter_search)
    #Make predictions using existing trained model
    if argu.make_pred:
        print(" ===================== ")
        print(" Starting to make predictions ... ")

        model = pickle.load(open(os.path.join(argu.model),"rb"))
        try:
            model.params_train
            argu.selected_features_diff = model.params_train.selected_features_diff
            argu.top_mean = model.params_train.top_mean
            argu.top_peaks = model.params_train.top_peaks
            argu.max_distance = model.params_train.max_distance
            argu.distance_bins = model.params_train.distance_bins
            argu.windowed_peak_picking = model.params_train.windowed_mode
            print(" ---------------------- ")
            print(" Loaded model settings: ")
            print(" Number of features selected  %s " % (len(argu.selected_features_diff)))
            print(" Maximum distance  %s " % (argu.max_distance))
            print(" Distance bins  %s " % (argu.distance_bins))
            print(" Used windows  %s " % (argu.windowed_peak_picking))
            print(" Top peak per MS2 %s" % (argu.top_peaks))
            print(" ---------------------- ")
        except:
            # Model parameters were not saved with trained model
            pass

        apply_model(argu.file_pred,argu.model,argu.selected_features_diff,out_dir=argu.out_dir,
                    windowed_peak_picking=argu.windowed_peak_picking,
                    top_peaks=argu.top_peaks,
                    max_distance=argu.max_distance,
                    distance_bins=argu.distance_bins)

        print(" ===================== ")

if __name__ == "__main__":
    main()


# Just a test
import audiovisualizer
import numpy
import numpy.fft as nfft
import pylab
import scipy
import scipy.io as sio

# The sample rate of the input matrix (Hz)
SAMPLE_RATE=44100
# Frequency range to display (audible is 16-16384Hz)
DISPLAY_FREQ=(16, 2000)
# FPS of output (Hz)
OUT_FPS = 30
# Size of the moving average (s)
MV_AVG_SIZE=0.5

def bin_approx_search(lst, tg):
    """
    Find the index of the element in lst which is closest to the number tg
    """
    top = len(lst) - 1
    bottom = 0
    while top > bottom:
        curri = (top - bottom)//2 + bottom
        if lst[curri] < tg:
            bottom = curri
        else:
            top = curri
        if top - bottom == 1:
            if abs(lst[top] - tg) < abs(lst[bottom] - tg):
                return top
            else:
                return bottom
    return top

def plotSpectrum(y,samplerate):
    """
    Plots a Single-Sided Amplitude Spectrum of y(t)
    """
    n = len(y) # length of the signal
    k = scipy.arange(n)
    T = n/samplerate
    frq = audiovisualizer.movingfft.get_x_axis(samplerate, n)

    Y = nfft.fft(y)/n # fft computing and normalization
    Y = Y[range(n//2)]
    print(Y.shape)

    bottomindex = bin_approx_search(frq, DISPLAY_FREQ[0])
    topindex = bin_approx_search(frq, DISPLAY_FREQ[1])

    frq = frq[bottomindex:topindex]
    Y = Y[bottomindex:topindex]
    print(Y.shape)

    # We want to plot frequencies in the audible range

    pylab.clf()
    pylab.plot(frq,abs(Y),'r') # plotting the spectrum
    pylab.xlabel('Freq (Hz)')
    pylab.ylabel('|Y(freq)|')
    pylab.ylim((0, 0.018))

audio = sio.loadmat("song1.mat")['x']

frame_increment = SAMPLE_RATE // OUT_FPS
avg_len = int(MV_AVG_SIZE * SAMPLE_RATE)
frame = 0
for start in range(0, audio.shape[0] - avg_len, frame_increment):
    print(start, start+avg_len)
    plotSpectrum(audio[start:start + avg_len, 0], SAMPLE_RATE)
    pylab.show()
    break
    pylab.savefig('frame_'+str(frame).zfill(4)+'.png', bbox_inches='tight')
    frame += 1


<reponame>msmbuilder/mdentropy
from ..utils import unique_row_count

from numpy import (array, atleast_1d, digitize, empty, floor, linspace, log2,
                   histogramdd, hstack, ndarray, sqrt, vstack)
from scipy.stats import skew

__all__ = ['hist', 'symbolic', 'doanes_rule']

def doanes_rule(x):
    """Convenience function for choosing an optimal number of bins using Doane's Rule.

    Parameters
    ----------
    x : numpy.ndarray or list of floats
        Data to be binned.

    Returns
    -------
    n_bins : int
    """
    if not isinstance(x, ndarray):
        x = array(x)

    n = x.shape[0]
    g1 = atleast_1d(skew(x))
    sg1 = sqrt(6 * (n - 2) / ((n + 1) * (n + 3)))

    return min(floor(1 + log2(n) + log2(1 + abs(g1)/sg1)))

def hist(n_bins, rng, *args):
    """Convenience function for histogramming N-dimentional data

    Parameters
    ----------
    n_bins : int
        Number of bins.
    rng : list of lists
        List of min/max values to bin data over.
    args : array_like, shape = (n_samples, )
        Data of which to histogram.

    Returns
    -------
    bins : array_like, shape = (n_bins, )
    """
    data = vstack((args)).T

    if n_bins is None:
        n_bins = doanes_rule(data)

    return histogramdd(data, bins=n_bins, range=rng)[0].flatten()

def symbolic(n_bins, rng, *args):
    """Symbolic binning of data

    Parameters
    ----------
    rng : list of lists
        List of min/max values for each dimention.
    n_bins : int
        Number of bins to use.
    args : array_like, shape = (n_samples, )
        Data of which to calculate entropy. Each array must have the same
        number of samples.

    Returns
    -------
    counts : float
    """
    labels = empty(0).reshape(args[0].shape[0], 0)
    if n_bins is None:
        n_bins = min(map(doanes_rule, args))

    for i, arg in enumerate(args):

        partitions = linspace(rng[i][0], rng[i][1], n_bins + 1)
        label = digitize(arg, partitions).reshape(-1, 1)

        labels = hstack((labels, label))

    return unique_row_count(labels)


<filename>training/buyer_utils.py
'''
Functions to evaluate buyer's fixed strategies
'''

import math
import numpy as np
import time
from scipy.optimize import minimize

def getBuyerExperience(sellers, buyer_info):
    # get the buyer experience with sellers based on previous purchases
    num_sellers = len(sellers)
    num_buyers = buyer_info.count
    cumulativeBuyerExperience = np.zeros((num_buyers, num_sellers))
    for i in range(0, num_buyers):
        for j in range(0, num_sellers):
            cumulativeBuyerExperience[i][j] = sellers[j].getBuyerExperience(i)
    return cumulativeBuyerExperience

def get_buyer_rewards(X, ys, probAll, cumulativeBuyerExperience, buyer_info):
    # get buyer utilities
    buyer_utilities = buyerUtilitiesCalculator(X, ys, buyer_info.V, buyer_info.a_val, probAll,
                                                  buyer_info.count,
                                                  cumulativeBuyerExperience, buyer_info.unfinished_task_penalty)

    # get buyer penalties
    buyer_penalties = buyerPenaltiesCalculator(X, ys, buyer_info.V, buyer_info.a_val, buyer_info.count,
                                                  cumulativeBuyerExperience, buyer_info.unfinished_task_penalty)

    return buyer_utilities, buyer_penalties

def getPurchases(buyer_info, cumulativeBuyerExperience, ys, probAll):
    # get the amount of resources purchased by each device based on y
    X = []
    for i in range(0, buyer_info.count):
        X_i = buyerPurchaseCalculator(cumulativeBuyerExperience[i, :], ys, buyer_info.V[i], buyer_info.a_val[i]
                                      , probAll, buyer_info.unfinished_task_penalty)
        X.append(X_i)
    X = np.array(X).T
    return X

# Callback to stop optimization after a set time limit
class TookTooLong(Warning):
    pass

# class MinimizeStopper(object):
#     def __init__(self, max_sec=60):
#         self.max_sec = max_sec
#         self.start = time.time()
#     def __call__(self, xk=None):
#         elapsed = time.time() - self.start
#         if elapsed > self.max_sec:
#             warnings.warn("Terminating optimization: time limit reached",
#                           TookTooLong)
        # else:
        #     # you might want to report other stuff here
        #     print("Elapsed: %.3f sec" % elapsed)

# Buyer Purchase Calculator
def buyerPurchaseCalculator(cumulativeBuyerExperience, yAll, V_i, a_i, y_prob, consumer_penalty_coeff):
    # get singleBuyer utility function to maximize
    N = len(y_prob)

    def singleBuyerUtilityFunction(x_i):
        buyerUtility = 0.
        for j in range(0, N):
            buyerUtility += (V_i * math.log(x_i[j] - a_i + np.e) \
                             - x_i[j] / yAll[j]) * y_prob[j] \
                            - consumer_penalty_coeff * (cumulativeBuyerExperience[j] - x_i[j]) ** 2
        #     buyerUtility += (V_i * math.log(x_i[j] - a_i + np.e) - (x_i[j] / yAll[j])) * y_prob[j]
        # buyerUtility -= consumer_penalty_coeff*(np.sum(cumulativeBuyerExperience) - np.sum(x_i))**2
        return -1 * buyerUtility

    # solve optimization function for each buyer try for two seconds
    x_init = cumulativeBuyerExperience
    # x_init = np.full_like(cumulativeBuyerExperience, 100)
    if np.random.uniform()>0.5:
        xi_opt_sol = minimize(singleBuyerUtilityFunction, x_init, bounds=[(0, 100)] * N,options={'maxiter':1000})
        # xi_opt_sol = minimize(singleBuyerUtilityFunction, x_init, options={'maxiter': 1000})
        x_opt = xi_opt_sol.x
    else:
        x_opt = x_init*np.random.uniform(low=1,high=1.2)
    # x_opt[x_opt < 0] = 0
    return x_opt

# Buyer Utilities Calculator
def buyerUtilitiesCalculator(X, yAll, V, a, y_prob, M, cumulativeBuyerExperience, consumer_penalty_coeff):
    N = len(y_prob)
    buyerUtilities = []
    for i in range(0, M):
        buyerUtility = 0
        for j in range(0, N):
            buyerUtility += (V[i] * math.log(X[j][i] - a[i] + np.e) \
                             - X[j][i] / yAll[j]) * y_prob[j]
            # todo: Add the regularizer based on Z values
        buyerUtilities.append(buyerUtility)
    buyerUtilities = np.array(buyerUtilities)
    return buyerUtilities

# Buyer Penalties Calculator
def buyerPenaltiesCalculator(X, yAll, V, a, M, cumulativeBuyerExperience, consumer_penalty_coeff):
    N = len(yAll)
    buyerPenalties = []
    for i in range(0, M):
        buyerPenalty = 0
        for j in range(0, N):
            buyerPenalty += consumer_penalty_coeff * (cumulativeBuyerExperience[i][j] - X[j][i]) ** 2
            # todo: Add the regularizer based on Z values
        buyerPenalties.append(buyerPenalty)
    buyerPenalties = np.array(buyerPenalties)
    return buyerPenalties