a
    Dff                  5   @   s  d dl Zd dlZd dlmZmZmZmZmZm	Z	m
Z
mZ ejdejdZejdejdZejdejdZeejjZeejjZejdd edD ejdZejd	d
dd Zejdejejjejjddd	dejjejjddd	dgd	ejjejjejj ejj!dddd Z"ejd	d
efddZ#ejd	d
dd Z$ejd	d
dd Z%ejd	d
dyddZ&ejd	d
edfddZ'ejd	d
efdd Z(ejd	d
d!d" Z)ejd	d
d#d$ Z*ejd	d
d%d& Z+ejd	d
d'd( Z,ejdejejjejjddd	dejjejjddd	dgd	ejjejjejjejj-ejj-ejj ejj!d)dd*d+ Z.ej/d	d
d,d- Z0ejd	d
d.d/ Z1ejd	d
d0d1 Z2ejd	d
d2d3 Z3ejd	d
d4d5 Z4ejd	d
d6d7 Z5ejd	d
d8d9 Z6ejd	d
d:d; Z7ejd	d
d<d= Z8ejd	d
d>d? Z9ejd	d
d@dA Z:ejdejejjejjddd	dejjejjddd	dgd	ejjejjejjejj ejj!dBddCdD Z;ejdd	ejjejj ejj!dEddFdG Z<ejdejejjejjddd	dejjejjddd	dgd	ejjejj ejj!dEddHdI Z=ej/d	d
dJdK Z>ejd	d
dLdM Z?ejd	d
dNdO Z@ej/d	d
dPdQ ZAejd	d
dRdS ZBejdejejjejjddd	dejjejjddd	dgd	ejjejjejjejj ejj!dTddUdV ZCejdejejjejjddd	dejjejjddd	dgd	ejjejjejjejj ejj!dTddWdX ZDej/d	d
dYdZ ZEe dzd\d]ZFejd	d
d^d_ ZGejd	d`edafdbdcZHejd	d
eddfdedfZIe dgdh ZJe d{didjZKe d|dkdlZLe dmdn ZMejdoejejjejj-ddd	dejjejj-ddd	dgd	ejjejj-ejj ejj!dpddqdr ZNejdoejejjejj-ddd	dejjejj-ddd	dgd	ejjejjejj-ejj-ejj ejj!dsddtdu ZOeee"e$e$e$e%e%e%e%e&e#e#e'e'e(e*e:e<eBe8e+eGe?e@eCeHeHeKeKeKeKeLeLeIeJeJeMeMeMe)e,e2e1e3e4e5e7e6e9eNeOdv4ZPe"ejQdwe"ejQdwe;e>dwe=e>dwe;eAdweDeEdwe.e0dwdxZRdS )}    N)allocate_graph_structuresinitialize_graph_structuresinitialize_supplyinitialize_costnetwork_simplex_core
total_costProblemStatussinkhorn_transport_plan   dtype)r
   r
   c                 C   s   g | ]}t |d qS )1)bincount).0i r   b/nfs/NAS7/SABIOD/METHODE/ermites/ermites_venv/lib/python3.9/site-packages/pynndescent/distances.py
<listcomp>       r      T)fastmathc                 C   s:   d}t | jd D ]}|| | ||  d 7 }qt|S )z_Standard euclidean distance.

    .. math::
        D(x, y) = \\sqrt{\sum_i (x_i - y_i)^2}
            r   r
   rangeshapenpsqrtxyresultr   r   r   r   	euclidean   s    r"   zf4(f4[::1],f4[::1])   C)readonly)r!   diffdimr   )r   localsc                 C   s<   d}| j d }t|D ] }| | ||  }||| 7 }q|S )zVSquared euclidean distance.

    .. math::
        D(x, y) = \sum_i (x_i - y_i)^2
    r   r   r   r   )r   r    r!   r'   r   r&   r   r   r   squared_euclidean,   s    
r*   c                 C   sB   d}t | jd D ]$}|| | ||  d ||  7 }qt|S )zEuclidean distance standardised against a vector of standard
    deviations per coordinate.

    .. math::
        D(x, y) = \sqrt{\sum_i \frac{(x_i - y_i)**2}{v_i}}
    r   r   r
   r   )r   r    sigmar!   r   r   r   r   standardised_euclideanK   s    "r,   c                 C   s6   d}t | jd D ]}|t| | ||  7 }q|S )z\Manhattan, taxicab, or l1 distance.

    .. math::
        D(x, y) = \sum_i |x_i - y_i|
    r   r   r   r   r   absr   r   r   r   	manhattanZ   s    r/   c                 C   s8   d}t | jd D ] }t|t| | ||  }q|S )zZChebyshev or l-infinity distance.

    .. math::
        D(x, y) = \max_i |x_i - y_i|
    r   r   )r   r   maxr   r.   r   r   r   r   	chebyshevh   s    r1   c                 C   sB   d}t | jd D ]"}|t| | ||  | 7 }q|d|  S )ah  Minkowski distance.

    .. math::
        D(x, y) = \left(\sum_i |x_i - y_i|^p\right)^{\frac{1}{p}}

    This is a general distance. For p=1 it is equivalent to
    manhattan distance, for p=2 it is Euclidean distance, and
    for p=infinity it is Chebyshev distance. In general it is better
    to use the more specialised functions for those distances.
    r   r         ?r-   )r   r    pr!   r   r   r   r   	minkowskiv   s     r4   c                 C   sJ   d}t | jd D ]*}||| t| | ||  |  7 }q|d|  S )aW  A weighted version of Minkowski distance.

    .. math::
        D(x, y) = \left(\sum_i w_i |x_i - y_i|^p\right)^{\frac{1}{p}}

    If weights w_i are inverse standard deviations of graph_data in each dimension
    then this represented a standardised Minkowski distance (and is
    equivalent to standardised Euclidean distance for p=1).
    r   r   r2   r-   )r   r    wr3   r!   r   r   r   r   weighted_minkowski   s    (r6   c                 C   s   d}t j| jd t jd}t| jd D ]}| | ||  ||< q(t| jd D ]D}d}t| jd D ]}||||f ||  7 }qf||||  7 }qPt |S )Nr   r   r   )r   emptyr   float32r   r   )r   r    Zvinvr!   r&   r   tmpjr   r   r   mahalanobis   s    r;   c                 C   sB   d}t | jd D ]}| | || kr|d7 }qt|| jd  S Nr   r   r2   r   r   floatr   r   r   r   hamming   s
    
r?   c                 C   s^   d}t | jd D ]F}t| | t||  }|dkr|t| | ||  | 7 }q|S Nr   r   r-   )r   r    r!   r   denominatorr   r   r   canberra   s     rB   c                 C   sl   d}d}t | jd D ]8}|t| | ||  7 }|t| | ||  7 }q|dkrdt|| S dS d S r@   )r   r   r   r.   r>   )r   r    	numeratorrA   r   r   r   r   bray_curtis   s    rD   c                 C   sl   d}d}t | jd D ]4}| | dk}|| dk}||p:|7 }||oF|7 }q|dkrXdS t|| | S d S r@   r=   )r   r    num_non_zero	num_equalr   x_truey_truer   r   r   jaccard   s    rI   )r!   rE   rF   rG   rH   r'   r   c                 C   sp   d}d}| j d }t|D ]4}| | dk}|| dk}||p>|7 }||oJ|7 }q|dkr\dS t||  S d S r@   )r   r   r   log2)r   r    rE   rF   r'   r   rG   rH   r   r   r   alternative_jaccard   s    
rK   c                 C   s   dt d|   S Nr2          @pow)vr   r   r   correct_alternative_jaccard  s    rQ   c                 C   sN   d}t | jd D ](}| | dk}|| dk}|||k7 }qt|| jd  S r@   r=   r   r    num_not_equalr   rG   rH   r   r   r   matching  s    rT   c                 C   sl   d}d}t | jd D ]4}| | dk}|| dk}||o:|7 }|||k7 }q|dkrXdS |d| |  S d S Nr   r   rM   r   r   r   r    num_true_truerS   r   rG   rH   r   r   r   dice  s    rY   c                 C   s   d}d}t | jd D ]4}| | dk}|| dk}||o:|7 }|||k7 }q|dkrXdS t|| | jd  || jd   S d S r@   r=   rW   r   r   r   	kulsinski#  s    rZ   c                 C   sR   d}t | jd D ](}| | dk}|| dk}|||k7 }qd| | jd |  S rU   rV   rR   r   r   r   rogers_tanimoto5  s    r[   c                 C   s   d}t | jd D ](}| | dk}|| dk}||o6|7 }q|t| dkkrd|t|dkkrddS t| jd | | jd  S d S r@   )r   r   r   sumr>   )r   r    rX   r   rG   rH   r   r   r   
russellrao@  s    $r]   c                 C   sR   d}t | jd D ](}| | dk}|| dk}|||k7 }qd| | jd |  S rU   rV   rR   r   r   r   sokal_michenerN  s    r^   c                 C   sl   d}d}t | jd D ]4}| | dk}|| dk}||o:|7 }|||k7 }q|dkrXdS |d| |  S d S Nr   r         ?rV   rW   r   r   r   sokal_sneathY  s    ra   c                 C   s   | j d dkrtdtd| d |d   }td| d |d   }t|d t| d t|d  |d   }dt| S )Nr   r
   z6haversine is only defined for 2 dimensional graph_datar`   r#   rM   )r   
