Model

def apply_updates(
    self,
    updates: Dict[str, np.ndarray],
) -> None:
    """Apply incoming updates to the wrapped model's parameters.

    Args:
        updates: Model parameters' updates to add (apply) to existing
            parameters' values.
    """
    self._assert_dict_inputs(updates)
    for key, val in updates.items():
        weights = getattr(self.model, key)
        setattr(self.model, key, weights + val)

def disable_internal_optimizer(self) -> None:
    """Disable the scikit-learn internal optimizer.

    Calling this method alters the wrapped model so that raw gradients are
    computed and attached to it (rather than relying on scikit-learn to
    apply a learning rate that may be scheduled to vary along time).

    ''' warning "Call it only if using an external optimizer"
    """
    # Record initial params, then override optimizer ones.
    self._optim_params = self.get_params()
    self.set_params(**self._null_optim_params)
    # Warn about overridden values.
    changed_params: List[str] = []
    for key, val in self._null_optim_params.items():
        param = self._optim_params.get(key)
        if param is not None and param != val:
            changed_params.append(key)
    if changed_params:
        changed = ",\n\t".join(changed_params)
        logger.warning(
            "The following non-default model parameters were overridden "
            f"due to the disabling of the scikit-learn internal optimizer:\n\t{changed}",
            broadcast=True
        )

def enable_internal_optimizer(self) -> None:
    """Enable the scikit-learn internal optimizer.

    Calling this method restores any model parameter previously overridden
    due to calling the counterpart `disable_internal_optimizer` method.
    """
    if self._optim_params:
        self.set_params(**self._optim_params)
        logger.debug("Internal Optimizer restored")

def export(self, filename: str) -> None:
    """Export the wrapped model to a dump file.

    Args:
        filename: path to the file where the model will be saved.

    !!! info "Notes":
        This method is designed to save the model to a local dump
        file for easy re-use by the same user, possibly outside of
        Fed-BioMed. It is not designed to produce trustworthy data
        dumps and is not used to exchange models and their weights
        as part of the federated learning process.

    !!! warning "Warning":
        This method uses `joblib.dump`, which relies on pickle and
        is therefore hard to trust by third-party loading methods.
    """
    with open(filename, "wb") as file:
        joblib.dump(self.model, file)

def flatten(self,
            only_trainable: bool = False,
            exclude_buffers: bool = True) -> List[float]:
    """Gets weights as flatten vector

    Args:
        only_trainable: Unused for scikit-learn models. (Whether to ignore
            non-trainable model parameters.)
        exclude_buffers: Unused for scikit-learn models. (Whether to ignore
            buffers.)

    Returns:
        to_list: Convert np.ndarray to a list if it is True.
    """

    weights = self.get_weights()
    flatten = []
    for _, w in weights.items():
        w_: List[float] = list(w.flatten().astype(float))
        flatten.extend(w_)

    return flatten

def get_gradients(
    self,
) -> Dict[str, np.ndarray]:
    """Return computed gradients attached to the model.

    Raises:
        FedbiomedModelError: If no gradients have been computed yet
            (i.e. the model has not been trained).

    Returns:
        Gradients, as a dict mapping parameters' names to their
            gradient's numpy array.
    """
    if not self._gradients:
        raise FedbiomedModelError(
            f"{ErrorNumbers.FB622.value}. Cannot get gradients if the "
            "model has not been trained beforehand."
        )
    gradients = self._gradients
    return gradients

@abstractmethod
def get_learning_rate(self) -> List[float]:
    """Retrieves learning rate of the model. Method implementation will
    depend on the attribute used to set up these arbitrary arguments

    Returns:
        Initial learning rate value(s); a single value if only on learning rate has been used, and
            a list of several learning rates, one for each layer of the model.
    """

def get_params(self, value: Any = None) -> Dict[str, Any]:
    """Return the wrapped scikit-learn model's hyperparameters.

    Please refer to [`baseEstimator documentation`]
    [https://scikit-learn.org/stable/modules/generated/sklearn.base.BaseEstimator.html] `get_params` method
    for further details.

    Args:
        value: if specified, returns a specific hyperparameter, otherwise, returns a dictionary
            with all the hyperparameters. Defaults to None.

    Returns:
        Dictionary mapping model hyperparameter names to their values
    """
    if value is not None:
        return self.model.get_params().get(value)
    return self.model.get_params()

