cosmoglint 1.0.0__py3-none-any.whl

cosmoglint/model/transformer_nf.py

@@ -0,0 +1,368 @@
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .transformer import Transformer1, Transformer2, TransformerWithGlobalCond, Transformer1WithAttn, Transformer2WithAttn
+
+from nflows.flows import Flow
+from nflows.distributions import StandardNormal, ConditionalDiagonalNormal
+from nflows.transforms import CompositeTransform, RandomPermutation, AffineCouplingTransform, ActNorm, PiecewiseRationalQuadraticCDF, AffineTransform
+from nflows.nn.nets import ResidualNet
+from nflows.utils import torchutils
+
+from torch.distributions import Normal, Categorical, MixtureSameFamily, Independent
+
+
+def transformer_nf_model(args, **kwargs):
+    
+    ### Transformer model ###
+    if args.model_name == "transformer1":
+        model_class = Transformer1
+    elif args.model_name == "transformer2":
+        model_class = Transformer2
+    elif args.model_name == "transformer1_with_global_cond":
+        model_class = TransformerWithGlobalCond
+        transformer_cls = Transformer1
+    elif args.model_name == "transformer2_with_global_cond":
+        model_class = TransformerWithGlobalCond
+        transformer_cls = Transformer2
+    elif args.model_name == "transformer1_with_attn":
+        model_class = Transformer1WithAttn
+    elif args.model_name == "transformer2_with_attn":
+        model_class = Transformer2WithAttn
+    else:
+        raise ValueError(f"Invalid model: {args.model_name}")
+    
+    common_args = dict(
+            num_condition = args.num_features_cond, 
+            d_model = args.d_model, 
+            num_layers = args.num_layers, 
+            num_heads = args.num_heads, 
+            max_length = args.max_length, 
+            num_features_in = args.num_features_in, 
+            num_features_out = args.num_context, 
+            num_token_types = 1,
+            last_activation = nn.Tanh(), 
+            pred_prob = False, 
+            **kwargs
+        )
+
+    if "with_global_cond" in args.model_name:
+        common_args["num_features_global"] = args.num_features_global
+        common_args["transformer_cls"] = transformer_cls
+    
+    model = model_class(**common_args)
+
+    ### Flow model ###
+    transforms = []
+
+    alpha = 0.05 # value in RealNVP
+    transforms.append(
+    AffineTransform(
+        shift=torch.full([args.num_features_in], fill_value=alpha),
+        scale=torch.full([args.num_features_in], fill_value=1.0 - 2.0*alpha),
+    )
+    )
+    transforms.append(Logit())
+
+    for ilayer in range(args.num_flows):
+        #transforms.append(ActNorm(args.num_features_out)) # for stabilyzing training and quick convergence
+        transforms.append(
+            AffineCouplingTransform(
+                #mask=torch.arange(args.num_features_out) % 2,
+                mask = torch.arange(args.num_features_in) % 2 if ilayer % 2 == 0 else (torch.arange(args.num_features_in) + 1) % 2,
+                transform_net_create_fn=lambda in_features, out_features: ResidualNet(
+                    in_features=in_features,
+                    out_features=out_features,
+                    hidden_features=args.hidden_dim,
+                    context_features=args.num_context,
+                    num_blocks=2,
+                    activation=nn.ReLU()
+                )
+            )
+        )
+        transforms.append(PiecewiseRationalQuadraticCDF(shape=[args.num_features_in], num_bins=8, tails='linear', tail_bound=3.0))
+        transforms.append(RandomPermutation(args.num_features_in))
+        
+    transform = CompositeTransform(transforms)
+    
+    if args.base_dist == "normal":
+        base_dist = StandardNormal(shape=[args.num_features_in])
+    elif args.base_dist == "conditional_normal":
+        base_dist = ConditionalDiagonalNormal(shape=[args.num_features_in], 
+            context_encoder=nn.Sequential(nn.Linear(args.num_context, args.hidden_dim), 
+                                          nn.LeakyReLU(), 
+                                          nn.Linear(args.hidden_dim, 2 * args.num_features_in)
+            )
+        )
+
+    elif args.base_dist == "bimodal":
+        base_dist = BimodalNormal(shape=[args.num_features_in], offset=2.0)
+    elif args.base_dist == "conditional_bimodal":
+        base_dist = ConditionalBimodal(context_dim=args.num_context, latent_dim=args.num_features_in, hidden_dim=args.hidden_dim)
+    else:
+        raise ValueError(f"Invalid base distribution: {args.base_dist}")
+    flow = Flow(transform, base_dist)
+
+    return model, flow
+    
+
+
+def calculate_transformer_nf_loss(transformer, flow, batch, stop=None, stop_predictor=None):
+    device = next(transformer.parameters()).device
+
+    condition = batch["context"].to(device)
+    seq = batch["target"].to(device)
+    mask = batch["mask"].to(device) # (batch, max_length, num_context)
+    global_cond = batch["global_context"].to(device)
+
+    input_seq = seq[:, :-1] # (batch, max_length-1, num_features)
+    output = transformer(condition, input_seq, global_cond=global_cond) # (batch, max_length, num_context)
+
+    mask = mask[:,:,0].reshape(-1) # (batch * max_length)
+
+    _, _, num_features = seq.shape
+    _, _, num_context = output.shape
+
+    context_for_flow = output.reshape(-1, num_context)[mask]  # flow context (num_galaxies, num_context)
+    target_for_flow = seq.reshape(-1, num_features)[mask] # flow target (num_galaxies, num_features)
+    log_prob = flow.log_prob(target_for_flow, context=context_for_flow)
+    loss = - log_prob.mean() / num_features
+
+    if stop is not None and stop_predictor is not None:
+        stop = stop.to(device) # (batch, max_length)
+        stop = stop.reshape(-1)[mask] # (num_galaxies,)
+        stop_pred = stop_predictor(context_for_flow)
+        loss_stop = torch.nn.BCELoss()(stop_pred, stop)  # (num_galaxies,)
+        return loss, loss_stop
+    else:
+        return loss
+
+def generate_with_transformer_nf(transformer, flow, x_cond, stop_predictor=None, prob_threshold=0, stop_threshold=None):
+    transformer.eval()
+    flow.eval()
+    if stop_predictor is not None:
+        stop_predictor.eval()
+
+    batch_size = len(x_cond)
+    max_length = transformer.max_length
+    num_features = transformer.num_features_in
+
+    x_seq = torch.zeros(batch_size, max_length, num_features).to(x_cond.device)
+    stop_flags = torch.zeros(batch_size, dtype=torch.bool).to(x_cond.device)
+
+    for t in range(max_length):
+        with torch.no_grad():
+            context = transformer(x_cond, x_seq[:,:t,:]) # (batch, t+1, num_context)
+            context = context[:, -1, :] # (batch, num_context)
+            if prob_threshold > 0:
+                raise ValueError("# prob_threshold > 0 is not yet implemented")
+                num_samples = 1
+                embedded_context = flow._embedding_net(context)
+
+                noise, log_prob = flow._distribution.sample_and_log_prob(num_samples, context)
+                
+                if embedded_context is not None:
+                    # Merge the context dimension with sample dimension in order to apply the transform.
+                    noise = torchutils.merge_leading_dims(noise, num_dims=2)
+                    embedded_context = torchutils.repeat_rows(embedded_context, num_reps=1)
+
+                samples, _ = flow._transform.inverse(noise, context=embedded_context)
+
+                if embedded_context is not None:
+                    # Split the context dimension from sample dimension.
+                    samples = torchutils.split_leading_dim(samples, shape=[-1, num_samples])
+                
+            else:
+                samples = flow.sample(1, context=context) # (batch, 1, num_features)
+
+            x_seq[:, t, :] = samples.squeeze(1)
+               
+            if stop_predictor is not None:  
+                stop_prob = stop_predictor(context) # (batch, )
+                stop_now = stop_prob > stop_threshold
+                stop_flags = stop_now
+            
+            elif stop_threshold is not None:
+                if t > 0:
+                    stop_now = x_seq[:,t,0] < stop_threshold # (batch, )
+                    stop_flags |= stop_now
+
+            if t == 0:
+                if num_features > 1:
+                    x_seq[:,0,1] = torch.randn(batch_size).to(x_cond.device) * 1e-3 # Random distance for central
+                if num_features > 2:
+                    x_seq[:,0,2] = torch.randn(batch_size).to(x_cond.device) * 1e-3
+            
+            if stop_flags.all():
+                break
+
+    return x_seq
+
+def my_stop_predictor(args):
+    return StopPredictor(context_dim=args.num_context, hidden_dim=args.hidden_dim_stop)
+
+
+class StopPredictor(nn.Module):
+    def __init__(self, context_dim, hidden_dim=64):
+        super().__init__()
+        self.layers = nn.Sequential(
+            nn.Linear(context_dim, hidden_dim),
+            nn.LeakyReLU(),
+            nn.Linear(hidden_dim, 1),
+            nn.Sigmoid()
+        )
+    
+    def forward(self, context):
+        x = self.layers(context)
+        return x.squeeze(-1)  # (batch_size,) in [0, 1]
+
+class BimodalNormal(StandardNormal):
+    """A simple mixture of two Gaussians (bimodal) with fixed means."""
+
+    def __init__(self, shape, offset=2.0):
+        super().__init__(shape)
+        self.offset = offset
+
+    def _sample(self, num_samples, context=None):
+        if context is None:
+            device = self._log_z.device
+            half = num_samples // 2
+
+            samples1 = torch.randn(half, *self._shape, device=device) - self.offset
+            samples2 = torch.randn(num_samples - half, *self._shape, device=device) + self.offset
+            samples = torch.cat([samples1, samples2], dim=0)
+            return samples[torch.randperm(num_samples)]  
+
+        else:
+            context_size = context.shape[0]
+            total = context_size * num_samples
+            device = context.device
+            half = total // 2
+
+            samples1 = torch.randn(half, *self._shape, device=device) - self.offset
+            samples2 = torch.randn(total - half, *self._shape, device=device) + self.offset
+            samples = torch.cat([samples1, samples2], dim=0)
+            samples = samples[torch.randperm(total)]
+
+            return torchutils.split_leading_dim(samples, [context_size, num_samples])
+
+    def _log_prob(self, inputs, context=None):
+        if inputs.shape[1:] != self._shape:
+            raise ValueError(
+                "Expected input of shape {}, got {}".format(
+                    self._shape, inputs.shape[1:]
+                )
+            )
+
+        x = inputs.view(inputs.size(0), -1)
+
+        logp1 = -0.5 * torch.sum((x + self.offset)**2, dim=1) - self._log_z
+        logp2 = -0.5 * torch.sum((x - self.offset)**2, dim=1) - self._log_z
+
+        log_prob = torch.logaddexp(logp1, logp2) - np.log(2.0)
+        return log_prob
+
+    def _mean(self, context):
+        if context is None:
+            return self._log_z.new_zeros(self._shape)
+        else:
+            return context.new_zeros(context.shape[0], *self._shape)
+
+class ConditionalBimodal(nn.Module):
+    def __init__(self, context_dim, latent_dim, hidden_dim=64):
+        super().__init__()
+        self.latent_dim = latent_dim
+
+        # context → offset (1,), logit(weight) (1,), scale (1,)
+        self.net = nn.Sequential(
+            nn.Linear(context_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, 2 * latent_dim + 2)  # mean1, mean2, scale, weight_logits
+        )
+
+    def get_distribution(self, context):
+        batch_size = context.shape[0]
+        params = self.net(context)  # (batch, 2*latent_dim + 2)
+
+        # 分解
+        mean1 = params[:, :self.latent_dim]
+        mean2 = params[:, self.latent_dim:2*self.latent_dim]
+        scale_raw = params[:, 2*self.latent_dim:2*self.latent_dim+1]  # shape (batch, 1)
+        weight_logits = params[:, 2*self.latent_dim+1:]  # shape (batch, 1)
+
+        scale = F.softplus(scale_raw) + 1e-3
+
+        mix_logits = torch.cat([weight_logits, -weight_logits], dim=1)  # (batch, 2)
+        mix_dist = Categorical(logits=mix_logits)
+
+        means = torch.stack([mean1, mean2], dim=1)  # (batch, 2, latent_dim)
+        scales = scale.unsqueeze(1).expand_as(means)  # (batch, 2, latent_dim)
+
+        comp_dist = Independent(Normal(means, scales), 1)  # treat latent_dim as event dim
+        mixture = MixtureSameFamily(mix_dist, comp_dist)
+        return mixture
+
+    def sample(self, num_samples, context):
+        """
+        context: (batch, context_dim)
+        return: (batch, num_samples, latent_dim)
+        """
+        dist = self.get_distribution(context)  
+        samples = dist.sample((num_samples,))  # (num_samples, batch, latent_dim)
+        return samples.permute(1, 0, 2)  # (batch, num_samples, latent_dim)
+
+    def log_prob(self, x, context):
+        """
+        x: (batch, latent_dim)
+        context: (batch, context_dim)
+        return: (batch,)
+        """
+        dist = self.get_distribution(context)
+        return dist.log_prob(x)
+    
+
+from nflows.transforms.base import InputOutsideDomain, Transform, InverseTransform
+    
+class Sigmoid(Transform):
+    def __init__(self, temperature=1, eps=1e-6, learn_temperature=False):
+        super().__init__()
+        self.eps = eps
+        if learn_temperature:
+            self.temperature = nn.Parameter(torch.Tensor([temperature]))
+        else:
+            temperature = torch.Tensor([temperature])
+            self.register_buffer('temperature', temperature)
+
+    def forward(self, inputs, context=None):
+        temperature = self.temperature.to(inputs.device)
+
+        inputs = temperature * inputs
+        outputs = torch.sigmoid(inputs)
+
+        logabsdet = torchutils.sum_except_batch(
+            torch.log(temperature) - F.softplus(-inputs) - F.softplus(inputs)
+        )
+        return outputs, logabsdet
+
+    def inverse(self, inputs, context=None):
+        if torch.min(inputs) < 0 or torch.max(inputs) > 1:
+            raise InputOutsideDomain()
+
+        inputs = torch.clamp(inputs, self.eps, 1 - self.eps)
+
+        temperature = self.temperature.to(inputs.device)
+        outputs = (1 / temperature) * (torch.log(inputs) - torch.log1p(-inputs))
+        logabsdet = -torchutils.sum_except_batch(
+            torch.log(temperature)
+            - F.softplus(-temperature * outputs)
+            - F.softplus(temperature * outputs)
+        )
+        return outputs, logabsdet
+
+
+class Logit(InverseTransform):
+    def __init__(self, temperature=1, eps=1e-6):
+        super().__init__(Sigmoid(temperature=temperature, eps=eps))