ValueErrorr   sinr   cosZarcsin)r   r    Zsin_latZsin_longr!   r   r   r   	haversinei  s    2re   c           	      C   s   d}d}d}t | jd D ]D}| | dk}|| dk}||o>|7 }||oL| 7 }|| oZ|7 }q| jd | | | }|dks|dkrdS d| | || ||   S d S rU   rV   )	r   r    rX   Znum_true_falseZnum_false_truer   rG   rH   Znum_false_falser   r   r   yules  s    
rf   c                 C   s   d}d}d}t | jd D ]8}|| | ||  7 }|| | d 7 }||| d 7 }q|dkrh|dkrhdS |dksx|dkr|dS d|t||   S d S Nr   r   r
   r2   r   )r   r    r!   norm_xnorm_yr   r   r   r   cosine  s    rj   )r!   rh   ri   r'   r   c                 C   s   d}d}d}| j d }t|D ]@}|| | ||  7 }|| | | |  7 }||| ||  7 }q|dkrt|dkrtdS |dks|dkrtS |dkrtS t|| | }t|S d S r@   )r   r   FLOAT32_MAXr   r   rJ   r   r    r!   rh   ri   r'   r   r   r   r   alternative_cosine  s     
rm   )r!   r'   r   c                 C   sH   d}| j d }t|D ]}|| | ||  7 }q|dkr<dS d| S d S r<   r)   r   r    r!   r'   r   r   r   r   dot  s    