def get_weights(
    self,
    only_trainable: bool = False,
    exclude_buffers: bool = True
) -> Dict[str, np.ndarray]:
    """Return a copy of the model's trainable weights.

    Args:
        only_trainable: Unused for scikit-learn models. (Whether to ignore
            non-trainable model parameters.)
        exclude_buffers: Unused for scikit-learn models. (Whether to ignore
            buffers.)

    Raises:
        FedbiomedModelError: If the model parameters are not initialized.

    Returns:
        Model weights, as a dictionary mapping parameters' names to their
            numpy array, or as a declearn NumpyVector wrapping such a dict.
    """
    if not self.param_list:
        raise FedbiomedModelError(
            f"{ErrorNumbers.FB622.value}. Attribute `param_list` is empty. You should "
            f"have initialized the model beforehand (try calling `set_init_params`)"
        )
    # Gather copies of the model weights.
    weights = {}  # type: Dict[str, np.ndarray]
    try:
        for key in self.param_list:
            val = getattr(self.model, key)
            if not isinstance(val, np.ndarray):
                raise FedbiomedModelError(
                    f"{ErrorNumbers.FB622.value}: SklearnModel parameter is not a numpy array."
                )
            weights[key] = val.copy()
    except AttributeError as err:
        raise FedbiomedModelError(
            f"{ErrorNumbers.FB622.value}. Unable to access weights of BaseEstimator "
            f"model {self.model} (details {err})"
        ) from err
    return weights

def init_training(self):
    """Initialises the training by setting up attributes.

    Raises:
        FedbiomedModelError: raised if `param_list` has not been defined
    """
    if not self.param_list:
        raise FedbiomedModelError(
            f"{ErrorNumbers.FB622.value}. Attribute `param_list` is empty. You should "
            f"have initialized the model beforehand (try calling `set_init_params`)"
        )

def predict(
    self,
    inputs: np.ndarray,
) -> np.ndarray:
    """Computes prediction given input data.

    Args:
        inputs: input data

    Returns:
        Model predictions
    """
    return self.model.predict(inputs)

def reload(self, filename: str) -> None:
    """Import and replace the wrapped model from a dump file.

    Args:
        filename: path to the file where the model has been exported.

    !!! info "Notes":
        This method is designed to load the model from a local dump
        file, that might not be in a trustworthy format. It should
        therefore only be used to re-load data exported locally and
        not received from someone else, including other FL peers.

    Raises:
        FedbiomedModelError: if the reloaded instance is of unproper type.
    """
    model = self._reload(filename)
    if not isinstance(model, self._model_type):
        err_msg = (
            f"{ErrorNumbers.FB622.value}: unproper type for imported model"
            f": expected '{self._model_type}', but 'got {type(model)}'."
        )
        logger.critical(err_msg)
        raise FedbiomedModelError(err_msg)
    self.model = model

def set_gradients(self, gradients: Dict[str, np.ndarray]) -> None:
    # TODO: either document or remove this (useless) method
    self._gradients = gradients

@abstractmethod
def set_init_params(self, model_args: Dict) -> None:
    """Zeroes scikit learn model parameters.

    Should be used before any training, as it sets the scikit learn model parameters
    and makes them accessible through the use of attributes. Model parameter attribute names
    will depend on the scikit learn model wrapped.

    Args:
        model_args: dictionary that contains specifications for setting initial model
    """

def set_params(self, **params: Any) -> Dict[str, Any]:
    """Assign some hyperparameters to the wrapped scikit-learn model.

    Please refer to [BaseEstimator][sklearn.base.BaseEstimator]
    [https://scikit-learn.org/stable/modules/generated/sklearn.base.BaseEstimator.html] `set_params` method
    for further details.

    Args:
        params: new hyperparameters to assign to the model.

    Returns:
        Dict[str, Any]: dictionary containing new hyperparameter values.
    """
    self.model.set_params(**params)
    return params

