Repository: pyro-ppl/numpyro

Issue:
[FR] Support for different supports in component distributions for mixture models
# Description
I am submitting a feature request to enhance the capabilities of Numpyro in handling mixture models with component distributions having different supports. Currently, when attempting to construct a mixture model using `numpyro.distributions.MixtureGeneral`, all component distributions are required to have the same support.

In my specific use case, I am implementing a truncated power law mixture model where one of the component distributions utilizes a truncated power law model with a high-mass cutoff represented by the Heaviside step function. The other component distribution involves a Gaussian distribution.

The formulation of the model is as follows:

$$
p(m_1\mid\theta) \propto (1-\lambda)A(\theta)m_1^{-\alpha}\Theta(m_{\text{max}}-m_1) + \lambda B(\theta)\exp\left(-\frac{(m_1-\mu_m)^2}{2\sigma_m^2}\right)
$$

Where $\theta$ represents the parameter vector $\left(\lambda,\alpha,m_{\text{min}},\delta_m,m_{\text{max}},\mu_m,\sigma_m\right)$, $A(\cdot)$ and $B(\cdot)$ are normalization factors.

To implement this model efficiently, it is essential to have the flexibility to use component distributions with varying supports within the same mixture model.

Therefore, I kindly request the implementation of this feature to enable users to construct mixture models with component distributions having different supports. If there is any solution which does not requires new feature, will be great.


