github.com/livekit/protocol@v1.16.1-0.20240517185851-47e4c6bba773/utils/welford.go (about)

     1  // Copyright 2023 LiveKit, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package utils
    16  
    17  import (
    18  	"math"
    19  )
    20  
    21  // Welford implements Welford's online algorithm for variance
    22  // SEE: https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Welford's_online_algorithm
    23  type Welford struct {
    24  	count float64
    25  	mean  float64
    26  	m2    float64
    27  }
    28  
    29  func (w *Welford) Reset() {
    30  	w.count = 0
    31  	w.mean = 0
    32  	w.m2 = 0
    33  }
    34  
    35  func (w *Welford) Update(v float64) {
    36  	w.count++
    37  	d := v - w.mean
    38  	w.mean += d / w.count
    39  	d2 := v - w.mean
    40  	w.m2 += d * d2
    41  }
    42  
    43  func (w Welford) Value() (mean, variance, sampleVariance float64) {
    44  	if w.count < 2 {
    45  		return w.mean, 0, 0
    46  	}
    47  	return w.mean, w.m2 / w.count, w.m2 / (w.count - 1)
    48  }
    49  
    50  func (w Welford) Count() float64 {
    51  	return w.count
    52  }
    53  
    54  func (w Welford) Mean() float64 {
    55  	return w.mean
    56  }
    57  
    58  func (w Welford) Variance() float64 {
    59  	return w.m2 / (w.count - 1)
    60  }
    61  
    62  func (w Welford) StdDev() float64 {
    63  	return math.Sqrt(w.Variance())
    64  }
    65  
    66  func WelfordMerge(ws ...Welford) Welford {
    67  	switch len(ws) {
    68  	case 0:
    69  		return Welford{}
    70  	case 1:
    71  		return ws[0]
    72  	case 2:
    73  		if ws[0].count == 0 {
    74  			return ws[1]
    75  		}
    76  		if ws[1].count == 0 {
    77  			return ws[0]
    78  		}
    79  		count := ws[0].count + ws[1].count
    80  		delta := ws[1].mean - ws[0].mean
    81  		return Welford{
    82  			count: count,
    83  			mean:  (ws[0].mean*ws[0].count + ws[1].mean*ws[1].count) / count,
    84  			m2:    ws[0].m2 + ws[1].m2 + delta*delta*ws[0].count*ws[1].count/count,
    85  		}
    86  	default:
    87  		n := len(ws) >> 1
    88  		return WelfordMerge(WelfordMerge(ws[:n]...), WelfordMerge(ws[n:]...))
    89  	}
    90  }