dfm-python 0.4.6__py3-none-any.whl → 0.4.7__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
dfm_python/__init__.py CHANGED
@@ -6,7 +6,6 @@ This package implements a comprehensive Dynamic Factor Model framework with supp
6
6
  - Tent kernel aggregation for low-to-high frequency mapping
7
7
  - Expectation-Maximization (EM) algorithm for parameter estimation
8
8
  - Kalman filtering and smoothing for factor extraction
9
- - News decomposition for nowcasting
10
9
  - Deep Dynamic Factor Models (DDFM) with nonlinear encoders (requires PyTorch)
11
10
 
12
11
  The package implements a clock-based approach to mixed-frequency DFMs, where all latent
@@ -23,7 +22,6 @@ Key Features:
23
22
  - Flexible block structure for factor modeling
24
23
  - Robust handling of missing data (internal spline interpolation)
25
24
  - Automatic standardization and data clipping
26
- - News decomposition for forecast updates
27
25
 
28
26
  Example (Standard Lightning Pattern):
29
27
  >>> from dfm_python import DFM, DFMDataModule, DFMTrainer
@@ -67,7 +65,7 @@ Note: DFMConfig and SeriesConfig are internal implementation details.
67
65
  For detailed documentation, see the README.md file and the tutorial notebooks/scripts.
