Skip to content

tinkoff-ai/etna

master
Switch branches/tags
Code

Files

Permalink
Failed to load latest commit information.
Type
Name
Latest commit message
Commit time
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Predict your time series the easiest way

PyPI Version Python versions Downloads

Coverage Test passing Docs publish License

Telegram Slack GitHub Discussions Contributors Stars

Homepage | Documentation | Tutorials | Contribution Guide | Release Notes

ETNA is an easy-to-use time series forecasting framework. It includes built in toolkits for time series preprocessing, feature generation, a variety of predictive models with unified interface - from classic machine learning to SOTA neural networks, models combination methods and smart backtesting. ETNA is designed to make working with time series simple, productive, and fun.

ETNA is the first python open source framework of Tinkoff.ru Artificial Intelligence Center. The library started as an internal product in our company - we use it in over 10+ projects now, so we often release updates. Contributions are welcome - check our Contribution Guide.

Get started

Let's load and prepare the data.

import pandas as pd
from etna.datasets import TSDataset

# Read the data
df = pd.read_csv("examples/data/example_dataset.csv")

# Create a TSDataset
df = TSDataset.to_dataset(df)
ts = TSDataset(df, freq="D")

# Choose a horizon
HORIZON = 14

# Make train/test split
train_ts, test_ts = ts.train_test_split(test_size=HORIZON)

Define transformations and model:

from etna.models import CatBoostModelMultiSegment
from etna.transforms import DateFlagsTransform
from etna.transforms import DensityOutliersTransform
from etna.transforms import FourierTransform
from etna.transforms import LagTransform
from etna.transforms import LinearTrendTransform
from etna.transforms import MeanTransform
from etna.transforms import SegmentEncoderTransform
from etna.transforms import TimeSeriesImputerTransform
from etna.transforms import TrendTransform

# Prepare transforms
transforms = [
    DensityOutliersTransform(in_column="target", distance_coef=3.0),
    TimeSeriesImputerTransform(in_column="target", strategy="forward_fill"),
    LinearTrendTransform(in_column="target"),
    TrendTransform(in_column="target", out_column="trend"),
    LagTransform(in_column="target", lags=list(range(HORIZON, 122)), out_column="target_lag"),
    DateFlagsTransform(week_number_in_month=True, out_column="date_flag"),
    FourierTransform(period=360.25, order=6, out_column="fourier"),
    SegmentEncoderTransform(),
    MeanTransform(in_column=f"target_lag_{HORIZON}", window=12, seasonality=7),
    MeanTransform(in_column=f"target_lag_{HORIZON}", window=7),
]

# Prepare model
model = CatBoostModelMultiSegment()

Fit Pipeline and make a prediction:

from etna.pipeline import Pipeline

# Create and fit the pipeline
pipeline = Pipeline(model=model, transforms=transforms, horizon=HORIZON)
pipeline.fit(train_ts)

# Make a forecast
forecast_ts = pipeline.forecast()

Let's plot the results:

from etna.analysis import plot_forecast

plot_forecast(forecast_ts=forecast_ts, test_ts=test_ts, train_ts=train_ts, n_train_samples=50)

Print the metric value across the segments:

from etna.metrics import SMAPE

metric = SMAPE(mode="macro")
metric_value = metric(y_true=test_ts, y_pred=forecast_ts)
>>> {'segment_b': 3.3017151519000967, 'segment_c': 5.270557433427279, 'segment_a': 5.272811627335398, 'segment_d': 4.689085450895735}

Installation

ETNA is available on PyPI, so you can use pip to install it.

Install default version:

pip install --upgrade pip
pip install etna

The default version doesn't contain all the dependencies, because some of them are needed only for specific models, e.g. Prophet, PyTorch. Available user extensions are the following:

  • prophet
  • torch
  • wandb

Install extension:

pip install etna[extension-name]

Install all extensions:

pip install etna[all]

There are also developer extensions. All the extensions are listed in pyproject.toml.

Without the appropriate extension you will get an ImportError trying to import the model that needs it. For example, etna.models.ProphetModel needs prophet extension and can't be used without it.

Configuration

ETNA supports configuration files. It means that library will check that all the specified packages are installed prior to script start and NOT during runtime.

To set up a configuration for your project you should create a .etna file at the project's root. To see the available options look at Settings. There is an example of configuration file.

Tutorials

We have also prepared a set of tutorials for an easy introduction:

Notebook Interactive launch
Get started Binder
Backtest Binder
EDA Binder
Outliers Binder
Clustering Binder
Deep learning models Binder
Ensembles Binder

Documentation

ETNA documentation is available here.

Community

To ask the questions or discuss the library you can join our telegram chat. Discussions section on github is also open for this purpose.

Resources

Acknowledgments

ETNA.Team

Andrey Alekseev, Nikita Barinov, Dmitriy Bunin, Aleksandr Chikov, Vladislav Denisov, Martin Gabdushev, Sergey Kolesnikov, Artem Makhin, Ivan Mitskovets, Albina Munirova, Nikolay Romantsov, Julia Shenshina

ETNA.Contributors

Artem Levashov, Aleksey Podkidyshev, Carlosbogo

License

Feel free to use our library in your commercial and private applications.

ETNA is covered by Apache 2.0. Read more about this license here

Please note that etna[prophet] is covered by GPL 2.0 due to pystan package.