"""Utility functions for simulations."""

import numpy as np

###################################################################################################
###################################################################################################

def normalized_sum(sig1, sig2, ratio):
    """Combine two signals, after transforming to have the desired variance ratio.

    Parameters
    ----------
    sig1, sig2 : 1d array
        Vectors of data to normalize variance with respect to each other.
    ratio : float
        Desired ratio of sig1 variance to sig2 variance.
        If > 1 - sig1 is stronger, if < 1 - sig2 is stronger.

    Returns
    -------
    1d array
        New vector with the combined input signals.
    """

    return sum(normalize_variance(sig1, sig2, ratio))


def normalize_variance(sig1, sig2, ratio):
    """Normalize the variance across two signals to reflect a specified ratio.

    Parameters
    ----------
    sig1, sig2 : 1d array
        Vectors of data to normalize variance with respect to each other.
    ratio : float
        Desired ratio of sig1 variance to sig2 variance.
        If > 1 - sig1 is stronger, if < 1 - sig2 is stronger.

    Returns
    -------
    1d array, 1d array
        New sig1 & sig2, where sig2 has been normalized so they have the desired variance ratio.
    """

    return sig1, np.sqrt(sig2**2 * np.var(sig1) / (np.var(sig2) * ratio)) * np.sign(sig2)

@@ -284,32 +284,53 @@
 ig1, sig2, ratio
+, select_nonzero=True
 ):%0A    %22%22%22Combin
@@ -658,32 +658,169 @@
 ig2 is stronger.
+%0A    select_nonzero : boolean, optional, default=True%0A        Whether to calculate the variance of sig1 across only non-zero data points.
 %0A%0A    Returns%0A  
@@ -954,16 +954,32 @@
 2, ratio
+, select_nonzero
 ))%0A%0A%0Adef
@@ -1015,16 +1015,37 @@
 2, ratio
+, select_nonzero=True
 ):%0A    %22
@@ -1385,16 +1385,153 @@
 tronger.
+%0A    select_nonzero : boolean, optional, default=True%0A        Whether to calculate the variance of sig1 across only non-zero data points.
 %0A%0A    Re
@@ -1672,64 +1672,469 @@
 io.%0A
+%0A
     
-%22%22%22%0A%0A    return sig1, np.sqrt(sig2**2 * np.var(sig1)
+Notes%0A    -----%0A    Only sig2 is actually modified by this procedure, relative to the variance in sig1.%0A    If sig1 is a potentially non-continuous periodic signal (for example, bursty), %60select_nonzero%60%0A    allows for normalizing the variance across only segments of the signal with signal present.%0A    %22%22%22%0A%0A    if select_nonzero:%0A        sig1_var = np.var(np.nonzero(sig1))%0A    else:%0A        sig1_var = np.var(sig1)%0A%0A    sig2 = np.sqrt(sig2**2 * sig1_var
  / (
@@ -2172,8 +2172,31 @@
 n(sig2)%0A
+%0A    return sig1, sig2%0A

Extend variance normalization to deal with non-continuous signals