Base

def fit(
    self,
    x: Tensor,
    y: Tensor,
    epochs: int = 200,
    learning_rate: float = 1e-3,
    criterion: Callable[[Tensor, Tensor], Tensor] | None = None,
    optimizer: torch.optim.Optimizer | None = None,
    batch_size: int | None = None,
    shuffle: bool = True,
    ur_weight: float = 0.0,
    ur_target: float | None = None,
    verbose: bool | int = False,
    x_val: Tensor | None = None,
    y_val: Tensor | None = None,
    patience: int | None = 20,
    restore_best: bool = True,
    weight_decay: float = 1e-8,
) -> dict[str, Any]:
    """Train the model with optional early stopping.

    When *x_val* and *y_val* are provided the model evaluates a
    task-specific metric (via :meth:`_evaluate_validation`) after every
    epoch and applies early stopping when the metric has not improved for
    *patience* consecutive epochs.
    By default the best model weights from validation are restored when
    ``restore_best=True``.

    Args:
        x: Training features of shape ``(N, n_inputs)``.
        y: Training targets of shape ``(N,)``.
        epochs: Maximum number of training epochs.
        learning_rate: Learning rate for the default AdamW optimizer.
        criterion: Optional loss function.  Defaults to
            :meth:`_default_criterion`.
        optimizer: Optional pre-built optimizer.  When ``None``, AdamW
            is constructed with separate parameter groups for antecedent
            (no weight decay) and consequent (*weight_decay*) layers.
        batch_size: Mini-batch size.  ``None`` uses the full dataset.
        shuffle: If ``True``, reshuffle sample indices each epoch.
        ur_weight: Non-negative weight for the uniform rule
            regularization term.  ``0.0`` disables it.
        ur_target: Target uniform activation for UR.  Must be in
            ``(0, 1]`` when provided.  ``None`` defaults to
            ``1 / n_rules``.
        verbose: Verbosity level. ``0`` = quiet, ``1`` = progress bar,
            ``2`` = per-epoch summary logging, ``3`` = per-epoch detailed
            logging. ``True`` is accepted as an alias for ``2``.
        x_val: Optional validation features of shape
            ``(M, n_inputs)``.
        y_val: Optional validation targets of shape ``(M,)``.
        patience: Number of consecutive epochs without improvement
            before early stopping.  Set to ``None`` to disable early
            stopping.  Only active when *x_val* and *y_val* are given.
        restore_best: If ``True`` (default), restore the model weights
            from the best validation epoch when early stopping is used.
        weight_decay: L2 weight decay applied to consequent parameters
            by the default AdamW optimizer.

    Returns:
        A dictionary with keys ``"train"``, ``"ur"``, and ``"val"``
        containing per-epoch loss lists.

    Raises:
        ValueError: If shapes of *x*, *y*, *x_val*, or *y_val* are
            incompatible, or if *ur_weight* < 0 or *ur_target* is
            outside ``(0, 1]``.
    """
    if x.ndim != 2 or x.shape[1] != self.n_inputs:
        raise ValueError(f"expected x shape (batch, {self.n_inputs}), got {tuple(x.shape)}")
    if y.ndim != 1:
        raise ValueError("expected y shape (batch,)")
    if ur_weight < 0.0:
        raise ValueError("ur_weight must be >= 0")
    if ur_target is not None and not (0.0 < ur_target <= 1.0):
        raise ValueError("ur_target must be in (0, 1] when provided")

    has_val = x_val is not None and y_val is not None
    if has_val:
        if x_val is None or y_val is None:  # pragma: no cover
            raise ValueError("x_val and y_val must both be provided")
        if x_val.ndim != 2 or x_val.shape[1] != self.n_inputs:
            raise ValueError(f"expected x_val shape (batch, {self.n_inputs}), got {tuple(x_val.shape)}")
        if y_val.ndim != 1:
            raise ValueError("expected y_val shape (batch,)")

    train_criterion = criterion or self._default_criterion()
    if optimizer is not None:
        train_optimizer = optimizer
    else:
        ante_params = list(self.membership_layer.parameters())
        rule_params = list(self.rule_layer.parameters())
        cons_params = list(self.consequent_layer.parameters())
        if self.consequent_bn is not None:
            cons_params.extend(self.consequent_bn.parameters())
        train_optimizer = torch.optim.AdamW(
            [
                {"params": ante_params, "weight_decay": 0.0},
                {"params": rule_params, "weight_decay": 0.0},
                {"params": cons_params, "weight_decay": weight_decay},
            ],
            lr=learning_rate,
        )

    history: dict[str, Any] = {"train": [], "ur": [], "val": []}
    best_metric = float("-inf")
    epochs_no_improve = 0
    best_state: dict[str, Any] | None = None
    verbose_level = self._resolve_verbose(verbose)

    self.train()
    pbar = None
    if verbose_level == 1:
        pbar = trange(epochs, desc="Training", leave=False)
        epoch_iterator = pbar
    else:
        epoch_iterator = range(epochs)

    for epoch in epoch_iterator:
        batch_losses: list[float] = []
        batch_ur_losses: list[float] = []
        for batch_idx in _iter_minibatch_indices(x.shape[0], batch_size=batch_size, shuffle=shuffle):
            x_b = x.index_select(0, batch_idx.to(device=x.device))
            y_b = y.index_select(0, batch_idx.to(device=y.device))

            train_optimizer.zero_grad(set_to_none=True)
            output, norm_w = self._forward_train(x_b)
            main_loss = self._compute_loss(train_criterion, output, y_b)

            ur_loss = _uniform_regularization_loss(norm_w, target=ur_target)
            loss = main_loss + (float(ur_weight) * ur_loss)
            loss.backward()
            train_optimizer.step()

            batch_losses.append(float(loss.detach().item()))
            batch_ur_losses.append(float(ur_loss.detach().item()))

        epoch_train_loss = float(sum(batch_losses) / max(len(batch_losses), 1))
        history["train"].append(epoch_train_loss)
        history["ur"].append(float(sum(batch_ur_losses) / max(len(batch_ur_losses), 1)))

        # --- validation & early stopping ---
        if has_val and x_val is not None and y_val is not None:
            self.eval()
            val_info = self._evaluate_validation(train_criterion, x_val, y_val)
            history["val"].append(val_info.get("val_loss", 0.0))
            # Store any extra keys (e.g. val_acc) in history
            for k, v in val_info.items():
                if k not in ("val_loss", "metric"):
                    history.setdefault(k, []).append(v)
            self.train()

            metric = val_info["metric"]
            if metric > best_metric:
                best_metric = metric
                epochs_no_improve = 0
                best_state = copy.deepcopy(self.state_dict())
            else:
                epochs_no_improve += 1

            if verbose_level == 1:
                if pbar is None:  # pragma: no cover
                    raise RuntimeError("progress bar unavailable for verbose level 1")
                postfix = [
                    f"train={epoch_train_loss:.4f}",
                    f"val={val_info.get('val_loss', 0.0):.4f}",
                ]
                pbar.set_postfix_str(" ".join(postfix))
            if verbose_level >= 2 and (
                verbose_level == 3 or ((epoch + 1) % max(epochs // 10, 1) == 0 or epoch == 0)
            ):
                log_parts = [
                    f"epoch={epoch + 1}/{epochs}",
                    f"train_loss={epoch_train_loss:.6f}",
                ]
                for k, v in val_info.items():
                    if k != "metric":
                        log_parts.append(f"{k}={v:.6f}" if isinstance(v, float) else f"{k}={v}")
                self._log(" ".join(log_parts), verbose=verbose_level)

            if patience is not None and epochs_no_improve >= patience:
                if verbose_level >= 2:
                    self._log(
                        "early stopping at epoch %s (patience=%s)",
                        epoch + 1,
                        patience,
                        verbose=verbose_level,
                    )
                break
        else:
            if verbose_level == 1:
                if pbar is None:  # pragma: no cover
                    raise RuntimeError("progress bar unavailable for verbose level 1")
                pbar.set_postfix_str(f"loss={epoch_train_loss:.4f}")
            if verbose_level >= 2 and (
                verbose_level == 3 or ((epoch + 1) % max(epochs // 10, 1) == 0 or epoch == 0)
            ):
                self._log(
                    "epoch=%s/%s loss=%.6f",
                    epoch + 1,
                    epochs,
                    epoch_train_loss,
                    verbose=verbose_level,
                )

    if pbar is not None:
        pbar.close()

    if restore_best and best_state is not None:
        self.load_state_dict(best_state)

    history["stopped_epoch"] = epoch + 1  # type: ignore[possibly-undefined]

    return history

def forward(self, x: Tensor) -> Tensor:
    """Full forward pass through the TSK pipeline."""
    output, _ = self._forward_train(x)
    return output

def forward_antecedents(self, x: Tensor) -> Tensor:
    """Compute normalized rule strengths from model antecedents."""
    mu = self.membership_layer(x)
    w = self.rule_layer(mu)
    return cast(Tensor, self.defuzzifier(w))