linear_regression.go

// Copyright (c) 2018 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package temporal

import (
	"fmt"
	"math"
	"time"

	"github.com/m3db/m3/src/query/executor/transform"
	"github.com/m3db/m3/src/query/ts"
)

const (
	// PredictLinearType predicts the value of time series t seconds from now, based on the input series, using simple linear regression.
	// PredictLinearType should only be used with gauges.
	PredictLinearType = "predict_linear"

	// DerivType calculates the per-second derivative of the time series, using simple linear regression.
	// DerivType should only be used with gauges.
	DerivType = "deriv"
)

type linearRegressionProcessor struct {
	fn      linearRegFn
	isDeriv bool
}

func (l linearRegressionProcessor) Init(op baseOp, controller *transform.Controller, opts transform.Options) Processor {
	return &linearRegressionNode{
		op:         op,
		controller: controller,
		timeSpec:   opts.TimeSpec,
		fn:         l.fn,
		isDeriv:    l.isDeriv,
	}
}

type linearRegFn func(float64, float64) float64

// NewLinearRegressionOp creates a new base temporal transform for linear regression functions
func NewLinearRegressionOp(args []interface{}, optype string) (transform.Params, error) {
	var (
		fn      linearRegFn
		isDeriv bool
	)

	switch optype {
	case PredictLinearType:
		if len(args) != 2 {
			return emptyOp, fmt.Errorf("invalid number of args for %s: %d", PredictLinearType, len(args))
		}

		duration, ok := args[1].(float64)
		if !ok {
			return emptyOp, fmt.Errorf("unable to cast to scalar argument: %v for %s", args[1], PredictLinearType)
		}

		fn = func(slope, intercept float64) float64 {
			return slope*duration + intercept
		}

	case DerivType:
		fn = func(slope, _ float64) float64 {
			return slope
		}

		isDeriv = true

	default:
		return nil, fmt.Errorf("unknown linear regression type: %s", optype)
	}

	l := linearRegressionProcessor{
		fn:      fn,
		isDeriv: isDeriv,
	}

	return newBaseOp(args, optype, l)
}

type linearRegressionNode struct {
	op         baseOp
	controller *transform.Controller
	timeSpec   transform.TimeSpec
	fn         linearRegFn
	isDeriv    bool
}

func (l linearRegressionNode) Process(dps ts.Datapoints, evaluationTime time.Time) float64 {
	if dps.Len() < 2 {
		return math.NaN()
	}

	slope, intercept := linearRegression(dps, evaluationTime, l.isDeriv)
	return l.fn(slope, intercept)
}

// linearRegression performs a least-square linear regression analysis on the
// provided datapoints. It returns the slope, and the intercept value at the
// provided time. The algorithm we use comes from https://en.wikipedia.org/wiki/Simple_linear_regression.
func linearRegression(dps ts.Datapoints, interceptTime time.Time, isDeriv bool) (float64, float64) {
	var (
		n                                   float64
		sumTimeDiff, sumVals                float64
		sumTimeDiffVals, sumTimeDiffSquared float64
		valueCount                          int
	)

	for _, dp := range dps {
		if math.IsNaN(dp.Value) {
			continue
		}

		if valueCount == 0 && isDeriv {
			// set interceptTime as timestamp of first non-NaN dp
			interceptTime = dp.Timestamp
		}

		valueCount++
		timeDiff := dp.Timestamp.Sub(interceptTime).Seconds()
		n += 1.0
		sumVals += dp.Value
		sumTimeDiff += timeDiff
		sumTimeDiffVals += timeDiff * dp.Value
		sumTimeDiffSquared += timeDiff * timeDiff
	}

	// need at least 2 non-NaN values to calculate slope and intercept
	if valueCount == 1 {
		return math.NaN(), math.NaN()
	}

	covXY := sumTimeDiffVals - sumTimeDiff*sumVals/n
	varX := sumTimeDiffSquared - sumTimeDiff*sumTimeDiff/n

	slope := covXY / varX
	intercept := sumVals/n - slope*sumTimeDiff/n

	return slope, intercept
}