From 37a29ee25bc97b2b09f6f9cd9c4cd7f53a671aa4 Mon Sep 17 00:00:00 2001
From: Icaro Costa <icarosadero@gmail.com>
Date: Mon, 18 Sep 2023 23:45:46 -0300
Subject: [PATCH 01/14] first sketch

---
 requirements.txt       |  18 ++++++
 synthia/copulas/gan.py | 131 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 149 insertions(+)
 create mode 100644 requirements.txt
 create mode 100644 synthia/copulas/gan.py

diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000..b5702ed
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,18 @@
+pytest
+numpy
+scipy
+xarray
+bottleneck
+matplotlib
+seaborn
+jupyter
+pip
+setuptools
+sphinx
+sphinx-autobuild
+sphinx_rtd_theme
+myst-parser
+nbsphinx
+sphinxcontrib-bibtex==1
+sphinx-copybutton
+torch
\ No newline at end of file
diff --git a/synthia/copulas/gan.py b/synthia/copulas/gan.py
new file mode 100644
index 0000000..7511166
--- /dev/null
+++ b/synthia/copulas/gan.py
@@ -0,0 +1,131 @@
+from typing import Literal, Union
+from .copula import Copula
+import torch
+from torch.nn import Module
+import numpy as np
+
+class GANCopula(Copula, Module):
+    """
+    Learns Copula from data using Generative Adversarial Networks.
+    """
+    def __init__(
+            self,
+            generator_deep_layers: int = 2,
+            discriminator_deep_layers: int = 2,
+            device: Literal['cpu', 'cuda', 'auto'] = 'auto',
+            generator_optimizer: Literal['adam', 'sgd'] = 'adam',
+            discriminator_optimizer: Literal['adam', 'sgd'] = 'adam',
+        )->None:
+        """
+        Args:
+            generator_deep_layers (int): Number of layers for the generator.
+            discriminator_deep_layers (int): Number of layers for the discriminator.
+            device (Literal['cpu', 'cuda', 'auto']): Device to use for training. Use 'auto' to automatically select the device.
+            generator_optimizer (Literal['adam', 'sgd']): Optimizer to use for training the generator.
+            discriminator_optimizer (Literal['adam', 'sgd']): Optimizer to use for training the discriminator.
+        
+        Returns:
+            None
+        """
+        Module.__init__(self)
+        Copula.__init__(self)
+
+        self.device = None
+        match device:
+            case 'auto':
+                self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
+            case 'cpu':
+                self.device = 'cpu'
+            case 'cuda':
+                self.device = 'cuda'
+            case _:
+                raise ValueError(f"Invalid device {device}. Use 'cpu', 'cuda' or 'auto'.")
+
+        self.n_gen = generator_deep_layers
+        self.n_dis = discriminator_deep_layers
+        self.g_opt = generator_optimizer
+        self.d_opt = discriminator_optimizer
+        self.loss = torch.nn.BCELoss()
+    
+    def get_optimizer(self, optimizer_type: Literal['adam', 'sgd'], lr) -> torch.optim.Optimizer:
+        """
+        Get the optimizer for the generator or discriminator.
+
+        Args:
+            optimizer (Literal['adam', 'sgd']): The optimizer to get.
+
+        Returns:
+            torch.optim.Optimizer: The optimizer.
+        """
+        match optimizer_type:
+            case 'adam':
+                return torch.optim.Adam(self.generator.parameters(), lr=lr)
+            case 'sgd':
+                return torch.optim.SGD(self.generator.parameters(), lr=lr)
+            case _:
+                raise ValueError(f"Invalid optimizer {optimizer_type}. Use 'adam' or 'sgd'.")
+
+    
+    def fit(
+            self,
+            X: Union[torch.Tensor, np.array],
+            epochs: int = 1,
+            lr: float = 0.001,
+            batch_size: int = 32,
+        ) -> None:
+        """
+        Fits the copula to data.
+        
+        Args:
+            X (Union[torch.Tensor, np.array]): Input data in the shape (n_samples, n_features).
+            epochs (int): Number of epochs to train the GAN.
+            lr (float): Learning rate for the optimizer.
+            batch_size (int): Batch size for training.
+        
+        Returns:
+            None
+        """
+        with self.device:
+            if isinstance(X, np.ndarray):
+                X = torch.from_numpy(X).float()
+            elif isinstance(X, torch.Tensor):
+                X = X.float()
+            else:
+                raise TypeError(f"Invalid type {type(X)} for X. Use torch.Tensor or np.array.")
+
+            #Create Layers
+            self.generator = torch.nn.Sequential(
+                torch.nn.Linear(1, X.shape[1]),
+                torch.nn.ReLU(),
+                *[layer for _ in range(self.n_gen) for layer in [torch.nn.Linear(X.shape[1], X.shape[1]), torch.nn.ReLU()]],
+                torch.nn.Linear(X.shape[1], X.shape[1])
+            )
+
+            self.discriminator = torch.nn.Sequential(
+                torch.nn.Linear(X.shape[1], X.shape[1]),
+                torch.nn.ReLU(),
+                *[layer for _ in range(self.n_dis) for layer in [torch.nn.Linear(X.shape[1], X.shape[1]), torch.nn.ReLU()]],
+                torch.nn.Linear(X.shape[1], 1)
+            )
+
+            self.generator_optimizer = self.get_optimizer(self.g_opt, lr)
+            self.discriminator_optimizer = self.get_optimizer(self.d_opt, lr)
+
+            #Train GAN
+            for epoch in range(epochs):
+                for i in range(0, X.shape[0], batch_size):
+                    #Train Discriminator
+                    self.discriminator.zero_grad()
+                    real = X[i:i+batch_size]
+                    fake = self.generator(torch.rand(batch_size, 1))
+                    loss = self.loss(self.discriminator(real), torch.ones(batch_size, 1)) +\
+                          self.loss(self.discriminator(fake), torch.zeros(batch_size, 1))
+                    loss.backward()
+                    self.generator_optimizer.step()
+
+                    #Train Generator
+                    self.generator.zero_grad()
+                    fake = self.generator(torch.rand(batch_size, 1))
+                    loss = self.loss(self.discriminator(fake), torch.ones(batch_size, 1))
+                    loss.backward()
+                    self.discriminator_optimizer.step()
\ No newline at end of file

From 427098d2ab296148af28ef4b0592d4e337459a4e Mon Sep 17 00:00:00 2001
From: Icaro Costa <icarosadero@gmail.com>
Date: Mon, 18 Sep 2023 23:54:36 -0300
Subject: [PATCH 02/14] Add generate method

---
 synthia/copulas/gan.py | 15 ++++++++++++++-
 1 file changed, 14 insertions(+), 1 deletion(-)