def set_weights(
    self,
    weights: Dict[str, np.ndarray],
) -> None:
    """Assign new values to the model's trainable weights.

    Args:
        weights: Model weights, as a dict mapping parameters' names
            to their numpy array.
    """
    self._assert_dict_inputs(weights)
    for key, val in weights.items():
        setattr(self.model, key, val.copy())

def train(
    self,
    inputs: np.ndarray,
    targets: np.ndarray,
    stdout: Optional[List[List[str]]] = None,
    **kwargs,
) -> None:
    """Run a training step, and record associated gradients.

    Args:
        inputs: inputs data.
        targets: targets, to be fit with inputs data.
        stdout: list of console outputs that have been collected
            during training, that contains losses values.
            Used to plot model losses. Defaults to None.

    Raises:
        FedbiomedModelError: if training has not been initialized.
    """
    batch_size = inputs.shape[0]
    w_init = self.get_weights()
    w_updt = {key: np.zeros_like(val) for key, val in w_init.items()}
    # Iterate over the batch; accumulate sample-wise gradients (and loss).
    for idx in range(batch_size):
        # Compute updated weights based on the sample. Capture loss prints.
        with capture_stdout() as console:
            self.model.partial_fit(inputs[idx : idx + 1], targets[idx])
        if stdout is not None:
            stdout.append(console)
        # Accumulate updated weights (weights + sum of gradients).
        # Reset the model's weights and iteration counter.
        for key in self.param_list:
            w_updt[key] += getattr(self.model, key)
            setattr(self.model, key, w_init[key].copy())
        self.model.n_iter_ -= 1
    # Compute the batch-averaged, learning-rate-scaled gradients.
    # Note: w_init: {w_t}, w_updt: {w_t - eta_t * sum_{s=1}^B(grad_s)}
    #       hence eta_t * avg(grad_s) = w_init - (w_updt / B)

    self._gradients = {
        key: w_init[key] - (w_updt[key] / batch_size)
        for key in self.param_list
    }

    # ------------------------------ WARNINGS ----------------------------------
    #
    # Warning 1: if `disable_internal_optimizer` has not been called before, gradients won't be scaled
    # (you will get un-scaled gradients, that need to be scaled back by dividing gradients by the learning rate)
    # here is a way to do so (with `lrate` as the learning rate): 
    # ```python
    # for key, val in self._gradients.items():
    #        val /= lrate
    # ````
    # Warning 2:  `_gradients` has different meanings, when using `disable_internal_optimizer`
    # if it is not called (ie when using native sklearn optimizer), it is not plain gradients, 
    # but rather the quantity `lr * grads`

    # Finally, increment the model's iteration counter.
    self.model.n_iter_ += 1

def unflatten(
        self,
        weights_vector: List[float],
        only_trainable: bool = False,
        exclude_buffers: bool = True
) -> Dict[str, np.ndarray]:
    """Unflatten vectorized model weights

    Args:
        weights_vector: Vectorized model weights to convert dict
        only_trainable: Unused for scikit-learn models. (Whether to ignore
            non-trainable model parameters.)
        exclude_buffers: Unused for scikit-learn models. (Whether to ignore
            buffers.)

    Returns:
        Model dictionary
    """

    super().unflatten(weights_vector, only_trainable, exclude_buffers)

    weights_vector = np.array(weights_vector)
    weights = self.get_weights()
    pointer = 0

    params = {}
    for key, w in weights.items():
        num_param = w.size
        params[key] = weights_vector[pointer: pointer + num_param].reshape(w.shape)

        pointer += num_param

    return params

def get_learning_rate(self) -> List[float]:
    return [self.model.learning_rate_init]

@abstractmethod
def apply_updates(self, updates: Dict[str, DT]):
    """Applies updates to the model.

    Args:
        updates: model updates.
    """

@abstractmethod
def export(self, filename: str) -> None:
    """Export the wrapped model to a dump file.

    Args:
        filename: path to the file where the model will be saved.

    !!! info "Notes":
        This method is designed to save the model to a local dump
        file for easy re-use by the same user, possibly outside of
        Fed-BioMed. It is not designed to produce trustworthy data
        dumps and is not used to exchange models and their weights
        as part of the federated learning process.
    """

