Issues with transforming of a Mixture of Generalized Pareto using Bijector

[Pranavesh-Panakkal]

I would like to transform a mixture of distributions into another distribution and then apply the learned transformation to the individual components for a specific engineering application. While I am able to accomplish this for a Normal distribution, it does not work for the Generalized Pareto distribution.

The code below illustrates the workflow.

import tensorflow_probability as tfp
from scipy.optimize import brentq
import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf

tfd = tfp.distributions
tfb = tfp.bijectors

normal1 = tfd.Normal(loc=1.5, scale=0.5)
normal2 = tfd.Normal(loc=1.0, scale=1.0)

source_dist = tfd.Mixture(cat=tfd.Categorical(probs=[0.4, 0.6]), components=[normal1, normal2])
target_dist = tfd.Exponential(rate=1.0)

# Define the variable transformer
transformer = ApproxTransform(source_dist=source_dist, target_dist=target_dist)

# Transform the source distribution
transformed_dist = tfd.TransformedDistribution(distribution=source_dist, bijector=transformer)

# Now apply the transformation to individual distributions
transformed_normal1 = tfd.TransformedDistribution(distribution=normal1, bijector=transformer)
transformed_normal2 = tfd.TransformedDistribution(distribution=normal2, bijector=transformer)
recombined = tfd.Mixture(cat=tfd.Categorical(probs=[0.4, 0.6]), components=[transformed_normal1, transformed_normal2])

Here

class ApproxTransform(tfb.Bijector):
    def __init__(self, source_dist, target_dist, min_support=-1e9, max_support=1e9, name="ApproxTransform"):
        super().__init__(forward_min_event_ndims=0, name=name)
        self.source = source_dist
        self.target = target_dist
        self.min_support = min_support
        self.max_support = max_support

    def approx_quantile(self, dist, qi):
        
        if np.ndim(qi) == 0:
            # scalar  
            def objective(x):
                return dist.cdf(x) - qi
            return brentq(objective, self.min_support, self.max_support)
        else:  
            # Vector over qi elements
            quantiles = [self.approx_quantile(dist, q) for q in qi]
            return np.array(quantiles)
        
    def _forward(self, x):
        x = tf.convert_to_tensor(x)
        x = tf.cast(x, self.source.dtype)
        u = self.source.cdf(x)
        return self.approx_quantile(self.target, u)

    def _inverse(self, y):
        y = tf.convert_to_tensor(y)
        y = tf.cast(y, self.source.dtype)
        u = self.target.cdf(y)
        return self.approx_quantile(self.source, u)

    def _forward_log_det_jacobian(self, x):
        """log|dy/dx| = log f_X(x) - log f_Y(y)"""
        desired_dtype = self.source.dtype
        x = tf.cast(x, desired_dtype)
        y = self._forward(x)
        y = tf.cast(y, desired_dtype)
        return self.source.log_prob(x) - self.target.log_prob(y)

    def _is_increasing(self):
        """Transformation is monotonic increasing."""
        return True
        

This works, as illustrated in the image below.

Now, I want to perform this process for a mixture of Generalized Pareto Distributions. Specifically, I aim to combine two Pareto distributions, transform the mixture into another distribution (for example, a Generalized Extreme Value distribution), and then recover the transformed components.

However, I am encountering an issue where TensorFlow Probability (TFP) reports PDF and CDF values below the distribution’s support, making my workflow impractical. I attempted to resolve this by implementing a custom class that wraps around tfd.GeneralizedPareto, but the resulting transformed distribution produces only NaN values. Is there any way to transform the mixture into a GEV distribution successfully?

class SafeGeneralizedPareto(tfd.GeneralizedPareto):
    def __init__(self, *args, **kwargs):
        # Force disable built-in validation to avoid assertion errors
        kwargs['validate_args'] = False
        super().__init__(*args, **kwargs)

    def prob(self, x):
        base_prob = super().prob(x)
        return tf.where(x < self.loc, tf.zeros_like(base_prob), base_prob)

    def cdf(self, x):
        base_cdf = super().cdf(x)
        return tf.where(x < self.loc, tf.zeros_like(base_cdf), base_cdf)

    def log_prob(self, x):
        # raise NotImplemented
        base_log_prob = super().log_prob(x)
        return tf.where(x < self.loc, tf.constant(-float('inf'), dtype=base_log_prob.dtype), base_log_prob)

gpd_1 = SafeGeneralizedPareto(loc=17.402, scale=37.613, concentration=-0.118)
gpd_2 = SafeGeneralizedPareto(loc=21.007, scale=9.453, concentration=0.136)

source_dist = tfd.Mixture(cat=tfd.Categorical(probs=[0.4, 0.6]), components=[gpd_1, gpd_2])
target_dist = tfd.GeneralizedExtremeValue(loc=3.930, scale=0.747, concentration=-0.2357)

# Define the variable transformer
transformer = ApproxTransform(source_dist=source_dist, target_dist=target_dist)

# Transform the source distribution
transformed_dist = tfd.TransformedDistribution(distribution=source_dist, bijector=transformer)