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

dfm_python/__init__.py

@@ -63,7 +63,7 @@ Note: DFMConfig uses frequency dict to specify series (column names -> frequenci
 For detailed documentation, see the README.md file and the tutorial notebooks/scripts.
 """
 
-__version__ = "0.6.13"
+__version__ = "0.7.0"
 
 # ============================================================================
 # PUBLIC API DEFINITION
@@ -98,6 +98,7 @@ from .utils.metric import calculate_rmse
 # Users can import these directly from dfm_python
 from .dataset.dfm_dataset import DFMDataset
 from .dataset.ddfm_dataset import DDFMDataset
+from .dataset.afm_dataset import AFMDataset
 
 # Model implementations
 from .models.base import BaseFactorModel
@@ -109,6 +110,12 @@ try:
 except ImportError:
     DDFM = None
 
+# AFM high-level API (Attention Factor Model; PyTorch is mandatory)
+try:
+    from .models.afm.afm import AFM
+except ImportError:
+    AFM = None
+
 __all__ = [
     # Core classes
     'DFM',
@@ -121,14 +128,16 @@ __all__ = [
     'BaseResult', 'DFMResult', 'DDFMResult', 'calculate_rmse',
 ]
 
-# DDFM high-level API (PyTorch is mandatory)
+# DDFM / AFM high-level API (PyTorch is mandatory)
 __all__.extend([
     'DDFM',  # High-level API class
+    'AFM',   # Attention Factor Model
 ])
 
 # Dataset classes
 __all__.extend([
     'DFMDataset',  # DFM Dataset class
     'DDFMDataset',  # DDFM Dataset class
+    'AFMDataset',  # AFM Dataset class
 ])
 

dfm_python/config/__init__.py

@@ -8,9 +8,9 @@ Note: Series are specified via frequency dict mapping column names to frequencie
 """
 
 from .schema import (
-    BaseModelConfig, DFMConfig, DDFMConfig,
-    BaseResult, DFMResult, DDFMResult,
-    DFMStateSpaceParams, DDFMStateSpaceParams,
+    BaseModelConfig, DFMConfig, DDFMConfig, AFMConfig,
+    BaseResult, DFMResult, DDFMResult, AFMResult,
+    DFMStateSpaceParams, DDFMStateSpaceParams, AFMModelState,
 )
 # DEFAULT_BLOCK_NAME is imported lazily where needed to avoid circular imports
 from .constants import (
@@ -97,11 +97,11 @@ __all__ = [
     'BaseModelConfig', 'BaseResult',
     # 'DEFAULT_BLOCK_NAME',  # Removed to avoid circular import
     # Model-specific configs (from schema.py)
-    'DFMConfig', 'DDFMConfig',
+    'DFMConfig', 'DDFMConfig', 'AFMConfig',
     # State-space parameters
-    'DFMStateSpaceParams', 'DDFMStateSpaceParams',
+    'DFMStateSpaceParams', 'DDFMStateSpaceParams', 'AFMModelState',
     # Model-specific results (from results.py)
-    'DFMResult', 'DDFMResult',
+    'DFMResult', 'DDFMResult', 'AFMResult',
     # Utilities
     'validate_frequency',
     # IO

dfm_python/config/constants.py

@@ -348,18 +348,23 @@ DEFAULT_KDFM_MA_ORDER = 0  # Default MA order (MA lag order q) for KDFM (0 = pur
 # iVDFM Defaults
 # ============================================================================
 
-# iVDFM sequence and dimension defaults
-DEFAULT_IVDFM_SEQUENCE_LENGTH = 100  # Default sequence length for sliding windows
+# iVDFM dimension defaults
 DEFAULT_IVDFM_LATENT_DIM = 3  # Default number of factors
 DEFAULT_IVDFM_AUX_DIM = 1  # Default auxiliary variable dimension
 
+# iVDFM regime (K=1 = baseline iVDFM, no regime structure)
+DEFAULT_IVDFM_NUM_REGIMES = 1  # Number of discrete regimes; 1 = baseline iVDFM
+DEFAULT_IVDFM_REGIME_TEMPERATURE = 1.0  # Softmax temperature for π; τ < 1 sharpens (local commitment)
+DEFAULT_IVDFM_MIXING = False  # If True, apply global post-factor mixing f -> M f (identity init) before decode
+
 # 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_ENCODER_N_HIDDEN_LAYERS = 2  # Default encoder hidden layers
+DEFAULT_IVDFM_DECODER_TYPE = "mlp"  # Default: 'linear', 'residual', or 'mlp' (mlp = current behavior)
 DEFAULT_IVDFM_DECODER_HIDDEN_DIM = 200  # Default decoder hidden dimension
-DEFAULT_IVDFM_DECODER_N_LAYERS = 3  # Default decoder layers
+DEFAULT_IVDFM_DECODER_N_HIDDEN_LAYERS = 2  # Default decoder hidden layers
 DEFAULT_IVDFM_PRIOR_HIDDEN_DIM = 100  # Default prior network hidden dimension
-DEFAULT_IVDFM_PRIOR_N_LAYERS = 2  # Default prior network layers
+DEFAULT_IVDFM_PRIOR_N_HIDDEN_LAYERS = 1  # Default prior network hidden layers
 
 # iVDFM dynamics and distribution defaults
 DEFAULT_IVDFM_FACTOR_ORDER = 1  # Default AR order for factors
@@ -384,6 +389,7 @@ 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
+DEFAULT_IVDFM_EARLY_STOP_PATIENCE = 20  # Default patience for early stopping (epochs without improvement)
 
 # ============================================================================
 # Tutorial Defaults
@@ -451,11 +457,35 @@ FREQUENCY_HIERARCHY: Dict[str, int] = {
 # For frequency gaps larger than this, the missing data approach is used instead
 MAX_TENT_SIZE: int = 12
 
+# ============================================================================
+# AFM (Attention Factor Model) defaults
+# ============================================================================
+DEFAULT_AFM_NUM_FACTORS = 5          # K latent factors
+DEFAULT_AFM_FACTOR_MODEL = 'attention'  # 'attention' | 'pca'
+DEFAULT_AFM_TRADING = 'longconv'     # 'longconv' | 'ou'
+DEFAULT_AFM_EMBED_DIM = 32           # characteristic embedding dimension d
+DEFAULT_AFM_HIST_LEN = 20            # residual-history length s for the filter
+DEFAULT_AFM_N_KERNELS = 32           # number of long convolutions
+DEFAULT_AFM_RIDGE = 1e-2             # ridge penalty for closed-form loadings
+DEFAULT_AFM_SQUASH_LAMBDA = 1e-3     # LongConv Squash kernel regularizer strength
+DEFAULT_AFM_PCA_WINDOW = 252         # trailing window for the PCA benchmark
+DEFAULT_AFM_REESTIM = 21             # PCA re-estimation stride
+DEFAULT_AFM_LAMBDA_VAR = 0.1         # explained-variance regularizer weight
+DEFAULT_AFM_TURNOVER_COST = 5e-4     # 5bps per unit turnover (paper)
+DEFAULT_AFM_SHORT_COST = 1e-4        # 1bp shorting cost (paper)
+DEFAULT_AFM_MAX_EPOCHS = 50          # training epochs
+
 # ============================================================================
 # Export all constants
 # ============================================================================
 
 __all__ = [
+    # AFM
+    'DEFAULT_AFM_NUM_FACTORS', 'DEFAULT_AFM_FACTOR_MODEL', 'DEFAULT_AFM_TRADING',
+    'DEFAULT_AFM_EMBED_DIM', 'DEFAULT_AFM_HIST_LEN', 'DEFAULT_AFM_N_KERNELS',
+    'DEFAULT_AFM_RIDGE', 'DEFAULT_AFM_SQUASH_LAMBDA', 'DEFAULT_AFM_PCA_WINDOW',
+    'DEFAULT_AFM_REESTIM', 'DEFAULT_AFM_LAMBDA_VAR', 'DEFAULT_AFM_TURNOVER_COST',
+    'DEFAULT_AFM_SHORT_COST', 'DEFAULT_AFM_MAX_EPOCHS',
     # Convergence
     'DEFAULT_CONVERGENCE_THRESHOLD',
     'DEFAULT_EM_THRESHOLD',
@@ -606,15 +636,16 @@ __all__ = [
     '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_ENCODER_N_HIDDEN_LAYERS',
+    'DEFAULT_IVDFM_NUM_REGIMES',
+    'DEFAULT_IVDFM_DECODER_TYPE',
     'DEFAULT_IVDFM_DECODER_HIDDEN_DIM',
-    'DEFAULT_IVDFM_DECODER_N_LAYERS',
+    'DEFAULT_IVDFM_DECODER_N_HIDDEN_LAYERS',
     'DEFAULT_IVDFM_PRIOR_HIDDEN_DIM',
-    'DEFAULT_IVDFM_PRIOR_N_LAYERS',
+    'DEFAULT_IVDFM_PRIOR_N_HIDDEN_LAYERS',
     'DEFAULT_IVDFM_FACTOR_ORDER',
     'DEFAULT_IVDFM_INNOVATION_DIST',
     'DEFAULT_IVDFM_DECODER_VAR',

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, iVDFMConfig
-from .results import BaseResult, DFMResult, DDFMResult, iVDFMResult
-from .params import DFMStateSpaceParams, DDFMStateSpaceParams, iVDFMModelState
+from .model import BaseModelConfig, DFMConfig, DDFMConfig, iVDFMConfig, AFMConfig
+from .results import BaseResult, DFMResult, DDFMResult, iVDFMResult, AFMResult
+from .params import DFMStateSpaceParams, DDFMStateSpaceParams, iVDFMModelState, AFMModelState
 
 __all__ = [
-    'BaseModelConfig', 'DFMConfig', 'DDFMConfig', 'iVDFMConfig',
-    'BaseResult', 'DFMResult', 'DDFMResult', 'iVDFMResult',
-    'DFMStateSpaceParams', 'DDFMStateSpaceParams', 'iVDFMModelState',
+    'BaseModelConfig', 'DFMConfig', 'DDFMConfig', 'iVDFMConfig', 'AFMConfig',
+    'BaseResult', 'DFMResult', 'DDFMResult', 'iVDFMResult', 'AFMResult',
+    'DFMStateSpaceParams', 'DDFMStateSpaceParams', 'iVDFMModelState', 'AFMModelState',
 ]
 

dfm_python/config/schema/model.py

@@ -56,6 +56,7 @@ from ..constants import (
     DEFAULT_IVDFM_SCHEDULER_PATIENCE,
     DEFAULT_IVDFM_SCHEDULER_FACTOR,
     DEFAULT_IVDFM_SCHEDULER_MIN_LR,
+    DEFAULT_IVDFM_EARLY_STOP_PATIENCE,
     DEFAULT_REGULARIZATION_SCALE,
     DEFAULT_STRUCTURAL_REG_WEIGHT,
     DEFAULT_CONVERGENCE_THRESHOLD,
@@ -76,17 +77,34 @@ from ..constants import (
     DEFAULT_CLOCK_FREQUENCY,
     DEFAULT_KDFM_AR_ORDER,
     DEFAULT_KDFM_MA_ORDER,
+    DEFAULT_AFM_NUM_FACTORS,
+    DEFAULT_AFM_FACTOR_MODEL,
+    DEFAULT_AFM_TRADING,
+    DEFAULT_AFM_EMBED_DIM,
+    DEFAULT_AFM_HIST_LEN,
+    DEFAULT_AFM_N_KERNELS,
+    DEFAULT_AFM_RIDGE,
+    DEFAULT_AFM_SQUASH_LAMBDA,
+    DEFAULT_AFM_PCA_WINDOW,
+    DEFAULT_AFM_REESTIM,
+    DEFAULT_AFM_LAMBDA_VAR,
+    DEFAULT_AFM_TURNOVER_COST,
+    DEFAULT_AFM_SHORT_COST,
+    DEFAULT_AFM_MAX_EPOCHS,
     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_ENCODER_N_HIDDEN_LAYERS,
+    DEFAULT_IVDFM_NUM_REGIMES,
+    DEFAULT_IVDFM_REGIME_TEMPERATURE,
+    DEFAULT_IVDFM_MIXING,
+    DEFAULT_IVDFM_DECODER_TYPE,
     DEFAULT_IVDFM_DECODER_HIDDEN_DIM,
-    DEFAULT_IVDFM_DECODER_N_LAYERS,
+    DEFAULT_IVDFM_DECODER_N_HIDDEN_LAYERS,
     DEFAULT_IVDFM_PRIOR_HIDDEN_DIM,
-    DEFAULT_IVDFM_PRIOR_N_LAYERS,
+    DEFAULT_IVDFM_PRIOR_N_HIDDEN_LAYERS,
     DEFAULT_IVDFM_FACTOR_ORDER,
     DEFAULT_IVDFM_INNOVATION_DIST,
     DEFAULT_IVDFM_DECODER_VAR,
@@ -676,28 +694,35 @@ class iVDFMConfig(BaseModelConfig):
     # 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.
+    window: Optional[int] = None  # Sliding window length for training. If None, uses full T (full series as one sequence). Default: None.
+    stride: int = 1  # Step size for sliding windows. stride=1 means overlapping windows, stride=window means non-overlapping.
+    context: Optional[Union[List[str], List[int]]] = None  # Column names (DataFrame) or indices (array) for exogenous context variables. These are concatenated with time features.
+    time_context: int = 1  # Dimension of time-based context features. Always included. time_context=1: time step only, time_context>1: adds periodic sine features.
     
     # ========================================================================
     # 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
+    encoder_n_hidden_layers: int = DEFAULT_IVDFM_ENCODER_N_HIDDEN_LAYERS  # Number of encoder hidden layers
+    decoder_type: str = DEFAULT_IVDFM_DECODER_TYPE  # 'linear', 'residual', or 'mlp'
+    decoder_hidden_dim: Union[int, List[int]] = DEFAULT_IVDFM_DECODER_HIDDEN_DIM  # Decoder MLP (residual/mlp)
+    decoder_n_hidden_layers: int = DEFAULT_IVDFM_DECODER_N_HIDDEN_LAYERS  # Decoder hidden layers (residual/mlp)
     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
+    prior_n_hidden_layers: int = DEFAULT_IVDFM_PRIOR_N_HIDDEN_LAYERS  # Number of prior network hidden layers
     activation: str = DEFAULT_IVDFM_ACTIVATION  # Activation function
     slope: float = DEFAULT_IVDFM_SLOPE  # Leaky ReLU slope
     
     # ========================================================================
     # Dynamics and Distribution Parameters
     # ========================================================================
+    num_regimes: int = DEFAULT_IVDFM_NUM_REGIMES  # K; 1 = baseline iVDFM (one global dynamics, one decoder)
+    regime_temperature: float = DEFAULT_IVDFM_REGIME_TEMPERATURE  # Softmax τ for π; τ < 1 = sharper (local commitment)
     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
+    beta_kl: float = 1.0  # Weight for KL term in ELBO: ELBO = recon_loss + beta_kl * kl_loss. β<1 reduces KL pressure.
+    use_layer_norm: bool = False  # Whether to use layer normalization in encoder/decoder networks
+    mixing: bool = DEFAULT_IVDFM_MIXING  # If True, global post-factor mixing f -> M f before decode (identity init)
     
     # ========================================================================
     # Training Parameters
@@ -709,6 +734,7 @@ class iVDFMConfig(BaseModelConfig):
     batch_size: int = DEFAULT_IVDFM_BATCH_SIZE  # Batch size
     max_epochs: int = DEFAULT_IVDFM_MAX_EPOCHS  # Maximum epochs
     tolerance: float = DEFAULT_TOLERANCE  # Convergence tolerance
+    patience: Optional[int] = DEFAULT_IVDFM_EARLY_STOP_PATIENCE  # Early stopping patience (epochs without improvement). None to disable.
     seed: Optional[int] = None  # Random seed
     
     # ========================================================================
@@ -727,16 +753,22 @@ class iVDFMConfig(BaseModelConfig):
         
         from ...utils.errors import ConfigurationError
         
-        # Validate sequence_length
-        if self.sequence_length < 1:
+        # Validate window
+        if self.window is not None and self.window < 1:
             raise ConfigurationError(
-                f"sequence_length must be >= 1, got {self.sequence_length}"
+                f"window must be >= 1, got {self.window}"
             )
         
-        # Validate context_dim
-        if self.context_dim < 1:
+        # Validate time_context
+        if self.time_context < 1:
             raise ConfigurationError(
-                f"context_dim must be >= 1, got {self.context_dim}"
+                f"time_context must be >= 1, got {self.time_context}"
+            )
+        
+        # Validate patience if provided
+        if self.patience is not None and self.patience < 1:
+            raise ConfigurationError(
+                f"patience must be >= 1 or None (to disable), got {self.patience}"
             )
         
         # Validate num_factors if provided
@@ -758,6 +790,22 @@ class iVDFMConfig(BaseModelConfig):
                 f"innovation_distribution must be one of {valid_dists}, got '{self.innovation_distribution}'"
             )
         
+        # Validate num_regimes
+        if self.num_regimes < 1:
+            raise ConfigurationError(
+                f"num_regimes must be >= 1, got {self.num_regimes}"
+            )
+        # Validate regime_temperature
+        if self.regime_temperature <= 0.0:
+            raise ConfigurationError(
+                f"regime_temperature must be > 0, got {self.regime_temperature}"
+            )
+        # Validate decoder_type
+        if self.decoder_type not in ("linear", "residual", "mlp"):
+            raise ConfigurationError(
+                f"decoder_type must be 'linear', 'residual', or 'mlp', got '{self.decoder_type}'"
+            )
+        
         # Validate activation
         valid_activations = {'relu', 'lrelu', 'tanh', 'sigmoid'}
         if self.activation not in valid_activations:
@@ -765,6 +813,12 @@ class iVDFMConfig(BaseModelConfig):
                 f"activation must be one of {valid_activations}, got '{self.activation}'"
             )
         
+        # Validate beta_kl
+        if self.beta_kl < 0:
+            raise ConfigurationError(
+                f"beta_kl must be >= 0, got {self.beta_kl}"
+            )
+        
         # Validate optimizer
         from ..constants import VALID_OPTIMIZERS
         if self.optimizer not in VALID_OPTIMIZERS:
@@ -804,15 +858,18 @@ class iVDFMConfig(BaseModelConfig):
     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_ENCODER_N_HIDDEN_LAYERS,
+            DEFAULT_IVDFM_NUM_REGIMES,
+            DEFAULT_IVDFM_REGIME_TEMPERATURE,
+            DEFAULT_IVDFM_MIXING,
+            DEFAULT_IVDFM_DECODER_TYPE,
             DEFAULT_IVDFM_DECODER_HIDDEN_DIM,
-            DEFAULT_IVDFM_DECODER_N_LAYERS,
+            DEFAULT_IVDFM_DECODER_N_HIDDEN_LAYERS,
             DEFAULT_IVDFM_PRIOR_HIDDEN_DIM,
-            DEFAULT_IVDFM_PRIOR_N_LAYERS,
+    DEFAULT_IVDFM_PRIOR_N_HIDDEN_LAYERS,
             DEFAULT_IVDFM_FACTOR_ORDER,
             DEFAULT_IVDFM_INNOVATION_DIST,
             DEFAULT_IVDFM_DECODER_VAR,
@@ -833,20 +890,27 @@ class iVDFMConfig(BaseModelConfig):
         
         ivdfm_params = cls._extract_params(data, {
             'num_factors': None,
-            'sequence_length': DEFAULT_IVDFM_SEQUENCE_LENGTH,
+            'window': None,  # Default: None (uses full T - full series as one sequence)
+            'stride': 1,  # Default: stride=1 (overlapping windows)
             'context': None,
-            'context_dim': DEFAULT_IVDFM_AUX_DIM,
+            'time_context': 1,  # Default: time step only
             'encoder_hidden_dim': DEFAULT_IVDFM_ENCODER_HIDDEN_DIM,
-            'encoder_n_layers': DEFAULT_IVDFM_ENCODER_N_LAYERS,
+            'encoder_n_hidden_layers': DEFAULT_IVDFM_ENCODER_N_HIDDEN_LAYERS,
+            'num_regimes': DEFAULT_IVDFM_NUM_REGIMES,
+            'regime_temperature': DEFAULT_IVDFM_REGIME_TEMPERATURE,
+            'decoder_type': DEFAULT_IVDFM_DECODER_TYPE,
             'decoder_hidden_dim': DEFAULT_IVDFM_DECODER_HIDDEN_DIM,
-            'decoder_n_layers': DEFAULT_IVDFM_DECODER_N_LAYERS,
+            'decoder_n_hidden_layers': DEFAULT_IVDFM_DECODER_N_HIDDEN_LAYERS,
             'prior_hidden_dim': DEFAULT_IVDFM_PRIOR_HIDDEN_DIM,
-            'prior_n_layers': DEFAULT_IVDFM_PRIOR_N_LAYERS,
+            'prior_n_hidden_layers': DEFAULT_IVDFM_PRIOR_N_HIDDEN_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,
+            'beta_kl': 1.0,
+            'use_layer_norm': False,
+            'mixing': DEFAULT_IVDFM_MIXING,
             'learning_rate': DEFAULT_LEARNING_RATE,
             'optimizer': 'Adam',
             'optimizer_weight_decay': DEFAULT_IVDFM_OPTIMIZER_WEIGHT_DECAY,
@@ -854,6 +918,7 @@ class iVDFMConfig(BaseModelConfig):
             'batch_size': DEFAULT_IVDFM_BATCH_SIZE,
             'max_epochs': DEFAULT_IVDFM_MAX_EPOCHS,
             'tolerance': DEFAULT_TOLERANCE,
+            'patience': DEFAULT_IVDFM_EARLY_STOP_PATIENCE,
             'seed': None,
             'scheduler_type': DEFAULT_IVDFM_SCHEDULER_TYPE,
             'scheduler_step_size': DEFAULT_IVDFM_SCHEDULER_STEP_SIZE,
@@ -926,6 +991,79 @@ def validate_frequency(frequency: str) -> str:
             f"Common frequencies: 'd' (daily), 'w' (weekly), 'm' (monthly), "
             f"'q' (quarterly), 'sa' (semi-annual), 'a' (annual)."
         )
-    
+
     return frequency
 
+
+@dataclass
+class AFMConfig(BaseModelConfig):
+    """Attention Factor Model configuration (Epstein et al., 2025).
+
+    Cross-sectional conditional factor model + residual trading policy. Like
+    iVDFM, AFM does not use block structure and is constructed directly via
+    ``AFMConfig.from_dict(OmegaConf.to_container(cfg))`` rather than through the
+    block-requiring YAML adapter.
+    """
+    num_factors: Optional[int] = None      # K latent factors
+    factor_model: str = DEFAULT_AFM_FACTOR_MODEL  # 'attention' (learned) | 'pca'
+    trading: str = DEFAULT_AFM_TRADING     # 'longconv' (learned) | 'ou' (Avellaneda-Lee)
+    embed_dim: int = DEFAULT_AFM_EMBED_DIM  # characteristic embedding dim d
+    hist_len: int = DEFAULT_AFM_HIST_LEN   # residual-history length s for the filter
+    n_kernels: int = DEFAULT_AFM_N_KERNELS  # number of long convolutions
+    ridge: float = DEFAULT_AFM_RIDGE       # ridge penalty for closed-form loadings
+    squash_lambda: float = DEFAULT_AFM_SQUASH_LAMBDA  # LongConv Squash regularizer
+    pca_window: int = DEFAULT_AFM_PCA_WINDOW  # trailing window for the PCA benchmark
+    reestim: int = DEFAULT_AFM_REESTIM     # PCA re-estimation stride
+    lambda_var: float = DEFAULT_AFM_LAMBDA_VAR  # explained-variance regularizer weight
+    turnover_cost: float = DEFAULT_AFM_TURNOVER_COST  # 5bps per unit turnover (paper)
+    short_cost: float = DEFAULT_AFM_SHORT_COST  # 1bp shorting cost (paper)
+    risk_free: float = 0.0
+    learning_rate: float = DEFAULT_LEARNING_RATE
+    max_epochs: int = DEFAULT_AFM_MAX_EPOCHS
+    patience: Optional[int] = None
+    seed: Optional[int] = None
+
+    def __post_init__(self):
+        super().__post_init__()
+        from ...utils.errors import ConfigurationError
+        if self.num_factors is not None and self.num_factors < 1:
+            raise ConfigurationError(f"num_factors must be >= 1, got {self.num_factors}")
+        if self.factor_model not in ('attention', 'pca'):
+            raise ConfigurationError(
+                f"factor_model must be 'attention' or 'pca', got '{self.factor_model}'")
+        if self.trading not in ('longconv', 'ou'):
+            raise ConfigurationError(
+                f"trading must be 'longconv' or 'ou', got '{self.trading}'")
+        if self.hist_len < 1:
+            raise ConfigurationError(f"hist_len must be >= 1, got {self.hist_len}")
+        if self.embed_dim < 1:
+            raise ConfigurationError(f"embed_dim must be >= 1, got {self.embed_dim}")
+
+    @classmethod
+    def _extract_afm(cls, data: Dict[str, Any]) -> Dict[str, Any]:
+        """Extract AFM-specific parameters from a config dict."""
+        return cls._extract_params(data, {
+            'num_factors': None,
+            'factor_model': DEFAULT_AFM_FACTOR_MODEL,
+            'trading': DEFAULT_AFM_TRADING,
+            'embed_dim': DEFAULT_AFM_EMBED_DIM,
+            'hist_len': DEFAULT_AFM_HIST_LEN,
+            'n_kernels': DEFAULT_AFM_N_KERNELS,
+            'ridge': DEFAULT_AFM_RIDGE,
+            'squash_lambda': DEFAULT_AFM_SQUASH_LAMBDA,
+            'pca_window': DEFAULT_AFM_PCA_WINDOW,
+            'reestim': DEFAULT_AFM_REESTIM,
+            'lambda_var': DEFAULT_AFM_LAMBDA_VAR,
+            'turnover_cost': DEFAULT_AFM_TURNOVER_COST,
+            'short_cost': DEFAULT_AFM_SHORT_COST,
+            'risk_free': 0.0, 'learning_rate': DEFAULT_LEARNING_RATE,
+            'max_epochs': DEFAULT_AFM_MAX_EPOCHS, 'patience': None, 'seed': None,
+        })
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> 'AFMConfig':
+        """Create AFMConfig from a dictionary."""
+        base_params = cls._extract_base(data)
+        afm_params = cls._extract_afm(data)
+        return cls(**base_params, **afm_params)
+

dfm_python/config/schema/params.py

@@ -514,3 +514,38 @@ DDFMStateSpaceParams = DDFMModelState
 # Backward compatibility alias
 DFMStateSpaceParams = DFMModelState
 
+
+@dataclass
+class AFMModelState:
+    """Attention Factor Model training state for checkpointing."""
+    num_iter: int = 0
+    converged: bool = False
+    loss_now: Optional[float] = None
+    sharpe: Optional[float] = None
+    explained_variance: Optional[float] = None
+    residuals: Optional[np.ndarray] = None       # (B, T, N) traded residual portfolios
+    book_returns: Optional[np.ndarray] = None    # (B, n_eval) after-cost book returns
+
+    @classmethod
+    def from_model(cls, model: Any) -> 'AFMModelState':
+        """Create state from a trained AFM instance."""
+        return cls(
+            num_iter=getattr(model, '_num_iter', 0),
+            converged=getattr(model, '_converged', False),
+            loss_now=getattr(model, 'loss_now', None),
+            sharpe=getattr(model, '_sharpe', None),
+            explained_variance=getattr(model, '_ev', None),
+            residuals=getattr(model, '_residuals', None),
+            book_returns=getattr(model, '_book_returns', None),
+        )
+
+    def apply_to_model(self, model: Any) -> None:
+        """Restore training diagnostics onto a model instance."""
+        model._num_iter = self.num_iter
+        model._converged = self.converged
+        model.loss_now = self.loss_now
+        model._sharpe = self.sharpe
+        model._ev = self.explained_variance
+        model._residuals = self.residuals
+        model._book_returns = self.book_returns
+

dfm_python/config/schema/results.py

@@ -458,6 +458,13 @@ class iVDFMResult(BaseResult):
     reconstructions: Optional[np.ndarray] = None      # decoded factors, (T, N) or (batch, T, N)
     full_state: Optional[np.ndarray] = None           # Augmented state for companion form (T x r*p) for p>1, (T x r) for p=1
 
+    # Regime and mixing (when num_regimes > 1 or mixing=True)
+    num_regimes: Optional[int] = None                # K; 1 = baseline (no regime structure)
+    regime_temperature: Optional[float] = None        # Softmax τ for π (τ < 1 = sharper commitment)
+    regime_weights: Optional[np.ndarray] = None      # (T, K) or (num_windows, K) π_t from regime_net; for Z = [factors, regime] in stability
+    mixing: Optional[bool] = None                    # Whether post-factor mixing f -> M f was used
+    mixing_matrix: Optional[np.ndarray] = None       # (r, r) learned M when mixing=True (for inspection)
+
     # Training diagnostics
     training_elbo: Optional[float] = None
     training_loss: Optional[float] = None
@@ -498,6 +505,14 @@ class iVDFMResult(BaseResult):
             lines.append(f"Reconstructions shape: {getattr(self.reconstructions, 'shape', None)}")
         if self.full_state is not None:
             lines.append(f"Full state (augmented) shape: {getattr(self.full_state, 'shape', None)}")
+        if self.num_regimes is not None:
+            lines.append(f"Num regimes: {self.num_regimes}")
+        if self.regime_temperature is not None:
+            lines.append(f"Regime temperature: {self.regime_temperature}")
+        if self.mixing is not None:
+            lines.append(f"Mixing: {self.mixing}")
+        if self.mixing_matrix is not None:
+            lines.append(f"Mixing matrix shape: {getattr(self.mixing_matrix, 'shape', None)}")
         lines.append("")
         lines.append("Training:")
         lines.append(f"  Converged: {self.converged}")
@@ -513,3 +528,46 @@ class iVDFMResult(BaseResult):
 
 # FitParams moved to config.schema.params as DFMParams
 
+
+@dataclass
+class AFMResult(BaseResult):
+    """Attention Factor Model results (statistical-arbitrage residual trading).
+
+    AFM is not a linear state-space model; its primary artifacts are the traded
+    residual portfolios, the after-cost book returns, and the objective values.
+    ``Z``/``x_sm`` alias the residuals so the common accessors still work.
+    """
+    residual_portfolios: Optional[np.ndarray] = None   # (B, T, N) traded residuals
+    book_returns: Optional[np.ndarray] = None          # (B, n_eval) after-cost returns
+    sharpe: Optional[float] = None                     # after-cost Sharpe (objective)
+    explained_variance: Optional[float] = None         # EV regularizer value
+    factor_model: Optional[str] = None                 # 'attention' | 'pca'
+    trading: Optional[str] = None                      # 'longconv' | 'ou'
+    num_factors: Optional[int] = None
+
+    @property
+    def factors(self) -> Optional[np.ndarray]:
+        return self.Z
+
+    @property
+    def reconstructed(self) -> Optional[np.ndarray]:
+        return self.x_sm
+
+    @property
+    def x_hat(self) -> Optional[np.ndarray]:
+        return self.x_sm
+
+    def summary(self) -> str:
+        lines = ["=" * 80, "AFM Result Summary", "=" * 80, ""]
+        lines.append(f"Factor model: {self.factor_model}, trading: {self.trading}")
+        if self.num_factors is not None:
+            lines.append(f"Num factors: {self.num_factors}")
+        if self.sharpe is not None:
+            lines.append(f"After-cost Sharpe (objective): {self.sharpe:.4f}")
+        if self.explained_variance is not None:
+            lines.append(f"Explained variance: {self.explained_variance:.4f}")
+        if self.book_returns is not None:
+            lines.append(f"Book returns shape: {getattr(self.book_returns, 'shape', None)}")
+        lines += ["", "=" * 80]
+        return "\n".join(lines)
+

dfm_python/dataset/__init__.py

@@ -11,6 +11,7 @@ This module provides:
 from .ddfm_dataset import DDFMDataset
 from .dfm_dataset import DFMDataset
 from .ivdfm_dataset import iVDFMDataset
+from .afm_dataset import AFMDataset
 from .time import TimeIndex
 
 __all__ = [
@@ -18,6 +19,7 @@ __all__ = [
     'DDFMDataset',
     'DFMDataset',
     'iVDFMDataset',
+    'AFMDataset',
     # Time utilities
     'TimeIndex',
 ]