@abstractmethod
def flatten(self,
            only_trainable: bool = False,
            exclude_buffers: bool = True) -> List[float]:
    """Flattens model weights

    Args:
        only_trainable: Whether to ignore non-trainable model parameters
            from outputs (e.g. frozen neural network layers' parameters),
            or include all model parameters (the default).
        exclude_buffers: Whether to ignore buffers (the default), or 
            include them.

    Returns:
        List of model weights as float.
    """

@abstractmethod
def get_gradients(self) -> Dict[str, DT]:
    """Return computed gradients attached to the model.

    Returns:
        Gradients, as a dict mapping parameters' names to their
            gradient's value.
    """

@abstractmethod
def get_weights(self, only_trainable: bool = False, exclude_buffers: bool = True) -> Dict[str, DT]:
    """Return a copy of the model's trainable weights.

    Args:
        only_trainable: Whether to ignore non-trainable model parameters
            from outputs (e.g. frozen neural network layers' parameters),
            or include all model parameters (the default).
        exclude_buffers: Whether to ignore buffers (the default), or 
            include them.

    Returns:
        Model weights, as a dict mapping parameters' names to their value.
    """

@abstractmethod
def init_training(self):
    """Initialize parameters before model training."""

@abstractmethod
def predict(self, inputs: Any) -> Any:
    """Return model predictions given input values.

    Args:
        inputs: input values.

    Returns:
        Any: predictions.
    """

@abstractmethod
def reload(self, filename: str) -> None:
    """Import and replace the wrapped model from a dump file.

    Args:
        filename: path to the file where the model has been exported.

    !!! info "Notes":
        This method is designed to load the model from a local dump
        file, that might not be in a trustworthy format. It should
        therefore only be used to re-load data exported locally and
        not received from someone else, including other FL peers.

    Raises:
        FedbiomedModelError: if the reloaded instance is of unproper type.
    """

def set_model(self, model: _MT) -> None:
    """Replace the wrapped model with a new one.

    Args:
        model: New model instance that needs assignment as the `model`
            attribute.
    """
    self._validate_model_type(model)
    self.model = model

@abstractmethod
def set_weights(self, weights: Dict[str, DT]) -> None:
    """Assign new values to the model's trainable weights.

    Args:
        weights: Model weights, as a dict mapping parameters' names
            to their value.
    """

@abstractmethod
def train(self, inputs: Any, targets: Any, **kwargs) -> None:
    """Perform a training step given inputs and targets data.

    !!! warning "Warning"
        Please run `init_training` method before running `train` method,
        so to initialize parameters needed for model training"

    !!! warning "Warning"
        This function usually does not update weights. You need to call
        `apply_updates` to ensure updates are applied to the model.

    Args:
        inputs: input (training) data.
        targets: target values.
    """

@abstractmethod
def unflatten(
        self,
        weights_vector: List[float],
        only_trainable: bool = False,
        exclude_buffers: bool = True
) -> None:
    """Revert flatten model weights back model-dict form.

    Args:
        weights_vector: Vectorized model weights to convert dict
        only_trainable: Whether to ignore non-trainable model parameters
            from outputs (e.g. frozen neural network layers' parameters),
            or include all model parameters (the default).
        exclude_buffers: Whether to ignore buffers (the default), or 
            include them.

    Returns:
        Model dictionary
    """

    if not isinstance(weights_vector, list) or not all([isinstance(w, float) for w in weights_vector]):
        raise FedbiomedModelError(
            f"{ErrorNumbers.FB622} `weights_vector should be 1D list of float containing flatten model parameters`"
        )

def set_init_params(self, model_args: Dict[str, Any]) -> None:
    """Initialize the model's trainable parameters."""
    # Set up zero-valued start weights, for binary of multiclass classif.
    n_classes = model_args["n_classes"]
    if n_classes == 2:
        init_params = {
            "intercept_": np.zeros((1,)),
            "coef_": np.zeros((1, model_args["n_features"])),
        }
    else:
        init_params = {
            "intercept_": np.zeros((n_classes,)),
            "coef_": np.zeros((n_classes, model_args["n_features"])),
        }
    # Assign these initialization parameters and retain their names.
    self.param_list = list(init_params)
    for key, val in init_params.items():
        setattr(self.model, key, val)
    # Also initialize the "classes_" slot with unique predictable labels.
    # FIXME: this assumes target values are integers in range(n_classes).
    setattr(self.model, "classes_", np.arange(n_classes))