Discussion:
Hi @Qazalbash, you can follow examples in the Discrete Latent Variables section e.g. [Gaussian Mixture Model](https://num.pyro.ai/en/stable/tutorials/gmm.html#). There we use enumeration for mixture models.
@fehiepsi Apologies if I'm missing something obvious, but is there a class of models where the enumeration trick can't be used?

As a toy example, let's say our data generation process looks like:

```python
def data_generation_process(mu):
    x = np.random.normal(loc=mu, scale=1, size=10_000)

    for _ in range(5):
        x = np.where((x > 0.4), x - 0.1, x)
    
    return x
```

and then we want to infer the mu parameter from data. I tried this using NumPyro:

```python
def difficult_without_mixtures(y=None):
    x = dist.Normal(loc=numpyro.param("mu", 0.0), scale=1)

    for _ in range(5):
        eligible_for_shift = (1 - x.cdf(0.4))
        x_eligible = dist.TruncatedDistribution(x, low=0.4)
        x_ineligible = dist.TruncatedDistribution(x, high=0.4)
        x = dist.Mixture(
            dist.Categorical([eligible_for_shift, 1 - eligible_for_shift]),
            [
                dist.TransformedDistribution(x_eligible, dist.transforms.AffineTransform(loc=-0.1, scale=1)),
                x_ineligible
            ]
        )
    
    numpyro.sample("obs", x, obs=y)
```

but got the error: `ValueError: All component distributions must have the same support.` Yet, I see no inherent reason why my code here couldn't work in principle; for example, PyMC allows mixtures with components with different support. I could implement a custom distribution class (like [this](https://num.pyro.ai/en/stable/tutorials/truncated_distributions.html#5.-Building-your-own-truncated-distribution)) but at that point I'd be writing the same code NumPyro would need to add this feature, I think.
I think we can allow `MixtureGeneral` take a `support` argument - we will skip the check if `support` is True. In this case, we need to enable `validate_args` of each component so that its log_prob will be -inf for out-of-domain samples. Do you want to make a PR for it?
@fehiepsi I can try!

Resolving patch:
diff --git a/numpyro/distributions/mixtures.py b/numpyro/distributions/mixtures.py
--- a/numpyro/distributions/mixtures.py
+++ b/numpyro/distributions/mixtures.py
@@ -271,6 +271,9 @@ class MixtureGeneral(_MixtureBase):
         ``mixture_size``.
     :param component_distributions: A list of ``mixture_size``
         :class:`~numpyro.distributions.Distribution` objects.
+    :param support: A :class:`~numpyro.distributions.constraints.Constraint`
+        object specifying the support of the mixture distribution. If not
+        provided, the support will be inferred from the component distributions.
 
     **Example**
 
@@ -288,13 +291,36 @@ class MixtureGeneral(_MixtureBase):
        >>> mixture = dist.MixtureGeneral(mixing_dist, component_dists)
        >>> mixture.sample(jax.random.PRNGKey(42)).shape
        ()
+
+    .. doctest::
+
+       >>> import jax
+       >>> import jax.numpy as jnp
+       >>> import numpyro.distributions as dist
+       >>> mixing_dist = dist.Categorical(probs=jnp.ones(2) / 2.)
+       >>> component_dists = [
+       ...     dist.Normal(loc=0.0, scale=1.0),
+       ...     dist.HalfNormal(scale=0.3),
+       ... ]
+       >>> mixture = dist.MixtureGeneral(mixing_dist, component_dists, support=dist.constraints.real)
+       >>> mixture.sample(jax.random.PRNGKey(42)).shape
+       ()
     """
 
-    pytree_data_fields = ("_mixing_distribution", "_component_distributions")
+    pytree_data_fields = (
+        "_mixing_distribution",
+        "_component_distributions",
+        "_support",
+    )
     pytree_aux_fields = ("_mixture_size",)
 
     def __init__(
-        self, mixing_distribution, component_distributions, *, validate_args=None
+        self,
+        mixing_distribution,
+        component_distributions,
+        *,
+        support=None,
+        validate_args=None,
     ):
         _check_mixing_distribution(mixing_distribution)
 
@@ -308,7 +334,7 @@ def __init__(
         for d in component_distributions:
             if not isinstance(d, Distribution):
                 raise ValueError(
-                    "All elements of 'component_distributions' must be instaces of "
+                    "All elements of 'component_distributions' must be instances of "
                     "numpyro.distributions.Distribution subclasses"
                 )
         if len(component_distributions) != self.mixture_size:
@@ -320,11 +346,19 @@ def __init__(
         # TODO: It would be good to check that the support of all the component
         # distributions match, but for now we just check the type, since __eq__
         # isn't consistently implemented for all support types.
-        support_type = type(component_distributions[0].support)
-        if any(
-            type(d.support) is not support_type for d in component_distributions[1:]
-        ):
-            raise ValueError("All component distributions must have the same support.")
+        self._support = support
+        if support is None:
+            support_type = type(component_distributions[0].support)
+            if any(
+                type(d.support) is not support_type for d in component_distributions[1:]
+            ):
+                raise ValueError(
+                    "All component distributions must have the same support."
+                )
+        else:
+            assert isinstance(
+                support, constraints.Constraint
+            ), "support must be a Constraint object"
 
         self._mixing_distribution = mixing_distribution
         self._component_distributions = component_distributions
@@ -357,6 +391,8 @@ def component_distributions(self):
 
     @constraints.dependent_property
     def support(self):
+        if self._support is not None:
+            return self._support
         return self.component_distributions[0].support
 
     @property
@@ -389,9 +425,14 @@ def component_sample(self, key, sample_shape=()):
         return jnp.stack(samples, axis=self.mixture_dim)
 
     def component_log_probs(self, value):
-        component_log_probs = jnp.stack(
-            [d.log_prob(value) for d in self.component_distributions], axis=-1
-        )
+        component_log_probs = []
+        for d in self.component_distributions:
+            log_prob = d.log_prob(value)
+            if (self._support is not None) and (not d._validate_args):
+                mask = d.support(value)
+                log_prob = jnp.where(mask, log_prob, -jnp.inf)
+            component_log_probs.append(log_prob)
+        component_log_probs = jnp.stack(component_log_probs, axis=-1)
         return jax.nn.log_softmax(self.mixing_distribution.logits) + component_log_probs
 
 


Test patch:
diff --git a/test/test_distributions_mixture.py b/test/test_distributions_mixture.py
--- a/test/test_distributions_mixture.py
+++ b/test/test_distributions_mixture.py
@@ -19,6 +19,13 @@ def get_normal(batch_shape):
     return normal
 
 
+def get_half_normal(batch_shape):
+    """Get parameterized HalfNormal with given batch shape."""
+    scale = jnp.ones(batch_shape)
+    half_normal = dist.HalfNormal(scale=scale)
+    return half_normal
+
+
 def get_mvn(batch_shape):
     """Get parameterized MultivariateNormal with given batch shape."""
     dimensions = 2
@@ -78,6 +85,44 @@ def test_mixture_broadcast_batch_shape(
     _test_mixture(mixing_distribution, component_distribution)
 
 
+@pytest.mark.parametrize("batch_shape", [(), (1,), (7,), (2, 5)])
+@pytest.mark.filterwarnings(
+    "ignore:Out-of-support values provided to log prob method."
+    " The value argument should be within the support.:UserWarning"
+)
+def test_mixture_with_different_support(batch_shape):
+    mixing_probabilities = jnp.ones(2) / 2
+    mixing_distribution = dist.Categorical(probs=mixing_probabilities)
+    component_distribution = [
+        get_normal(batch_shape),
+        get_half_normal(batch_shape),
+    ]
+    mixture = dist.MixtureGeneral(
+        mixing_distribution=mixing_distribution,
+        component_distributions=component_distribution,
+        support=dist.constraints.real,
+    )
+    assert mixture.batch_shape == batch_shape
+    sample_shape = (11,)
+    component_distribution[0]._validate_args = True
+    component_distribution[1]._validate_args = True
+    xx = component_distribution[0].sample(rng_key, sample_shape)
+    log_prob_0 = component_distribution[0].log_prob(xx)
+    log_prob_1 = component_distribution[1].log_prob(xx)
+    expected_log_prob = jax.scipy.special.logsumexp(
+        jnp.stack(
+            [
+                log_prob_0 + jnp.log(mixing_probabilities[0]),
+                log_prob_1 + jnp.log(mixing_probabilities[1]),
+            ],
+            axis=-1,
+        ),
+        axis=-1,
+    )
+    result = mixture.log_prob(xx)
+    assert jnp.allclose(result, expected_log_prob)
+
+
 def _test_mixture(mixing_distribution, component_distribution):
     # Create mixture
     mixture = dist.Mixture(