68
66
  """
69
67
 
70
- __version__ = "0.4.6"
68
+ __version__ = "0.4.7"
71
69
 
72
70
  # ============================================================================
73
71
  # PUBLIC API DEFINITION
@@ -81,7 +79,7 @@ __version__ = "0.4.6"
81
79
  # 2. High-level API: DFM, DDFM, module-level convenience functions
82
80
  # 3. Core utilities: TimeIndex, diagnostics
83
81
  # 4. Models: BaseFactorModel, DDFM (low-level)
84
- # 5. Nowcasting: Nowcast, result classes, para_const
82
+ # 5. Data & Results: DFMResult
85
83
  # 6. Data & Results: DFMResult
86
84
  # ============================================================================
87
85
 
@@ -90,8 +88,9 @@ from .config import (
90
88
  DEFAULT_BLOCK_NAME,
91
89
  ConfigSource, YamlSource, DictSource, HydraSource,
92
90
  MergedConfigSource, make_config_source,
93
- DFMConfig, SeriesConfig, # Internal use - prefer Hydra YAML configs
94
91
  )
92
+ # Internal imports (for backward compatibility, but not recommended)
93
+ from .config import DFMConfig, SeriesConfig # Internal use only
95
94
 
96
95
  # Results
97
96
  from .config.results import DFMResult, DDFMResult, BaseResult
@@ -101,23 +100,15 @@ from .utils.diagnostics import diagnose_series, print_series_diagnosis
101
100
  from .utils.time import calculate_rmse
102
101
 
103
102
  # PyTorch Lightning modules (mandatory dependency)
103
+ # Users can import these directly from dfm_python
104
104
  from .lightning import (
105
105
  DFMDataModule,
106
- DFMDataset,
107
- KalmanFilter,
108
- EMAlgorithm,
106
+ DFMDataset, # Dataset class (usually not needed directly)
107
+ KalmanFilter, # Module class
108
+ EMAlgorithm, # Module class
109
109
  )
110
110
 
111
- # Nowcasting (from nowcast/ subpackage)
112
- from .nowcast import (
113
- Nowcast,
114
- NowcastResult,
115
- )
116
- from .nowcast.helpers import (
117
- para_const,
118
- NewsDecompResult,
119
- BacktestResult,
120
- )
111
+ # Nowcasting removed - use src.nowcast instead
121
112
 
122
113
  # Model implementations
123
114
  from .models.base import BaseFactorModel
@@ -128,11 +119,9 @@ from .models.ddfm import DDFM, DDFMModel
128
119
 
129
120
  __all__ = [
130
121
  # Core classes
131
- 'DFM', 'Nowcast',
122
+ 'DFM',
132
123
  # Model base and implementations
133
124
  'BaseFactorModel',
134
- # Nowcast result classes
135
- 'NowcastResult', 'NewsDecompResult', 'BacktestResult',
136
125
  # Constants
137
126
  'DEFAULT_BLOCK_NAME',
138
127
  # Config sources
@@ -140,7 +129,6 @@ __all__ = [
140
129
  'MergedConfigSource', 'make_config_source',
141
130
  # Low-level API (functional interface - advanced usage)
142
131
  'BaseResult', 'DFMResult', 'DDFMResult', 'calculate_rmse', 'diagnose_series', 'print_series_diagnosis',
143
- 'para_const',
144
132
  ]
145
133
 
146
134
  # DDFM high-level API (PyTorch is mandatory)
@@ -152,9 +140,9 @@ __all__.extend([
152
140
  # Lightning modules (mandatory dependency)
153
141
  __all__.extend([
154
142
  'DFMDataModule',
155
- 'DFMDataset',
156
- 'KalmanFilter',
157
- 'EMAlgorithm',
143
+ 'DFMDataset', # Dataset class (usually not needed directly)
144
+ 'KalmanFilter', # Module class
145
+ 'EMAlgorithm', # Module class
158
146
  ])
159
147
 
160
148
  # Trainer classes (mandatory dependency)
dfm_python/models/base.py CHANGED
@@ -91,7 +91,6 @@ class BaseFactorModel(pl.LightningModule):
91
91
  self._result: Optional[BaseResult] = None
92
92
  self.training_state: Optional[Any] = None
93
93
  self._data_module: Optional[Any] = None
94
- self._nowcast: Optional[Any] = None
95
94
 
96
95
  @property
97
96
  def config(self) -> DFMConfig:
@@ -494,7 +493,6 @@ class BaseFactorModel(pl.LightningModule):
494
493
  self._config = None
495
494
  self._data_module = None
496
495
  self._result = None
497
- self._nowcast = None
498
496
  if hasattr(self, 'training_state'):
499
497
  self.training_state = None
500
498
  return self
dfm_python/models/ddfm.py CHANGED
@@ -31,9 +31,8 @@ from ..config.results import DDFMResult
31
31
  from ..config.utils import get_periods_per_year
32
32
  from ..encoder.vae import Decoder, Encoder, extract_decoder_params
33
33
  from ..logger import get_logger
34
- from ..nowcast.dataview import DataView
35
- from ..nowcast.nowcast import Nowcast
36
- from ..utils.data import create_data_view, rem_nans_spline
34
+ # Nowcasting removed - use src.nowcast instead
35
+ from ..utils.data import rem_nans_spline
37
36
  from ..utils.helpers import (
38
37
  find_series_index,
39
38
  get_clock_frequency,
@@ -291,131 +290,8 @@ class DDFMModel:
291
290
  return X_forecast
292
291
  return Z_forecast
293
292
 
294
- def generate_dataset(
295
- self,
296
- target_series: str,
297
- periods: List[datetime],
298
- backward: int = 0,
299
- forward: int = 0,
300
- dataview: Optional['DataView'] = None
301
- ) -> Dict[str, Any]:
302
- """Generate dataset for DFM evaluation.
303
-
304
- Note: Requires data to be stored during fit(). Data is automatically
305
- stored from DataModule during fit().
306
- """
307
- if not hasattr(self, '_data') or self._data is None:
308
- error_msg = format_error_message(
309
- model_type="DDFM",
310
- operation="generate_dataset",
311
- reason="model has not been fitted with DataModule yet",
312
- guidance="Please call trainer.fit(model, data_module) first to store data"
313
- )
314
- raise ValueError(error_msg)
315
-
316
- i_series = find_series_index(self._config, target_series)
317
- X_features, y_baseline, y_actual, metadata, backward_results = [], [], [], [], []
318
-
319
- if dataview is not None:
320
- dataview_factory = dataview
321
- else:
322
- # Convert data to numpy if needed
323
- if hasattr(self._data, 'to_numpy'):
324
- X_data = self._data.to_numpy()
325
- else:
326
- X_data = np.asarray(self._data)
327
-
328
- dataview_factory = DataView.from_arrays(
329
- X=X_data, Time=self._time,
330
- Z=self._original_data, config=self._config,
331
- X_frame=self._data_frame
332
- )
333
- if dataview_factory.config is None:
334
- dataview_factory.config = self._config
335
-
336
- for period in periods:
337
- view_obj = dataview_factory.with_view_date(period)
338
- X_view, Time_view, _ = view_obj.materialize()
339
-
340
- if backward > 0:
341
- nowcasts, data_view_dates = [], []
342
- for weeks_back in range(backward, -1, -1):
343
- data_view_date = period - timedelta(weeks=weeks_back)
344
- view_past = dataview_factory.with_view_date(data_view_date)
345
- X_view_past, Time_view_past, _ = view_past.materialize()
346
- # Access nowcast through _nowcast_ref (set by high-level DDFM class)
347
- nowcast_obj = getattr(self, '_nowcast_ref', None)
348
- if nowcast_obj is None:
349
- error_msg = format_error_message(
350
- model_type="DDFM",
351
- operation="nowcast",
352
- reason="requires high-level DDFM instance",
353
- guidance="Please call nowcast() from DDFM class, not DDFMModel"
354
- )
355
- raise ValueError(error_msg)
356
- nowcast_val = nowcast_obj(
357
- target_series=target_series,
358
- view_date=view_past.view_date or data_view_date,
359
- target_period=period
360
- )
361
- nowcasts.append(nowcast_val)
362
- data_view_dates.append(view_past.view_date or data_view_date)
363
- baseline_nowcast = nowcasts[-1]
364
- backward_results.append({
365
- 'nowcasts': np.array(nowcasts),
366
- 'data_view_dates': data_view_dates,
367
- 'target_date': period
368
- })
369
- else:
370
- # Access nowcast through _nowcast_ref (set by high-level DDFM class)
371
- nowcast_obj = getattr(self, '_nowcast_ref', None)
372
- if nowcast_obj is None:
373
- error_msg = format_error_message(
374
- model_type="DDFM",
375
- operation="nowcast",
376
- reason="requires high-level DDFM instance",
377
- guidance="Please call nowcast() from DDFM class, not DDFMModel"
378
- )
379
- raise ValueError(error_msg)
380
- baseline_nowcast = nowcast_obj(
381
- target_series=target_series,
382
- view_date=view_obj.view_date or period,
383
- target_period=period
384
- )
385
-
386
- y_baseline.append(baseline_nowcast)
387
- t_idx = find_time_index(self._time, period)
388
- actual_val = np.nan
389
- # Convert data to numpy for indexing
390
- if hasattr(self._data, 'to_numpy'):
391
- data_array = self._data.to_numpy()
392
- else:
393
- data_array = np.asarray(self._data)
394
- if t_idx is not None and t_idx < data_array.shape[0] and i_series < data_array.shape[1]:
395
- actual_val = data_array[t_idx, i_series]
396
- y_actual.append(actual_val)
397
-
398
- # Extract features
399
- if self._result is not None and hasattr(self._result, 'Z'):
400
- latest_factors = self._result.Z[-1, :] if self._result.Z.shape[0] > 0 else np.zeros(self._result.Z.shape[1])
401
- else:
402
- latest_factors = np.array([])
403
- if X_view.shape[0] > 0:
404
- mean_residual = np.nanmean(X_view[-1, :]) if X_view.shape[0] > 0 else 0.0
405
- else:
406
- mean_residual = 0.0
407
- features = np.concatenate([latest_factors, [mean_residual]])
408
- X_features.append(features)
409
- metadata.append({'period': period, 'target_series': target_series})
410
-
411
- return {
412
- 'X': np.array(X_features),
413
- 'y_baseline': np.array(y_baseline),
414
- 'y_actual': np.array(y_actual),
415
- 'y_target': np.array(y_actual) - np.array(y_baseline),
416
- 'metadata': metadata,
417
- 'backward_results': backward_results if backward > 0 else []
418
- }
293
+ # generate_dataset() method removed - application-specific functionality
294
+ # Use src.nowcast for nowcasting-related dataset generation
419
295
 
420
296
  def get_state(
421
297
  self,
@@ -450,22 +326,12 @@ class DDFMModel:
450
326
  else:
451
327
  X_data = np.asarray(self._data)
452
328
 
453
- X_view, Time_view, _ = create_data_view(
454
- X=X_data, Time=self._time,
455
- Z=self._original_data, config=self._config, view_date=t
456
- )
329
+ # create_data_view removed - use src.nowcast for data view functionality
330
+ # For get_state, we use the raw data directly
331
+ X_view = X_data
457
332
 
458
- # Access nowcast through _nowcast_ref (set by high-level DDFM class)
459
- nowcast_obj = getattr(self, '_nowcast_ref', None)
460
- if nowcast_obj is None:
461
- error_msg = format_error_message(
462
- model_type="DDFM",
463
- operation="nowcast",
464
- reason="requires high-level DDFM instance",
465
- guidance="Please call nowcast() from DDFM class, not DDFMModel"
466
- )
467
- raise ValueError(error_msg)
468
- baseline_nowcast = nowcast_obj(target_series=target_series, view_date=t, target_period=None)
333
+ # Nowcast removed - use src.nowcast for nowcasting functionality
334
+ baseline_nowcast = np.nan # Placeholder - use src.nowcast for actual nowcast values
469
335
 
470
336
  baseline_forecast, actual_history, residuals, factors_history = [], [], [], []
471
337
  t_idx = find_time_index(self._time, t)
@@ -480,7 +346,11 @@ class DDFMModel:
480
346
 
481
347
  for i in range(max(0, t_idx - lookback + 1), t_idx + 1):
482
348
  if i < X_data.shape[0]:
483
- forecast_val = baseline_nowcast
349
+ # Use smoothed value from result if available, otherwise NaN
350
+ if self._result is not None and hasattr(self._result, 'X_sm') and i < self._result.X_sm.shape[0]:
351
+ forecast_val = self._result.X_sm[i, i_series] if i_series < self._result.X_sm.shape[1] else np.nan
352
+ else:
353
+ forecast_val = np.nan
484
354
  baseline_forecast.append(forecast_val)
485
355
  actual_val = X_data[i, i_series] if i_series < X_data.shape[1] else np.nan
486
356
  actual_history.append(actual_val)
@@ -505,7 +375,6 @@ class DDFMModel:
505
375
  # ============================================================================
506
376
 
507
377
  if TYPE_CHECKING:
508
- from ..nowcast import Nowcast
509
378
  from ..lightning import DFMDataModule
510
379
 
511
380
  class DDFM(BaseFactorModel):
@@ -874,9 +743,11 @@ class DDFM(BaseFactorModel):
874
743
  f"Possible causes: (1) learning rate too high, (2) gradient explosion, "
875
744
  f"(3) extreme input values, (4) unstable activation functions."
876
745
  )
877
- # Return a large but finite loss to signal the issue, but don't update weights
878
- # This prevents NaN from propagating to decoder weights
879
- loss = torch.tensor(1e6, device=data.device, dtype=data.dtype, requires_grad=False)
746
+ # Return a large but finite loss to signal the issue
747
+ # Must have requires_grad=True for Lightning backward() to work
748
+ # Create loss from model parameters to ensure it's in the computation graph
749
+ dummy_loss = sum(p.sum() * 0.0 for p in self.parameters() if p.requires_grad)
750
+ loss = dummy_loss + torch.tensor(1e6, device=data.device, dtype=data.dtype, requires_grad=True)
880
751
  self.log('train_loss', loss, on_step=True, on_epoch=True, prog_bar=True)
881
752
  return loss
882
753
 
@@ -2604,113 +2475,8 @@ class DDFM(BaseFactorModel):
2604
2475
  return self
2605
2476
 
2606
2477
 
2607
- def generate_dataset(
2608
- self,
2609
- target_series: str,
2610
- periods: List[datetime],
2611
- backward: int = 0,
2612
- forward: int = 0,
2613
- dataview: Optional['DataView'] = None
2614
- ) -> Dict[str, Any]:
2615
- """Generate dataset for DFM evaluation.
2616
-
2617
- Note: Requires data to be stored during fit(). Data is automatically
2618
- stored from DataModule during fit().
2619
- """
2620
- if not hasattr(self, '_data') or self._data is None:
2621
- error_msg = format_error_message(
2622
- model_type="DDFM",
2623
- operation="generate_dataset",
2624
- reason="model has not been fitted with DataModule yet",
2625
- guidance="Please call trainer.fit(model, data_module) first to store data"
2626
- )
2627
- raise ValueError(error_msg)
2628
-
2629
- i_series = find_series_index(self._config, target_series)
2630
- X_features, y_baseline, y_actual, metadata, backward_results = [], [], [], [], []
2631
-
2632
- if dataview is not None:
2633
- dataview_factory = dataview
2634
- else:
2635
- # Convert data to numpy if needed
2636
- if hasattr(self._data, 'to_numpy'):
2637
- X_data = self._data.to_numpy()
2638
- else:
2639
- X_data = np.asarray(self._data)
2640
-
2641
- dataview_factory = DataView.from_arrays(
2642
- X=X_data, Time=self._time,
2643
- Z=self._original_data, config=self._config,
2644
- X_frame=self._data_frame
2645
- )
2646
- if dataview_factory.config is None:
2647
- dataview_factory.config = self._config
2648
-
2649
- for period in periods:
2650
- view_obj = dataview_factory.with_view_date(period)
2651
- X_view, Time_view, _ = view_obj.materialize()
2652
-
2653
- if backward > 0:
2654
- nowcasts, data_view_dates = [], []
2655
- for weeks_back in range(backward, -1, -1):
2656
- data_view_date = period - timedelta(weeks=weeks_back)
2657
- view_past = dataview_factory.with_view_date(data_view_date)
2658
- X_view_past, Time_view_past, _ = view_past.materialize()
2659
- # Use nowcast property directly
2660
- nowcast_val = self.nowcast(
2661
- target_series=target_series,
2662
- view_date=view_past.view_date or data_view_date,
2663
- target_period=period
2664
- )
2665
- nowcasts.append(nowcast_val)
2666
- data_view_dates.append(view_past.view_date or data_view_date)
2667
- baseline_nowcast = nowcasts[-1]
2668
- backward_results.append({
2669
- 'nowcasts': np.array(nowcasts),
2670
- 'data_view_dates': data_view_dates,
2671
- 'target_date': period
2672
- })
2673
- else:
2674
- # Use nowcast property directly
2675
- baseline_nowcast = self.nowcast(
2676
- target_series=target_series,
2677
- view_date=view_obj.view_date or period,
2678
- target_period=period
2679
- )
2680
-
2681
- y_baseline.append(baseline_nowcast)
2682
- t_idx = find_time_index(self._time, period)
2683
- actual_val = np.nan
2684
- # Convert data to numpy for indexing
2685
- if hasattr(self._data, 'to_numpy'):
2686
- data_array = self._data.to_numpy()
2687
- else:
2688
- data_array = np.asarray(self._data)
2689
- if t_idx is not None and t_idx < data_array.shape[0] and i_series < data_array.shape[1]:
2690
- actual_val = data_array[t_idx, i_series]
2691
- y_actual.append(actual_val)
2692
-
2693
- # Extract features
2694
- if self._result is not None and hasattr(self._result, 'Z'):
2695
- latest_factors = self._result.Z[-1, :] if self._result.Z.shape[0] > 0 else np.zeros(self._result.Z.shape[1])
2696
- else:
2697
- latest_factors = np.array([])
2698
- if X_view.shape[0] > 0:
2699
- mean_residual = np.nanmean(X_view[-1, :]) if X_view.shape[0] > 0 else 0.0
2700
- else:
2701
- mean_residual = 0.0
2702
- features = np.concatenate([latest_factors, [mean_residual]])
2703
- X_features.append(features)
2704
- metadata.append({'period': period, 'target_series': target_series})
2705
-
2706
- return {
2707
- 'X': np.array(X_features),
2708
- 'y_baseline': np.array(y_baseline),
2709
- 'y_actual': np.array(y_actual),
2710
- 'y_target': np.array(y_actual) - np.array(y_baseline),
2711
- 'metadata': metadata,
2712
- 'backward_results': backward_results if backward > 0 else []
2713
- }
2478
+ # generate_dataset() method removed - application-specific functionality
2479
+ # Use src.nowcast for nowcasting-related dataset generation
2714
2480
 
2715
2481
  def get_state(
2716
2482
  self,
@@ -2720,7 +2486,7 @@ class DDFM(BaseFactorModel):
2720
2486
  ) -> Dict[str, Any]:
2721
2487
  """Get DFM state at time t.