def set_init_params(self, model_args: Dict[str, Any]):
    """Initialize the model's trainable parameters."""
    init_params = {
        "intercept_": np.array([0.0]),
        "coef_": np.array([0.0] * model_args["n_features"]),
    }
    self.param_list = list(init_params)
    for key, val in init_params.items():
        setattr(self.model, key, val)

def add_corrections_to_gradients(
    self,
    corrections: Dict[str, torch.Tensor],
) -> None:
    """Add values to the gradients currently attached to the model.

    Args:
        corrections: corrections to be added to the model's gradients.
    """
    self._assert_dict_inputs(corrections)
    for name, update in corrections.items():
        param = self.model.get_parameter(name)
        if param.grad is not None:
            param.grad.add_(update.to(param.grad.device))

def apply_updates(
    self,
    updates: Dict[str, torch.Tensor],
) -> None:
    """Apply incoming updates to the wrapped model's parameters.

    Args:
        updates: model updates to be added to the model.
    """
    self._assert_dict_inputs(updates)
    with torch.no_grad():
        for name, update in updates.items():
            param = self.model.get_parameter(name)
            param.add_(update.to(param.device))

def export(self, filename: str) -> None:
    """Export the wrapped model to a dump file.

    For PyTorch only export the model weights.

    Args:
        filename: path to the file where the model will be saved.

    !!! info "Notes":
        This method is designed to save the model to a local dump
        file for easy re-use by the same user, possibly outside of
        Fed-BioMed. It is not designed to produce trustworthy data
        dumps and is not used to exchange models and their weights
        as part of the federated learning process.

    !!! warning "Warning":
        This method uses `torch.save`, which relies on pickle and
        is therefore hard to trust by third-party loading methods.
    """
    torch.save(self.model.state_dict(), filename)

def flatten(self,
            only_trainable: bool = False,
            exclude_buffers: bool = True) -> List[float]:
    """Gets weights as flatten vector

    Args:
        only_trainable: Whether to ignore non-trainable model parameters
            from outputs (e.g. frozen neural network layers' parameters),
            or include all model parameters (the default).
        exclude_buffers: Whether to ignore buffers (the default), or 
            include them.

    Returns:
        to_list: Convert np.ndarray to a list if it is True.
    """
    params: List[float] = torch.nn.utils.parameters_to_vector(
        self.get_weights(only_trainable=only_trainable, exclude_buffers=exclude_buffers).values()
    ).tolist()

    return params

def get_gradients(
    self,
) -> Dict[str, torch.Tensor]:
    """Return the gradients attached to the model.

    Returns:
        Gradients, as a dict mapping parameters' names to their gradient's
            torch tensor.
    """
    gradients = {
        name: param.grad.detach().clone()
        for name, param in self.model.named_parameters()
        if (param.requires_grad and param.grad is not None)
    }
    if len(gradients) < len(list(self.model.named_parameters())):
        # FIXME: this will be triggered when having some frozen weights
        #        even if training was properly conducted
        logger.warning(
            "Warning: can not retrieve all gradients from the model. "
            "Are you sure you have trained the model beforehand?"
        )
    return gradients

def get_weights(
    self,
    only_trainable: bool = False,
    exclude_buffers: bool = True
) -> Dict[str, torch.Tensor]:
    """Return the model's parameters.

    Args:
        only_trainable: Whether to ignore non-trainable model parameters
            from outputs (e.g. frozen neural network layers' parameters),
            or include all model parameters (the default).
        exclude_buffers: Whether to ignore buffers (the default), or 
            include them.

    Returns:
        Model weights, as a dictionary mapping parameters' names to their
            torch tensor.
    """
    param_iterator = self.model.named_parameters() if exclude_buffers else self.model.state_dict().items()
    parameters = {
        name: param.detach().clone()
        for name, param in param_iterator
        if param.requires_grad or not only_trainable
    }
    return parameters

