pybroker.model module

Contains model related functionality.

class CachedModel(model: Any, input_cols: tuple[str] | None)

Bases: NamedTuple

Stores cached model data.

model

Trained model instance.

Type:

Any

input_cols

Names of the columns to be used as input for the model when making predictions.

Type:

tuple[str] | None

input_cols: tuple[str] | None

Alias for field number 1

model: Any

Alias for field number 0

class ModelLoader(name: str, load_fn: Callable[[...], Any | tuple[Any, Iterable[str]]], indicator_names: Iterable[str], input_data_fn: Callable[[DataFrame], DataFrame] | None, predict_fn: Callable[[Any, DataFrame], ndarray[Any, dtype[_ScalarType_co]]] | None, kwargs: dict[str, Any])

Bases: ModelSource

Loads a pre-trained model.

Parameters:

• name – Name of model.

• load_fn – Callable[[symbol: str, train_start_date: datetime, train_end_date: datetime, ...], DataFrame] used to load and return a pre-trained model. This is expected to return either a trained model instance, or a tuple containing a trained model instance and a Iterable of column names to to be used as input for the model when making predictions.

• indicator_names – Iterable of names of pybroker.indicator.Indicators used as features of the model.

• input_data_fn – Callable[[DataFrame], DataFrame] for preprocessing input data passed to the model when making predictions. If set, input_data_fn will be called with a pandas.DataFrame containing all test data.

• predict_fn – Callable[[Model, DataFrame], ndarray] that overrides calling the model’s default predict function. If set, predict_fn will be called with the trained model and a pandas.DataFrame containing all test data.

• kwargs – dict of kwargs to pass to load_fn.

__call__(symbol: str, train_start_date: datetime, train_end_date: datetime) → Any | tuple[Any, Iterable[str]]

Loads pre-trained model.

Parameters:

• symbol – Ticker symbol for loading the pre-trained model.

• train_start_date – Start date of training window.

• train_end_date – End date of training window.

Returns:

Pre-trained model.

class ModelSource(name: str, indicator_names: Iterable[str], input_data_fn: Callable[[DataFrame], DataFrame] | None, predict_fn: Callable[[Any, DataFrame], ndarray[Any, dtype[_ScalarType_co]]] | None, kwargs: dict[str, Any])

Bases: object

Base class of a model source. A model source provides a model instance either by training one or by loading a pre-trained model.

Parameters:

• name – Name of model.

• indicator_names – Iterable of names of pybroker.indicator.Indicators used as features of the model.

• input_data_fn – Callable[[DataFrame], DataFrame] for preprocessing input data passed to the model when making predictions. If set, input_data_fn will be called with a pandas.DataFrame containing all test data.

• predict_fn – Callable[[Model, DataFrame], ndarray] that overrides calling the model’s default predict function. If set, predict_fn will be called with the trained model and a pandas.DataFrame containing all test data.

• kwargs – dict of additional kwargs.

prepare_input_data(df: DataFrame) → DataFrame

Prepares a pandas.DataFrame of input data for passing to a model when making predictions. If set, the input_data_fn is used to preprocess the input data. If False, then indicator columns in df are used as input features.

class ModelTrainer(name: str, train_fn: Callable[[...], Any | tuple[Any, Iterable[str]]], indicator_names: Iterable[str], input_data_fn: Callable[[DataFrame], DataFrame] | None, predict_fn: Callable[[Any, DataFrame], ndarray[Any, dtype[_ScalarType_co]]] | None, kwargs: dict[str, Any])

Bases: ModelSource

Trains a model.

Parameters:

• name – Name of model.

• train_fn – Callable[[symbol: str, train_data: DataFrame, test_data: DataFrame, ...], DataFrame] used to train and return a model. This is expected to return either a trained model instance, or a tuple containing a trained model instance and a Iterable of column names to to be used as input for the model when making predictions.

• indicator_names – Iterable of names of pybroker.indicator.Indicators used as features of the model.

• input_data_fn – Callable[[DataFrame], DataFrame] for preprocessing input data passed to the model when making predictions. If set, input_data_fn will be called with a pandas.DataFrame containing all test data.

• predict_fn – Callable[[Model, DataFrame], ndarray] that overrides calling the model’s default predict function. If set, predict_fn will be called with the trained model and a pandas.DataFrame containing all test data.

• kwargs – dict of kwargs to pass to train_fn.

__call__(symbol: str, train_data: DataFrame, test_data: DataFrame) → Any | tuple[Any, Iterable[str]]

Trains model.

Parameters:

• symbol – Ticker symbol of model (models are trained per symbol).

• train_data – Train data.

• test_data – Test data.

Returns:

Trained model.

class ModelsMixin

Bases: object

Mixin implementing model related functionality.

train_models(model_syms: Iterable[ModelSymbol], train_data: DataFrame, test_data: DataFrame, indicator_data: Mapping[IndicatorSymbol, Series], cache_date_fields: CacheDateFields) → dict[ModelSymbol, TrainedModel]

Trains models for the provided pybroker.common.ModelSymbol pairs.

Parameters:

• model_syms – Iterable of pybroker.common.ModelSymbol pairs of models to train.

• train_data – pandas.DataFrame of training data.

• test_data – pandas.DataFrame of test data.

• indicator_data – Mapping of pybroker.common.IndicatorSymbol pairs to pandas.Series of pybroker.indicator.Indicator values.

• cache_date_fields – Date fields used to key cache data.

Returns:

dict mapping each pybroker.common.ModelSymbol pair to a pybroker.common.TrainedModel.

model(name: str, fn: Callable[[...], Any | tuple[Any, Iterable[str]]], indicators: Iterable[Indicator] | None = None, input_data_fn: Callable[[DataFrame], DataFrame] | None = None, predict_fn: Callable[[Any, DataFrame], ndarray[Any, dtype[_ScalarType_co]]] | None = None, pretrained: bool = False, **kwargs) → ModelSource

Creates a ModelSource instance and registers it globally with name.

Parameters:

• name – Name for referencing the model globally.

• fn – Callable used to either train or load a model instance. If for training, then fn has signature Callable[[symbol: str, train_data: DataFrame, test_data: DataFrame, ...], DataFrame]. If for loading, then fn has signature Callable[[symbol: str, train_start_date: datetime, train_end_date: datetime, ...], DataFrame]. This is expected to return either a trained model instance, or a tuple containing a trained model instance and a Iterable of column names to to be used as input for the model when making predictions.

• indicators – Iterable of pybroker.indicator.Indicators used as features of the model.

• input_data_fn – Callable[[DataFrame], DataFrame] for preprocessing input data passed to the model when making predictions. If set, input_data_fn will be called with a pandas.DataFrame containing all test data.

• predict_fn – Callable[[Model, DataFrame], ndarray] that overrides calling the model’s default predict function. If set, predict_fn will be called with the trained model and a pandas.DataFrame containing all test data.

• pretrained – If True, then fn is used to load and return a pre-trained model. If False, fn is used to train and return a new model. Defaults to False.

• **kwargs – Additional arguments to pass to fn.

Returns:

ModelSource instance.