2722
2488
 
2723
- Returns state information including baseline nowcast, actual history,
2489
+ Returns state information including baseline forecast, actual history,
2724
2490
  residuals, and factor history for the specified time period.
2725
2491
  """
2726
2492
  if not hasattr(self, '_data') or self._data is None:
@@ -2743,8 +2509,8 @@ class DDFM(BaseFactorModel):
2743
2509
  else:
2744
2510
  X_data = np.asarray(self._data)
2745
2511
 
2746
- # Use nowcast property directly
2747
- baseline_nowcast = self.nowcast(target_series=target_series, view_date=t, target_period=None)
2512
+ # Nowcast removed - use src.nowcast for nowcasting functionality
2513
+ baseline_nowcast = np.nan # Placeholder - use src.nowcast for actual nowcast values
2748
2514
 
2749
2515
  baseline_forecast, actual_history, residuals, factors_history = [], [], [], []
2750
2516
  t_idx = find_time_index(self._time, t)
@@ -2759,7 +2525,11 @@ class DDFM(BaseFactorModel):
2759
2525
 
2760
2526
  for i in range(max(0, t_idx - lookback + 1), t_idx + 1):
2761
2527
  if i < X_data.shape[0]:
2762
- forecast_val = baseline_nowcast
2528
+ # Use smoothed value from result if available, otherwise NaN
2529
+ if self._result is not None and hasattr(self._result, 'X_sm') and i < self._result.X_sm.shape[0]:
2530
+ forecast_val = self._result.X_sm[i, i_series] if i_series < self._result.X_sm.shape[1] else np.nan
2531
+ else:
2532
+ forecast_val = np.nan
2763
2533
  baseline_forecast.append(forecast_val)
2764
2534
  actual_val = X_data[i, i_series] if i_series < X_data.shape[1] else np.nan
2765
2535
  actual_history.append(actual_val)
dfm_python/models/dfm.py CHANGED
@@ -219,6 +219,10 @@ class DFM(BaseFactorModel):
219
219
  For DFM, each "step" is actually one EM iteration. The batch contains
220
220
  the full time series data.
221
221
 
222
+ **Important**: If `fit_em()` has already completed training (in `on_train_start()`),
223
+ this method skips additional EM iterations to prevent double training and log-likelihood
224
+ deterioration. The model will use the results from `fit_em()` instead.
225
+
222
226
  Parameters
223
227
  ----------
224
228
  batch : torch.Tensor or tuple
@@ -231,6 +235,19 @@ class DFM(BaseFactorModel):
231
235
  loss : torch.Tensor
232
236
  Negative log-likelihood (to minimize)
233
237
  """
