rm: rm legacy CategoricalNet

Author: jnsbck

sbi/neural_nets/estimators/__init__.py

@@ -2,7 +2,6 @@
 from sbi.neural_nets.estimators.categorical_net import (
     CategoricalMADE,
     CategoricalMassEstimator,
-    CategoricalNet,
 )
 from sbi.neural_nets.estimators.flowmatching_estimator import FlowMatchingEstimator
 from sbi.neural_nets.estimators.mixed_density_estimator import MixedDensityEstimator

sbi/neural_nets/estimators/categorical_net.py

@@ -8,7 +8,6 @@
 from nflows.utils import torchutils
 from torch import Tensor, nn
 from torch.distributions import Categorical
-from torch.nn import Sigmoid, Softmax
 from torch.nn import functional as F
 
 from sbi.neural_nets.estimators.base import ConditionalDensityEstimator
@@ -87,7 +86,7 @@ def forward(self, inputs: Tensor, context: Optional[Tensor] = None) -> Tensor:
             condition: Conditioning variable. (batch_size, *condition_shape)
 
         Returns:
-            Predicted categorical probabilities. (batch_size, *input_shape,
+            Predicted categorical logits. (batch_size, *input_shape,
                 num_categories)
         """
         embedded_context = self.embedding_net.forward(context)
@@ -149,119 +148,19 @@ def _initialize(self):
         pass
 
 
-class CategoricalNet(nn.Module):
-    """Conditional density (mass) estimation for a categorical random variable.
-
-    Takes as input parameters theta and learns the parameters p of a Categorical.
-
-    Defines log prob and sample functions.
-    """
-
-    def __init__(
-        self,
-        num_input: int,
-        num_categories: int,
-        num_hidden: int = 20,
-        num_layers: int = 2,
-        embedding_net: Optional[nn.Module] = None,
-    ):
-        """Initialize the neural net.
-
-        Args:
-            num_input: number of input units, i.e., dimensionality of the features.
-            num_categories: number of output units, i.e., number of categories.
-            num_hidden: number of hidden units per layer.
-            num_layers: number of hidden layers.
-            embedding_net: emebedding net for input.
-        """
-        super().__init__()
-
-        self.num_hidden = num_hidden
-        self.num_input = num_input
-        self.activation = Sigmoid()
-        self.softmax = Softmax(dim=1)
-        self.num_categories = num_categories
-        self.num_variables = 1
-
-        # Maybe add embedding net in front.
-        if embedding_net is not None:
-            self.input_layer = nn.Sequential(
-                embedding_net, nn.Linear(num_input, num_hidden)
-            )
-        else:
-            self.input_layer = nn.Linear(num_input, num_hidden)
-
-        # Repeat hidden units hidden layers times.
-        self.hidden_layers = nn.ModuleList()
-        for _ in range(num_layers):
-            self.hidden_layers.append(nn.Linear(num_hidden, num_hidden))
-
-        self.output_layer = nn.Linear(num_hidden, num_categories)
-
-    def forward(self, condition: Tensor) -> Tensor:
-        """Return categorical probability predicted from a batch of inputs.
-
-        Args:
-            condition: batch of context parameters for the net.
-
-        Returns:
-            Tensor: batch of predicted categorical probabilities.
-        """
-        # forward path
-        condition = self.activation(self.input_layer(condition))
-
-        # iterate n hidden layers, input condition and calculate tanh activation
-        for layer in self.hidden_layers:
-            condition = self.activation(layer(condition))
-
-        return self.softmax(self.output_layer(condition))
-
-    def log_prob(self, input: Tensor, condition: Tensor) -> Tensor:
-        """Return categorical log probability of categories input, given condition.
-
-        Args:
-            input: categories to evaluate.
-            condition: parameters.
-
-        Returns:
-            Tensor: log probs with shape (input.shape[0],)
-        """
-        # Predict categorical ps and evaluate.
-        ps = self.forward(condition)
-        # Squeeze the last dimension (event dim) because `Categorical` has
-        # `event_shape=()` but our data usually has an event_shape of `(1,)`.
-        return Categorical(probs=ps).log_prob(input.squeeze(dim=-1))
-
-    def sample(self, sample_shape: torch.Size, condition: Tensor) -> Tensor:
-        """Returns samples from categorical random variable with probs predicted from
-        the neural net.
-
-        Args:
-            sample_shape: number of samples to obtain.
-            condition: batch of parameters for prediction.
-
-        Returns:
-            Tensor: Samples with shape (num_samples, 1)
-        """
-
-        # Predict Categorical ps and sample.
-        ps = self.forward(condition)
-        return Categorical(probs=ps).sample(sample_shape=sample_shape)
-
-
 class CategoricalMassEstimator(ConditionalDensityEstimator):
     """Conditional density (mass) estimation for a categorical random variable.
 
     The event_shape of this class is `()`.
     """
 
     def __init__(
-        self, net: CategoricalNet, input_shape: torch.Size, condition_shape: torch.Size
+        self, net: CategoricalMADE, input_shape: torch.Size, condition_shape: torch.Size
     ) -> None:
         """Initialize the mass estimator.
 
         Args:
-            net: CategoricalNet.
+            net: CategoricalMADE.
             input_shape: Shape of the input data.
             condition_shape: Shape of the condition data
         """

sbi/neural_nets/net_builders/categorial.py

@@ -9,7 +9,6 @@
 from sbi.neural_nets.estimators import (
     CategoricalMADE,
     CategoricalMassEstimator,
-    CategoricalNet,
 )
 from sbi.neural_nets.estimators.mixed_density_estimator import _is_discrete
 from sbi.utils.nn_utils import get_numel
@@ -18,55 +17,6 @@
 
 
 def build_categoricalmassestimator(
-    batch_x: Tensor,
-    batch_y: Tensor,
-    z_score_x: Optional[str] = "none",
-    z_score_y: Optional[str] = "independent",
-    num_hidden: int = 20,
-    num_layers: int = 2,
-    embedding_net: nn.Module = nn.Identity(),
-):
-    """Returns a density estimator for a categorical random variable.
-
-    Args:
-        batch_x: A batch of input data.
-        batch_y: A batch of condition data.
-        z_score_x: Whether to z-score the input data.
-        z_score_y: Whether to z-score the condition data.
-        num_hidden: Number of hidden units per layer.
-        num_layers: Number of hidden layers.
-        embedding_net: Embedding net for y.
-    """
-
-    if z_score_x != "none":
-        raise ValueError("Categorical input should not be z-scored.")
-
-    check_data_device(batch_x, batch_y)
-    if batch_x.shape[1] > 1:
-        raise NotImplementedError("CategoricalMassEstimator only supports 1D input.")
-    num_categories = unique(batch_x).numel()
-    dim_condition = get_numel(batch_y, embedding_net=embedding_net)
-
-    z_score_y_bool, structured_y = z_score_parser(z_score_y)
-    if z_score_y_bool:
-        embedding_net = nn.Sequential(
-            standardizing_net(batch_y, structured_y), embedding_net
-        )
-
-    categorical_net = CategoricalNet(
-        num_input=dim_condition,
-        num_categories=num_categories,
-        num_hidden=num_hidden,
-        num_layers=num_layers,
-        embedding_net=embedding_net,
-    )
-
-    return CategoricalMassEstimator(
-        categorical_net, input_shape=batch_x[0].shape, condition_shape=batch_y[0].shape
-    )
-
-
-def build_autoregressive_categoricalmassestimator(
     batch_x: Tensor,
     batch_y: Tensor,
     z_score_x: Optional[str] = "none",

sbi/neural_nets/net_builders/mnle.py

@@ -13,7 +13,6 @@
     _separate_input,
 )
 from sbi.neural_nets.net_builders.categorial import (
-    build_autoregressive_categoricalmassestimator,
     build_categoricalmassestimator,
 )
 from sbi.neural_nets.net_builders.flow import (
@@ -59,7 +58,6 @@ def build_mnle(
     z_score_x: Optional[str] = "independent",
     z_score_y: Optional[str] = "independent",
     flow_model: str = "nsf",
-    categorical_model: str = "mlp",
     num_categorical_columns: Optional[Tensor] = None,
     embedding_net: nn.Module = nn.Identity(),
     combined_embedding_net: Optional[nn.Module] = None,
@@ -157,32 +155,16 @@ def build_mnle(
     combined_condition = torch.cat([disc_x, embedded_batch_y], dim=-1)
 
     # Set up a categorical RV neural net for modelling the discrete data.
-    if categorical_model == "made":
-        discrete_net = build_autoregressive_categoricalmassestimator(
-            disc_x,
-            batch_y,
-            z_score_x="none",  # discrete data should not be z-scored.
-            z_score_y="none",  # y-embedding net already z-scores.
-            num_hidden=hidden_features,
-            num_layers=hidden_layers,
-            embedding_net=embedding_net,
-            num_categories=num_categorical_columns,
-        )
-    elif categorical_model == "mlp":
-        assert num_disc == 1, "MLP only supports 1D input."
-        discrete_net = build_categoricalmassestimator(
-            disc_x,
-            batch_y,
-            z_score_x="none",  # discrete data should not be z-scored.
-            z_score_y="none",  # y-embedding net already z-scores.
-            num_hidden=hidden_features,
-            num_layers=hidden_layers,
-            embedding_net=embedding_net,
-        )
-    else:
-        raise ValueError(
-            f"Unknown categorical net {categorical_model}. Must be 'made' or 'mlp'."
-        )
+    discrete_net = build_categoricalmassestimator(
+        disc_x,
+        batch_y,
+        z_score_x="none",  # discrete data should not be z-scored.
+        z_score_y="none",  # y-embedding net already z-scores.
+        num_hidden=hidden_features,
+        num_layers=hidden_layers,
+        embedding_net=embedding_net,
+        num_categories=num_categorical_columns,
+    )
 
     if combined_embedding_net is None:
         # set up linear embedding net for combining discrete and continuous