ro   c                 C   sL   d}| j d }t|D ]}|| | ||  7 }q|dkr<tS t| S d S r@   )r   r   rk   r   rJ   rn   r   r   r   alternative_dot  s    
rp   c                 C   s   dt d|   S rL   rN   dr   r   r   correct_alternative_cosine  s    rs   c                 C   s   d}d}d}d}| j d }t|D ]\}| | ||  }||| 7 }|| | ||  7 }|| | | |  7 }||| ||  7 }q"t|}t|}t|| }	|||  }t|td }
t||	 d |
 }|| t|
 d }|| S )Nr   r   
   r
   rM   )r   r   r   r   r.   arccosradiansrc   )r   r    Zd_euc_squaredZd_cosrh   ri   r'   r   r&   Zmagnitude_differencethetaZsectorZtriangler   r   r   tsss  s&    


rx   c                 C   s   d}d}d}| j d }t|D ]@}|| | ||  7 }|| | | |  7 }||| ||  7 }q|dkrt|dkrtdS |dks|dkrtS |dkrtS |t||  }dt|tj  S d S r<   )r   r   rk   r   r   ru   pirl   r   r   r   true_angular  s     
rz   c                 C   s   dt td|  t j  S rL   )r   ru   rO   ry   rq   r   r   r   true_angular_from_alt_cosine&  s    r{   c           
      C   s   d}d}d}d}d}t | jd D ]}|| | 7 }||| 7 }q"|| jd  }|| jd  }t | jd D ]@}| | | }|| | }	||d 7 }||	d 7 }|||	 7 }qj|dkr|dkrdS |dkrdS d|t||   S d S rg   r   )