238
+ # CRITICAL FIX: Skip additional EM iterations if fit_em() already completed training
239
+ # This prevents double training that causes log-likelihood deterioration.
240
+ # fit_em() runs in on_train_start() and completes the full EM algorithm.
241
+ # Lightning's training loop should not run additional EM iterations after that.
242
+ if self.training_state is not None:
243
+ # fit_em() has already run - use its results instead of doing more EM iterations
244
+ loglik = self.training_state.loglik
245
+ self.log('loglik', loglik, on_step=True, on_epoch=True, prog_bar=True)
246
+ self.log('em_iteration', float(self.training_state.num_iter), on_step=True, on_epoch=True)
247
+ # Return negative log-likelihood as loss (consistent with normal training_step)
248
+ device = next(self.parameters()).device if len(list(self.parameters())) > 0 else torch.device('cpu')
249
+ return -torch.tensor(loglik, device=device, dtype=torch.float32)
250
+
234
251
  # Handle both tuple and single tensor batches
235
252
  if isinstance(batch, tuple):
236
253
  data, _ = batch
@@ -306,7 +323,10 @@ class DFM(BaseFactorModel):
306
323
 
307
324
  if converged:
308
325
  self.training_state.converged = True
309
- _logger.info(f"EM algorithm converged at iteration {self.current_epoch}")
326
+ # Only log convergence once, not on every training_step call
327
+ if not hasattr(self, '_convergence_logged'):
328
+ _logger.info(f"EM algorithm converged at iteration {self.training_state.num_iter}")
329
+ self._convergence_logged = True
310
330
 
