Issues with transforming of a Mixture of Generalized Pareto using Bijector

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)

Start by running the supplied ApproxTransform and SafeGeneralizedPareto reproducer, focusing on _forward, _inverse, and the GeneralizedPareto support behavior. Compare the Normal case with the mixture of Generalized Pareto components and the Generalized Extreme Value target. Done means the transformed mixture produces valid finite results without unsupported PDF or CDF values.