Add lc endpoint: predict_modelwise_probas

choice_learn/models/base_model.py

@@ -444,7 +444,7 @@ def fit(
         self.callbacks.on_train_end(logs=temps_logs)
         return losses_history
 
-    @tf.function
+    @tf.function(reduce_retracing=True)
     def batch_predict(
         self,
         shared_features_by_choice,
@@ -731,7 +731,6 @@ def f(params_1d):
             # calculate gradients and convert to 1D tf.Tensor
             grads = tape.gradient(loss_value, self.trainable_weights)
             grads = tf.dynamic_stitch(idx, grads)
-            # print out iteration & loss
             f.iter.assign_add(1)
 
             # store loss value so we can retrieve later

choice_learn/models/latent_class_base_model.py

@@ -104,12 +104,18 @@ def instantiate(self, **kwargs):
             name="Latent-Logits",
         )
         self.latent_logits = init_logit
-        self.models = [self.model_class(**mp) for mp in self.model_parameters]
-        for model in self.models:
-            model.instantiate(**kwargs)
 
+        self.models = self.instantiate_latent_models(**kwargs)
         self.instantiated = True
 
+    def instantiate_latent_models(self, **kwargs):
+        """Instantiate latent models."""
+        models = [self.model_class(**mp) for mp in self.model_parameters]
+        for model in models:
+            model.instantiate(**kwargs)
+
+        return models
+
     # @tf.function
     def batch_predict(
         self,
@@ -230,7 +236,7 @@ def compute_batch_utility(
             utilities.append(model_utilities)
         return utilities
 
-    def fit(self, choice_dataset, sample_weight=None, verbose=0):
+    def fit(self, choice_dataset, sample_weight=None, val_dataset=None, verbose=0):
         """Fit the model on a ChoiceDataset.
 
         Parameters
@@ -239,6 +245,8 @@ def fit(self, choice_dataset, sample_weight=None, verbose=0):
             Dataset to be used for coefficients estimations
         sample_weight : np.ndarray, optional
             sample weights to apply, by default None
+        val_dataset: ChoiceDataset
+            Validation dataset for MLE Gradient Descent Optimization
         verbose : int, optional
             print level, for debugging, by default 0
 
@@ -249,7 +257,6 @@ def fit(self, choice_dataset, sample_weight=None, verbose=0):
         """
         if self.fit_method.lower() == "em":
             self.minf = np.log(1e-3)
-            print("Expectation-Maximization estimation algorithm not well implemented yet.")
             return self._em_fit(
                 choice_dataset=choice_dataset, sample_weight=sample_weight, verbose=verbose
             )
@@ -272,7 +279,10 @@ def fit(self, choice_dataset, sample_weight=None, verbose=0):
                     self.optimizer = tf.keras.optimizers.Adam(self.lr)
 
             return self._fit_with_gd(
-                choice_dataset=choice_dataset, sample_weight=sample_weight, verbose=verbose
+                choice_dataset=choice_dataset,
+                sample_weight=sample_weight,
+                verbose=verbose,
+                val_dataset=val_dataset,
             )
 
         raise ValueError(f"Fit method not implemented: {self.fit_method}")
@@ -757,45 +767,6 @@ def _fit_with_gd(
         # self.callbacks.on_train_end(logs=temps_logs)
         return losses_history
 
-    def _nothing(self, inputs):
-        """_summary_.
-
-        Parameters
-        ----------
-        inputs : _type_
-            _description_
-
-        Returns
-        -------
-        _type_
-            _description_
-        """
-        latent_probas = tf.clip_by_value(
-            self.latent_logits - tf.reduce_max(self.latent_logits), self.minf, 0
-        )
-        latent_probas = tf.math.exp(latent_probas)
-        # latent_probas = tf.math.abs(self.logit_latent_probas)  # alternative implementation
-        latent_probas = latent_probas / tf.reduce_sum(latent_probas)
-        proba_list = []
-        avail = inputs[4]
-        for q in range(self.n_latent_classes):
-            combined = self.models[q].compute_batch_utility(*inputs)
-            combined = tf.clip_by_value(
-                combined - tf.reduce_max(combined, axis=1, keepdims=True), self.minf, 0
-            )
-            combined = tf.keras.layers.Activation(activation=tf.nn.softmax)(combined)
-            # combined = tf.keras.layers.Softmax()(combined)
-            combined = combined * avail
-            combined = latent_probas[q] * tf.math.divide(
-                combined, tf.reduce_sum(combined, axis=1, keepdims=True)
-            )
-            combined = tf.expand_dims(combined, -1)
-            proba_list.append(combined)
-            # print(combined.get_shape()) # it is useful to print the shape of tensors for debugging
-
-        proba_final = tf.keras.layers.Concatenate(axis=2)(proba_list)
-        return tf.math.reduce_sum(proba_final, axis=2, keepdims=False)
-
     def _expectation(self, choice_dataset):
         predicted_probas = [model.predict_probas(choice_dataset) for model in self.models]
         latent_probabilities = self.get_latent_classes_weights()
@@ -824,7 +795,7 @@ def _expectation(self, choice_dataset):
         )
 
         return tf.clip_by_value(
-            predicted_probas / np.sum(predicted_probas, axis=1, keepdims=True), 1e-10, 1
+            predicted_probas / np.sum(predicted_probas, axis=1, keepdims=True), 1e-6, 1
         ), loss
 
     def _maximization(self, choice_dataset, verbose=0):
@@ -842,10 +813,17 @@ def _maximization(self, choice_dataset, verbose=0):
         np.ndarray
             latent probabilities resulting of maximization step
         """
-        self.models = [self.model_class(**mp) for mp in self.model_parameters]
+        # models = [self.model_class(**mp) for mp in self.model_parameters]
+        # for i in range(len(models)):
+        #     for j, var in enumerate(self.models[i].trainable_weights):
+        #         models[i]._trainable_weights[j] = var
+        # self.instantiate_latent_models(choice_dataset)
+
         # M-step: MNL estimation
         for q in range(self.n_latent_classes):
-            self.models[q].fit(choice_dataset, sample_weight=self.weights[:, q], verbose=verbose)
+            self.models[q].fit(
+                choice_dataset, sample_weight=self.weights[:, q].numpy(), verbose=verbose
+            )
 
         # M-step: latent probability estimation
         latent_probas = np.sum(self.weights, axis=0)
@@ -876,7 +854,9 @@ def _em_fit(self, choice_dataset, sample_weight=None, verbose=0):
 
         # Initialization
         init_sample_weight = np.random.rand(self.n_latent_classes, len(choice_dataset))
-        init_sample_weight = init_sample_weight / np.sum(init_sample_weight, axis=0, keepdims=True)
+        init_sample_weight = np.clip(
+            init_sample_weight / np.sum(init_sample_weight, axis=0, keepdims=True), 1e-6, 1
+        )
         for i, model in enumerate(self.models):
             # model.instantiate()
             model.fit(choice_dataset, sample_weight=init_sample_weight[i], verbose=verbose)
@@ -888,7 +868,7 @@ def _em_fit(self, choice_dataset, sample_weight=None, verbose=0):
             if np.sum(np.isnan(self.latent_logits)) > 0:
                 print("Nan in logits")
                 break
-        return hist_logits, hist_loss
+        return hist_loss, hist_logits
 
     def predict_probas(self, choice_dataset, batch_size=-1):
         """Predicts the choice probabilities for each choice and each product of a ChoiceDataset.
@@ -922,6 +902,43 @@ def predict_probas(self, choice_dataset, batch_size=-1):
 
         return tf.concat(stacked_probabilities, axis=0)
 
+    def predict_modelwise_probas(self, choice_dataset, batch_size=-1):
+        """Predicts the choice probabilities for each choice and each product of a ChoiceDataset.
+
+        Stacks each model probability.
+
+        Parameters
+        ----------
+        choice_dataset : ChoiceDataset
+            Dataset on which to apply to prediction
+        batch_size : int, optional
+            Batch size to use for the prediction, by default -1
+
+        Returns
+        -------
+        np.ndarray (n_choices, n_items)
+            Choice probabilties for each choice and each product
+        """
+        modelwise_probabilities = []
+        for model in self.models:
+            stacked_probabilities = []
+            for (
+                shared_features,
+                items_features,
+                available_items,
+                choices,
+            ) in choice_dataset.iter_batch(batch_size=batch_size):
+                _, probabilities = model.batch_predict(
+                    shared_features_by_choice=shared_features,
+                    items_features_by_choice=items_features,
+                    available_items_by_choice=available_items,
+                    choices=choices,
+                )
+                stacked_probabilities.append(probabilities)
+            modelwise_probabilities.append(tf.concat(stacked_probabilities, axis=0))
+
+        return tf.stack(modelwise_probabilities, axis=0)
+
     def get_latent_classes_weights(self):
         """Return the latent classes weights / probabilities from logits.
 

choice_learn/models/latent_class_mnl.py

@@ -4,6 +4,8 @@
 
 import tensorflow as tf
 
+import choice_learn.tf_ops as tf_ops
+
 from .conditional_logit import ConditionalLogit, MNLCoefficients
 from .latent_class_base_model import BaseLatentClassModel
 from .simple_mnl import SimpleMNL
@@ -23,6 +25,7 @@ def __init__(
         intercept=None,
         optimizer="Adam",
         lr=0.001,
+        epochs_maximization=1000,
         **kwargs,
     ):
         """Initialize model.
@@ -56,7 +59,7 @@ def __init__(
             "batch_size": batch_size,
             "lbfgs_tolerance": lbfgs_tolerance,
             "lr": lr,
-            "epochs": 1000,
+            "epochs": epochs_maximization,
         }
 
         super().__init__(
@@ -88,6 +91,15 @@ def instantiate_latent_models(self, n_items, n_shared_features, n_items_features
             model.indexes, model.weights = model.instantiate(
                 n_items, n_shared_features, n_items_features
             )
+            model.exact_nll = tf_ops.CustomCategoricalCrossEntropy(
+                from_logits=False,
+                label_smoothing=0.0,
+                sparse=False,
+                axis=-1,
+                epsilon=1e-25,
+                name="exact_categorical_crossentropy",
+                reduction="sum_over_batch_size",
+            )
             model.instantiated = True
 
     def instantiate(self, n_items, n_shared_features, n_items_features):