Fix cat distrib

.gitignore

@@ -129,4 +129,5 @@ dmypy.json
 .pyre/
 
 /Results*
-/Dataset/Downloaded/*
+/Dataset/Downloaded/*
+.DS_STORE

Model/Energy/EnergyForDistribution/__init__.py

@@ -1,6 +1,6 @@
 from .categorical import EnergyCategoricalDistrib
 from .conv import ConvEnergy
-from .ising import EnergyIsing
+from .ising import ErdosRenyiEnergyIsing
 from .linear import fc_energy
 from .poisson import EnergyPoissonDistribution
 from .rbm import EnergyRBM

Model/Energy/EnergyForDistribution/ising.py

@@ -1,13 +1,16 @@
 from math import prod
 from typing import Tuple
 
+import igraph as ig
+import numpy as np
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from jaxtyping import Float
+from torch.distributions import bernoulli
 
 
-class EnergyIsing(nn.Module):
+class ErdosRenyiEnergyIsing(nn.Module):
     """Implement the energy of an Ising model. C.f. Oops I took a gradient from Grathwohl et al.
 
     According to table 1 of the paper the energy is defined as:
@@ -25,7 +28,8 @@ class EnergyIsing(nn.Module):
 
     Attributes:
         W: nn.Linear (input_size, hidden_dim), the parameters of the energy.
-            W is initialized with a Bernoulli distribution with p=0.5.
+            W is initialized as the adjacency matrix of an Erdos-Renyi graph with probability p=4/prod(input_size).
+            So each node has an average degree of 4.
         b: torch.Tensor of size (hidden_dim), the parameters of the energy.
             b is initialized as a tensor of ones.
 
@@ -34,28 +38,48 @@ class EnergyIsing(nn.Module):
     def __init__(
         self,
         input_size: Tuple[int],
-        learn_W: bool = True,
-        learn_b: bool = True,
+        n_node: int,
+        average_degree: int = 4,
+        init_bias: float = 0.0,
+        learn_G: bool = False,
+        learn_bias: bool = False,
     ) -> None:
-        super().__init__()
-        self.W = nn.parameter.Parameter(
-            torch.ones(prod(input_size), prod(input_size)) * 0.5, requires_grad=learn_W
-        )
-        self.b = nn.parameter.Parameter(
-            torch.ones(prod(input_size)), requires_grad=learn_b
+        super(ErdosRenyiEnergyIsing, self).__init__()
+        # Code from Oops I took a gradient
+        # g = ig.Graph.Erdos_Renyi(n_node, float(average_degree) / float(n_node))
+        # A = np.asarray(g.get_adjacency().data)  # g.get_sparse_adjacency()
+        # A = torch.tensor(A).float()
+        A = torch.randn((n_node, n_node)) * 0.01
+        weights = torch.randn_like(A) * ((1.0 / average_degree) ** 0.5)
+        weights = weights * (1 - torch.tril(torch.ones_like(weights)))
+        weights = weights + weights.t()
+
+        self.G = nn.Parameter(A * weights, requires_grad=learn_G)
+        self.bias = nn.Parameter(
+            torch.ones((n_node,)).float() * init_bias, requires_grad=learn_bias
         )
+        self.data_dim = n_node
+
+    @property
+    def J(self):
+        return self.G
 
     def forward(
-        self, x: Float[torch.Tensor, "batch_size *dim"]
+        self, x: Float[torch.Tensor, "batch_size nb_point_in_graph"]
     ) -> Float[torch.Tensor, "batch_size"]:
         """Compute the energy of the Ising model.
 
         Args:
-            x: Float[torch.Tensor, "batch_size *dim"], batch input of the energy.
+            x: Float[torch.Tensor, "batch_size nb_point_in_graph"], batch input of the energy.
+            x is a vector of zeros and ones.
 
         Returns:
             Float[torch.Tensor, "batch_size"], E(x), the energy of the Ising model.
         """
-        x = x.flatten(1)
-        Wx = torch.matmul(x, self.W.T)
-        return -torch.sum(x * Wx, dim=1) - x @ self.b
+        # code from Oops I took a gradient
+
+        x = 2 * x - 1  # convert 0/1 to -1/1
+        xg = x @ self.J
+        xgx = (xg * x).sum(-1)
+        b = (self.bias[None, :] * x).sum(-1)
+        return -xgx - b

Model/Energy/EnergyForDistribution/poisson.py

@@ -55,6 +55,4 @@ def forward(
         Returns:
             Float[torch.Tensor, "batch_size 1"], E(x), the energy of the poisson distribution
         """
-        # print(f"self.lambda_: {self.lambda_}")
-        # print(f"x: {x.shape}")
         return torch.lgamma(x + 1) - x * torch.log(self.lambda_)

Model/Energy/energy_getter.py

@@ -1,27 +1,31 @@
+import numpy as np
+
 from .EnergyForDistribution import (
     ConvEnergy,
     EnergyCategoricalDistrib,
-    EnergyIsing,
     EnergyPoissonDistribution,
     EnergyRBM,
+    ErdosRenyiEnergyIsing,
     fc_energy,
 )
-from .EnergyForRegression import EnergyNetworkRegression_Large, EnergyNetworkRegression_Toy
+from .EnergyForRegression import (
+    EnergyNetworkRegression_Large,
+    EnergyNetworkRegression_Toy,
+)
 from .FeatureExtractor import Resnet18_FeatureExtractor, ToyFeatureNet
 
-import numpy as np
 dic_energy = {
     "fc": fc_energy,
     "conv": ConvEnergy,
     "rbm": EnergyRBM,
     "categorical": EnergyCategoricalDistrib,
     "poisson": EnergyPoissonDistribution,
-    "ising": EnergyIsing,
+    "ising": ErdosRenyiEnergyIsing,
 }
 
 dic_energy_regression = {
-    'fc': EnergyNetworkRegression_Large,
-    'toy': EnergyNetworkRegression_Toy,
+    "fc": EnergyNetworkRegression_Large,
+    "toy": EnergyNetworkRegression_Toy,
 }
 
 
@@ -56,8 +60,8 @@ def get_energy_regression(input_size_x, input_size_y, args_dict):
 
 
 dic_feature_extractor = {
-    'resnet' : Resnet18_FeatureExtractor,
-    'toy' : ToyFeatureNet,
+    "resnet": Resnet18_FeatureExtractor,
+    "toy": ToyFeatureNet,
 }
 
 
@@ -69,37 +73,48 @@ def get_feature_extractor(
         return None
     if args_dict["feature_extractor_name"] not in dic_feature_extractor:
         raise ValueError("Feature extractor name not valid")
-
-    feature_extractor = dic_feature_extractor[args_dict['feature_extractor_name']]
-    if 'feature_extractor_params' not in args_dict.keys():
-        args_dict['feature_extractor_params'] = {}
-    feature_extractor = feature_extractor(input_dim=input_size_x, **args_dict['feature_extractor_params'])
-    print(args_dict['train_feature_extractor'])
-    if args_dict['train_feature_extractor'] == False :
+
+    feature_extractor = dic_feature_extractor[args_dict["feature_extractor_name"]]
+    if "feature_extractor_params" not in args_dict.keys():
+        args_dict["feature_extractor_params"] = {}
+    feature_extractor = feature_extractor(
+        input_dim=input_size_x, **args_dict["feature_extractor_params"]
+    )
+    print(args_dict["train_feature_extractor"])
+    if args_dict["train_feature_extractor"] == False:
         for param in feature_extractor.parameters():
             param.requires_grad = False
     return feature_extractor
 
+
 from .ExplicitBiasForRegression import Layer1FC, Layer2FC, Layer3FC
+
 dic_explicit_bias_regression = {
-    '1_layer_fc' : Layer1FC,
-    '2_layer_fc' : Layer2FC,
-    '3_layer_fc' : Layer3FC,
-    'none' : None,
+    "1_layer_fc": Layer1FC,
+    "2_layer_fc": Layer2FC,
+    "3_layer_fc": Layer3FC,
+    "none": None,
 }
 
-def get_explicit_bias_regression(args_dict,
-                      input_size_x,
-                      ):
-    if 'explicit_bias_name' not in args_dict:
+
+def get_explicit_bias_regression(
+    args_dict,
+    input_size_x,
+):
+    if "explicit_bias_name" not in args_dict:
         return None
-    if args_dict['explicit_bias_name'] not in dic_explicit_bias_regression:
-        raise ValueError('Explicit bias name not valid')
-    if args_dict['explicit_bias_name'] is None or args_dict['explicit_bias_name'] == 'none' :
-            return None
-    explicit_bias = dic_explicit_bias_regression[args_dict['explicit_bias_name']]
-    if 'explicit_bias_params' not in args_dict.keys():
-        args_dict['explicit_bias_params'] = {}
-
-    explicit_bias = explicit_bias(input_size_x=input_size_x, **args_dict['explicit_bias_params'])
-    return explicit_bias
+    if args_dict["explicit_bias_name"] not in dic_explicit_bias_regression:
+        raise ValueError("Explicit bias name not valid")
+    if (
+        args_dict["explicit_bias_name"] is None
+        or args_dict["explicit_bias_name"] == "none"
+    ):
+        return None
+    explicit_bias = dic_explicit_bias_regression[args_dict["explicit_bias_name"]]
+    if "explicit_bias_params" not in args_dict.keys():
+        args_dict["explicit_bias_params"] = {}
+
+    explicit_bias = explicit_bias(
+        input_size_x=input_size_x, **args_dict["explicit_bias_params"]
+    )
+    return explicit_bias

Model/Proposals/ProposalForDistributionEstimation/__init__.py

@@ -3,4 +3,6 @@
 from .gaussian_mixture import GaussianMixtureProposal
 from .gaussian_mixture_adaptive import GaussianMixtureAdaptiveProposal
 from .noise_gradation_adaptive import NoiseGradationAdaptiveProposal
-from .student import StudentProposal
+from .student import StudentProposal
+from .ising_proposal import IsingProposal
+from .ising_proposal_adaptive import IsingProposalAdaptive

Model/Proposals/ProposalForDistributionEstimation/categorical.py

@@ -28,7 +28,7 @@ def sample(self, nb_sample: int = 1):
 
         return samples_one_hot
 
-    def log_prob(self, x: Float[torch.Tensor, "batch_size"]):
+    def log_prob(self, x: Float[torch.Tensor, "batch_size 1"]):
         return (
             categorical.Categorical(self.logit_parameters).log_prob(x).flatten(1).sum(1)
         )

Model/Proposals/ProposalForDistributionEstimation/ising_proposal.py

@@ -0,0 +1,74 @@
+import numpy as np
+import torch
+import torch.nn as nn
+from jaxtyping import Float
+from torch.nn.parameter import Parameter
+
+
+class IsingProposal(nn.Module):
+    """Proposal for Ising model
+
+    Attributes:
+        dataset: torch.utils.data.Dataset, dataset of the Ising model
+        p: float, probability of not flipping the spin
+    """
+
+    def __init__(self, input_size, dataset, centers = None, p: float = 0.9):
+        super(IsingProposal, self).__init__()
+        self.dataset = dataset
+        self.p = p
+        self.dummy_param = Parameter(torch.Tensor([0.0]), requires_grad=False)
+
+        shape = self.dataset[0][0].shape
+        len_dataset = len(self.dataset)
+        if centers is not None:
+            # For the adaptive proposal, we can pass the centers directly
+            self.centers = centers
+            self.len_centers = len(centers)
+        else :
+            self.len_centers = len_dataset
+            # if np.prod(shape)*len(dataset)<1e5:
+                # Checking the full size of storing everything :
+                # self.centers = torch.stack([self.dataset[i][0] for i in range(len(self.dataset))])
+            # else :
+            self.centers = None
+
+    def get_centers(self, index):
+        # Might be worth it time wise to store the samples in memory rather than recalculating them everytime
+        if self.centers is not None :
+            return self.centers[index]
+        else :
+            return torch.stack([self.dataset[i][0] for i in index])
+
+    def sample(self, nb_sample: int = 1) -> Float[torch.Tensor, "nb_sample nb_nodes"]:
+        with torch.no_grad(): # Lower memory size
+            if nb_sample < self.len_centers:
+                index = np.random.choice(self.len_centers, nb_sample)
+            else:
+                index = np.random.choice(self.len_centers, nb_sample, replace=True)
+
+            center = self.get_centers(index).to(self.dummy_param.device)
+            bernoulli_keep = torch.distributions.Bernoulli(
+                torch.full_like(center, self.p)
+            ).sample()
+
+            samples = center * bernoulli_keep + (1 - center) * (1 - bernoulli_keep)
+        return samples.detach()
+
+    def log_prob(
+        self, x: Float[torch.Tensor, "batch_size nb_nodes"]
+    ) -> Float[torch.Tensor, "batch_size"]:
+
+        data = self.get_centers(range(self.len_centers)).unsqueeze(1).to(x.device)
+        x_expanded = x.unsqueeze(0)  # 1, batch_size, nb_nodes
+
+        dependency = (data - x_expanded).abs()  # len(dataset), batch_size, nb_nodes
+        log_prob = (
+            torch.distributions.Bernoulli(torch.full_like(dependency, self.p))
+            .log_prob(dependency)
+            .sum(-1)
+        )  # len(dataset), batch_size
+        log_prob = log_prob.logsumexp(0) - torch.log(
+            torch.tensor(self.len_centers)
+        )  # batch_size
+        return log_prob

Model/Proposals/ProposalForDistributionEstimation/ising_proposal_adaptive.py

@@ -0,0 +1,43 @@
+import numpy as np
+import torch
+import torch.nn as nn
+from jaxtyping import Float
+from torch.nn.parameter import Parameter
+from .ising_proposal import IsingProposal
+
+class IsingProposalAdaptive(IsingProposal):
+    """Proposal for Ising model
+
+    Attributes:
+        dataset: torch.utils.data.Dataset, dataset of the Ising model
+        p: float, probability of not flipping the spin
+    """
+
+    def __init__(self, default_proposal, input_size, dataset, p: float = 0.9):
+        super(IsingProposalAdaptive, self).__init__(input_size =input_size, dataset=dataset, centers=None, p = p)
+        self.x = None
+    def set_x(self, x):
+        self.x = x
+
+    def get_center(self, index):
+        # Might be worth it time wise to store the samples in memory rather than recalculating them everytime
+        if self.center is not None :
+            return self.center[index]
+        else :
+            return torch.stack([self.dataset[i][0] for i in index])
+
+    def sample(self, nb_sample: int = 1) -> Float[torch.Tensor, "nb_sample nb_nodes"]:
+        if self.x is not None :
+            aux_ising = IsingProposal(self.x.shape[1], self.dataset, p = self.p, centers=self.x)
+            return aux_ising.sample(nb_sample).detach()
+        else :
+            return super().sample(nb_sample).detach()
+
+    def log_prob(
+        self, x: Float[torch.Tensor, "batch_size nb_nodes"]
+    ) -> Float[torch.Tensor, "batch_size"]:
+        if self.x is not None :
+            aux_ising = IsingProposal(self.x.shape[1], self.dataset, p = self.p, centers=self.x)
+            return aux_ising.log_prob(x)
+        else :
+            return super().log_prob(x)