311
331
  def fit_em(
312
332
  self,
@@ -363,6 +383,8 @@ class DFM(BaseFactorModel):
363
383
 
364
384
  # Initialize state
365
385
  previous_loglik = float('-inf')
386
+ best_loglik = float('-inf')
387
+ best_params = None
366
388
  num_iter = 0
367
389
  converged = False
368
390
 
@@ -422,6 +444,41 @@ class DFM(BaseFactorModel):
422
444
  self.Z_0.data = Z_0_new
423
445
  self.V_0.data = V_0_new
424
446
 
447
+ # Track best log-likelihood and parameters (for early stopping on divergence)
448
+ if loglik > best_loglik:
449
+ best_loglik = loglik
450
+ # Store best parameters in case we need to revert
451
+ best_params = {
452
+ 'A': self.A.data.clone(),
453
+ 'C': self.C.data.clone(),
454
+ 'Q': self.Q.data.clone(),
455
+ 'R': self.R.data.clone(),
456
+ 'Z_0': self.Z_0.data.clone(),
457
+ 'V_0': self.V_0.data.clone()
458
+ }
459
+
460
+ # Check for significant log-likelihood deterioration (divergence detection)
461
+ # Stop early if log-likelihood worsens by more than 1000 from best (indicates divergence)
462
+ if num_iter > 10 and best_loglik != float('-inf') and loglik < best_loglik - 1000:
463
+ _logger.warning(
464
+ f"EM algorithm: Log-likelihood deteriorated significantly from best value. "
465
+ f"Best: {best_loglik:.4f}, Current: {loglik:.4f}, Deterioration: {best_loglik - loglik:.4f}. "
466
+ f"Stopping early at iteration {num_iter + 1} to prevent further divergence. "
467
+ f"Reverting to best parameters from iteration {num_iter}."
468
+ )
469
+ # Revert to best parameters
470
+ if best_params is not None:
471
+ with torch.no_grad():
472
+ self.A.data = best_params['A']
473
+ self.C.data = best_params['C']
474
+ self.Q.data = best_params['Q']
475
+ self.R.data = best_params['R']
476
+ self.Z_0.data = best_params['Z_0']
477
+ self.V_0.data = best_params['V_0']
478
+ # Use best log-likelihood for final state
479
+ loglik = best_loglik
480
+ break
481
+
425
482
  # Check convergence - use self.em.check_convergence() instead of em_converged()
426
483
  if num_iter > 2:
427
484
  converged, change = self.em.check_convergence(