utils for smoothing or generating smooth noise

smooth.py

def smooth_noise(shp, dim=1, smoothing=30):
    """
    Generate smoothed random noise that looks like a random walk.
    
    - shp: shape of noise
    - dim: dimension to smooth along
    - smooth_steps: smoothing steps
    
    might want to just use use hhttps://en.wikipedia.org/wiki/Ornstein%E2%80%93Uhlenbeck_process
    """
    shp = list(shp)
    shp[dim] += smoothing-1 # add this dim to counter smoothing
    r = torch.randn(size=shp)
    # r = np.random.normal(size=shp)
    fudge_factor = 2**(1/3)
    n = r.unfold(dim, smoothing, 1).mean(-1) * np.sqrt(smoothing) * fudge_factor
    return n.numpy()


def smooth_simple(x: np.array, n=4):
    return pd.DataFrame(x).rolling(window=n, axis=0, center=True, min_periods=1).mean().values


def moving_average_pd(x, axis=1, n=3, pad=True, retain_amp=1):
    """we want to keep peaks in same place, and retain mean"""
    x = np.swapaxes(x, axis, 0)
    assert np.isfinite(x).all()

    x = pd.DataFrame(x).rolling(window=n, axis=0, center=True, min_periods=1).mean()

    if retain_amp > 0:
        x = x * np.sqrt(n)
    
    if pad:
        x = x.bfill().ffill().values
        assert np.isfinite(x).all()
    else:
        x = x.dropna().values
        assert np.isfinite(x).all()
    x = np.swapaxes(x, 0, axis)
    return x

def smooth(x: np.array, dim=1, smoothing=30, pad=False, retain_amp=1) -> np.array:
    """
    retain_amp=0: will smooth peaks
    retain_amp=1: will retain peaks
    """
    assert x.shape[dim]>smoothing
    m = x.mean(dim, keepdims=True)
    s = x.std(dim, keepdims=True) + 1e-3
    x = x.copy()
    x = x - m
    x = x / s
    x = moving_average_pd(x, axis=dim, n=smoothing, pad=pad)
    x = x * s
    x = x + m
    return x