
    zj]                        d Z ddlZddlZddlmZ dZ ej        ddgddgddgg          Z	d Z
ddZ ed	d	
          d             Zd ZddZ ed	d	
          d             Zd Zd Z ed	d	
          d             ZddZd ZdS )a   
utils.py

Mathematical and signal processing utilities for acoustic tracking.

(c) Walter M.X. Zimmer (2025)
Original code produced for the CIAN Tool Kit, associated with the manuscript: 
Berkenbaum, L., Glotin, H., Sarano, F., Zimmer, W.M.X., Sarano, V., Adam, O., Giraudet, P. (2026). "Audiovisual diarization of overlapping click trains in sperm whale (Physeter macrocephalus) vocal sparring using a three-hydrophone array," J. Acoust. Soc. Am.

Includes shared constants for the OPALE array configuration.

    N)jitg     @      c                 d   t          j        |           5 }|j        }|j        }|j        }|                    t          ||z                       |                    t          ||z                      }|j        dk    r|	                    dd          }d d d            n# 1 swxY w Y   |||fS )Nr   )
sf	SoundFile
sampleratechannelsframesseekintreadndimreshape)filenamet1dtsoundfsnchnddatas           [/Users/laraberkenbaum/Desktop/PhD_thesis/JASA/JASA_VocalSparring_2026/02_notebooks/utils.pywav_readr      s    	h		 &5n\

3r"u::zz#be**%%9>><<1%%D& & & & & & & & & & & & & & & R<s   BB""B&)B&c                 0   |j         dk    r|                    dd          }| j         dk    r|                     dd          } | j        \  }}|j        \  }}t          ||          }d|dz
                                  z  }	|	|z
  dz   }
||	k     rt          d          }nt          d||
          }t          j        |           st          j        |          rHt          j        j        }t          j        j	        }t          j
        ||	z   |ft          j                  }n;t          j        j        }t          j        j        }t          j
        ||	z   |f          } || |	d          }|D ]H} |||||
z   d d f         |	d          }||z  }||||	z   d d fxx          ||d          z  cc<   I|Q|d |j        d         d d f         |z   |d |j        d         d d f<   ||d d d f         }||||z   d d f         |fS ||||z   d d f         S )Nr   r   r   r   dtypenaxisr!   )r   r   shapemax
bit_lengthrangenpiscomplexobjfftifftzeros
complex128rfftirfft)bxzinhL_IN_IL_sigN_sigN_chanL_FL_Soffsetsfft_func	ifft_funcresFDirr    u1u2s                      r   fft_filtrA   &   sE   v{{		"QAv{{		"QAwHC7LE5__F
c!e!!
!C
)a-Cs{{((5#&&	q ,R_Q// ,6:FK	hc	6*"-@@@6;FL	hc	6*++8A1%%%D 0 0Xa!C%	lc222$YAaeGAAAI))BQ////	~!,28A;,qqq.1B6LRXa[LN]2bh;qqq=!2%%2bh;qqq=!!    T)nopythoncachec                 6    d}| dk    rdS | r| dz  } |dz  }| |S Nr   r    )r    bitss     r   r%   r%   L   s?    DAvvaa
 	a	   KrB   c                     | t          j        | j        d         df          t          | j        d                   z  z   S rF   )r'   onesr#   r&   )uus    r   stackrL   U   s4    !Q((5!+=+====rB   c                    d}t          j        |           d         }t          j        |          d         }t          j        ||f          }|d k    r|}t          |          D ]}| d d ||df         f         }| d d ||df         f         }	t           j                            ||z  |d          }
t           j                            |	|z  |d          }|t          j        |
          z  }|t          j        t          j        |                    z  }t           j        	                    |d          }t           j        
                    |          d |         |d d |f<   |S )Nr   r   r   r"   )r'   r#   r+   r&   r)   r-   	conjugatesqrtabsr.   fftshift)zzhselnfftwinnzncvvjju0r?   v0v1wowws                 r   fftCorrr_   X   sC   
C	"aB	$	B	2r(		Bt||BTBii - -41:41:V[[C4a[00V[[C4a[00",r""""'"&**%%%V\\"1\%%FOOB'',111b5		IrB   c                 p   t          j        |          d         }t          j        d|z   |z   |f          }t          |          D ]}| d d ||df         f                                         }| d d ||df         f                                         }t          |          D ]3}t          j        |d ||z
            |||z
  d                    |||f<   4t          j        ||          |||f<   t          d|dz             D ]0}t          j        ||d          |d |                    |||z   |f<   1|S rF   )r'   r#   r+   r&   copydot)	rK   rS   mcrW   rX   rY   rZ   r?   iis	            r   xCorrre   i   sI   	$	B	1R47B-	 	 BBii 6 641:##%%41:##%%)) 	= 	=B9beH9r2b5((|<<Br2vJJVB^^2r6
2a4.. 	6 	6BF2bcc7BttH55Br"ubyMM	6 IrB   c                    | j         \  }}t          |          }||k    rd S t          j        |dz
  t          j        d|                    }| |t          |          f         }| |dz
  t          |          f         }| |dz   t          |          f         }||z   d|z  z
  dz  }||z
  d|z  z  }	|||	dz  z  z
  }
|                    d          |	z   }|
}||fS )Nr   r      float)r#   lenr'   minimummaximumr&   astype)xxidlnxrW   niuoumupr/   xoyojdladls                r   quadIntrx   y   s    XFB	SB	Bwwtt
*RT2:a--
.
.C	CrN	B	CE599	B	CE599	B	b1R41A
r'ac	B	aAgB
**W


"C
C8OrB   c                 t   t          j        ||z  |z                                t                    }t          j        |t
                    }t          j        |          |z  t          j        dt           j        z  dz  |z  t          j	        |          z  |z            z  }t          || d d          }|S )Nr   y              ?r   )r1   r2   )r'   ceilrl   r   r+   complexhanningexppiarangerA   )rm   r   fkQNWyys          r   baseband_filterr      s    
2	!!#&&A
'"""A

1BF2:?Q#61#Ea#GHHHA	!RDQ	'	'	'BIrB   c                    t          j        t          j        |          |z            }| j        d         }t	          d|          D ]:}| |dz
  d d f         |z  }t          j        || |d d f                   | |d d f<   ;| S rF   )r'   r~   logr#   r&   rk   )rR   zoaabbnyrd   axs          r   leakr      s    	r

2		B	!BArll , ,1QQQZ"_:b"RT(++2aaa4JrB     c                 2    ||z   }||z   }| ||d d f         S NrG   )rK   rd   i0n0i1i2s         r   extractr      s*    	bB	bBbeAAAg;rB   c                 V   |                                  }| d d df         t          j        |          z  | d d df         t          j        |          z  z
  |d d df<   | d d df         t          j        |          z  | d d df         t          j        |          z  z   |d d df<   |S )Nr   r   )ra   r'   cossin)XaYs      r   rotyr      s    	AqqqsVbfQii!AAAaC&26!99"44AaaacFqqqsVbfQii!AAAaC&26!99"44AaaacFHrB   )Nr   r   )r   r   )__doc__numpyr'   	soundfiler   numbar   C_SOUNDarrayHSELr   rA   r%   rL   r_   re   rx   r   r   r   r   rG   rB   r   <module>r      sy                 
 rx!Q!Q!Q())
	 	 	$" $" $" $"L d$   > > >   " d$     $   d$      
    rB   