r   r    Zmu_xZmu_yrh   ri   Zdot_productr   Z	shifted_xZ	shifted_yr   r   r   correlation+  s*    r|   )r!   	l1_norm_x	l1_norm_yr'   r   c                 C   s   d}d}d}| j d }t|D ]6}|t| | ||  7 }|| | 7 }||| 7 }q|dkrj|dkrjdS |dksz|dkr~dS td|t||   S d S )Nr   r   r2   r#   )r   r   r   r   r   r    r!   r}   r~   r'   r   r   r   r   	hellingerI  s    
r   c                 C   s   d}d}d}| j d }t|D ]6}|t| | ||  7 }|| | 7 }||| 7 }q|dkrj|dkrjdS |dksz|dkr~tS |dkrtS t|| | }t|S d S r@   )r   r   r   r   rk   rJ   r   r   r   r   alternative_hellingerm  s     
r   c                 C   s   t dtd|   S rL   )r   r   rO   rq   r   r   r   correct_alternative_hellinger  s    r   averagec           	      C   sH  t t | }|dkr&|jdd}n|jdd}t j|jt jd}t |j||< |dkrl|d t j	S || }t 
|jt j}|dd  |d d k|dd < | | }|dkr|t j	S t |d	 }t |t t|g|jf}|d
kr
|| t j	S |dkr,||d  d t j	S d|| ||d   d  S )NZordinalZ	mergesort)kindZ	quicksortr   r#   denser   r0   minr`   )r   ZravelZasarrayZargsortr7   sizeintpZarangeastypefloat64onesZbool_ZcumsumnonzeroZconcatenatearraylenr   )	amethodZarrZsorterinvZobsr   r   r   r   r   r   rankdata  s*     

r   c                 C   s   t | }t |}t||S )N)r   r|   )r   r    Zx_rankZy_rankr   r   r   	spearmanr  s    r   )Znogili c                 C   s  | dk}|dk}| |  tj}||  tj}| }| }	|| }||	 }||d d f d d |f }
t|jd |jd d\}}}t|| ||j t|
||j	 t
|||}|dkrtdt||||}|tjkrtdn|tjkrtdt|j|j	}|S )Nr   FzDKantorovich distance inputs must be valid probability distributions.z>Optimal transport problem was INFEASIBLE. Please check inputs.z=Optimal transport problem was UNBOUNDED. Please check inputs.)r   r   r   r\   r   r   r   Zsupplyr   costr   rb   r   r   Z
INFEASIBLEZ	UNBOUNDEDr   Zflow)r   r    r   Zmax_iterrow_maskcol_maskr   ba_sumb_sumsub_costZnode_arc_dataZspanning_treegraphZinit_statusZsolve_statusr!   r   r   r   kantorovich  s<    


r   r2   c                 C   s   | dk}|dk}| |  tj}||  tj}| }| }	|| }||	 }||d d f d d |f }
t| ||
|d}|jd }|jd }d}t|D ].}t|D ] }||||f |||f  7 }qq|S )Nr   )r   regularizationr#   r   )r   r   r   r\   r	   r   r   )r   r    r   r   r   r   r   r   r   r   r   Ztransport_planZdim_iZdim_jr!   r   r:   r   r   r   sinkhorn  s&    

 r   c           
   
   C   s   d}d}d}| j d }t|D ]}|| | 7 }||| 7 }q|t| 7 }|t| 7 }| t | }|t | }d||  }	t|D ]H}|d|| t|| |	|   || t|| |	|     7 }q|S r_   r   r   FLOAT32_EPSr   log)
