diff --git a/rbms/bm/__init__.py b/rbms/bm/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/rbms/bm/bernoulli.py b/rbms/bm/bernoulli.py new file mode 100644 index 0000000..e69de29 diff --git a/rbms/bm/implement.py b/rbms/bm/implement.py new file mode 100644 index 0000000..41dc210 --- /dev/null +++ b/rbms/bm/implement.py @@ -0,0 +1,65 @@ +import torch +from torch import Tensor + +from rbms.custom_fn import one_hot + + +@torch.jit.script +def _get_freq_single_point( + data: Tensor, + weights: Tensor, + pseudo_count: float, +) -> torch.Tensor: + _, _, q = data.shape + frequencies = (data * weights).sum(dim=0) + # Set to zero the negative frequencies. Used for the reintegration. + torch.clamp_(frequencies, min=0.0) + + return (1.0 - pseudo_count) * frequencies + (pseudo_count / q) + + +@torch.jit.script +def _get_freq_two_points( + data: torch.Tensor, + weights: torch.Tensor, + pseudo_count: float, +) -> torch.Tensor: + + M, L, q = data.shape + data_oh = data.reshape(M, q * L) + + fij = (data_oh * weights).T @ data_oh + # Set to zero the negative frequencies. Used for the reintegration. + torch.clamp_(fij, min=0.0) + # Apply the pseudo count + fij = (1.0 - pseudo_count) * fij + (pseudo_count / q**2) + # Diagonal terms must represent the single point frequencies + fij_diag = _get_freq_single_point( + data, weights.reshape(M, 1, 1), pseudo_count + ).ravel() + # Set the diagonal terms of fij to the single point frequencies + fij = torch.diagonal_scatter(fij, fij_diag, dim1=0, dim2=1) + + return fij.reshape(L, q, L, q) + + +@torch.jit.script +def _sample_one_visible_potts( + v: Tensor, weight_matrix: Tensor, bias: Tensor, beta: float +) -> Tensor: + device = v.device + dtype = weight_matrix.dtype + N, L, q = v.shape + idx = torch.randint(0, L, (1,), device=device)[0] + couplings_residue = weight_matrix[idx].reshape(q, L * q) + logit_residue = beta * ( + bias[idx].unsqueeze(0) + v.reshape(N, L * q).to(dtype) @ couplings_residue.T + ) # (N, q) + new_residues = one_hot( + torch.multinomial(torch.softmax(logit_residue, dim=-1), num_samples=1).squeeze( + -1 + ), + num_classes=q, + ) + v[:, idx] = new_residues + return v diff --git a/rbms/bm/potts.py b/rbms/bm/potts.py new file mode 100644 index 0000000..ff7e5d1 --- /dev/null +++ b/rbms/bm/potts.py @@ -0,0 +1,360 @@ +from __future__ import annotations + +import numpy as np +import torch +from torch import Tensor + +from rbms.bm.implement import _sample_one_visible_potts +from rbms.bm.utils import get_freq_single_point, get_freq_two_points +from rbms.classes import EBM +from rbms.custom_fn import check_keys_dict, one_hot +from rbms.dataset.dataset_class import RBMDataset + + +class PBM(EBM): + """Potts Boltzmann Machine (bmDCA)""" + + visible_type: str = "categorical" + + def __init__( + self, + weight_matrix: Tensor, + bias: Tensor, + device: torch.device | str | None = None, + dtype: torch.dtype | None = None, + ): + if device is None: + device = weight_matrix.device + if dtype is None: + dtype = weight_matrix.dtype + self.device = device + self.dtype = dtype + self.weight_matrix = weight_matrix.to(device=self.device, dtype=self.dtype) + self.bias = bias.to(device=self.device, dtype=self.dtype) + self.name = "PBM" + self.flags = [] + + def __add__(self, other): + """Add the parameters of two EBMs. Useful for interpolation""" + return PBM( + weight_matrix=self.weight_matrix + other.weight_matrix, + bias=self.bias + other.bias, + ) + + def __mul__(self, other): + """Multiplies the p ofsy a float.""" + return PBM( + weight_matrix=self.weight_matrix * other, + bias=self.bias * other, + ) + + def sample_one_visible(self, chains: dict[str, Tensor], beta: float = 1.0): + one_hot_gen = one_hot(chains["visible"].long(), self.num_states) + res = _sample_one_visible_potts(one_hot_gen, self.weight_matrix, self.bias, beta) + chains["visible"] = res.argmax(-1).to(self.dtype) + return chains + + def sample_visibles( + self, chains: dict[str, Tensor], beta: float = 1.0 + ) -> dict[str, Tensor]: + """Sample L randomly selected . + + Args: + chains (dict[str, Tensor]): The parallel chains used for sampling. + beta (float, optional): The inverse temperature. Defaults to 1.0. + + Returns: + dict[str, Tensor]: The updated chains with sampled hidden states. + """ + one_hot_gen = one_hot(chains["visible"].long(), self.num_states) + for _ in range(self.num_visibles): + one_hot_gen = _sample_one_visible_potts( + one_hot_gen, self.weight_matrix, self.bias, beta + ) + chains["visible"] = one_hot_gen.argmax(-1).to(self.dtype) + return chains + + def compute_energy_visibles(self, v: Tensor) -> Tensor: + """Returns the marginalized energy of the model computed on the visible configurations + + Args: + v (Tensor): Visible configurations + + Returns: + Tensor: The computed energy. + """ + L, q = self.bias.shape + batch_size = v.shape[0] + x_flat = ( + one_hot(v.long(), num_classes=self.num_states) + .to(self.dtype) + .view(batch_size, -1) + ) + bias_flat = self.bias.view(-1) + couplings_flat = self.weight_matrix.reshape(L * q, L * q) + bias_term = x_flat @ bias_flat + coupling_term = torch.sum(x_flat * (x_flat @ couplings_flat), dim=1) + energy = -bias_term - 0.5 * coupling_term + + return energy + + def init_chains( + self, + num_samples: int, + weights: Tensor | None = None, + start_v: Tensor | None = None, + ) -> dict[str, Tensor]: + """Initialize a Markov chain for the EBM by sampling a uniform distribution on the visible layer + and sampling the hidden layer according to the visible one. + + Args: + num_samples (int): The number of samples to initialize. + start_v (Tensor, optional): The initial visible states. Defaults to None. + + Returns: + dict[str, Tensor]: The initialized Markov chain. + + Notes: + - If start_v is specified, its number of samples will override the num_samples argument. + """ + if num_samples <= 0: + if start_v is not None: + num_samples = start_v.shape[0] + else: + raise ValueError(f"Got negative num_samples arg: {num_samples}") + + if start_v is None: + # Dummy mean visible + mv = ( + torch.ones( + size=(num_samples, self.num_visibles), + device=self.device, + dtype=self.dtype, + ) + / 2 + ) + v = torch.bernoulli(mv) + else: + # Dummy mean visible + mv = torch.ones_like(start_v, device=self.device, dtype=self.dtype) / 2 + v = start_v.to(device=self.device, dtype=self.dtype) + + # Initialize chains + if weights is None: + weights = torch.ones(v.shape[0], device=v.device, dtype=v.dtype) + return dict( + visible=v, + visible_mag=mv, + weights=weights, + ) + + def compute_gradient( + self, + data: dict[str, Tensor], + chains: dict[str, Tensor], + centered: bool = True, + ) -> None: + """Compute the gradient for each of the parameters and attach it. + + Args: + data (dict[str, Tensor]): The data state. + chains (dict[str, Tensor]): The parallel chains used for gradient computation. + centered (bool, optional): Whether to use centered gradients. Defaults to True. + lambda_l1 (float, optional): factor for the L1 regularization. Defaults to 0. + lambda_l2 (float, optional): factor for the L2 regularization. Defaults to 0. + """ + pseudo_count = 1e-4 + one_hot_data = one_hot(data["visible"].long(), self.num_states).to(self.dtype) + one_hot_gen = one_hot(chains["visible"].long(), self.num_states).to(self.dtype) + fi_data = get_freq_single_point(one_hot_data, data["weights"], pseudo_count) + fi_gen = get_freq_single_point(one_hot_gen, chains["weights"], pseudo_count) + fij_data = get_freq_two_points(one_hot_data, data["weights"], pseudo_count) + fij_gen = get_freq_two_points(one_hot_gen, chains["weights"], pseudo_count) + self.bias.grad = fi_data - fi_gen + self.weight_matrix.grad = fij_data - fij_gen + + def parameters(self) -> list[Tensor]: + """Returns a list containing the parameters of the RBM. + + Returns: + List[Tensor]: A list containing the weight matrix, visible bias, and hidden bias. + """ + return [self.weight_matrix, self.bias] + + def named_parameters(self) -> dict[str, np.ndarray]: + return { + "weight_matrix": self.weight_matrix.cpu().numpy(), + "bias": self.bias.cpu().numpy(), + } + + @staticmethod + def set_named_parameters( + named_params: dict[str, np.ndarray], + device: torch.device | str, + dtype: torch.dtype, + ) -> EBM: + names = ["bias", "weight_matrix"] + check_keys_dict(d=named_params, names=names) + params = PBM( + weight_matrix=torch.from_numpy(named_params.pop("weight_matrix")).to( + device=device, dtype=dtype + ), + bias=torch.from_numpy(named_params.pop("bias")).to( + device=device, dtype=dtype + ), + ) + if len(named_params.keys()) > 0: + raise ValueError( + f"Too many keys in params dictionary. Remaining keys: {named_params.keys()}" + ) + return params + + def to( + self, + device: torch.device | str | None = None, + dtype: torch.dtype | None = None, + ): + """Move the parameters to the specified device and/or convert them to the specified data type. + + Args: + device (Optional[torch.device], optional): The device to move the parameters to. + Defaults to None. + dtype (Optional[torch.dtype], optional): The data type to convert the parameters to. + Defaults to None. + + Returns: + RBM: The modified RBM instance. + """ + if device is not None: + self.device = device + if dtype is not None: + self.dtype = dtype + self.weight_matrix = self.weight_matrix.to(device=self.device, dtype=self.dtype) + self.bias = self.bias.to(device=self.device, dtype=self.dtype) + return self + + def clone( + self, device: torch.device | str | None = None, dtype: torch.dtype | None = None + ) -> EBM: + """Create a clone of the RBM instance. + + Args: + device (Optional[torch.device], optional): The device for the cloned parameters. + Defaults to the current device. + dtype (Optional[torch.dtype], optional): The data type for the cloned parameters. + Defaults to the current data type. + + Returns: + RBM: A new RBM instance with cloned parameters. + """ + if device is None: + device = self.device + if dtype is None: + dtype = self.dtype + + return PBM( + self.weight_matrix.clone(), + self.bias.clone(), + device=device, + dtype=dtype, + ) + + @staticmethod + def init_parameters( + num_hiddens: int, + dataset: RBMDataset, + device: torch.device | str, + dtype: torch.dtype, + var_init: float = 1e-4, + ) -> PBM: + """Initialize the parameters of the RBM. + + Args: + num_hiddens (int): Number of hidden units. + dataset (RBMDataset): Training dataset. + device (torch.device): PyTorch device for the parameters. + dtype (torch.dtype): PyTorch dtype for the parameters. + var_init (float, optional): Variance of the weight matrix. Defaults to 1e-4. + + Notes: + - The number of visible units is induced from the dataset provided. + - Hidden biases are set to 0. + - Visible biases are set to the frequencies of the dataset. + - The weight matrix is initialized with a Gaussian distribution of variance `var_init`. + """ + + fi = get_freq_single_point( + one_hot(dataset.data.long(), dataset.get_num_states()), dataset.weights, 1e-4 + ) + return PBM( + weight_matrix=torch.zeros( + ( + dataset.get_num_visibles(), + dataset.get_num_states(), + dataset.get_num_visibles(), + dataset.get_num_states(), + ), + device=fi.device, + dtype=fi.dtype, + ), + bias=torch.log(fi), + ) + + @property + def num_visibles(self) -> int: + """Number of visible units""" + return self.weight_matrix.shape[0] + + @property + def ref_log_z(self) -> float: + """Reference log partition function with weights set to 0 (except for the visible bias).""" + return torch.logsumexp(self.bias, dim=1).sum().item() + + def independent_model(self) -> PBM: + """Independent model where only local fields are preserved.""" + return PBM( + torch.zeros_like(self.weight_matrix), + self.bias.clone(), + device=self.device, + dtype=self.dtype, + ) + + def sample_state( + self, chains: dict[str, Tensor], n_steps: int, beta: float = 1.0 + ) -> dict[str, Tensor]: + """Sample the model for n_steps + + Args: + chains (): The starting position of the chains. + n_steps (int): The number of sampling steps. + beta (float, optional): The inverse temperature. Defaults to 1.0 + + Returns: + dict[str, Tensor]: The updated chains after n_steps of sampling. + """ + chains_mutate = { + "visible": chains["visible"].clone(), + "weights": chains["weights"].clone(), + } # avoids to modify the chains inplace + + for _ in torch.arange(n_steps): + chains_mutate = self.sample_visibles(chains_mutate, beta) + + return chains_mutate + + def get_metrics(self, metrics: dict[str, float]) -> dict[str, float]: + return metrics + + def pre_grad_update(self) -> None: + pass + + def post_grad_update(self) -> None: + pass + + @property + def effective_number_variables(self) -> float: + return 1 + + @property + def num_states(self) -> int: + return self.weight_matrix.shape[1] diff --git a/rbms/bm/utils.py b/rbms/bm/utils.py new file mode 100644 index 0000000..39cb43c --- /dev/null +++ b/rbms/bm/utils.py @@ -0,0 +1,67 @@ +import torch +from torch import Tensor + +from rbms.bm.implement import _get_freq_single_point, _get_freq_two_points + + +def get_freq_two_points( + data: Tensor, + weights: Tensor | None = None, + pseudo_count: float = 0.0, +) -> Tensor: + """ + Computes the 2-points statistics of the input MSA. + + Args: + data (torch.Tensor): One-hot encoded data array. + weights (Optional[torch.Tensor], optional): Array of weights to assign to the sequences of shape. + pseudo_count (float, optional): Pseudo count for the single and two points statistics. Acts as a regularization. Defaults to 0.0. + + Raises: + ValueError: If the input data is not a 3D tensor. + + Returns: + torch.Tensor: Matrix of two-point frequencies of shape (L, q, L, q). + """ + if data.dim() != 3: + raise ValueError( + f"Expected data to be a 3D tensor, but got {data.dim()}D tensor instead" + ) + + M = len(data) + if weights is not None: + norm_weights = weights.reshape(M, 1) / weights.sum() + else: + norm_weights = torch.ones((M, 1), device=data.device, dtype=data.dtype) / M + + return _get_freq_two_points(data, norm_weights, pseudo_count) + + +def get_freq_single_point( + data: Tensor, + weights: Tensor | None = None, + pseudo_count: float = 0.0, +) -> Tensor: + """Computes the single point frequencies of the input MSA. + Args: + data (torch.Tensor): One-hot encoded data array. + weights (Optional[torch.Tensor], optional): Weights of the sequences. + pseudo_count (float, optional): Pseudo count to be added to the frequencies. Defaults to 0.0. + + Raises: + ValueError: If the input data is not a 3D tensor. + + Returns: + torch.Tensor: Single point frequencies. + """ + if data.dim() != 3: + raise ValueError( + f"Expected data to be a 3D tensor, but got {data.dim()}D tensor instead" + ) + M = len(data) + if weights is not None: + norm_weights = weights.reshape(M, 1, 1) / weights.sum() + else: + norm_weights = torch.ones((M, 1, 1), device=data.device, dtype=data.dtype) / M + + return _get_freq_single_point(data, norm_weights, pseudo_count) diff --git a/rbms/custom_fn.py b/rbms/custom_fn.py index 7bc1fcf..62370a1 100644 --- a/rbms/custom_fn.py +++ b/rbms/custom_fn.py @@ -2,9 +2,10 @@ import numpy as np import torch from torch import Tensor +from torch.nn.functional import one_hot as one_hot -def one_hot( +def one_hot_old( x: Tensor, num_classes: int = -1, dtype: torch.dtype = torch.float32 ) -> Tensor: """A one-hot encoding function faster than the PyTorch one working with torch.int32 and returning a float Tensor diff --git a/rbms/dataset/__init__.py b/rbms/dataset/__init__.py index bf96690..0a04be7 100644 --- a/rbms/dataset/__init__.py +++ b/rbms/dataset/__init__.py @@ -16,6 +16,7 @@ def load_dataset( use_weights: bool = False, alphabet="protein", remove_duplicates: bool = False, + shuffle: bool = False, device: torch.device | str = "cpu", dtype: torch.dtype = torch.float32, ) -> tuple[RBMDataset, RBMDataset | None]: @@ -61,13 +62,14 @@ def load_dataset( else: unique_ind = np.arange(data.shape[0]) - idx = torch.randperm(unique_ind.shape[0]) if unique_ind.shape[0] < data.shape[0]: print(f"N_samples: {data.shape[0]} -> {unique_ind.shape[0]}") - data = data[unique_ind[idx]] - labels = labels[unique_ind[idx]] - weights = weights[unique_ind[idx]] - names = names[unique_ind[idx]] + if shuffle: + unique_ind = unique_ind[torch.randperm(unique_ind.shape[0])] + data = data[unique_ind] + labels = labels[unique_ind] + weights = weights[unique_ind] + names = names[unique_ind] return_datasets.append( RBMDataset( diff --git a/rbms/dataset/fasta_utils.py b/rbms/dataset/fasta_utils.py index 2f17295..e16c9b0 100644 --- a/rbms/dataset/fasta_utils.py +++ b/rbms/dataset/fasta_utils.py @@ -152,7 +152,7 @@ def compute_weights( """ device = torch.device(device) data_tensor = torch.from_numpy(data).to(device=device) - assert len(data_tensor) == 2, "'data' must be a 2-dimensional array" + assert len(data_tensor.shape) == 2, "'data' must be a 2-dimensional array" _, L = data_tensor.shape def get_sequence_weight(s: torch.Tensor, data: torch.Tensor, L: int, th: float): @@ -160,7 +160,9 @@ def get_sequence_weight(s: torch.Tensor, data: torch.Tensor, L: int, th: float): n_clust = torch.sum(seq_id >= th) return 1.0 / n_clust - weights = torch.vstack([get_sequence_weight(s, data_tensor, L, th) for s in data]) + weights = torch.vstack( + [get_sequence_weight(s, data_tensor, L, th) for s in data_tensor] + ) return weights.cpu().numpy() diff --git a/rbms/dataset/load_h5.py b/rbms/dataset/load_h5.py index 9d26444..f9d8535 100644 --- a/rbms/dataset/load_h5.py +++ b/rbms/dataset/load_h5.py @@ -39,10 +39,13 @@ def load_HDF5( variable_type = f["variable_type"][()].decode() weights = np.ones(dataset.shape[0]) if use_weights: - if variable_type != "categorical": - print("Ignoring compute weights since data is not categorical") + if "weights" in f.keys(): + weights = f["weights"][()] else: - weights = compute_weights(data=dataset, device=device) + if variable_type != "categorical": + print("Ignoring compute weights since data is not categorical") + else: + weights = compute_weights(data=dataset, device=device) if "labels" in f.keys(): labels = np.array(f["labels"][()]) diff --git a/rbms/dataset/utils.py b/rbms/dataset/utils.py index 8dee55f..50d0c8b 100644 --- a/rbms/dataset/utils.py +++ b/rbms/dataset/utils.py @@ -111,8 +111,11 @@ def get_covariance_matrix( num_data = num_extract if num_classes != 2: - data = data.to(device=device, dtype=torch.int32) - data_oh = one_hot(data, num_classes=num_classes).reshape(num_data, -1) + data_oh = ( + one_hot(data.long().to(device), num_classes=num_classes) + .reshape(num_data, -1) + .to(dtype) + ) else: data_oh = data.to(device=device, dtype=dtype) diff --git a/rbms/map_model.py b/rbms/map_model.py index f339dfd..be75930 100644 --- a/rbms/map_model.py +++ b/rbms/map_model.py @@ -1,5 +1,6 @@ from rbms.bernoulli_bernoulli.classes import BBRBM from rbms.bernoulli_gaussian.classes import BGRBM +from rbms.bm.potts import PBM from rbms.classes import EBM from rbms.ising_gaussian.classes import IGRBM from rbms.ising_ising.classes import IIRBM @@ -11,4 +12,5 @@ "BGRBM": BGRBM, "IGRBM": IGRBM, "IIRBM": IIRBM, + "PBM": PBM, } diff --git a/rbms/optim.py b/rbms/optim.py index e2d708e..8ab62a1 100644 --- a/rbms/optim.py +++ b/rbms/optim.py @@ -55,7 +55,7 @@ def step(self, closure=None): def setup_optim(optim: str, args: dict, params: EBM) -> list[Optimizer]: - match args["optim"]: + match optim: case "sgd": optim_class = SGD case "cossim": diff --git a/rbms/potts_bernoulli/classes.py b/rbms/potts_bernoulli/classes.py index 0e506f3..dbdbf73 100644 --- a/rbms/potts_bernoulli/classes.py +++ b/rbms/potts_bernoulli/classes.py @@ -161,6 +161,7 @@ def init_parameters(num_hiddens, dataset, device, dtype, var_init=0.0001): vbias, hbias, weight_matrix = _init_parameters( num_hiddens=num_hiddens, data=data, + weights=dataset.weights, device=device, dtype=dtype, var_init=var_init, diff --git a/rbms/potts_bernoulli/implement.py b/rbms/potts_bernoulli/implement.py index b4be369..47e6309 100644 --- a/rbms/potts_bernoulli/implement.py +++ b/rbms/potts_bernoulli/implement.py @@ -2,7 +2,9 @@ from torch import Tensor from torch.nn.functional import softmax +from rbms.bm.utils import get_freq_single_point from rbms.custom_fn import one_hot +from rbms.dataset.utils import convert_data def _sample_hiddens( @@ -11,8 +13,10 @@ def _sample_hiddens( dtype = weight_matrix.dtype num_visibles, num_states, num_hiddens = weight_matrix.shape weight_matrix_oh = weight_matrix.view(num_visibles * num_states, num_hiddens) - v_oh = one_hot(v.to(torch.int32), num_classes=num_states, dtype=dtype).view( - -1, num_visibles * num_states + v_oh = ( + one_hot(v.long(), num_classes=num_states) + .to(dtype=dtype) + .view(-1, num_visibles * num_states) ) mh = torch.sigmoid(beta * (hbias + v_oh @ weight_matrix_oh)) h = torch.bernoulli(mh).to(weight_matrix.dtype) @@ -40,8 +44,10 @@ def _compute_energy( ): dtype = weight_matrix.dtype num_visibles, num_states, num_hiddens = weight_matrix.shape - v_oh = one_hot(v.to(torch.int32), num_classes=num_states, dtype=dtype).view( - -1, num_visibles * num_states + v_oh = ( + one_hot(v.long(), num_classes=num_states) + .to(dtype=dtype) + .view(-1, num_visibles * num_states) ) vbias_oh = vbias.flatten() weight_matrix_oh = weight_matrix.view(num_visibles * num_states, num_hiddens) @@ -55,8 +61,10 @@ def _compute_energy_visibles( ): dtype = weight_matrix.dtype num_visibles, num_states, num_hiddens = weight_matrix.shape - v_oh = one_hot(v.to(torch.int32), num_classes=num_states, dtype=dtype).view( - -1, num_visibles * num_states + v_oh = ( + one_hot(v.long(), num_classes=num_states) + .to(dtype=dtype) + .view(-1, num_visibles * num_states) ) vbias_oh = vbias.flatten() @@ -90,17 +98,18 @@ def _compute_gradient( ): w_data = w_data.view(-1, 1, 1) w_chain = w_chain.view(-1, 1, 1) - int_dtype = torch.int32 dtype = weight_matrix.dtype num_states = weight_matrix.shape[1] # One-hot representation of the data - v_data_one_hot = one_hot(v_data.to(int_dtype), num_classes=num_states, dtype=dtype) - v_gen_one_hot = one_hot(v_chain.to(int_dtype), num_classes=num_states, dtype=dtype) + v_data_one_hot = one_hot(v_data.long(), num_classes=num_states).to(dtype) + v_gen_one_hot = one_hot(v_chain.long(), num_classes=num_states).to(dtype=dtype) # Turn the weights of the chains into normalized weights chain_weights = softmax(-w_chain, dim=0) w_chain_norm = chain_weights.sum() + # The weights should be normalized on the batch by dividing with the sum + # of the weights of the batch w_data_norm = w_data.sum() # Averages over data and generated samples v_data_mean = (v_data_one_hot * w_data).sum(0) / w_data_norm @@ -184,9 +193,11 @@ def _init_chains( else: v = start_v.to(weight_matrix.dtype) weight_matrix_oh = weight_matrix.view(num_visibles * num_states, num_hiddens) - v_oh = one_hot( - v.to(torch.int32), num_classes=num_states, dtype=weight_matrix_oh.dtype - ).view(-1, num_visibles * num_states) + v_oh = ( + one_hot(v.long(), num_classes=num_states) + .to(dtype=weight_matrix_oh.dtype) + .view(-1, num_visibles * num_states) + ) mv = torch.zeros(v.shape[0], v.shape[1], num_states) mh = torch.sigmoid(hbias + v_oh @ weight_matrix_oh) h = torch.bernoulli(mh) @@ -196,6 +207,7 @@ def _init_chains( def _init_parameters( num_hiddens: int, data: Tensor, + weights: Tensor, device: torch.device, dtype: torch.dtype, var_init: float = 1e-4, @@ -205,6 +217,14 @@ def _init_parameters( num_states = int(torch.max(data) + 1) all_states = torch.arange(num_states).reshape(-1, 1, 1).to(data.device) frequencies = (data == all_states).type(torch.float32).mean(1).to(device) + frequencies = get_freq_single_point( + convert_data["categorical"]["bernoulli"](data).view( + data.shape[0], data.shape[1], num_states + ), + weights / weights.sum(), + 1e-4, + ).T + frequencies = torch.clamp(frequencies, min=eps, max=(1.0 - eps)) vbias = ( (torch.log(frequencies) - 1.0 / num_states * torch.sum(torch.log(frequencies), 0))