diff --git a/synthia/copulas/gan.py b/synthia/copulas/gan.py
index 7511166..2cf1aa3 100644
--- a/synthia/copulas/gan.py
+++ b/synthia/copulas/gan.py
@@ -128,4 +128,17 @@ def fit(
                     fake = self.generator(torch.rand(batch_size, 1))
                     loss = self.loss(self.discriminator(fake), torch.ones(batch_size, 1))
                     loss.backward()
-                    self.discriminator_optimizer.step()
\ No newline at end of file
+                    self.discriminator_optimizer.step()
+        
+    def generate(self, n_samples: int) -> np.ndarray:
+        """
+        Generates samples from the copula.
+        
+        Args:
+            n_samples (int): Number of samples to generate.
+        
+        Returns:
+            np.ndarray: Samples from the copula.
+        """
+        with self.device:
+            return self.generator(torch.rand(n_samples, 1)).detach().numpy()
\ No newline at end of file

From b082eed22a7c00710a2becaf80f9bbddebdc4722 Mon Sep 17 00:00:00 2001
From: Icaro Costa <icarosadero@gmail.com>
Date: Sat, 23 Sep 2023 21:29:59 -0300
Subject: [PATCH 03/14] Fix for batch size

---
 synthia/copulas/gan.py | 59 ++++++++++++++++++++++++------------------
 1 file changed, 34 insertions(+), 25 deletions(-)

diff --git a/synthia/copulas/gan.py b/synthia/copulas/gan.py
index 2cf1aa3..6adf510 100644
--- a/synthia/copulas/gan.py
+++ b/synthia/copulas/gan.py
@@ -3,6 +3,8 @@
 import torch
 from torch.nn import Module
 import numpy as np
+import os
+os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
 
 class GANCopula(Copula, Module):
     """
@@ -10,16 +12,16 @@ class GANCopula(Copula, Module):
     """
     def __init__(
             self,
-            generator_deep_layers: int = 2,
-            discriminator_deep_layers: int = 2,
+            generator_deep_layers: list[int] = [32, 32],
+            discriminator_deep_layers: list[int] = [32, 32],
             device: Literal['cpu', 'cuda', 'auto'] = 'auto',
             generator_optimizer: Literal['adam', 'sgd'] = 'adam',
             discriminator_optimizer: Literal['adam', 'sgd'] = 'adam',
         )->None:
         """
         Args:
-            generator_deep_layers (int): Number of layers for the generator.
-            discriminator_deep_layers (int): Number of layers for the discriminator.
+            generator_deep_layers (list[int]): Number of deep layers for the generator. (Default: [32, 32])
+            discriminator_deep_layers (list[int]): Number of deep layers for the discriminator. (Default: [32, 32])
             device (Literal['cpu', 'cuda', 'auto']): Device to use for training. Use 'auto' to automatically select the device.
             generator_optimizer (Literal['adam', 'sgd']): Optimizer to use for training the generator.
             discriminator_optimizer (Literal['adam', 'sgd']): Optimizer to use for training the discriminator.
@@ -33,11 +35,11 @@ def __init__(
         self.device = None
         match device:
             case 'auto':
-                self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
+                self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
             case 'cpu':
-                self.device = 'cpu'
+                self.device = torch.device('cpu')
             case 'cuda':
-                self.device = 'cuda'
+                self.device = torch.device('cuda')
             case _:
                 raise ValueError(f"Invalid device {device}. Use 'cpu', 'cuda' or 'auto'.")
 
@@ -72,7 +74,7 @@ def fit(
             epochs: int = 1,
             lr: float = 0.001,
             batch_size: int = 32,
-        ) -> None:
+        ) -> tuple[torch.Tensor, torch.Tensor]:
         """
         Fits the copula to data.
         
@@ -83,8 +85,9 @@ def fit(
             batch_size (int): Batch size for training.
         
         Returns:
-            None
+            tuple[torch.Tensor, torch.Tensor]: Generator and discriminator loss.
         """
+        generator_loss, discriminator_loss = None, None
         with self.device:
             if isinstance(X, np.ndarray):
                 X = torch.from_numpy(X).float()
@@ -92,20 +95,22 @@ def fit(
                 X = X.float()
             else:
                 raise TypeError(f"Invalid type {type(X)} for X. Use torch.Tensor or np.array.")
+            X = X.to(self.device)
+            self.n_features = X.shape[1]
 
-            #Create Layers
+            #Create Generator
+            self.deep_gen_layers = [2] + self.n_gen + [self.n_features]
+            self.deep_gen_layers = zip(self.deep_gen_layers[:-1], self.deep_gen_layers[1:])
             self.generator = torch.nn.Sequential(
-                torch.nn.Linear(1, X.shape[1]),
-                torch.nn.ReLU(),
-                *[layer for _ in range(self.n_gen) for layer in [torch.nn.Linear(X.shape[1], X.shape[1]), torch.nn.ReLU()]],
-                torch.nn.Linear(X.shape[1], X.shape[1])
+                *[x for i, j in self.deep_gen_layers for x in [torch.nn.Linear(i, j), torch.nn.ReLU()]],
             )
 
+            #Create Discriminator
+            self.deep_dis_layers = [self.n_features] + self.n_dis + [1]
+            self.deep_dis_layers = zip(self.deep_dis_layers[:-1], self.deep_dis_layers[1:])
             self.discriminator = torch.nn.Sequential(
-                torch.nn.Linear(X.shape[1], X.shape[1]),
-                torch.nn.ReLU(),
-                *[layer for _ in range(self.n_dis) for layer in [torch.nn.Linear(X.shape[1], X.shape[1]), torch.nn.ReLU()]],
-                torch.nn.Linear(X.shape[1], 1)
+                *[x for i,j in self.deep_dis_layers for x in [torch.nn.Linear(i, j), torch.nn.ReLU()]],
+                torch.nn.Sigmoid(),
             )
 
             self.generator_optimizer = self.get_optimizer(self.g_opt, lr)
@@ -117,19 +122,23 @@ def fit(
                     #Train Discriminator
                     self.discriminator.zero_grad()
                     real = X[i:i+batch_size]
-                    fake = self.generator(torch.rand(batch_size, 1))
-                    loss = self.loss(self.discriminator(real), torch.ones(batch_size, 1)) +\
-                          self.loss(self.discriminator(fake), torch.zeros(batch_size, 1))
-                    loss.backward()
+                    actual_batch_size = real.shape[0]
+                    #At the end the batch size might be smaller than the specified batch size
+                    fake = self.generator(torch.rand(actual_batch_size, 2))
+                    generator_loss = self.loss(self.discriminator(real), torch.ones(actual_batch_size, 1)) +\
+                          self.loss(self.discriminator(fake), torch.zeros(actual_batch_size, 1))
+                    generator_loss.backward()
                     self.generator_optimizer.step()
 
                     #Train Generator
                     self.generator.zero_grad()
-                    fake = self.generator(torch.rand(batch_size, 1))
-                    loss = self.loss(self.discriminator(fake), torch.ones(batch_size, 1))
-                    loss.backward()
+                    fake = self.generator(torch.rand(batch_size, 2))
+                    discriminator_loss = self.loss(self.discriminator(fake), torch.ones(batch_size, 1))
+                    discriminator_loss.backward()
                     self.discriminator_optimizer.step()
         
+        return generator_loss, discriminator_loss
+        
     def generate(self, n_samples: int) -> np.ndarray:
         """
         Generates samples from the copula.

From b7eb0339ef115ec0ec59197a8f1cbd4fc783e5ef Mon Sep 17 00:00:00 2001
From: Icaro Costa <icarosadero@gmail.com>
Date: Sat, 23 Sep 2023 22:36:32 -0300
Subject: [PATCH 04/14] debug generator

---
 synthia/copulas/gan.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/synthia/copulas/gan.py b/synthia/copulas/gan.py
index 6adf510..19e427d 100644
--- a/synthia/copulas/gan.py
+++ b/synthia/copulas/gan.py
@@ -150,4 +150,4 @@ def generate(self, n_samples: int) -> np.ndarray:
             np.ndarray: Samples from the copula.
         """
         with self.device:
-            return self.generator(torch.rand(n_samples, 1)).detach().numpy()
\ No newline at end of file
+            return self.generator(torch.rand(n_samples, 2)).detach().cpu().numpy()
\ No newline at end of file

From c704dc8e640f12622e7c04691069866cd20c5c09 Mon Sep 17 00:00:00 2001
From: Icaro Costa <icarosadero@gmail.com>
Date: Sat, 23 Sep 2023 22:56:38 -0300
Subject: [PATCH 05/14] patch

---
 synthia/copulas/gan.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/synthia/copulas/gan.py b/synthia/copulas/gan.py
index 19e427d..4391baf 100644
--- a/synthia/copulas/gan.py
+++ b/synthia/copulas/gan.py
@@ -133,7 +133,7 @@ def fit(
                     #Train Generator
                     self.generator.zero_grad()
                     fake = self.generator(torch.rand(batch_size, 2))
-                    discriminator_loss = self.loss(self.discriminator(fake), torch.ones(batch_size, 1))
+                    discriminator_loss = self.loss(self.discriminator(fake), torch.zeros(batch_size, 1))
                     discriminator_loss.backward()
                     self.discriminator_optimizer.step()
         

From 3b3d56fadcc9e0e38f4991304678ebea1f15b3f9 Mon Sep 17 00:00:00 2001
From: Icaro Costa <icarosadero@gmail.com>
Date: Sat, 23 Sep 2023 23:30:16 -0300
Subject: [PATCH 06/14] Updated GAN training method

---
 synthia/copulas/gan.py | 61 ++++++++++++++++++++++++------------------
 1 file changed, 35 insertions(+), 26 deletions(-)

diff --git a/synthia/copulas/gan.py b/synthia/copulas/gan.py
index 4391baf..0939f7f 100644
--- a/synthia/copulas/gan.py
+++ b/synthia/copulas/gan.py
@@ -15,16 +15,14 @@ def __init__(
             generator_deep_layers: list[int] = [32, 32],
             discriminator_deep_layers: list[int] = [32, 32],
             device: Literal['cpu', 'cuda', 'auto'] = 'auto',
-            generator_optimizer: Literal['adam', 'sgd'] = 'adam',
-            discriminator_optimizer: Literal['adam', 'sgd'] = 'adam',
+            optimizer: Literal['adam', 'sgd'] = 'adam',
         )->None:
         """
         Args:
             generator_deep_layers (list[int]): Number of deep layers for the generator. (Default: [32, 32])
             discriminator_deep_layers (list[int]): Number of deep layers for the discriminator. (Default: [32, 32])
             device (Literal['cpu', 'cuda', 'auto']): Device to use for training. Use 'auto' to automatically select the device.
-            generator_optimizer (Literal['adam', 'sgd']): Optimizer to use for training the generator.
-            discriminator_optimizer (Literal['adam', 'sgd']): Optimizer to use for training the discriminator.
+            optimizer (Literal['adam', 'sgd']): Optimizer to use for training.
         
         Returns:
             None
@@ -45,16 +43,16 @@ def __init__(
 
         self.n_gen = generator_deep_layers
         self.n_dis = discriminator_deep_layers
-        self.g_opt = generator_optimizer
-        self.d_opt = discriminator_optimizer
-        self.loss = torch.nn.BCELoss()
+        self.opt = optimizer
+        self.loss = torch.nn.BCEWithLogitsLoss()
     
-    def get_optimizer(self, optimizer_type: Literal['adam', 'sgd'], lr) -> torch.optim.Optimizer:
+    def get_optimizer(self, optimizer_type: Literal['adam', 'sgd'], lr: float) -> torch.optim.Optimizer:
         """
         Get the optimizer for the generator or discriminator.
 
         Args:
             optimizer (Literal['adam', 'sgd']): The optimizer to get.
+            lr (float): The learning rate for the optimizer.
 
         Returns:
             torch.optim.Optimizer: The optimizer.
@@ -87,7 +85,7 @@ def fit(
         Returns:
             tuple[torch.Tensor, torch.Tensor]: Generator and discriminator loss.
         """
-        generator_loss, discriminator_loss = None, None
+        loss = None
         with self.device:
             if isinstance(X, np.ndarray):
                 X = torch.from_numpy(X).float()
@@ -109,35 +107,46 @@ def fit(
             self.deep_dis_layers = [self.n_features] + self.n_dis + [1]
             self.deep_dis_layers = zip(self.deep_dis_layers[:-1], self.deep_dis_layers[1:])
             self.discriminator = torch.nn.Sequential(
-                *[x for i,j in self.deep_dis_layers for x in [torch.nn.Linear(i, j), torch.nn.ReLU()]],
-                torch.nn.Sigmoid(),
+                *[x for i,j in self.deep_dis_layers for x in [torch.nn.Linear(i, j), torch.nn.ReLU()]]
             )
 
-            self.generator_optimizer = self.get_optimizer(self.g_opt, lr)
-            self.discriminator_optimizer = self.get_optimizer(self.d_opt, lr)
+            self.generator_optimizer = self.get_optimizer(self.opt, lr)
+            self.discriminator_optimizer = self.get_optimizer(self.opt, lr)
 
             #Train GAN
             for epoch in range(epochs):
                 for i in range(0, X.shape[0], batch_size):
-                    #Train Discriminator
-                    self.discriminator.zero_grad()
+
                     real = X[i:i+batch_size]
-                    actual_batch_size = real.shape[0]
                     #At the end the batch size might be smaller than the specified batch size
+                    actual_batch_size = real.shape[0]
                     fake = self.generator(torch.rand(actual_batch_size, 2))
-                    generator_loss = self.loss(self.discriminator(real), torch.ones(actual_batch_size, 1)) +\
-                          self.loss(self.discriminator(fake), torch.zeros(actual_batch_size, 1))
-                    generator_loss.backward()
-                    self.generator_optimizer.step()
 
-                    #Train Generator
-                    self.generator.zero_grad()
-                    fake = self.generator(torch.rand(batch_size, 2))
-                    discriminator_loss = self.loss(self.discriminator(fake), torch.zeros(batch_size, 1))
-                    discriminator_loss.backward()
+                    self.discriminator.zero_grad()
+                    disc_real = self.discriminator(real)
+                    disc_fake = self.discriminator(fake)
+
+                    loss_real = self.loss(disc_real, torch.ones(actual_batch_size))
+                    #Inserting 1 for y effectively calculates log(D(x)), as the second term on
+                    #the loss vanishes as it is proportional to (1-y)
+                    loss_real.backward()
+
+                    loss_fake = self.loss(disc_fake, torch.zeros(actual_batch_size))
+                    #In a similar fashion, inserting 0 for y effectively calculates log(1-D(G(z)))
+                    loss_fake.backward()
+                    loss_discriminator = loss_real + loss_fake
+
                     self.discriminator_optimizer.step()
+
+                    self.generator.zero_grad()
+                    new_disc_fake = self.discriminator(fake)
+                    loss_generator = self.loss(new_disc_fake, torch.ones(actual_batch_size))
+                    loss_generator.backward()
+                    self.generator_optimizer.step()
+                    
+
         
-        return generator_loss, discriminator_loss
+        return loss_discriminator, loss_generator
         
     def generate(self, n_samples: int) -> np.ndarray:
         """

From c7750561089eeb25d2d67136f94b102b88768b30 Mon Sep 17 00:00:00 2001
From: Icaro Costa <icarosadero@gmail.com>
Date: Sat, 23 Sep 2023 23:42:44 -0300
Subject: [PATCH 07/14] patch dimensions

---
 synthia/copulas/gan.py | 10 ++++------
 1 file changed, 4 insertions(+), 6 deletions(-)

diff --git a/synthia/copulas/gan.py b/synthia/copulas/gan.py
index 0939f7f..03585e5 100644
--- a/synthia/copulas/gan.py
+++ b/synthia/copulas/gan.py
@@ -126,21 +126,19 @@ def fit(
                     disc_real = self.discriminator(real)
                     disc_fake = self.discriminator(fake)
 
-                    loss_real = self.loss(disc_real, torch.ones(actual_batch_size))
+                    loss_real = self.loss(disc_real, torch.ones(actual_batch_size, 1))
                     #Inserting 1 for y effectively calculates log(D(x)), as the second term on
                     #the loss vanishes as it is proportional to (1-y)
-                    loss_real.backward()
 
-                    loss_fake = self.loss(disc_fake, torch.zeros(actual_batch_size))
+                    loss_fake = self.loss(disc_fake, torch.zeros(actual_batch_size, 1))
                     #In a similar fashion, inserting 0 for y effectively calculates log(1-D(G(z)))
-                    loss_fake.backward()
                     loss_discriminator = loss_real + loss_fake
-
+                    loss_discriminator.backward(retain_graph=True)
                     self.discriminator_optimizer.step()
 
                     self.generator.zero_grad()
                     new_disc_fake = self.discriminator(fake)
-                    loss_generator = self.loss(new_disc_fake, torch.ones(actual_batch_size))
+                    loss_generator = self.loss(new_disc_fake, torch.ones(actual_batch_size, 1))
                     loss_generator.backward()
                     self.generator_optimizer.step()
                     

From 6edd1ec73d06b8be5c49a2efd9a8f61259beef9e Mon Sep 17 00:00:00 2001
From: Icaro Costa <icarosadero@gmail.com>
Date: Sun, 24 Sep 2023 05:34:49 -0300
Subject: [PATCH 08/14] Updated training

---
 synthia/copulas/gan.py | 84 +++++++++++++++++++++++++++++++-----------
 1 file changed, 63 insertions(+), 21 deletions(-)

diff --git a/synthia/copulas/gan.py b/synthia/copulas/gan.py
index 03585e5..c92bd63 100644
--- a/synthia/copulas/gan.py
+++ b/synthia/copulas/gan.py
@@ -16,6 +16,8 @@ def __init__(
             discriminator_deep_layers: list[int] = [32, 32],
             device: Literal['cpu', 'cuda', 'auto'] = 'auto',
             optimizer: Literal['adam', 'sgd'] = 'adam',
+            latent_dim: int = 10,
+            generator_fake_size: int = 1000,
         )->None:
         """
         Args:
@@ -23,6 +25,8 @@ def __init__(
             discriminator_deep_layers (list[int]): Number of deep layers for the discriminator. (Default: [32, 32])
             device (Literal['cpu', 'cuda', 'auto']): Device to use for training. Use 'auto' to automatically select the device.
             optimizer (Literal['adam', 'sgd']): Optimizer to use for training.
+            latent_dim (int): Dimension of the latent space. (Default: 10).
+            generator_fake_size (int): Number of fake samples to generate for the generator. (Default: 1000)
         
         Returns:
             None
@@ -45,6 +49,22 @@ def __init__(
         self.n_dis = discriminator_deep_layers
         self.opt = optimizer
         self.loss = torch.nn.BCEWithLogitsLoss()
+        self.latent_dim = latent_dim
+        self.fake_batch_size = generator_fake_size
+    
+    def init_weights(self, m: torch.nn.Module) -> None:
+        """
+        Initializes weights for the generator and discriminator.
+        
+        Args:
+            m (torch.nn.Module): Module to initialize weights for.
+        
+        Returns:
+            None
+        """
+        if isinstance(m, torch.nn.Linear):
+            torch.nn.init.xavier_uniform_(m.weight)
+            m.bias.data.fill_(0.1)
     
     def get_optimizer(self, optimizer_type: Literal['adam', 'sgd'], lr: float) -> torch.optim.Optimizer:
         """
@@ -69,23 +89,34 @@ def get_optimizer(self, optimizer_type: Literal['adam', 'sgd'], lr: float) -> to
     def fit(
             self,
             X: Union[torch.Tensor, np.array],
-            epochs: int = 1,
+            global_iterations: int = 1,
+            discriminator_iterations: int = 1,
             lr: float = 0.001,
             batch_size: int = 32,
+            dropout_proba: float = 0.1,
         ) -> tuple[torch.Tensor, torch.Tensor]:
         """
         Fits the copula to data.
+
+        Loss calculation:
+
+        .. math::
+            \\begin{cases}
+            \min_D L_D(D, \mu_G) = -\mathbb{E}_{x\sim \mu_{G}}[\ln (1-D(x))] - \mathbb{E}_{x\sim \mu_{\\text{ref}}}[\ln (D(x))]\\
+            \min_G L_G(D, \mu_G) = -\mathbb{E}_{x\sim \mu_G}[\ln (D(x))]
+            \end{cases}
         
         Args:
             X (Union[torch.Tensor, np.array]): Input data in the shape (n_samples, n_features).
-            epochs (int): Number of epochs to train the GAN.
+            global_iterations (int): Number of iterations to train the GAN.
+            discriminator_iterations (int): Number of iterations to train the discriminator for each global iteration.
             lr (float): Learning rate for the optimizer.
             batch_size (int): Batch size for training.
+            dropout_proba (float): Dropout probability for the generator and discriminator.
         
         Returns:
             tuple[torch.Tensor, torch.Tensor]: Generator and discriminator loss.
         """
-        loss = None
         with self.device:
             if isinstance(X, np.ndarray):
                 X = torch.from_numpy(X).float()
@@ -97,50 +128,61 @@ def fit(
             self.n_features = X.shape[1]
 
             #Create Generator
-            self.deep_gen_layers = [2] + self.n_gen + [self.n_features]
-            self.deep_gen_layers = zip(self.deep_gen_layers[:-1], self.deep_gen_layers[1:])
+            self.deep_gen_layers = [self.latent_dim] + self.n_gen + [self.n_features]
+            self.deep_gen_layers = list(zip(self.deep_gen_layers[:-1], self.deep_gen_layers[1:]))
             self.generator = torch.nn.Sequential(
                 *[x for i, j in self.deep_gen_layers for x in [torch.nn.Linear(i, j), torch.nn.ReLU()]],
+                torch.nn.Dropout(dropout_proba),
+                torch.nn.Linear(self.deep_gen_layers[-1][1], self.deep_gen_layers[-1][1]),
             )
+            #Initialize weights
+            self.generator.apply(self.init_weights)
 
             #Create Discriminator
-            self.deep_dis_layers = [self.n_features] + self.n_dis + [1]
-            self.deep_dis_layers = zip(self.deep_dis_layers[:-1], self.deep_dis_layers[1:])
+            self.deep_dis_layers = [self.n_features] + self.n_dis + [2]
+            self.deep_dis_layers = list(zip(self.deep_dis_layers[:-1], self.deep_dis_layers[1:]))
             self.discriminator = torch.nn.Sequential(
-                *[x for i,j in self.deep_dis_layers for x in [torch.nn.Linear(i, j), torch.nn.ReLU()]]
+                *[x for i,j in self.deep_dis_layers for x in [torch.nn.Linear(i, j), torch.nn.ReLU()]],
+                torch.nn.Dropout(dropout_proba),
+                torch.nn.Linear(2,1)
             )
+            #Initialize weights
+            self.discriminator.apply(self.init_weights)
 
             self.generator_optimizer = self.get_optimizer(self.opt, lr)
             self.discriminator_optimizer = self.get_optimizer(self.opt, lr)
 
             #Train GAN
-            for epoch in range(epochs):
-                for i in range(0, X.shape[0], batch_size):
+            for _ in range(global_iterations):
+                for i in range(discriminator_iterations):
 
-                    real = X[i:i+batch_size]
+                    #Sample from X
+                    real = X[torch.randperm(X.shape[0])[:batch_size]]
                     #At the end the batch size might be smaller than the specified batch size
                     actual_batch_size = real.shape[0]
-                    fake = self.generator(torch.rand(actual_batch_size, 2))
+                    fake = self.generator(torch.rand(actual_batch_size, self.latent_dim))
 
                     self.discriminator.zero_grad()
                     disc_real = self.discriminator(real)
                     disc_fake = self.discriminator(fake)
 
                     loss_real = self.loss(disc_real, torch.ones(actual_batch_size, 1))
-                    #Inserting 1 for y effectively calculates log(D(x)), as the second term on
-                    #the loss vanishes as it is proportional to (1-y)
+                    #Inserting 1 for y effectively calculates log(D(x)), as the other term on
+                    #the loss vanishes as it is proportional to y
 
                     loss_fake = self.loss(disc_fake, torch.zeros(actual_batch_size, 1))
-                    #In a similar fashion, inserting 0 for y effectively calculates log(1-D(G(z)))
+                    #In a similar fashion, inserting 0 for y yields log(1-D(G(z)))
                     loss_discriminator = loss_real + loss_fake
                     loss_discriminator.backward(retain_graph=True)
                     self.discriminator_optimizer.step()
 
-                    self.generator.zero_grad()
-                    new_disc_fake = self.discriminator(fake)
-                    loss_generator = self.loss(new_disc_fake, torch.ones(actual_batch_size, 1))
-                    loss_generator.backward()
-                    self.generator_optimizer.step()
+                self.generator.zero_grad()
+                new_fake = self.generator(torch.rand(self.fake_batch_size, self.latent_dim))
+                new_disc_fake = self.discriminator(new_fake)
+                loss_generator = self.loss(new_disc_fake, torch.ones(self.fake_batch_size, 1))
+                #Inserting 0 for y calculates log(D(G(z)))
+                loss_generator.backward()
+                self.generator_optimizer.step()
                     
 
         
@@ -157,4 +199,4 @@ def generate(self, n_samples: int) -> np.ndarray:
             np.ndarray: Samples from the copula.
         """
         with self.device:
-            return self.generator(torch.rand(n_samples, 2)).detach().cpu().numpy()
\ No newline at end of file
+            return self.generator(torch.rand(n_samples, self.latent_dim)).detach().cpu().numpy()
\ No newline at end of file

From fbd29037642effac5139e5caa512c363b3d04db3 Mon Sep 17 00:00:00 2001
From: Icaro Costa <icarosadero@gmail.com>
Date: Sun, 24 Sep 2023 21:06:04 -0300
Subject: [PATCH 09/14] Applying GAN training best practices from
 https://github.com/soumith/ganhacks

---
 synthia/copulas/gan.py | 25 +++++++++++++++----------
 1 file changed, 15 insertions(+), 10 deletions(-)

diff --git a/synthia/copulas/gan.py b/synthia/copulas/gan.py
index c92bd63..fdeb342 100644
--- a/synthia/copulas/gan.py
+++ b/synthia/copulas/gan.py
@@ -107,7 +107,7 @@ def fit(
             \end{cases}
         
         Args:
-            X (Union[torch.Tensor, np.array]): Input data in the shape (n_samples, n_features).
+            X (Union[torch.Tensor, np.array]): Input data in the shape (n_samples, n_features). Values must be within [-1,1]
             global_iterations (int): Number of iterations to train the GAN.
             discriminator_iterations (int): Number of iterations to train the discriminator for each global iteration.
             lr (float): Learning rate for the optimizer.
@@ -117,6 +117,10 @@ def fit(
         Returns:
             tuple[torch.Tensor, torch.Tensor]: Generator and discriminator loss.
         """
+
+        #Check X values
+        assert X.max() <= 1 and X.min() >= -1, "Values must be within [-1,1]. Try using np.tanh first."
+
         with self.device:
             if isinstance(X, np.ndarray):
                 X = torch.from_numpy(X).float()
@@ -131,9 +135,10 @@ def fit(
             self.deep_gen_layers = [self.latent_dim] + self.n_gen + [self.n_features]
             self.deep_gen_layers = list(zip(self.deep_gen_layers[:-1], self.deep_gen_layers[1:]))
             self.generator = torch.nn.Sequential(
-                *[x for i, j in self.deep_gen_layers for x in [torch.nn.Linear(i, j), torch.nn.ReLU()]],
+                *[x for i, j in self.deep_gen_layers for x in [torch.nn.Linear(i, j), torch.nn.LeakyReLU()]],
                 torch.nn.Dropout(dropout_proba),
                 torch.nn.Linear(self.deep_gen_layers[-1][1], self.deep_gen_layers[-1][1]),
+                torch.nn.Tanh()
             )
             #Initialize weights
             self.generator.apply(self.init_weights)
@@ -142,7 +147,7 @@ def fit(
             self.deep_dis_layers = [self.n_features] + self.n_dis + [2]
             self.deep_dis_layers = list(zip(self.deep_dis_layers[:-1], self.deep_dis_layers[1:]))
             self.discriminator = torch.nn.Sequential(
-                *[x for i,j in self.deep_dis_layers for x in [torch.nn.Linear(i, j), torch.nn.ReLU()]],
+                *[x for i,j in self.deep_dis_layers for x in [torch.nn.Linear(i, j), torch.nn.LeakyReLU()]],
                 torch.nn.Dropout(dropout_proba),
                 torch.nn.Linear(2,1)
             )
@@ -160,27 +165,27 @@ def fit(
                     real = X[torch.randperm(X.shape[0])[:batch_size]]
                     #At the end the batch size might be smaller than the specified batch size
                     actual_batch_size = real.shape[0]
-                    fake = self.generator(torch.rand(actual_batch_size, self.latent_dim))
+                    fake = self.generator(torch.randn(actual_batch_size, self.latent_dim))
 
                     self.discriminator.zero_grad()
                     disc_real = self.discriminator(real)
                     disc_fake = self.discriminator(fake)
 
                     loss_real = self.loss(disc_real, torch.ones(actual_batch_size, 1))
-                    #Inserting 1 for y effectively calculates log(D(x)), as the other term on
+                    #Inserting 1 for y effectively calculates -log(D(x)), as the other term on
                     #the loss vanishes as it is proportional to y
 
                     loss_fake = self.loss(disc_fake, torch.zeros(actual_batch_size, 1))
-                    #In a similar fashion, inserting 0 for y yields log(1-D(G(z)))
+                    #In a similar fashion, inserting 0 for y yields -log(1-D(G(z)))
                     loss_discriminator = loss_real + loss_fake
-                    loss_discriminator.backward(retain_graph=True)
+                    loss_discriminator.backward()
                     self.discriminator_optimizer.step()
 
                 self.generator.zero_grad()
-                new_fake = self.generator(torch.rand(self.fake_batch_size, self.latent_dim))
+                new_fake = self.generator(torch.randn(self.fake_batch_size, self.latent_dim))
                 new_disc_fake = self.discriminator(new_fake)
                 loss_generator = self.loss(new_disc_fake, torch.ones(self.fake_batch_size, 1))
-                #Inserting 0 for y calculates log(D(G(z)))
+                #Inserting 0 for y calculates -log(D(G(z)))
                 loss_generator.backward()
                 self.generator_optimizer.step()
                     
@@ -199,4 +204,4 @@ def generate(self, n_samples: int) -> np.ndarray:
             np.ndarray: Samples from the copula.
         """
         with self.device:
-            return self.generator(torch.rand(n_samples, self.latent_dim)).detach().cpu().numpy()
\ No newline at end of file
+            return self.generator(torch.randn(n_samples, self.latent_dim)).detach().cpu().numpy()
\ No newline at end of file

From a066f9b7f90cede6e179475030acb6b931f7759d Mon Sep 17 00:00:00 2001
From: Icaro Costa <icarosadero@gmail.com>
Date: Sun, 24 Sep 2023 21:08:05 -0300
Subject: [PATCH 10/14] Dropout just in generator

---
 synthia/copulas/gan.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/synthia/copulas/gan.py b/synthia/copulas/gan.py
index fdeb342..e272333 100644
--- a/synthia/copulas/gan.py
+++ b/synthia/copulas/gan.py
@@ -112,7 +112,7 @@ def fit(
             discriminator_iterations (int): Number of iterations to train the discriminator for each global iteration.
             lr (float): Learning rate for the optimizer.
             batch_size (int): Batch size for training.
-            dropout_proba (float): Dropout probability for the generator and discriminator.
+            dropout_proba (float): Dropout probability for the generator.
         
         Returns:
             tuple[torch.Tensor, torch.Tensor]: Generator and discriminator loss.
@@ -148,7 +148,6 @@ def fit(
             self.deep_dis_layers = list(zip(self.deep_dis_layers[:-1], self.deep_dis_layers[1:]))
             self.discriminator = torch.nn.Sequential(
                 *[x for i,j in self.deep_dis_layers for x in [torch.nn.Linear(i, j), torch.nn.LeakyReLU()]],
-                torch.nn.Dropout(dropout_proba),
                 torch.nn.Linear(2,1)
             )
             #Initialize weights

From e4d6f938a831439b3265028d26fc49dc27144892 Mon Sep 17 00:00:00 2001
From: Icaro Costa <icarosadero@gmail.com>
Date: Sun, 24 Sep 2023 21:26:49 -0300
Subject: [PATCH 11/14] Moved GAN to generators

---
 synthia/{copulas => generators}/gan.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)
 rename synthia/{copulas => generators}/gan.py (98%)

diff --git a/synthia/copulas/gan.py b/synthia/generators/gan.py
similarity index 98%
rename from synthia/copulas/gan.py
rename to synthia/generators/gan.py
index e272333..c6f16da 100644
--- a/synthia/copulas/gan.py
+++ b/synthia/generators/gan.py
@@ -6,7 +6,7 @@
 import os
 os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
 
-class GANCopula(Copula, Module):
+class GAN(Module):
     """
     Learns Copula from data using Generative Adversarial Networks.
     """
@@ -31,8 +31,7 @@ def __init__(
         Returns:
             None
         """
-        Module.__init__(self)
-        Copula.__init__(self)
+        super.__init__()
 
         self.device = None
         match device:

From 86e0cd152cf1917cb8d970f3adbc733f42fdb041 Mon Sep 17 00:00:00 2001
From: Icaro Costa <icarosadero@gmail.com>
Date: Wed, 27 Sep 2023 18:56:07 -0300
Subject: [PATCH 12/14] Using GRU as discriminant

---
 requirements.txt              |   3 +-
 synthia/copulas/copula_gan.py | 213 ++++++++++++++++++++++++++++++++++
 synthia/generators/gan.py     | 126 ++++++++++++--------
 3 files changed, 292 insertions(+), 50 deletions(-)
 create mode 100644 synthia/copulas/copula_gan.py

diff --git a/requirements.txt b/requirements.txt
index b5702ed..75c6ed8 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -15,4 +15,5 @@ myst-parser
 nbsphinx
 sphinxcontrib-bibtex==1
 sphinx-copybutton
-torch
\ No newline at end of file
+torch
+derivative
\ No newline at end of file
diff --git a/synthia/copulas/copula_gan.py b/synthia/copulas/copula_gan.py
new file mode 100644
index 0000000..c64f2eb
--- /dev/null
+++ b/synthia/copulas/copula_gan.py
@@ -0,0 +1,213 @@
+from typing import Literal, Union, Annotated
+from .copula import Copula
+from ..generators.gan import GAN
+from derivative import dxdt
+import numpy as np
+import torch
+import zfit
+
+def derivative(
+        x: Annotated[torch.Tensor, 'CDF'],
+        t: Annotated[torch.Tensor, 'Support'] = None,
+        device: torch.device = torch.device('cpu')
+    ) -> torch.Tensor:
+    """
+    Derivative of the CDF with respect to support t.
+
+    Args:
+        x (Annotated[torch.Tensor, 'CDF']): CDF of the data.
+        t (Annotated[torch.Tensor, 'Support']): Support of the data. Will default to the interval [0,1].
+        device (torch.device): Device to use for computation. (Default: torch.device('cpu'))
+    """
+    if t is None:
+        t = torch.linspace(0,1,x.shape[0], device=device)
+    return dxdt(x.T, t).T
+
+def compute_cdf(kde_model, n_samples=1000):
+    """
+    Computes the CDF of a KDE model using the cumulative trapezoidal rule.
+
+    Args:
+        kde_model: KDE model to compute CDF of.
+        n_samples: Number of samples to use for computing the CDF. (Default: 1000)
+    """
+    lower_limit = float(kde_model.space.limits[0])
+    upper_limit = float(kde_model.space.limits[1])
+
+    support = torch.linspace(lower_limit, upper_limit, n_samples)
+    pdf_values = torch.Tensor([float(kde_model.pdf(x)) for x in support])
+
+    # Compute the CDF using cumulative trapezoidal rule
+    cdf_values = torch.cumulative_trapezoid(pdf_values, dx=(support[1] - support[0]))
+
+    # Normalize the CDF to [0, 1] by dividing by the final value
+    cdf_values /= cdf_values[-1].clone()
+    #Add zero on beginning
+    cdf_values = torch.cat((torch.zeros(1), cdf_values))
+
+    return cdf_values, support
+
+def map_cdf(cdf_values, support)->function:
+    """
+    Maps the data in the support to its corresponding CDF value.
+
+    Args:
+        cdf_values: CDF values to map to.
+        support: Support of the data.
+    """
+    return lambda x: np.interp(x, support, cdf_values)
+
+def inverse_map_cdf(cdf_values, support):
+    """
+    Computes the inverse CDF of a KDE model using linear interpolation.
+
+    Args:
+        cdf_values: CDF values to compute inverse of.
+        support: Support of the data.
+    """
+    return lambda x: np.interp(x, cdf_values, support)
+
+class CopulaGAN(Copula, GAN):
+    """
+    Learns Copula from data using Generative Adversarial Networks.
+    """
+    def __init__(
+            self,
+            generator_deep_layers: list[int] = [32, 32],
+            discriminator_deep_layers: list[int] = [32, 32],
+            device: Literal['cpu', 'cuda', 'auto'] = 'auto',
+            optimizer: Literal['adam', 'sgd'] = 'adam',
+            generator_fake_size: int = 1000,
+        ) -> None:
+        super().__init__(
+            generator_deep_layers,
+            discriminator_deep_layers,
+            device,
+            optimizer,
+            generator_fake_size
+        )
+
+    def initialize(self, X: torch.Tensor, **kwargs):
+        # Generate marginals and sample space
+        n_features = X.shape[1]
+
+        # Create target KDEs and CDFs
+        target_kdes = []
+        target_pdfs = []
+        cdfs = []
+        real_spaces = []
+        real_U = []
+        for i in range(n_features):
+            obs_space = zfit.Space('X' + str(i), limits=(X[:, i].min(), X[:, i].max()))
+            kde = zfit.pdf.KDE1DimGrid(
+                obs=obs_space,
+                data=zfit.Data.from_tensor(tensor=X[:, i], obs=obs_space)
+            )
+            cdf, support = compute_cdf(kde, **kwargs)
+            real_U.append(map_cdf(cdf, support)(X[:,i]))
+            pdf = kde.pdf(obs_space, torch.linspace(X[:,i].min(),X[:,i].max(),n_samples))
+            target_kdes.append(kde)
+            cdfs.append(cdf)
+            real_spaces.append(obs_space)
+            target_pdfs.append(pdf)
+
+        self.real_space = zfit.dimension.combine_spaces(*real_spaces)
+        self.target_kdes = target_kdes
+        self.target_cdf = torch.stack(cdfs, dim=1)
+        self.target_pdf = torch.stack(target_pdfs, dim=1)
+        self.real_U = torch.stack(real_U, dim=1)
+
+    def construct_discriminator(
+            self,
+            n_features: int,
+            hidden_size: int = 32,
+        ) -> torch.nn.Sequential:
+        """
+        Constructs the discriminator for the GAN. It is a GRU that receives samples from CDFs as input
+        and outputs a single value.
+
+        Input (batch_size, sequence_length, n_features)
+        Output (batch_size, 1)
+
+        Args:
+            n_features (int): Number of features for the discriminator.
+            hidden_size (int): Hidden size for the GRU. (Default: 32)
+        """
+        self.discriminator = torch.nn.Sequential(
+            torch.nn.GRU(
+                input_size = n_features,
+                hidden_size = hidden_size,
+                batch_first=True),
+            torch.nn.Linear(hidden_size, 1)
+        ).to(self.device)
+        self.discriminator.apply(self.init_weights)
+        return self.discriminator
+
+    def fit(
+            self,
+            X: Union[torch.Tensor, np.ndarray],
+            global_iterations: int,
+            discriminator_iterations: int = 1,
+            lr: float = 0.001,
+            batch_size: int = 32,
+            generator_kwargs: dict = {},
+            discriminator_kwargs: dict = {},
+            CDFN: int = 1000
+        ):
+        """
+        Uses GAN to learn a Copula flow from vector U sampled from the marginals of an arbitrary dataset.
+
+        If each variable $X_i$ in a dataset has a CDF denoted by $F_i$, then the marginals are defined as $U_i = F_i(X_i)$.
+
+        Args:
+            X (Union[torch.Tensor, np.ndarray]): Dataset in format (n_samples, n_features).
+            global_iterations (int): Number of global iterations to train for.
+            discriminator_iterations (int): Number of discriminator iterations per global iteration. (Default: 1)
+            lr (float): Learning rate for the optimizer. (Default: 0.001)
+            batch_size (int): Batch size for training. (Default: 32)
+            generator_kwargs (dict): Keyword arguments to pass to the generator. (Default: {})
+            discriminator_kwargs (dict): Keyword arguments to pass to the discriminator. (Default: {})
+            CDFN (int): Number of samples to use for computing the CDF. (Default: 1000)
+        """
+        if isinstance(X, np.ndarray):
+            X = torch.from_numpy(X).float()
+        elif isinstance(X, torch.Tensor):
+            X = X.float()
+        else:
+            raise TypeError(f"Invalid type {type(X)} for X. Use torch.Tensor or np.array.")
+
+        self.initialize(X, n_samples = CDFN)
+
+        self.construct_generator(X.shape[1], activation='sigmoid', **generator_kwargs)
+        self.construct_discriminator(X.shape[1], **discriminator_kwargs)
+
+        with self.device:
+
+            self.generator_optimizer = self.get_optimizer(self.opt, lr)
+            self.discriminator_optimizer = self.get_optimizer(self.opt, lr)
+
+            #Train
+            for _ in range(global_iterations):
+                for __ in range(discriminator_iterations):
+                    #Train Discriminator
+                    rand = torch.rand(batch_size, self.latent_dimension, device=self.device)
+                    fake_U = self.generator(rand)
+                    real_U = self.real_U[torch.randperm(self.real_U.shape[0])[:batch_size]]
+                    real_discriminant, _ = self.discriminator(real_U)
+                    fake_discriminant, _ = self.discriminator(fake_U)
+                    self.discriminator.zero_grad()
+                    (
+                        self.loss(real_discriminant, torch.ones_like(real_discriminant)) +\
+                        self.loss(fake_discriminant, torch.zeros_like(fake_discriminant))
+                    ).backward()
+                    self.discriminator_optimizer.step()
+                self.generator.zero_grad()
+                rand = torch.rand(batch_size, self.latent_dimension, device=self.device)
+                fake_U = self.generator(rand)
+                fake_discriminant, _ = self.discriminator(fake_U)
+                loss = self.loss(fake_discriminant, torch.ones_like(fake_discriminant))
+                loss.backward()
+                self.generator_optimizer.step()
+    
+    def generate(self, n_samples: int, **kws) -> np.ndarray:
+        return super().generate(n_samples, **kws)
\ No newline at end of file
diff --git a/synthia/generators/gan.py b/synthia/generators/gan.py
index c6f16da..6c49464 100644
--- a/synthia/generators/gan.py
+++ b/synthia/generators/gan.py
@@ -1,5 +1,4 @@
 from typing import Literal, Union
-from .copula import Copula
 import torch
 from torch.nn import Module
 import numpy as np
@@ -8,7 +7,7 @@
 
 class GAN(Module):
     """
-    Learns Copula from data using Generative Adversarial Networks.
+    PyTorch implementation of a Generative Adversarial Network.
     """
     def __init__(
             self,
@@ -16,7 +15,6 @@ def __init__(
             discriminator_deep_layers: list[int] = [32, 32],
             device: Literal['cpu', 'cuda', 'auto'] = 'auto',
             optimizer: Literal['adam', 'sgd'] = 'adam',
-            latent_dim: int = 10,
             generator_fake_size: int = 1000,
         )->None:
         """
@@ -25,13 +23,12 @@ def __init__(
             discriminator_deep_layers (list[int]): Number of deep layers for the discriminator. (Default: [32, 32])
             device (Literal['cpu', 'cuda', 'auto']): Device to use for training. Use 'auto' to automatically select the device.
             optimizer (Literal['adam', 'sgd']): Optimizer to use for training.
-            latent_dim (int): Dimension of the latent space. (Default: 10).
-            generator_fake_size (int): Number of fake samples to generate for the generator. (Default: 1000)
+            generator_fake_size (int): Number of fake samples to generate for the generator. (Default: 1000).
         
         Returns:
             None
         """
-        super.__init__()
+        super().__init__()
 
         self.device = None
         match device:
@@ -48,7 +45,6 @@ def __init__(
         self.n_dis = discriminator_deep_layers
         self.opt = optimizer
         self.loss = torch.nn.BCEWithLogitsLoss()
-        self.latent_dim = latent_dim
         self.fake_batch_size = generator_fake_size
     
     def init_weights(self, m: torch.nn.Module) -> None:
@@ -84,34 +80,87 @@ def get_optimizer(self, optimizer_type: Literal['adam', 'sgd'], lr: float) -> to
             case _:
                 raise ValueError(f"Invalid optimizer {optimizer_type}. Use 'adam' or 'sgd'.")
 
+    def construct_generator(
+            self,
+            n_features: int,
+            dropout: float = 0.1,
+            latent_dimension: int = 10,
+            activation: Literal['tanh', 'sigmoid'] = 'tanh'
+        ) -> torch.nn.Sequential:
+        """
+        Constructs the generator for the GAN.
+
+        Args:
+            n_features (int): Number of features for the generator.
+            dropout (float): Dropout probability for the generator.
+            latent_dimension (int): Dimension of the latent space.
+            activation (Literal['tanh', 'sigmoid']): Activation function for the generator.
+        """
+        self.latent_dimension = latent_dimension
+        match activation:
+            case 'tanh':
+                activation = torch.nn.Tanh()
+            case 'sigmoid':
+                activation = torch.nn.Sigmoid()
+            case _:
+                raise ValueError(f"Invalid activation {activation}. Use 'tanh' or 'sigmoid'.")
+        #Create Generator
+        self.deep_gen_layers = [latent_dimension] + self.n_gen + [n_features]
+        self.deep_gen_layers = list(zip(self.deep_gen_layers[:-1], self.deep_gen_layers[1:]))
+        self.generator = torch.nn.Sequential(
+            *[x for i, j in self.deep_gen_layers for x in [torch.nn.Linear(i, j), torch.nn.LeakyReLU()]],
+            torch.nn.Dropout(dropout),
+            torch.nn.Linear(self.deep_gen_layers[-1][1], self.deep_gen_layers[-1][1]),
+            activation
+        )
+        #Send to device
+        self.generator.to(self.device)
+        #Initialize weights
+        self.generator.apply(self.init_weights)
     
+    def construct_discriminator(
+            self,
+            n_features: int
+        )->torch.nn.Sequential:
+        """
+        Constructs the discriminator for the GAN.
+
+        Args:
+            n_features (int): Number of features for the discriminator.
+        """
+        #Create Discriminator
+        self.deep_dis_layers = [n_features] + self.n_dis + [2]
+        self.deep_dis_layers = list(zip(self.deep_dis_layers[:-1], self.deep_dis_layers[1:]))
+        self.discriminator = torch.nn.Sequential(
+            *[x for i,j in self.deep_dis_layers for x in [torch.nn.Linear(i, j), torch.nn.LeakyReLU()]],
+            torch.nn.Linear(2,1)
+        )
+        #Send to device
+        self.discriminator.to(self.device)
+        #Initialize weights
+        self.discriminator.apply(self.init_weights)
+
     def fit(
             self,
             X: Union[torch.Tensor, np.array],
-            global_iterations: int = 1,
+            global_iterations: int,
             discriminator_iterations: int = 1,
             lr: float = 0.001,
             batch_size: int = 32,
-            dropout_proba: float = 0.1,
+            generator_kwargs: dict = {},
+            discriminator_kwargs: dict = {}
         ) -> tuple[torch.Tensor, torch.Tensor]:
         """
-        Fits the copula to data.
-
-        Loss calculation:
+        Fits the model to data.
 
-        .. math::
-            \\begin{cases}
-            \min_D L_D(D, \mu_G) = -\mathbb{E}_{x\sim \mu_{G}}[\ln (1-D(x))] - \mathbb{E}_{x\sim \mu_{\\text{ref}}}[\ln (D(x))]\\
-            \min_G L_G(D, \mu_G) = -\mathbb{E}_{x\sim \mu_G}[\ln (D(x))]
-            \end{cases}
-        
         Args:
-            X (Union[torch.Tensor, np.array]): Input data in the shape (n_samples, n_features). Values must be within [-1,1]
-            global_iterations (int): Number of iterations to train the GAN.
-            discriminator_iterations (int): Number of iterations to train the discriminator for each global iteration.
+            X (Union[torch.Tensor, np.array]): Data to fit the model to.
+            global_iterations (int): Number of global iterations for training.
+            discriminator_iterations (int): Number of discriminator iterations for training.
             lr (float): Learning rate for the optimizer.
             batch_size (int): Batch size for training.
-            dropout_proba (float): Dropout probability for the generator.
+            generator_kwargs (dict): Keyword arguments for the generator. See :func:`construct_generator`.
+            discriminator_kwargs (dict): Keyword arguments for the discriminator. See :func:`construct_discriminator`.
         
         Returns:
             tuple[torch.Tensor, torch.Tensor]: Generator and discriminator loss.
@@ -120,6 +169,9 @@ def fit(
         #Check X values
         assert X.max() <= 1 and X.min() >= -1, "Values must be within [-1,1]. Try using np.tanh first."
 
+        self.construct_discriminator(X.shape[1], **discriminator_kwargs)
+        self.construct_generator(X.shape[1], **generator_kwargs)
+
         with self.device:
             if isinstance(X, np.ndarray):
                 X = torch.from_numpy(X).float()
@@ -130,28 +182,6 @@ def fit(
             X = X.to(self.device)
             self.n_features = X.shape[1]
 
-            #Create Generator
-            self.deep_gen_layers = [self.latent_dim] + self.n_gen + [self.n_features]
-            self.deep_gen_layers = list(zip(self.deep_gen_layers[:-1], self.deep_gen_layers[1:]))
-            self.generator = torch.nn.Sequential(
-                *[x for i, j in self.deep_gen_layers for x in [torch.nn.Linear(i, j), torch.nn.LeakyReLU()]],
-                torch.nn.Dropout(dropout_proba),
-                torch.nn.Linear(self.deep_gen_layers[-1][1], self.deep_gen_layers[-1][1]),
-                torch.nn.Tanh()
-            )
-            #Initialize weights
-            self.generator.apply(self.init_weights)
-
-            #Create Discriminator
-            self.deep_dis_layers = [self.n_features] + self.n_dis + [2]
-            self.deep_dis_layers = list(zip(self.deep_dis_layers[:-1], self.deep_dis_layers[1:]))
-            self.discriminator = torch.nn.Sequential(
-                *[x for i,j in self.deep_dis_layers for x in [torch.nn.Linear(i, j), torch.nn.LeakyReLU()]],
-                torch.nn.Linear(2,1)
-            )
-            #Initialize weights
-            self.discriminator.apply(self.init_weights)
-
             self.generator_optimizer = self.get_optimizer(self.opt, lr)
             self.discriminator_optimizer = self.get_optimizer(self.opt, lr)
 
@@ -163,7 +193,7 @@ def fit(
                     real = X[torch.randperm(X.shape[0])[:batch_size]]
                     #At the end the batch size might be smaller than the specified batch size
                     actual_batch_size = real.shape[0]
-                    fake = self.generator(torch.randn(actual_batch_size, self.latent_dim))
+                    fake = self.generator(torch.randn(actual_batch_size, self.latent_dimension))
 
                     self.discriminator.zero_grad()
                     disc_real = self.discriminator(real)
@@ -180,14 +210,12 @@ def fit(
                     self.discriminator_optimizer.step()
 
                 self.generator.zero_grad()
-                new_fake = self.generator(torch.randn(self.fake_batch_size, self.latent_dim))
+                new_fake = self.generator(torch.randn(self.fake_batch_size, self.latent_dimension))
                 new_disc_fake = self.discriminator(new_fake)
                 loss_generator = self.loss(new_disc_fake, torch.ones(self.fake_batch_size, 1))
                 #Inserting 0 for y calculates -log(D(G(z)))
                 loss_generator.backward()
                 self.generator_optimizer.step()
-                    
-
         
         return loss_discriminator, loss_generator
         
@@ -202,4 +230,4 @@ def generate(self, n_samples: int) -> np.ndarray:
             np.ndarray: Samples from the copula.
         """
         with self.device:
-            return self.generator(torch.randn(n_samples, self.latent_dim)).detach().cpu().numpy()
\ No newline at end of file
+            return self.generator(torch.randn(n_samples, self.latent_dimension)).detach().cpu().numpy()
\ No newline at end of file

From d245cea609b8fccea1021f622a9ad3b9321e1aaa Mon Sep 17 00:00:00 2001
From: Icaro Costa <icarosadero@gmail.com>
Date: Wed, 27 Sep 2023 23:25:47 -0300
Subject: [PATCH 13/14] debugging

---
 .gitignore                    |  1 +
 synthia/copulas/copula_gan.py | 61 +++++++++++++++++++++--------------
 synthia/generators/gan.py     |  9 +++---
 3 files changed, 42 insertions(+), 29 deletions(-)

diff --git a/.gitignore b/.gitignore
index bb7fb5c..ba09183 100644
--- a/.gitignore
+++ b/.gitignore
@@ -3,6 +3,7 @@
 docs/assets/
 *.nc
 *.pkl
+TODO.md
 
 # Byte-compiled / optimized / DLL files
 __pycache__/
diff --git a/synthia/copulas/copula_gan.py b/synthia/copulas/copula_gan.py
index c64f2eb..f17219f 100644
--- a/synthia/copulas/copula_gan.py
+++ b/synthia/copulas/copula_gan.py
@@ -47,7 +47,7 @@ def compute_cdf(kde_model, n_samples=1000):
 
     return cdf_values, support
 
-def map_cdf(cdf_values, support)->function:
+def map_cdf(cdf_values, support):
     """
     Maps the data in the support to its corresponding CDF value.
 
@@ -55,7 +55,7 @@ def map_cdf(cdf_values, support)->function:
         cdf_values: CDF values to map to.
         support: Support of the data.
     """
-    return lambda x: np.interp(x, support, cdf_values)
+    return lambda x: torch.from_numpy(np.interp(x, support, cdf_values))
 
 def inverse_map_cdf(cdf_values, support):
     """
@@ -65,7 +65,23 @@ def inverse_map_cdf(cdf_values, support):
         cdf_values: CDF values to compute inverse of.
         support: Support of the data.
     """
-    return lambda x: np.interp(x, cdf_values, support)
+    return lambda x: torch.from_numpy(np.interp(x, cdf_values, support))
+
+class GRU(torch.nn.Module):
+    def __init__(self, n_features: int, hidden_size: int = 1):
+        super().__init__()
+        self.gru = torch.nn.GRU(
+            input_size = n_features,
+            hidden_size = hidden_size,
+            batch_first=True
+        )
+        self.linear = torch.nn.Linear(hidden_size, 1)
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x, _ = self.gru(x)
+        x = x[:,-1,:]
+        x = self.linear(x)
+        return x
+
 
 class CopulaGAN(Copula, GAN):
     """
@@ -74,14 +90,13 @@ class CopulaGAN(Copula, GAN):
     def __init__(
             self,
             generator_deep_layers: list[int] = [32, 32],
-            discriminator_deep_layers: list[int] = [32, 32],
             device: Literal['cpu', 'cuda', 'auto'] = 'auto',
             optimizer: Literal['adam', 'sgd'] = 'adam',
             generator_fake_size: int = 1000,
         ) -> None:
         super().__init__(
             generator_deep_layers,
-            discriminator_deep_layers,
+            [],
             device,
             optimizer,
             generator_fake_size
@@ -105,7 +120,10 @@ def initialize(self, X: torch.Tensor, **kwargs):
             )
             cdf, support = compute_cdf(kde, **kwargs)
             real_U.append(map_cdf(cdf, support)(X[:,i]))
-            pdf = kde.pdf(obs_space, torch.linspace(X[:,i].min(),X[:,i].max(),n_samples))
+            pdf = torch.from_numpy(kde.pdf(
+                torch.linspace(X[:,i].min().item(),X[:,i].max().item(),len(support)),
+                obs_space
+            ).numpy())
             target_kdes.append(kde)
             cdfs.append(cdf)
             real_spaces.append(obs_space)
@@ -113,15 +131,15 @@ def initialize(self, X: torch.Tensor, **kwargs):
 
         self.real_space = zfit.dimension.combine_spaces(*real_spaces)
         self.target_kdes = target_kdes
-        self.target_cdf = torch.stack(cdfs, dim=1)
-        self.target_pdf = torch.stack(target_pdfs, dim=1)
-        self.real_U = torch.stack(real_U, dim=1)
+        self.target_cdf = torch.stack(cdfs, dim=1).float().to(self.device)
+        self.target_pdf = torch.stack(target_pdfs, dim=1).float().to(self.device)
+        self.real_U = torch.stack(real_U, dim=1).float().to(self.device)
 
     def construct_discriminator(
             self,
             n_features: int,
             hidden_size: int = 32,
-        ) -> torch.nn.Sequential:
+        ) -> torch.nn.Module:
         """
         Constructs the discriminator for the GAN. It is a GRU that receives samples from CDFs as input
         and outputs a single value.
@@ -133,13 +151,7 @@ def construct_discriminator(
             n_features (int): Number of features for the discriminator.
             hidden_size (int): Hidden size for the GRU. (Default: 32)
         """
-        self.discriminator = torch.nn.Sequential(
-            torch.nn.GRU(
-                input_size = n_features,
-                hidden_size = hidden_size,
-                batch_first=True),
-            torch.nn.Linear(hidden_size, 1)
-        ).to(self.device)
+        self.discriminator = GRU(n_features, hidden_size).to(self.device)
         self.discriminator.apply(self.init_weights)
         return self.discriminator
 
@@ -191,10 +203,10 @@ def fit(
                 for __ in range(discriminator_iterations):
                     #Train Discriminator
                     rand = torch.rand(batch_size, self.latent_dimension, device=self.device)
-                    fake_U = self.generator(rand)
-                    real_U = self.real_U[torch.randperm(self.real_U.shape[0])[:batch_size]]
-                    real_discriminant, _ = self.discriminator(real_U)
-                    fake_discriminant, _ = self.discriminator(fake_U)
+                    fake_U = self.generator(rand).unsqueeze(0)
+                    real_U = self.real_U[torch.randperm(self.real_U.shape[0], device=self.device)[:batch_size]].unsqueeze(0)
+                    real_discriminant = self.discriminator(real_U)
+                    fake_discriminant = self.discriminator(fake_U)
                     self.discriminator.zero_grad()
                     (
                         self.loss(real_discriminant, torch.ones_like(real_discriminant)) +\
@@ -203,11 +215,10 @@ def fit(
                     self.discriminator_optimizer.step()
                 self.generator.zero_grad()
                 rand = torch.rand(batch_size, self.latent_dimension, device=self.device)
-                fake_U = self.generator(rand)
-                fake_discriminant, _ = self.discriminator(fake_U)
+                fake_U = self.generator(rand).unsqueeze(0)
+                fake_discriminant = self.discriminator(fake_U)
                 loss = self.loss(fake_discriminant, torch.ones_like(fake_discriminant))
                 loss.backward()
                 self.generator_optimizer.step()
-    
-    def generate(self, n_samples: int, **kws) -> np.ndarray:
+    def generate(self, n_samples: int, **kws) -> torch.Tensor:
         return super().generate(n_samples, **kws)
\ No newline at end of file
diff --git a/synthia/generators/gan.py b/synthia/generators/gan.py
index 6c49464..a690eb7 100644
--- a/synthia/generators/gan.py
+++ b/synthia/generators/gan.py
@@ -219,7 +219,7 @@ def fit(
         
         return loss_discriminator, loss_generator
         
-    def generate(self, n_samples: int) -> np.ndarray:
+    def sample(self, n_samples: int) -> torch.Tensor:
         """
         Generates samples from the copula.
         
@@ -227,7 +227,8 @@ def generate(self, n_samples: int) -> np.ndarray:
             n_samples (int): Number of samples to generate.
         
         Returns:
-            np.ndarray: Samples from the copula.
+            torch.Tensor: Samples from the copula.
         """
-        with self.device:
-            return self.generator(torch.randn(n_samples, self.latent_dimension)).detach().cpu().numpy()
\ No newline at end of file
+        with torch.no_grad():
+            with self.device:
+                return self.generator(torch.randn(n_samples, self.latent_dimension))
\ No newline at end of file

From ecc1a062d442abd3b8b02d4746105df635804612 Mon Sep 17 00:00:00 2001
From: Icaro Costa <icarosadero@gmail.com>
Date: Thu, 28 Sep 2023 20:40:10 -0300
Subject: [PATCH 14/14] tested

---
 requirements.txt              |   3 +-
 synthia/copulas/copula_gan.py | 117 ++++++++++++++++++++++------------
 2 files changed, 79 insertions(+), 41 deletions(-)

diff --git a/requirements.txt b/requirements.txt
index 75c6ed8..93d9ca5 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -16,4 +16,5 @@ nbsphinx
 sphinxcontrib-bibtex==1
 sphinx-copybutton
 torch
-derivative
\ No newline at end of file
+derivative
+tqdm
\ No newline at end of file
diff --git a/synthia/copulas/copula_gan.py b/synthia/copulas/copula_gan.py
index f17219f..44a44d1 100644
--- a/synthia/copulas/copula_gan.py
+++ b/synthia/copulas/copula_gan.py
@@ -2,6 +2,7 @@
 from .copula import Copula
 from ..generators.gan import GAN
 from derivative import dxdt
+from tqdm import tqdm
 import numpy as np
 import torch
 import zfit
@@ -65,7 +66,7 @@ def inverse_map_cdf(cdf_values, support):
         cdf_values: CDF values to compute inverse of.
         support: Support of the data.
     """
-    return lambda x: torch.from_numpy(np.interp(x, cdf_values, support))
+    return lambda x: np.interp(x, cdf_values, support)
 
 class GRU(torch.nn.Module):
     def __init__(self, n_features: int, hidden_size: int = 1):
@@ -82,11 +83,12 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         x = self.linear(x)
         return x
 
-
 class CopulaGAN(Copula, GAN):
     """
     Learns Copula from data using Generative Adversarial Networks.
     """
+    _support = None
+    _initialized = False
     def __init__(
             self,
             generator_deep_layers: list[int] = [32, 32],
@@ -102,38 +104,54 @@ def __init__(
             generator_fake_size
         )
 
-    def initialize(self, X: torch.Tensor, **kwargs):
-        # Generate marginals and sample space
-        n_features = X.shape[1]
-
-        # Create target KDEs and CDFs
-        target_kdes = []
-        target_pdfs = []
-        cdfs = []
-        real_spaces = []
-        real_U = []
-        for i in range(n_features):
-            obs_space = zfit.Space('X' + str(i), limits=(X[:, i].min(), X[:, i].max()))
-            kde = zfit.pdf.KDE1DimGrid(
-                obs=obs_space,
-                data=zfit.Data.from_tensor(tensor=X[:, i], obs=obs_space)
-            )
-            cdf, support = compute_cdf(kde, **kwargs)
-            real_U.append(map_cdf(cdf, support)(X[:,i]))
-            pdf = torch.from_numpy(kde.pdf(
-                torch.linspace(X[:,i].min().item(),X[:,i].max().item(),len(support)),
-                obs_space
-            ).numpy())
-            target_kdes.append(kde)
-            cdfs.append(cdf)
-            real_spaces.append(obs_space)
-            target_pdfs.append(pdf)
-
-        self.real_space = zfit.dimension.combine_spaces(*real_spaces)
-        self.target_kdes = target_kdes
-        self.target_cdf = torch.stack(cdfs, dim=1).float().to(self.device)
-        self.target_pdf = torch.stack(target_pdfs, dim=1).float().to(self.device)
-        self.real_U = torch.stack(real_U, dim=1).float().to(self.device)
+    def initialize(self, X: torch.Tensor, verbose: bool = False, **kwargs):
+        if not self._initialized:
+            # Generate marginals and sample space
+            n_features = X.shape[1]
+
+            # Create target KDEs and CDFs
+            target_kdes = []
+            target_pdfs = []
+            cdfs = []
+            real_spaces = []
+            real_U = []
+            for i in range(n_features):
+                obs_space = zfit.Space('X' + str(i), limits=(X[:, i].min(), X[:, i].max()))
+                kde = zfit.pdf.KDE1DimGrid(
+                    obs=obs_space,
+                    data=zfit.Data.from_tensor(tensor=X[:, i], obs=obs_space)
+                )
+                cdf, support = compute_cdf(kde, **kwargs)
+                real_U.append(map_cdf(cdf, support)(X[:,i]))
+                pdf = torch.from_numpy(kde.pdf(
+                    torch.linspace(X[:,i].min().item(),X[:,i].max().item(),len(support)),
+                    obs_space
+                ).numpy())
+                target_kdes.append(kde)
+                cdfs.append(cdf)
+                real_spaces.append(obs_space)
+                target_pdfs.append(pdf)
+
+            self.real_space = zfit.dimension.combine_spaces(*real_spaces)
+            self.target_kdes = target_kdes
+            self.target_cdf = torch.stack(cdfs, dim=1).float().to(self.device)
+            self.target_pdf = torch.stack(target_pdfs, dim=1).float().to(self.device)
+            self.real_U = torch.stack(real_U, dim=1).float().to(self.device)
+            self._initialized = True
+        else:
+            tqdm.write("Already initialized") if verbose else None
+
+    @property
+    def support(self):
+        if hasattr(self, 'real_space'):
+            M_samples = self.target_cdf.shape[0]
+            self._support = torch.stack([
+                torch.linspace(float(x_0), float(x_1), M_samples)\
+                        for x_0, x_1 in zip(self.real_space.limits[0][0],self.real_space.limits[1][0])
+                ], dim=1).detach().cpu().numpy()
+        else:
+            raise AttributeError("Cannot compute support without real_space. Run initialize() first.")
+        return self._support
 
     def construct_discriminator(
             self,
@@ -164,7 +182,8 @@ def fit(
             batch_size: int = 32,
             generator_kwargs: dict = {},
             discriminator_kwargs: dict = {},
-            CDFN: int = 1000
+            CDFN: int = 1000,
+            verbose: bool = True
         ):
         """
         Uses GAN to learn a Copula flow from vector U sampled from the marginals of an arbitrary dataset.
@@ -179,7 +198,8 @@ def fit(
             batch_size (int): Batch size for training. (Default: 32)
             generator_kwargs (dict): Keyword arguments to pass to the generator. (Default: {})
             discriminator_kwargs (dict): Keyword arguments to pass to the discriminator. (Default: {})
-            CDFN (int): Number of samples to use for computing the CDF. (Default: 1000)
+            CDFN (int): Number of samples to use for computing the CDF. (Default: 1000).
+            verbose (bool): Whether to show progress bar and print status. (Default: True)
         """
         if isinstance(X, np.ndarray):
             X = torch.from_numpy(X).float()
@@ -188,7 +208,8 @@ def fit(
         else:
             raise TypeError(f"Invalid type {type(X)} for X. Use torch.Tensor or np.array.")
 
-        self.initialize(X, n_samples = CDFN)
+        tqdm.write("Initializing") if verbose else None
+        self.initialize(X, n_samples = CDFN, verbose = verbose)
 
         self.construct_generator(X.shape[1], activation='sigmoid', **generator_kwargs)
         self.construct_discriminator(X.shape[1], **discriminator_kwargs)
@@ -199,7 +220,7 @@ def fit(
             self.discriminator_optimizer = self.get_optimizer(self.opt, lr)
 
             #Train
-            for _ in range(global_iterations):
+            for _ in tqdm(range(global_iterations)):
                 for __ in range(discriminator_iterations):
                     #Train Discriminator
                     rand = torch.rand(batch_size, self.latent_dimension, device=self.device)
@@ -220,5 +241,21 @@ def fit(
                 loss = self.loss(fake_discriminant, torch.ones_like(fake_discriminant))
                 loss.backward()
                 self.generator_optimizer.step()
-    def generate(self, n_samples: int, **kws) -> torch.Tensor:
-        return super().generate(n_samples, **kws)
\ No newline at end of file
+    def sample_copula(self, n_samples: int, **kws) -> torch.Tensor:
+        return super().sample(n_samples, **kws)
+
+    def sample(self, n_samples: int) -> np.ndarray:
+        """
+        Samples from the learned Copula.
+
+        Args:
+            n_samples (int): Number of samples to generate.
+        """
+        U = self.sample_copula(n_samples).detach().cpu().numpy()
+
+        target_cdf = self.target_cdf.detach().cpu().numpy()
+
+        return np.vstack([inverse_map_cdf(target_cdf[:,i], self.support[:,i])(U[:,1]) for i in range (U.shape[1])]).T
+    
+    def generate(self, n_samples: int, **kws) -> np.ndarray:
+        return self.sample(n_samples, **kws)
\ No newline at end of file