r   r    r!   r}   r~   r'   r   pdf_xpdf_ymr   r   r   jensen_shannon_divergence  s"    
:r   c                 C   s   d}d}t | jd D ]}|| | 7 }||| 7 }q| | }|| }t d|jd D ]4}||  ||d  7  < ||  ||d  7  < qTt|||S )Nr   r   r#   )r   r   r4   )r   r    r3   x_sumy_sumr   x_cdfy_cdfr   r   r   wasserstein_1d-  s    r   c                 C   s  d}d}t | jd D ]}|| | 7 }||| 7 }q| | }|| }t d|jd D ]4}||  ||d  7  < ||  ||d  7  < qTt|| | }d}	|dkrt |jd D ]&}|	t|| ||  | | 7 }	q|	d|  S |dkr4t |jd D ]&}|| ||  | }
|	|
|
 7 }	qt|	S |dkrvt |jd D ]$}|	t|| ||  | 7 }	qL|	S tdd S )Nr   r   r#   r
   r2   z)Invalid p supplied to Kantorvich distance)r   r   r   Zmedianr.   r   rb   )r   r    r3   r   r   r   r   r   mur!   valr   r   r   circular_kantorovich?  s4    $


"r   c           	   	   C   s   d}d}d}| j d }t|D ]}|| | 7 }||| 7 }q|t| 7 }|t| 7 }| t | }|t | }t|D ]D}||| t|| ||   || t|| ||    7 }qt|S r@   r   )	r   r    r!   r}   r~   r'   r   r   r   r   r   r   symmetric_kl_divergencei  s     
(r   zf4(u1[::1],u1[::1]))r!   intersectionr'   r   c                 C   s<   d}| j d }t|D ] }| | || A }|t| 7 }q|S r@   )r   r   popcnt)r   r    r!   r'   r   r   r   r   r   bit_hamming  s    
r   )r!   denomand_or_r'   r   c                 C   sh   d}d}| j d }t|D ]<}| | || @ }| | || B }|t| 7 }|t| 7 }qt||  S r@   )r   r   r   r   r   )r   r    r!   r   r'   r   r   r   r   r   r   bit_jaccard  s    
r   )4r"   l2Zsqeuclideanr/   Ztaxicabl1r1   Z	linfinityZlinftyZlinfr4   Z
seuclideanr,   Z
wminkowskir6   r;   rB   rj   ro   r|   re   Z
braycurtisr   rx   rz   r   r   Zwassersteinr   zwasserstein-1dzkantorovich-1dZkantorovich_1dr   Zcircular_wassersteinr   zjensen-shannonZjensen_shannonzsymmetric-klZsymmetric_klZsymmetric_kullback_lieblerr?   rI   rY   rT   rZ   Zrogerstanimotor]   ZsokalsneathZsokalmichenerrf   r   r   )distZ
correction)r"   r   rj   ro   rz   r   rI   )r
   )r   )r#   )r#   )Snumpyr   ZnumbaZpynndescent.optimal_transportr   r   r   r   r   r   r   r	   Zeyer8   Z_mock_identityr   Z
_mock_oneszerosr   Z_dummy_costZfinfoZepsr   r0   rk   r   r   r   Znjitr"   typesArrayr   Zuint16r*   r,   r/   r1   r4   r6   r;   r?   rB   rD   rI   Zuint8rK   Z	vectorizerQ   rT   rY   rZ   r[   r]   r^   ra   re   rf   rj   rm   ro   rp   rs   rx   rz   r{   r|   r   r   r   r   r   r   r   r   r   r   r   r   r   Znamed_distancesr   Zfast_distance_alternativesr   r   r   r   <module>   s8  (












	




























	




	















#


0

)


C

