Feat : Augmented Normalizing Flows (ANF)

## Description I builded ANF introduced in #28. (See also : [Arxiv - ANF](https://arxiv.org/abs/2002.07101)) ### Checked list - [x] Code convention : **Ruff** : I checked `augmented.py`, `__init__.py`. - [ ] CI using GitHub Action : Unfortunately, my docker was broken. `act`, `docker` didn't work temporarily. :cry:   ## Testing module I tested using `test_anf.py` (following code). ```python import numpy as np import torch from zuko.flows import ANF, HierarchicalANF def toy_example(): # Set random seed for reproducibility torch.manual_seed(42) # Create a simple ANF model features, noise_features, context_size = 2, 2, 4 # 2D data, 2D noise, Optional context dimension # Create model with 3 autoencoding transform layers model = ANF( features=features, noise_features=noise_features, context=context_size, steps=3, hidden_features=[64, 64], # Hidden layer sizes for the networks ) # Generate dummy context context = torch.randn(context_size) # Sample from the model n_samples = 1000 samples_with_noise = model(context).sample((n_samples,)) # Extract only the data part (not the noise) samples = samples_with_noise[:, :features] print(f"Generated {n_samples} samples with shape: {samples.shape}") # Calculate log probability of the samples log_probs = [] for i in range(5): # Just check a few samples # Generate random noise for the sample noise = torch.randn(noise_features) sample_with_noise = torch.cat([samples[i], noise]) # Calculate log probability log_prob = model(context).log_prob(sample_with_noise) log_probs.append(log_prob.item()) def hierarchical_example(): """Example with hierarchical ANF.""" print("\nRunning hierarchical ANF example...") # Set random seed for reproducibility torch.manual_seed(42) # Create a hierarchical ANF model features = 2 # 2D data noise_features_list = [2, 3, 2] # Three levels of latent variables context_size = 4 # Optional context dimension # Create hierarchical model model = HierarchicalANF( features=features, noise_features_list=noise_features_list, context=context_size, hidden_features=[64, 64], # Hidden layer sizes for the networks ) print(f"Created hierarchical model with {len(noise_features_list)} levels") # Generate dummy context context = torch.randn(context_size) # Sample from the model n_samples = 1000 samples_with_noise = model(context).sample((n_samples,)) # Extract only the data part (not the noise) samples = samples_with_noise[:, :features] print(f"Generated {n_samples} samples with shape: {samples.shape}") def density_estimation_example(): """Demonstrate density estimation on a 2D toy dataset.""" print("\nRunning density estimation example...") # Set random seed for reproducibility torch.manual_seed(42) # Create a toy 2D mixture of Gaussians dataset def sample_mog(n_samples=1000, n_components=5): """Generate samples from a mixture of 2D Gaussians.""" # Fixed centers for the Gaussians centers = [ [-4, -4], [-4, 4], [4, -4], [4, 4], [0, 0] ][:n_components] # Generate samples samples = [] for _ in range(n_samples): # Choose a component k = np.random.randint(n_components) # Sample from that Gaussian x = np.random.normal(centers[k][0], 0.5) y = np.random.normal(centers[k][1], 0.5) samples.append([x, y]) return torch.tensor(samples, dtype=torch.float32) # Generate dataset n_samples = 1000 data = sample_mog(n_samples=n_samples) print(f"Generated {n_samples} samples for training") # Create ANF model for density estimation features, noise_features = 2, 2 model = ANF( features=features, noise_features=noise_features, steps=5, hidden_features=[128, 128], ) # Train the model (simplified training loop) optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) n_epochs, batch_size = 100, 64 print(f"Training ANF model for {n_epochs} epochs...") for epoch in range(n_epochs): # Shuffle data perm = torch.randperm(n_samples) running_loss = 0.0 # Train in batches for i in range(0, n_samples, batch_size): indices = perm[i:i+batch_size] x_batch = data[indices] # Generate random noise e_batch = torch.randn(len(indices), noise_features) # Combine data and noise x_e_batch = torch.cat([x_batch, e_batch], dim=1) # Zero the gradients optimizer.zero_grad() # Compute loss (negative log-likelihood) log_prob = model().log_prob(x_e_batch) loss = -log_prob.mean() # Backpropagation loss.backward() optimizer.step() running_loss += loss.item() * len(indices) epoch_loss = running_loss / n_samples if (epoch + 1) % 20 == 0: print(f"Epoch {epoch+1}/{n_epochs}, Loss: {epoch_loss:.4f}") print("Training completed!") # Generate samples from the trained model with torch.no_grad(): samples_with_noise = model().sample((n_samples,)) if __name__ == "__main__": toy_example() hierarchical_example() density_estimation_example() ``` And the result is :