Error in vae.train with Custom Data and Parameter Configuration

[sd68515]

Dear authors,

I encountered some issues when running the velovae package with my custom data, and I would like to report the following problems:

1. Cluster Key Specification in vv.preprocess
   When using vv.preprocess on my own dataset, the cluster_key parameter needs to explicitly specify the key name for clusters. It cannot always default to clusters. This behavior might not be clear to users and could benefit from clarification in the documentation.

2. Errors with vae.train
   While running the following code:

vae = vv.VAE(adata, 
             tmax=20, 
             dim_z=5, 
             device='cuda:0', 
             reverse_gene_mode=False,
             full_vb=True,
             discrete=True,
             count_distribution='auto'
            )

I am not sure whether the full_vb=True or discrete=True parameter caused the subsequent issue. When executing:

vae.train(adata,
          config=config,
          plot=False,
          gene_plot=gene_plot,
          figure_path=figure_path,
          embed='umap')

the code throws an error. Unfortunately, the specific error message is unclear, but I suspect it could be due to an incompatibility between these parameters and my dataset.
I found that modifying the file ./VeloVAE/velovae/model/vae.py enabled the code to execute successfully. Specifically, I added two additional input parameters u0=None and s0=None to the definition of the def _vae_risk_nb function. After making this change, the code worked as expected.

I am unsure if this modification aligns with the intended behavior of the model and whether it could have unintended side effects. I would appreciate your clarification or suggestions for a more robust solution.

Thank you for your help!

Kind regards,

Here are the specific error details I encountered:
Learning Rate based on Data Sparsity: 0.0005
--------------------------- Train a VeloVAE ---------------------------
********* Creating Training/Validation Datasets *********
********* Finished. *********
********* Creating optimizers *********
********* Finished. *********
********* Start training *********
********* Stage 1 *********
Total Number of Iterations Per Epoch: 100, test iteration: 198

TypeError Traceback (most recent call last)
Cell In[23], line 5
1 config = {
2 # You can change any hyperparameters here!
3 }
----> 5 vae.train(adata,
6 config=config,
7 plot=True,
8 gene_plot=gene_plot,
9 figure_path=figure_path,
10 embed='umap')

File ~\Documents\PythonScripts\VeloVAE-master\velovae\model\vae.py:1401, in VAE.train(self, adata, config, plot, gene_plot, cluster_key, figure_path, embed)
1395 stop_training = self._train_epoch(data_loader,
1396 test_set,
1397 optimizer_ode,
1398 optimizer,
1399 self.config["k_alt"])
1400 else:
-> 1401 stop_training = self._train_epoch(data_loader,
1402 test_set,
1403 optimizer,
1404 None,
1405 self.config["k_alt"])
1407 if plot and (epoch == 0 or (epoch+1) % self.config["save_epoch"] == 0):
1408 elbo_train = self._test(train_set,
1409 Xembed[self.train_idx],
1410 f"train{epoch+1}",
(...)
1414 plot,
1415 figure_path)

File ~\Documents\PythonScripts\VeloVAE-master\velovae\model\vae.py:1129, in VAE._train_epoch(self, train_loader, test_set, optimizer, optimizer2, K)
1127 for i, batch in enumerate(train_loader):
1128 if self.counter == 1 or self.counter % self.config["test_iter"] == 0:
-> 1129 elbo_test = self._test(test_set,
1130 None,
1131 self.counter,
1132 True)
1134 if len(self.loss_test) > 0: # update the number of epochs with dropping/converging ELBO
1135 if elbo_test - self.loss_test[-1] <= self.config["early_stop_thred"]:

File ~\Documents\PythonScripts\VeloVAE-master\velovae\model\vae.py:1821, in VAE._test(self, dataset, Xembed, testid, test_mode, gind, gene_plot, plot, path, **kwargs)
1819 if self.train_stage == 2:
1820 out_type.append("v")
-> 1821 out, elbo = self.pred_all(dataset.data, self.cell_labels, mode, out_type, gind)
1822 Uhat, Shat, t = out["uhat"], out["shat"], out["t"]
1824 G = dataset.data.shape[1]//2

File ~\Documents\PythonScripts\VeloVAE-master\velovae\model\vae.py:1639, in VAE.pred_all(self, data, cell_labels, mode, output, gene_idx)
1637 uhat_fw = uhat_fwlu_scale
1638 shat_fw = uhat_fwls_scale
-> 1639 loss = self.vae_risk((mu_tx, std_tx), p_t,
1640 (mu_zx, std_zx), p_z,
1641 data_in[:, :G], data_in[:, G:],
1642 uhatlu_scale, shatls_scale,
1643 vu, vs,
1644 u0, s0,
1645 uhat_fw, shat_fw,
1646 vu_fw, vs_fw,
1647 u1, s1,
1648 None)
1649 elbo = elbo - (B/N)*loss
1650 if "uhat" in output and gene_idx is not None:

TypeError: VAE._vae_risk_nb() takes from 9 to 19 positional arguments but 20 were given