Metadata-Version: 2.4
Name: aurorax-model
Version: 0.1.0
Summary: Aurora-X: a time series foundation model with native covariate support
Author-email: XJWu <1793410861@qq.com>
License-Expression: Apache-2.0
Project-URL: Homepage, https://huggingface.co/DecisionIntelligence/Aurora-X
Project-URL: Model weights, https://huggingface.co/DecisionIntelligence/Aurora-X
Keywords: time-series,forecasting,foundation-model,covariates,transformer
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Operating System :: OS Independent
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.4.0
Requires-Dist: transformers>=4.50.0
Requires-Dist: huggingface_hub>=0.16.0
Requires-Dist: safetensors>=0.4.0
Requires-Dist: numpy>=1.21.0
Requires-Dist: einops>=0.8.1
Dynamic: license-file

# Aurora-X

Aurora-X is a time series foundation model with **native covariate support**.
It forecasts univariate and multivariate series zero-shot, and can condition on
past-only and known-future covariates in a single forward pass.

Model weights: [DecisionIntelligence/Aurora-X](https://huggingface.co/DecisionIntelligence/Aurora-X)

## Installation

```bash
pip install aurorax-model
```

The architecture config ships with the package; the weights (~4 GB) are pulled
from the Hugging Face Hub the first time you load the model and cached locally.

## Quick start

```python
import numpy as np
from aurorax import load_pipeline

pipe = load_pipeline()          # downloads weights on first use

context = np.random.randn(512)                      # 512 historical steps
preds = pipe.predict(context, prediction_length=96) # list with one entry
print(preds[0].shape)                               # (1, 20, 96) = (n_targets, n_samples, horizon)
```

`predict` returns a list with one tensor per input case, shaped
`(n_targets, K, prediction_length)`, where `K` is `num_samples` in sampling mode
or the number of quantile levels in quantile mode.

### Point forecast + quantiles

```python
quantiles, mean = pipe.predict_quantiles(
    context,
    prediction_length=96,
    quantile_levels=[0.1, 0.5, 0.9],
)
mean[0].shape         # (1, 96)     -> point forecast
quantiles[0].shape    # (1, 96, 3)  -> one column per quantile level
```

### Multivariate

Wrap a `(n_variates, context_length)` array in a list to forecast its variates
jointly, so the model can exploit cross-variate structure.

```python
series = np.random.randn(3, 512)              # 3 variates, 512 steps

preds = pipe.predict([series], prediction_length=96)   # list of 2D -> one multivariate case
preds[0].shape        # (3, 20, 96)
```

> **Careful:** a bare 2D array is read as a *batch of univariate series*, not as
> one multivariate case. `pipe.predict(series, ...)` returns **three separate**
> univariate forecasts. Use `[series]` (or the 3D form `series[None]`) whenever
> the variates belong together.

### Covariates

Give each case a `target` plus optional covariates. Every covariate needs its
history in `past_covariates`; add it to `future_covariates` as well when the
future values are known ahead of time (calendar, weather forecast, promotions).

```python
preds = pipe.predict(
    [{
        "target": np.random.randn(512),
        # past-only covariate: history is used, future is unknown
        "past_covariates":   {"sales": np.random.randn(512),
                              "temp":  np.random.randn(512)},
        # known-future covariate: future values are fed to the model
        "future_covariates": {"temp":  np.random.randn(96)},
    }],
    prediction_length=96,
)
preds[0].shape        # (1, 20, 96) -> only target rows are returned
```

Covariate rows condition the forecast but are never scored, so the output only
covers the `target` rows.

### Batching

Pass a list to forecast many cases in one go. Contexts may have different
lengths; shorter ones are left-padded and masked automatically.

```python
preds = pipe.predict(
    [np.random.randn(300), np.random.randn(512), np.random.randn(1024)],
    prediction_length=96,
)
len(preds)            # 3
```

## Inputs

`predict` and `predict_quantiles` accept any of:

| Form | Meaning |
|---|---|
| 1D array `(T,)` | one univariate series |
| 2D array `(N, T)` | `N` **independent univariate** series |
| 3D array `(N, V, T)` | `N` multivariate cases, `V` variates each |
| list of 1D arrays | `N` univariate cases, context lengths may differ |
| list of 2D arrays `(V, T)` | `N` multivariate cases, context lengths may differ |
| list of dicts | cases with covariates (see above) |

NumPy arrays and PyTorch tensors are interchangeable everywhere; `NaN` marks
missing values and is masked out. Returned tensors are always on CPU.

## Key arguments

| Argument | Meaning | Default |
|---|---|---|
| `prediction_length` | forecast horizon | required |
| `num_samples` | number of sampled trajectories (`mode="sample"`) | 20 |
| `mode` | `"sample"` or `"quantile"` | `"sample"` |
| `quantile_levels` | levels to return when `mode="quantile"` | — |
| `inference_token_len` | patch size; lower it (8/16/32) for short series | model default (48) |
| `batch_size` | max variate rows per forward pass | 256 |
| `cross_learning` | share attention across all cases in the batch | `False` |

## Lower-level access

`load_model` returns the raw model when you want to drive `generate` yourself:

```python
import torch
from aurorax import load_model

model = load_model()            # eval mode, on cuda if available
preds = model.generate(
    inputs=torch.randn(4, 512),
    max_output_length=96,
    num_samples=20,
)                               # (4, 20, 96)
```

Both loaders take the same arguments:

```python
load_pipeline(
    repo_id="DecisionIntelligence/Aurora-X",  # or a local checkpoint directory
    cache_dir=None,                           # where to cache the weights
    force_download=False,
    device=None,                              # default: cuda if available
)
```
