FranSys

FranSys prediction framework: diagnosis/prognosis models and training utilities.

Diag_RNN

Diag_RNN(input_size: int, output_size: int, hidden_size: int = 100, rnn_layer: int = 1, linear_layer: int = 1, **kwargs)

Bases: Module

RNN-based diagnosis model that outputs flat state vectors.

Outputs [batch, seq_len, output_size] — a flat state estimate per timestep. Use RNN.state_size from the prognosis model to set output_size so the diagnosis output can be directly unflattened into the prognosis hidden state.

Parameters:

Name	Type	Description	Default
input_size	int	number of input features	required
output_size	int	flat state dimension (typically prognosis_rnn.state_size)	required
hidden_size	int	number of hidden units in the RNN	100
rnn_layer	int	number of RNN layers	1
linear_layer	int	number of linear layers in the output head	1

Source code in tsfast/prediction/fransys.py

def __init__(
    self,
    input_size: int,
    output_size: int,
    hidden_size: int = 100,
    rnn_layer: int = 1,
    linear_layer: int = 1,
    **kwargs,
):
    super().__init__()
    self.rnn = RNN(input_size, hidden_size, rnn_layer, ret_full_hidden=False, **kwargs)
    self.final = SeqLinear(hidden_size, output_size, hidden_layer=linear_layer - 1)

Diag_RNN_raw

Diag_RNN_raw(input_size: int, hidden_size: int = 100, rnn_layer: int = 1)

Bases: Module

Raw RNN diagnosis model that returns flattened hidden states directly.

Parameters:

Name	Type	Description	Default
input_size	int	number of input features	required
hidden_size	int	number of hidden units in the RNN	100
rnn_layer	int	number of RNN layers	1

Source code in tsfast/prediction/fransys.py

def __init__(
    self,
    input_size: int,
    hidden_size: int = 100,
    rnn_layer: int = 1,
):
    super().__init__()
    self.rnn = RNN(input_size, hidden_size, rnn_layer, ret_full_hidden=True)

Diag_TCN

Diag_TCN(input_size: int, output_size: int, hl_width: int, mlp_layers: int = 0, **kwargs)

Bases: Module

TCN-based diagnosis model with optional MLP output head.

Parameters:

Name	Type	Description	Default
input_size	int	number of input features	required
output_size	int	flat state dimension (typically prognosis_rnn.state_size)	required
hl_width	int	width of TCN hidden layers	required
mlp_layers	int	number of additional MLP layers after the TCN	0

Source code in tsfast/prediction/fransys.py

def __init__(self, input_size: int, output_size: int, hl_width: int, mlp_layers: int = 0, **kwargs):
    super().__init__()
    if mlp_layers > 0:
        self._model = TCN(input_size, hl_width, hl_width=hl_width, **kwargs)
        self.final = SeqLinear(hl_width, output_size, hidden_size=hl_width, hidden_layer=mlp_layers)
    else:
        self._model = TCN(input_size, output_size, hl_width=hl_width, **kwargs)
        self.final = nn.Identity()

ARProg_Init

ARProg_Init(n_u: int, n_y: int, init_sz: int, n_x: int = 0, hidden_size: int = 100, rnn_layer: int = 1, diag_model: Module | None = None, linear_layer: int = 1, final_layer: int = 0, init_sz_range: tuple[int, int] | None = None, **kwargs)

Bases: Module

Autoregressive prognosis model with diagnosis-based initialization.

Parameters:

Name	Type	Description	Default
n_u	int	number of input channels	required
n_y	int	number of output channels	required
init_sz	int	number of initial time steps used for diagnosis	required
n_x	int	number of external state channels	0
hidden_size	int	number of hidden units in the RNN	100
rnn_layer	int	number of RNN layers	1
diag_model	Module | None	custom diagnosis model outputting [batch, seq, state_dim]	None
linear_layer	int	number of linear layers in the diagnosis output head	1
final_layer	int	number of additional final layers (unused, reserved)	0
init_sz_range	tuple[int, int] | None	if set, randomize init_sz within (min, max) during training	None

Source code in tsfast/prediction/fransys.py

