dfm-python 0.6.13__py3-none-any.whl → 0.6.14__py3-none-any.whl

dfm_python/config/constants.py

@@ -344,6 +344,47 @@ COMPUTATION_ERROR_TYPES = (RuntimeError, ValueError, TypeError, AttributeError,
 DEFAULT_KDFM_AR_ORDER = 1  # Default AR order (VAR lag order p) for KDFM
 DEFAULT_KDFM_MA_ORDER = 0  # Default MA order (MA lag order q) for KDFM (0 = pure VAR)
 
+# ============================================================================
+# iVDFM Defaults
+# ============================================================================
+
+# iVDFM sequence and dimension defaults
+DEFAULT_IVDFM_SEQUENCE_LENGTH = 100  # Default sequence length for sliding windows
+DEFAULT_IVDFM_LATENT_DIM = 3  # Default number of factors
+DEFAULT_IVDFM_AUX_DIM = 1  # Default auxiliary variable dimension
+
+# iVDFM network architecture defaults
+DEFAULT_IVDFM_ENCODER_HIDDEN_DIM = 200  # Default encoder hidden dimension
+DEFAULT_IVDFM_ENCODER_N_LAYERS = 3  # Default encoder layers
+DEFAULT_IVDFM_DECODER_HIDDEN_DIM = 200  # Default decoder hidden dimension
+DEFAULT_IVDFM_DECODER_N_LAYERS = 3  # Default decoder layers
+DEFAULT_IVDFM_PRIOR_HIDDEN_DIM = 100  # Default prior network hidden dimension
+DEFAULT_IVDFM_PRIOR_N_LAYERS = 2  # Default prior network layers
+
+# iVDFM dynamics and distribution defaults
+DEFAULT_IVDFM_FACTOR_ORDER = 1  # Default AR order for factors
+DEFAULT_IVDFM_INNOVATION_DIST = 'laplace'  # Default innovation distribution
+DEFAULT_IVDFM_DECODER_VAR = 0.01  # Default decoder variance
+DEFAULT_IVDFM_ACTIVATION = 'lrelu'  # Default activation function
+DEFAULT_IVDFM_SLOPE = 0.1  # Default leaky ReLU slope
+
+# iVDFM training defaults
+DEFAULT_IVDFM_BATCH_SIZE = 32  # Default batch size
+DEFAULT_IVDFM_MAX_EPOCHS = 100  # Default maximum epochs
+DEFAULT_IVDFM_AUX_VARIABLE_TYPE = 'time'  # Default auxiliary variable type
+
+# iVDFM optimizer defaults
+DEFAULT_IVDFM_OPTIMIZER_WEIGHT_DECAY = 0.0  # Default weight decay for optimizer
+DEFAULT_IVDFM_OPTIMIZER_MOMENTUM = 0.9  # Default momentum for SGD
+
+# iVDFM scheduler defaults
+DEFAULT_IVDFM_SCHEDULER_TYPE = 'step'  # Default scheduler type: 'step', 'plateau', 'cosine', 'exponential', None
+DEFAULT_IVDFM_SCHEDULER_STEP_SIZE = None  # Default step_size for StepLR (auto: max_epochs // 3)
+DEFAULT_IVDFM_SCHEDULER_GAMMA = 0.5  # Default gamma for StepLR/ExponentialLR
+DEFAULT_IVDFM_SCHEDULER_PATIENCE = 10  # Default patience for ReduceLROnPlateau
+DEFAULT_IVDFM_SCHEDULER_FACTOR = 0.1  # Default factor for ReduceLROnPlateau
+DEFAULT_IVDFM_SCHEDULER_MIN_LR = 0.0  # Default min_lr for ReduceLROnPlateau
+
 # ============================================================================
 # Tutorial Defaults
 # ============================================================================
@@ -564,6 +605,24 @@ __all__ = [
     # KDFM defaults
     'DEFAULT_KDFM_AR_ORDER',
     'DEFAULT_KDFM_MA_ORDER',
+    # iVDFM defaults
+    'DEFAULT_IVDFM_SEQUENCE_LENGTH',
+    'DEFAULT_IVDFM_LATENT_DIM',
+    'DEFAULT_IVDFM_AUX_DIM',
+    'DEFAULT_IVDFM_ENCODER_HIDDEN_DIM',
+    'DEFAULT_IVDFM_ENCODER_N_LAYERS',
+    'DEFAULT_IVDFM_DECODER_HIDDEN_DIM',
+    'DEFAULT_IVDFM_DECODER_N_LAYERS',
+    'DEFAULT_IVDFM_PRIOR_HIDDEN_DIM',
+    'DEFAULT_IVDFM_PRIOR_N_LAYERS',
+    'DEFAULT_IVDFM_FACTOR_ORDER',
+    'DEFAULT_IVDFM_INNOVATION_DIST',
+    'DEFAULT_IVDFM_DECODER_VAR',
+    'DEFAULT_IVDFM_ACTIVATION',
+    'DEFAULT_IVDFM_SLOPE',
+    'DEFAULT_IVDFM_BATCH_SIZE',
+    'DEFAULT_IVDFM_MAX_EPOCHS',
+    'DEFAULT_IVDFM_AUX_VARIABLE_TYPE',
     # Tutorial defaults
     'TUTORIAL_MAX_PERIODS',
     'TUTORIAL_MAX_EPOCHS',

dfm_python/config/schema/__init__.py

@@ -7,13 +7,13 @@ This package contains:
 Note: Series are specified via frequency dict mapping column names to frequencies.
 """
 
-from .model import BaseModelConfig, DFMConfig, DDFMConfig
-from .results import BaseResult, DFMResult, DDFMResult
-from .params import DFMStateSpaceParams, DDFMStateSpaceParams
+from .model import BaseModelConfig, DFMConfig, DDFMConfig, iVDFMConfig
+from .results import BaseResult, DFMResult, DDFMResult, iVDFMResult
+from .params import DFMStateSpaceParams, DDFMStateSpaceParams, iVDFMModelState
 
 __all__ = [
-    'BaseModelConfig', 'DFMConfig', 'DDFMConfig',
-    'BaseResult', 'DFMResult', 'DDFMResult',
-    'DFMStateSpaceParams', 'DDFMStateSpaceParams',
+    'BaseModelConfig', 'DFMConfig', 'DDFMConfig', 'iVDFMConfig',
+    'BaseResult', 'DFMResult', 'DDFMResult', 'iVDFMResult',
+    'DFMStateSpaceParams', 'DDFMStateSpaceParams', 'iVDFMModelState',
 ]
 

dfm_python/config/schema/model.py

@@ -48,6 +48,14 @@ from ..constants import (
     DEFAULT_BATCH_SIZE,
     DEFAULT_DDFM_WINDOW_SIZE,
     DEFAULT_GRAD_CLIP_VAL,
+    DEFAULT_IVDFM_OPTIMIZER_WEIGHT_DECAY,
+    DEFAULT_IVDFM_OPTIMIZER_MOMENTUM,
+    DEFAULT_IVDFM_SCHEDULER_TYPE,
+    DEFAULT_IVDFM_SCHEDULER_STEP_SIZE,
+    DEFAULT_IVDFM_SCHEDULER_GAMMA,
+    DEFAULT_IVDFM_SCHEDULER_PATIENCE,
+    DEFAULT_IVDFM_SCHEDULER_FACTOR,
+    DEFAULT_IVDFM_SCHEDULER_MIN_LR,
     DEFAULT_REGULARIZATION_SCALE,
     DEFAULT_STRUCTURAL_REG_WEIGHT,
     DEFAULT_CONVERGENCE_THRESHOLD,
@@ -70,6 +78,23 @@ from ..constants import (
     DEFAULT_KDFM_MA_ORDER,
     FREQUENCY_HIERARCHY,
     DEFAULT_HIERARCHY_VALUE,
+    DEFAULT_IVDFM_SEQUENCE_LENGTH,
+    DEFAULT_IVDFM_LATENT_DIM,
+    DEFAULT_IVDFM_AUX_DIM,
+    DEFAULT_IVDFM_ENCODER_HIDDEN_DIM,
+    DEFAULT_IVDFM_ENCODER_N_LAYERS,
+    DEFAULT_IVDFM_DECODER_HIDDEN_DIM,
+    DEFAULT_IVDFM_DECODER_N_LAYERS,
+    DEFAULT_IVDFM_PRIOR_HIDDEN_DIM,
+    DEFAULT_IVDFM_PRIOR_N_LAYERS,
+    DEFAULT_IVDFM_FACTOR_ORDER,
+    DEFAULT_IVDFM_INNOVATION_DIST,
+    DEFAULT_IVDFM_DECODER_VAR,
+    DEFAULT_IVDFM_ACTIVATION,
+    DEFAULT_IVDFM_SLOPE,
+    DEFAULT_IVDFM_BATCH_SIZE,
+    DEFAULT_IVDFM_MAX_EPOCHS,
+    DEFAULT_IVDFM_AUX_VARIABLE_TYPE,
 )
 
 
@@ -136,6 +161,7 @@ class BaseModelConfig:
     # ========================================================================
     clock: str = 'm'  # Required: Base frequency for latent factors (global clock): 'd', 'w', 'm', 'q', 'sa', 'a' (defaults to 'm' for monthly)
     target_scaler: Optional[ScalerType] = None  # Fitted sklearn scaler instance (StandardScaler, RobustScaler, etc.) for target series only. Must be a fitted scaler object (call .fit() on target data first). Pass scaler object directly, not string. Feature series are assumed to be manually preprocessed. If None, target series are assumed to be already in the desired scale.
+    scaler: Optional[str] = None  # Scaler type as string: 'standard', 'robust', 'minmax', 'maxabs', 'quantile', or null. Applied to targets only (not context/auxiliary variables). For iVDFM, this is the preferred way to specify scaler (like DDFM pattern).
     
     def __post_init__(self):
         """Validate basic model structure.
@@ -626,6 +652,244 @@ class DDFMConfig(BaseModelConfig):
         )
 
 
+@dataclass
+class iVDFMConfig(BaseModelConfig):
+    """Identifiable Variational Dynamic Factor Model configuration.
+    
+    Uses 'scaler' (string) pattern like DDFM, not 'target_scaler' (instance).
+    Scaler is applied to targets only, not to context/auxiliary variables.
+    """
+    """Identifiable Variational Dynamic Factor Model configuration.
+    
+    This configuration class extends BaseModelConfig with parameters specific
+    to iVDFM models trained using variational inference with identifiable
+    innovation priors.
+    
+    Note: iVDFM does NOT use block structure. Use num_factors directly to specify
+    the number of factors.
+    
+    The configuration can be built from:
+    - Main settings (training parameters) from config files
+    - Series definitions via frequency dict (column names -> frequencies)
+    """
+    # ========================================================================
+    # Model Structure
+    # ========================================================================
+    num_factors: Optional[int] = None  # Number of factors (inferred from config if None)
+    sequence_length: int = DEFAULT_IVDFM_SEQUENCE_LENGTH  # Sequence length for training
+    context: Optional[Union[List[str], List[int]]] = None  # Column names (DataFrame) or indices (array) for context variables. If None, context_dim is used to generate time-based context.
+    context_dim: int = DEFAULT_IVDFM_AUX_DIM  # Dimension of context. Used when context is None (generates time-based context) or when context is multivariate.
+    
+    # ========================================================================
+    # Neural Network Architecture
+    # ========================================================================
+    encoder_hidden_dim: Union[int, List[int]] = DEFAULT_IVDFM_ENCODER_HIDDEN_DIM  # Encoder architecture
+    encoder_n_layers: int = DEFAULT_IVDFM_ENCODER_N_LAYERS  # Number of encoder layers
+    decoder_hidden_dim: Union[int, List[int]] = DEFAULT_IVDFM_DECODER_HIDDEN_DIM  # Decoder architecture
+    decoder_n_layers: int = DEFAULT_IVDFM_DECODER_N_LAYERS  # Number of decoder layers
+    prior_hidden_dim: Union[int, List[int]] = DEFAULT_IVDFM_PRIOR_HIDDEN_DIM  # Prior network architecture
+    prior_n_layers: int = DEFAULT_IVDFM_PRIOR_N_LAYERS  # Number of prior network layers
+    activation: str = DEFAULT_IVDFM_ACTIVATION  # Activation function
+    slope: float = DEFAULT_IVDFM_SLOPE  # Leaky ReLU slope
+    
+    # ========================================================================
+    # Dynamics and Distribution Parameters
+    # ========================================================================
+    factor_order: int = DEFAULT_IVDFM_FACTOR_ORDER  # AR order for factors
+    innovation_distribution: str = DEFAULT_IVDFM_INNOVATION_DIST  # Innovation distribution type
+    decoder_var: float = DEFAULT_IVDFM_DECODER_VAR  # Decoder variance
+    
+    # ========================================================================
+    # Training Parameters
+    # ========================================================================
+    learning_rate: float = DEFAULT_LEARNING_RATE  # Learning rate
+    optimizer: str = 'Adam'  # Optimizer type
+    optimizer_weight_decay: float = DEFAULT_IVDFM_OPTIMIZER_WEIGHT_DECAY  # Weight decay (L2 regularization)
+    optimizer_momentum: float = DEFAULT_IVDFM_OPTIMIZER_MOMENTUM  # Momentum for SGD
+    batch_size: int = DEFAULT_IVDFM_BATCH_SIZE  # Batch size
+    max_epochs: int = DEFAULT_IVDFM_MAX_EPOCHS  # Maximum epochs
+    tolerance: float = DEFAULT_TOLERANCE  # Convergence tolerance
+    seed: Optional[int] = None  # Random seed
+    
+    # ========================================================================
+    # Scheduler Parameters
+    # ========================================================================
+    scheduler_type: Optional[str] = DEFAULT_IVDFM_SCHEDULER_TYPE  # Scheduler type: 'step', 'plateau', 'cosine', 'exponential', None
+    scheduler_step_size: Optional[int] = DEFAULT_IVDFM_SCHEDULER_STEP_SIZE  # Step size for StepLR (None = auto: max_epochs // 3)
+    scheduler_gamma: float = DEFAULT_IVDFM_SCHEDULER_GAMMA  # Gamma for StepLR/ExponentialLR
+    scheduler_patience: int = DEFAULT_IVDFM_SCHEDULER_PATIENCE  # Patience for ReduceLROnPlateau
+    scheduler_factor: float = DEFAULT_IVDFM_SCHEDULER_FACTOR  # Factor for ReduceLROnPlateau
+    scheduler_min_lr: float = DEFAULT_IVDFM_SCHEDULER_MIN_LR  # Min learning rate for ReduceLROnPlateau
+    
+    def __post_init__(self):
+        """Validate iVDFM configuration."""
+        super().__post_init__()  # Validate base config
+        
+        from ...utils.errors import ConfigurationError
+        
+        # Validate sequence_length
+        if self.sequence_length < 1:
+            raise ConfigurationError(
+                f"sequence_length must be >= 1, got {self.sequence_length}"
+            )
+        
+        # Validate context_dim
+        if self.context_dim < 1:
+            raise ConfigurationError(
+                f"context_dim must be >= 1, got {self.context_dim}"
+            )
+        
+        # Validate num_factors if provided
+        if self.num_factors is not None and self.num_factors < 1:
+            raise ConfigurationError(
+                f"num_factors must be >= 1, got {self.num_factors}"
+            )
+        
+        # Validate factor_order
+        if self.factor_order < 1:
+            raise ConfigurationError(
+                f"factor_order must be >= 1, got {self.factor_order}"
+            )
+        
+        # Validate innovation_distribution
+        valid_dists = {'laplace', 'gaussian', 'student_t', 'gamma', 'beta', 'exponential'}
+        if self.innovation_distribution not in valid_dists:
+            raise ConfigurationError(
+                f"innovation_distribution must be one of {valid_dists}, got '{self.innovation_distribution}'"
+            )
+        
+        # Validate activation
+        valid_activations = {'relu', 'lrelu', 'tanh', 'sigmoid'}
+        if self.activation not in valid_activations:
+            raise ConfigurationError(
+                f"activation must be one of {valid_activations}, got '{self.activation}'"
+            )
+        
+        # Validate optimizer
+        from ..constants import VALID_OPTIMIZERS
+        if self.optimizer not in VALID_OPTIMIZERS:
+            raise ConfigurationError(
+                f"optimizer must be one of {VALID_OPTIMIZERS}, got '{self.optimizer}'"
+            )
+        
+        # Validate scheduler_type
+        valid_schedulers = {'step', 'plateau', 'cosine', 'exponential', None}
+        if self.scheduler_type not in valid_schedulers:
+            raise ConfigurationError(
+                f"scheduler_type must be one of {valid_schedulers}, got '{self.scheduler_type}'"
+            )
+        
+        # Validate scheduler parameters
+        if self.scheduler_gamma <= 0 or self.scheduler_gamma > 1:
+            raise ConfigurationError(
+                f"scheduler_gamma must be in (0, 1], got {self.scheduler_gamma}"
+            )
+        
+        if self.scheduler_patience < 0:
+            raise ConfigurationError(
+                f"scheduler_patience must be >= 0, got {self.scheduler_patience}"
+            )
+        
+        if self.scheduler_factor <= 0 or self.scheduler_factor > 1:
+            raise ConfigurationError(
+                f"scheduler_factor must be in (0, 1], got {self.scheduler_factor}"
+            )
+        
+        if self.scheduler_min_lr < 0:
+            raise ConfigurationError(
+                f"scheduler_min_lr must be >= 0, got {self.scheduler_min_lr}"
+            )
+    
+    @classmethod
+    def _extract_ivdfm(cls, data: Dict[str, Any]) -> Dict[str, Any]:
+        """Extract iVDFM-specific parameters from config dict."""
+        from ..constants import (
+            DEFAULT_IVDFM_SEQUENCE_LENGTH,
+            DEFAULT_IVDFM_LATENT_DIM,
+            DEFAULT_IVDFM_AUX_DIM,
+            DEFAULT_IVDFM_ENCODER_HIDDEN_DIM,
+            DEFAULT_IVDFM_ENCODER_N_LAYERS,
+            DEFAULT_IVDFM_DECODER_HIDDEN_DIM,
+            DEFAULT_IVDFM_DECODER_N_LAYERS,
+            DEFAULT_IVDFM_PRIOR_HIDDEN_DIM,
+            DEFAULT_IVDFM_PRIOR_N_LAYERS,
+            DEFAULT_IVDFM_FACTOR_ORDER,
+            DEFAULT_IVDFM_INNOVATION_DIST,
+            DEFAULT_IVDFM_DECODER_VAR,
+            DEFAULT_IVDFM_ACTIVATION,
+            DEFAULT_IVDFM_SLOPE,
+            DEFAULT_IVDFM_BATCH_SIZE,
+            DEFAULT_IVDFM_MAX_EPOCHS,
+            DEFAULT_IVDFM_AUX_VARIABLE_TYPE,
+            DEFAULT_IVDFM_OPTIMIZER_WEIGHT_DECAY,
+            DEFAULT_IVDFM_OPTIMIZER_MOMENTUM,
+            DEFAULT_IVDFM_SCHEDULER_TYPE,
+            DEFAULT_IVDFM_SCHEDULER_STEP_SIZE,
+            DEFAULT_IVDFM_SCHEDULER_GAMMA,
+            DEFAULT_IVDFM_SCHEDULER_PATIENCE,
+            DEFAULT_IVDFM_SCHEDULER_FACTOR,
+            DEFAULT_IVDFM_SCHEDULER_MIN_LR,
+        )
+        
+        ivdfm_params = cls._extract_params(data, {
+            'num_factors': None,
+            'sequence_length': DEFAULT_IVDFM_SEQUENCE_LENGTH,
+            'context': None,
+            'context_dim': DEFAULT_IVDFM_AUX_DIM,
+            'encoder_hidden_dim': DEFAULT_IVDFM_ENCODER_HIDDEN_DIM,
+            'encoder_n_layers': DEFAULT_IVDFM_ENCODER_N_LAYERS,
+            'decoder_hidden_dim': DEFAULT_IVDFM_DECODER_HIDDEN_DIM,
+            'decoder_n_layers': DEFAULT_IVDFM_DECODER_N_LAYERS,
+            'prior_hidden_dim': DEFAULT_IVDFM_PRIOR_HIDDEN_DIM,
+            'prior_n_layers': DEFAULT_IVDFM_PRIOR_N_LAYERS,
+            'activation': DEFAULT_IVDFM_ACTIVATION,
+            'slope': DEFAULT_IVDFM_SLOPE,
+            'factor_order': DEFAULT_IVDFM_FACTOR_ORDER,
+            'innovation_distribution': DEFAULT_IVDFM_INNOVATION_DIST,
+            'decoder_var': DEFAULT_IVDFM_DECODER_VAR,
+            'learning_rate': DEFAULT_LEARNING_RATE,
+            'optimizer': 'Adam',
+            'optimizer_weight_decay': DEFAULT_IVDFM_OPTIMIZER_WEIGHT_DECAY,
+            'optimizer_momentum': DEFAULT_IVDFM_OPTIMIZER_MOMENTUM,
+            'batch_size': DEFAULT_IVDFM_BATCH_SIZE,
+            'max_epochs': DEFAULT_IVDFM_MAX_EPOCHS,
+            'tolerance': DEFAULT_TOLERANCE,
+            'seed': None,
+            'scheduler_type': DEFAULT_IVDFM_SCHEDULER_TYPE,
+            'scheduler_step_size': DEFAULT_IVDFM_SCHEDULER_STEP_SIZE,
+            'scheduler_gamma': DEFAULT_IVDFM_SCHEDULER_GAMMA,
+            'scheduler_patience': DEFAULT_IVDFM_SCHEDULER_PATIENCE,
+            'scheduler_factor': DEFAULT_IVDFM_SCHEDULER_FACTOR,
+            'scheduler_min_lr': DEFAULT_IVDFM_SCHEDULER_MIN_LR,
+            'scaler': None,  # Scaler string: 'standard', 'robust', 'minmax', 'maxabs', 'quantile', or null
+        })
+        
+        return ivdfm_params
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> 'iVDFMConfig':
+        """Create iVDFMConfig from dictionary.
+        
+        Parameters
+        ----------
+        data : Dict[str, Any]
+            Configuration dictionary
+            
+        Returns
+        -------
+        iVDFMConfig
+            iVDFM configuration instance
+        """
+        # Extract base params
+        base_params = cls._extract_base(data)
+        
+        # Extract iVDFM-specific params
+        ivdfm_params = cls._extract_ivdfm(data)
+        
+        # Combine and create config
+        return cls(**base_params, **ivdfm_params)
+
+
 # ============================================================================
 # Validation Functions
 # ============================================================================

dfm_python/config/schema/params.py

@@ -5,6 +5,7 @@ This module defines dataclasses for storing state-space parameters and model str
 
 from dataclasses import dataclass
 import numpy as np
+import torch
 from typing import Optional, Any
 
 
@@ -326,6 +327,188 @@ class DDFMModelState:
 
 # Backward compatibility aliases
 DDFMTrainingState = DDFMModelState
+
+
+@dataclass
+class iVDFMModelState:
+    """iVDFM model state including training state and SSM parameters.
+    
+    Consolidates all iVDFM model state for checkpointing: training state (convergence,
+    loss, factors, innovations) and fitted state-space model parameters.
+    
+    Attributes
+    ----------
+    num_iter : int
+        Current training iteration/epoch.
+    loss_now : float, optional
+        Current training loss value.
+    elbo : float, optional
+        Final ELBO value.
+    converged : bool
+        Whether training has converged.
+    factors : np.ndarray, optional
+        Extracted factors (T x r) or (batch x T x r).
+    innovations : np.ndarray, optional
+        Innovations (T x r) or (batch x T x r).
+    factor_order : int
+        AR order for factor dynamics (p in AR(p)).
+    A : np.ndarray, optional
+        Transition matrix (r x r) for p=1, or AR coefficients (r x p) for p>1.
+    B : np.ndarray, optional
+        Innovation loading matrix (r x r).
+    f0 : np.ndarray, optional
+        Initial state (r,).
+    """
+    # Training state
+    num_iter: int = 0
+    loss_now: Optional[float] = None
+    elbo: Optional[float] = None
+    converged: bool = False
+    factors: Optional[np.ndarray] = None
+    innovations: Optional[np.ndarray] = None
+    full_state: Optional[np.ndarray] = None  # Augmented state for companion form (T x r*p) for p>1, (T x r) for p=1
+    
+    # SSM parameters
+    factor_order: int = 1
+    A: Optional[np.ndarray] = None  # Transition matrix or AR coefficients
+    B: Optional[np.ndarray] = None  # Innovation loading matrix
+    f0: Optional[np.ndarray] = None  # Initial state
+    
+    @classmethod
+    def from_model(cls, model: Any) -> 'iVDFMModelState':
+        """Create iVDFMModelState from iVDFM model instance.
+        
+        Parameters
+        ----------
+        model : iVDFM
+            iVDFM model instance
+            
+        Returns
+        -------
+        iVDFMModelState
+            Model state dataclass
+        """
+        # Get training state from model attributes
+        factors = getattr(model, 'factors', None)
+        innovations = getattr(model, 'innovations', None)
+        num_iter = getattr(model, '_num_iter', 0)
+        loss_now = getattr(model, 'loss_now', None)
+        elbo = getattr(model, '_elbo', None)
+        converged = getattr(model, '_converged', False)
+        
+        # Get SSM parameters
+        factor_order = getattr(model, 'factor_order', 1)
+        A = B = f0 = None
+        
+        # Extract from SSM if available
+        if hasattr(model, 'ssm') and model.ssm is not None:
+            ssm = model.ssm
+            if factor_order == 1:
+                # First-order: A is diagonal, stored as (r,)
+                if hasattr(ssm, 'A'):
+                    A = ssm.A.data.cpu().numpy() if hasattr(ssm.A, 'data') else ssm.A
+            else:
+                # Higher-order: AR coefficients stored as (r, p)
+                if hasattr(ssm, 'ar_coeffs'):
+                    A = ssm.ar_coeffs.data.cpu().numpy() if hasattr(ssm.ar_coeffs, 'data') else ssm.ar_coeffs
+            
+            if hasattr(ssm, 'B'):
+                B = ssm.B.data.cpu().numpy() if hasattr(ssm.B, 'data') else ssm.B
+            
+            if hasattr(ssm, 'f0'):
+                f0 = ssm.f0.data.cpu().numpy() if hasattr(ssm.f0, 'data') else ssm.f0
+        
+        # Fallback: get from model attributes directly
+        if A is None:
+            A = getattr(model, 'A', None)
+            if A is not None and hasattr(A, 'data'):
+                A = A.data.cpu().numpy()
+        if B is None:
+            B = getattr(model, 'B', None)
+            if B is not None and hasattr(B, 'data'):
+                B = B.data.cpu().numpy()
+        if f0 is None:
+            f0 = getattr(model, 'f0', None)
+            if f0 is not None and hasattr(f0, 'data'):
+                f0 = f0.data.cpu().numpy()
+        
+        # Compute full_state (augmented state for companion form)
+        full_state = None
+        if factors is not None:
+            factors_np = factors
+            if isinstance(factors_np, np.ndarray):
+                # Average over batch if needed
+                if factors_np.ndim == 3:
+                    factors_np = np.mean(factors_np, axis=0)  # (T, r)
+                
+                T, r = factors_np.shape
+                if factor_order == 1:
+                    # For p=1, full_state is just factors
+                    full_state = factors_np  # (T, r)
+                else:
+                    # For p>1, construct augmented state (T, r*p)
+                    # s_t[i*p : (i+1)*p] = [f_t[i], f_{t-1}[i], ..., f_{t-p+1}[i]]
+                    full_state = np.zeros((T, r * factor_order), dtype=factors_np.dtype)
+                    for t in range(T):
+                        for i in range(r):
+                            for lag in range(factor_order):
+                                idx = t - lag
+                                if idx >= 0:
+                                    full_state[t, i * factor_order + lag] = factors_np[idx, i]
+                                else:
+                                    # Use f0 for negative indices
+                                    if f0 is not None and i < len(f0):
+                                        full_state[t, i * factor_order + lag] = f0[i]
+        
+        return cls(
+            num_iter=num_iter,
+            loss_now=loss_now,
+            elbo=elbo,
+            converged=converged,
+            factors=factors,
+            innovations=innovations,
+            full_state=full_state,
+            factor_order=factor_order,
+            A=A,
+            B=B,
+            f0=f0
+        )
+    
+    def apply_to_model(self, model: Any) -> None:
+        """Apply this state to iVDFM model instance.
+        
+        Parameters
+        ----------
+        model : iVDFM
+            iVDFM model instance to update
+        """
+        # Apply training state
+        if self.factors is not None:
+            model.factors = self.factors
+        if self.innovations is not None:
+            model.innovations = self.innovations
+        model._num_iter = self.num_iter
+        model.loss_now = self.loss_now
+        if self.elbo is not None:
+            model._elbo = self.elbo
+        model._converged = self.converged
+        
+        # Apply SSM parameters if available
+        if self.A is not None and self.B is not None:
+            if hasattr(model, 'ssm') and model.ssm is not None:
+                ssm = model.ssm
+                if self.factor_order == 1:
+                    if hasattr(ssm, 'A'):
+                        ssm.A.data = torch.from_numpy(self.A).to(ssm.A.device)
+                else:
+                    if hasattr(ssm, 'ar_coeffs'):
+                        ssm.ar_coeffs.data = torch.from_numpy(self.A).to(ssm.ar_coeffs.device)
+                
+                if hasattr(ssm, 'B'):
+                    ssm.B.data = torch.from_numpy(self.B).to(ssm.B.device)
+                
+                if self.f0 is not None and hasattr(ssm, 'f0'):
+                    ssm.f0.data = torch.from_numpy(self.f0).to(ssm.f0.device)
 DDFMStateSpaceParams = DDFMModelState
 
 # Backward compatibility alias