Coverage for anfis_toolbox / optim / hybrid.py: 100%

1from __future__ import annotations 

2 

3import logging 

4from dataclasses import dataclass 

5 

6import numpy as np 

7 

8from ..losses import MSELoss 

9from ..model import TSKANFIS 

10from .base import BaseTrainer, ModelLike 

11 

12 

13@dataclass 

14class HybridTrainer(BaseTrainer): 

15 """Original Jang (1993) hybrid training: LSM for consequents + GD for antecedents. 

16 

17 Notes: 

18 This trainer assumes a single-output regression head. It is not compatible with 

19 :class:`~anfis_toolbox.model.TSKANFISClassifier` or the high-level 

20 :class:`~anfis_toolbox.classifier.ANFISClassifier` facade. 

21 """ 

22 

23 learning_rate: float = 0.01 

24 epochs: int = 100 

25 verbose: bool = False 

26 _loss_fn: MSELoss = MSELoss() 

27 

28 def init_state(self, model: ModelLike, X: np.ndarray, y: np.ndarray) -> None: 

29 """Hybrid trainer doesn't maintain optimizer state; returns None.""" 

30 self._require_regression_model(model) 

31 return None 

32 

33 def train_step(self, model: ModelLike, Xb: np.ndarray, yb: np.ndarray, state: None) -> tuple[float, None]: 

34 """Perform one hybrid step on a batch and return (loss, state). 

35 

36 Equivalent to one iteration of the hybrid algorithm on the given batch. 

37 """ 

38 model = self._require_regression_model(model) 

39 Xb, yb = self._prepare_training_data(model, Xb, yb) 

40 # Forward to get normalized weights 

41 normalized_weights = model.forward_antecedents(Xb) 

42 

43 # Build LSM system for batch 

44 ones_col = np.ones((Xb.shape[0], 1), dtype=float) 

45 x_bar = np.concatenate([Xb, ones_col], axis=1) 

46 A_blocks = [normalized_weights[:, j : j + 1] * x_bar for j in range(model.n_rules)] 

47 A = np.concatenate(A_blocks, axis=1) 

48 try: 

49 regularization = 1e-6 * np.eye(A.shape[1]) 

50 ATA_reg = A.T @ A + regularization 

51 theta = np.linalg.solve(ATA_reg, A.T @ yb.flatten()) 

52 except np.linalg.LinAlgError: 

53 logging.getLogger(__name__).warning("Matrix singular in LSM, using pseudo-inverse") 

54 theta = np.linalg.pinv(A) @ yb.flatten() 

55 model.consequent_layer.parameters = theta.reshape(model.n_rules, model.n_inputs + 1) 

56 

57 # Loss and backward for antecedents only 

58 y_pred = model.consequent_layer.forward(Xb, normalized_weights) 

59 loss = self._loss_fn.loss(yb, y_pred) 

60 dL_dy = self._loss_fn.gradient(yb, y_pred) 

61 dL_dnorm_w, _ = model.consequent_layer.backward(dL_dy) 

62 dL_dw = model.normalization_layer.backward(dL_dnorm_w) 

63 gradients = model.rule_layer.backward(dL_dw) 

64 model.membership_layer.backward(gradients) 

65 model.update_membership_parameters(self.learning_rate) 

66 return float(loss), state 

67 

68 def compute_loss(self, model: ModelLike, X: np.ndarray, y: np.ndarray) -> float: 

69 """Compute the hybrid MSE loss on prepared data without side effects.""" 

70 model = self._require_regression_model(model) 

71 X_arr, y_arr = self._prepare_validation_data(model, X, y) 

72 normalized_weights = model.forward_antecedents(X_arr) 

73 preds = model.consequent_layer.forward(X_arr, normalized_weights) 

74 return float(self._loss_fn.loss(y_arr, preds)) 

75 

76 @staticmethod 

77 def _require_regression_model(model: ModelLike) -> TSKANFIS: 

78 if not isinstance(model, TSKANFIS): 

79 raise TypeError("HybridTrainer supports TSKANFIS regression models only") 

80 return model