def __init__(
    self,
    n_u: int,
    n_y: int,
    init_sz: int,
    n_x: int = 0,
    hidden_size: int = 100,
    rnn_layer: int = 1,
    diag_model: nn.Module | None = None,
    linear_layer: int = 1,
    final_layer: int = 0,
    init_sz_range: tuple[int, int] | None = None,
    **kwargs,
):
    super().__init__()
    self.n_u = n_u
    self.n_y = n_y
    self.init_sz = init_sz
    self.init_sz_range = init_sz_range
    self.n_x = n_x
    self.hidden_size = hidden_size
    self.rnn_layer = rnn_layer
    self.diag_model = diag_model
    self.linear_layer = linear_layer
    self.final_layer = final_layer

    rnn_kwargs = dict(hidden_size=hidden_size, num_layers=rnn_layer)
    rnn_kwargs = dict(rnn_kwargs, **kwargs)

    self.rnn_prognosis = AR_Model(
        SimpleRNN(input_size=n_u + n_x + n_y, output_size=n_x + n_y, return_state=True, **rnn_kwargs),
        model_has_state=True,
        return_state=False,
        ar=True,
        out_sz=n_x + n_y,
    )
    self._prog_rnn = self.rnn_prognosis.model.rnn

    if diag_model is None:
        self.rnn_diagnosis = Diag_RNN(
            n_u + n_x + n_y,
            self._prog_rnn.state_size,
            hidden_size=hidden_size,
            rnn_layer=rnn_layer,
            linear_layer=linear_layer,
        )
    else:
        self.rnn_diagnosis = diag_model

FranSys

FranSys(n_u: int, n_y: int, init_sz: int, n_x: int = 0, hidden_size: int = 100, rnn_layer: int = 1, diag_model: Module | None = None, linear_layer: int = 1, init_diag_only: bool = False, final_layer: int = 0, init_sz_range: tuple[int, int] | None = None, **kwargs)

Bases: Module

Framework for Analysis of Systems: combined diagnosis/prognosis RNN model.

The diagnosis model can be any nn.Module that outputs [batch, seq_len, state_dim] (flat state vectors per timestep), where state_dim equals rnn_prognosis.state_size. The prognosis RNN converts flat vectors to its internal hidden-state format via unflatten_state.

Parameters:

Name	Type	Description	Default
n_u	int	number of input channels	required
n_y	int	number of output channels	required
init_sz	int	number of initial time steps used for diagnosis	required
n_x	int	number of external state channels	0
hidden_size	int	number of hidden units in the RNN	100
rnn_layer	int	number of RNN layers	1
diag_model	Module | None	custom diagnosis model outputting [batch, seq, state_dim]	None
linear_layer	int	number of linear layers in the diagnosis output head	1
init_diag_only	bool	if True, limit diagnosis to init_sz time steps during training	False
final_layer	int	number of additional layers in the shared output head	0
init_sz_range	tuple[int, int] | None	if set, randomize init_sz within (min, max) during training	None

Source code in tsfast/prediction/fransys.py

def __init__(
    self,
    n_u: int,
    n_y: int,
    init_sz: int,
    n_x: int = 0,
    hidden_size: int = 100,
    rnn_layer: int = 1,
    diag_model: nn.Module | None = None,
    linear_layer: int = 1,
    init_diag_only: bool = False,
    final_layer: int = 0,
    init_sz_range: tuple[int, int] | None = None,
    **kwargs,
):
    super().__init__()
    self.n_u = n_u
    self.n_y = n_y
    self.n_x = n_x
    self.init_sz = init_sz
    self.init_diag_only = init_diag_only
    self.init_sz_range = init_sz_range

    rnn_kwargs = dict(hidden_size=hidden_size, num_layers=rnn_layer, ret_full_hidden=True)
    rnn_kwargs = dict(rnn_kwargs, **kwargs)

    self.rnn_prognosis = RNN(n_u, **rnn_kwargs)

    if diag_model is None:
        self.rnn_diagnosis = Diag_RNN(
            n_u + n_x + n_y,
            self.rnn_prognosis.state_size,
            hidden_size=hidden_size,
            rnn_layer=rnn_layer,
            linear_layer=linear_layer,
            **kwargs,
        )
    else:
        self.rnn_diagnosis = diag_model

    self.final = SeqLinear(hidden_size, n_y, hidden_layer=final_layer)