def init_training(self) -> None:
    """Initializes and sets attributes before the training.

    Initializes `init_params` as a copy of the initial parameters of the model
    """
    # initial aggregated model parameters
    self.init_params = {
        key: param.data.detach().clone()
        for key, param in self.model.named_parameters()
    }
    self.model.train()  # pytorch switch for training
    self.model.zero_grad()

def predict(
    self,
    inputs: torch.Tensor,
) -> np.ndarray:
    """Computes prediction given input data.

    Args:
        inputs: input data

    Returns:
        Model predictions returned as a numpy array
    """
    self.model.eval()  # pytorch switch for model inference-mode
    with torch.no_grad():
        pred = self.model(inputs)
    return pred.cpu().numpy()

def reload(self, filename: str) -> None:
    """Import and replace the wrapped model from a dump file.

    For PyTorch, only import the model weights.

    Args:
        filename: path to the file where the model has been exported.

    !!! info "Notes":
        This method is designed to load the model from a local dump
        file, that might not be in a trustworthy format. It should
        therefore only be used to re-load data exported locally and
        not received from someone else, including other FL peers.

    Raises:
        FedbiomedModelError: if the reloaded instance is of unproper type.
    """
    weights = torch.load(filename)
    # check format of weights and apply them to the model
    self.set_weights(weights)

def send_to_device(
    self,
    device: torch.device,
) -> None:
    """Sends model to device

    Args:
        device: device set for using GPU or CPU.
    """
    self.model.to(device)

def set_weights(
    self,
    weights: Dict[str, torch.Tensor],
) -> None:
    """Sets model weights.

    Args:
        weights: Model weights, as a dict mapping parameters' names
            to their torch tensor.
    """
    self._assert_dict_inputs(weights)
    incompatible = self.model.load_state_dict(weights, strict=False)
    # Warn about (probably-)missing trainable weights.
    # Note: state_dict may include values that do not belong to the model's
    # parameters, and/or input weights may exclude non-trainable weights,
    # without requiring a warning.
    if incompatible.missing_keys:
        params = {key for key, prm in self.model.named_parameters() if prm.requires_grad}
        missing = params.intersection(incompatible.missing_keys)
        if missing:
            logger.warning(
                "'TorchModel.set_weights' received inputs that did not cover all"
                "trainable model parameters; missing weights: %s",
                missing
            )
    # Warn about invalid (hence, unused) inputs.
    if incompatible.unexpected_keys:
        logger.warning(
            "'TorchModel.set_weights' received inputs with unexpected names: %s",
            incompatible.unexpected_keys
        )

def train(
    self,
    inputs: torch.Tensor,
    targets: torch.Tensor,
    **kwargs,
) -> None:
    # TODO: should we pass loss function here? and do the backward prop?
    if not self.init_params:
        raise FedbiomedModelError(
            f"{ErrorNumbers.FB622.value}. Training has not been initialized, please initialize it beforehand"
        )

def unflatten(
        self,
        weights_vector: List[float],
        only_trainable: bool = False,
        exclude_buffers: bool = True
) -> Dict[str, torch.Tensor]:
    """Unflatten vectorized model weights using [`vector_to_parameters`][torch.nn.utils.vector_to_parameters]

    This method does not manipulate current model weights modify model parameters.

    Args:
        weights_vector: Vectorized model weights to convert dict
        only_trainable: Whether to ignore non-trainable model parameters
            from outputs (e.g. frozen neural network layers' parameters),
            or include all model parameters (the default).
        exclude_buffers: Whether to ignore buffers (the default), or 
            include them.

    Returns:
        Model dictionary
    """

    super().unflatten(weights_vector, only_trainable, exclude_buffers)

    # Copy model to make sure global model parameters won't be overwritten
    model = copy.deepcopy(self)
    vector = torch.as_tensor(weights_vector).type(torch.DoubleTensor)
    weights = model.get_weights(only_trainable=only_trainable, exclude_buffers=exclude_buffers)

    # Following operation updates model parameters of copied model object
    try:
        torch.nn.utils.vector_to_parameters(vector, weights.values())
    except TypeError as e:
        FedbiomedModelError(
            f"{ErrorNumbers.FB622.value} Can not unflatten model parameters. {e}"
        )
    return weights