FranSysLearner

FranSysLearner(dls, init_sz: int, attach_output: bool = False, loss_func=nn.L1Loss(), metrics: list | None = None, opt_func=torch.optim.Adam, lr: float = 0.003, transforms: list | None = None, augmentations: list | None = None, aux_losses: list | None = None, input_norm: type | None = StandardScaler, output_norm: type | None = None, **kwargs) -> Learner

Create a Learner configured for FranSys diagnosis/prognosis training.

Parameters:

Name	Type	Description	Default
dls		DataLoaders with norm_stats for input/output normalization	required
init_sz	int	number of initial time steps used for diagnosis	required
attach_output	bool	if True, use prediction_concat to concatenate output to input	False
loss_func		loss function	L1Loss()
metrics	list | None	metrics to track	None
opt_func		optimizer constructor	Adam
lr	float	learning rate	0.003
transforms	list | None	additional transforms (train + valid)	None
augmentations	list | None	additional augmentations (train only)	None
aux_losses	list | None	additional auxiliary losses	None
input_norm	type | None	scaler class for input normalization, None to disable	StandardScaler
output_norm	type | None	scaler class for output denormalization, None to disable	None

Source code in tsfast/prediction/fransys.py

def FranSysLearner(
    dls,
    init_sz: int,
    attach_output: bool = False,
    loss_func=nn.L1Loss(),
    metrics: list | None = None,
    opt_func=torch.optim.Adam,
    lr: float = 3e-3,
    transforms: list | None = None,
    augmentations: list | None = None,
    aux_losses: list | None = None,
    input_norm: type | None = StandardScaler,
    output_norm: type | None = None,
    **kwargs,
) -> Learner:
    """Create a Learner configured for FranSys diagnosis/prognosis training.

    Args:
        dls: DataLoaders with norm_stats for input/output normalization
        init_sz: number of initial time steps used for diagnosis
        attach_output: if True, use prediction_concat to concatenate output to input
        loss_func: loss function
        metrics: metrics to track
        opt_func: optimizer constructor
        lr: learning rate
        transforms: additional transforms (train + valid)
        augmentations: additional augmentations (train only)
        aux_losses: additional auxiliary losses
        input_norm: scaler class for input normalization, None to disable
        output_norm: scaler class for output denormalization, None to disable
    """
    if metrics is None:
        metrics = [fun_rmse]
    transforms = list(transforms) if transforms else []
    augmentations = list(augmentations) if augmentations else []
    aux_losses = list(aux_losses) if aux_losses else []

    _batch = dls.one_batch()
    inp = _batch[0].shape[-1]
    out = _batch[1].shape[-1]

    norm_u, norm_y = dls.norm_stats

    if attach_output:
        model = FranSys(inp, out, init_sz, **kwargs)

        # Add prediction_concat transform if not already present
        if not any(isinstance(t, prediction_concat) for t in transforms):
            transforms.insert(0, prediction_concat(t_offset=0))

        # Input will be [u, y] after prediction_concat
        combined_input_stats = norm_u + norm_y
    else:
        model = FranSys(inp - out, out, init_sz, **kwargs)

        # Input is [u, y] from prediction-mode dls
        combined_input_stats = norm_u + norm_y

    # Wrap model with input normalization and optional output denormalization
    if input_norm is not None:
        in_scaler = input_norm.from_stats(combined_input_stats)
        out_scaler = output_norm.from_stats(norm_y) if output_norm is not None else None
        model = ScaledModel(model, in_scaler, out_scaler)

    # For long sequences, add truncate_sequence augmentation
    seq_len = _batch[0].shape[1]
    LENGTH_THRESHOLD = 300
    if seq_len > init_sz + LENGTH_THRESHOLD:
        if not any(isinstance(a, truncate_sequence) for a in augmentations):
            INITIAL_SEQ_LEN = 100
            augmentations.append(truncate_sequence(init_sz + INITIAL_SEQ_LEN))

    return Learner(
        model,
        dls,
        loss_func=loss_func,
        metrics=metrics,
        n_skip=init_sz,
        opt_func=opt_func,
        lr=lr,
        transforms=transforms,
        augmentations=augmentations,
        aux_losses=aux_losses,
    )