Merge branch 'master' into harald/qT-tempering

src/sampling/laris.py

@@ -135,7 +135,8 @@ def sample_arborescence_from_weighted_graph(graph: nx.DiGraph,
 
 def _sample_feasible_arc(weighted_arcs):
     # weighted_arcs is a list of 3-tuples (u, v, weight)
-    unnorm_probs = torch.stack([w for u, v, w in weighted_arcs])
+    # weights are negative: need transformation
+    unnorm_probs = 1 / (-torch.stack([w for u, v, w in weighted_arcs]))
     probs = unnorm_probs / unnorm_probs.sum()
     c = np.random.choice(np.arange(len(weighted_arcs)), p=probs.numpy())
     return weighted_arcs[c][:2], probs[c]

src/simul.py

@@ -104,7 +104,7 @@ def make_anndata(obs, raw_counts, chr_dataframe, c, z, mu, tree, obs_names: list
     return adata
 
 
-def simulate_full_dataset(config: Config, eps_a=5., eps_b=50., mu0=1., lambda0=10.,
+def simulate_full_dataset(config: Config, eps_a=500., eps_b=50000., mu0=1., lambda0=1000.,
                           alpha0=500., beta0=50., dir_delta: [float | list[float]] = 1., tree=None, raw_reads=True,
                           chr_df: pd.DataFrame | None = None, nans: bool = False,
                           fixed_z:torch.Tensor = None, cne_length_factor: int = 0):

src/variational_distributions/var_dists.py

@@ -1217,6 +1217,8 @@ def _normalize_graph_weights(self, w_matrix):
             norm_w_matrix[...] = w_matrix - math_utils.nanlogsumexp(w_matrix, dim=1, keepdim=True)
         else:
             raise ValueError(f"normalization method {self._norm_method} not recognized")
+        # weight shouldn't be equal to 0 for stability
+        norm_w_matrix = torch.clamp(norm_w_matrix, max=-1e-8)
 
         return norm_w_matrix
 

tests/tests_utils/test_data_handling.py

@@ -6,7 +6,7 @@
 
 import simul
 from inference.victree import VICTree
-from utils.config import Config
+from utils.config import Config, set_seed
 from utils.data_handling import DataHandler
 from tests.data.generate_data import generate_2chr_adata
 from variational_distributions.joint_dists import VarTreeJointDist
@@ -18,6 +18,7 @@ class dataHandlingTestCase(unittest.TestCase):
 
     def setUp(self) -> None:
 
+        set_seed(0)
         self.output_dir = "./test_output"
         if not os.path.exists(self.output_dir):
             os.mkdir(self.output_dir)
@@ -39,7 +40,7 @@ def test_write_output(self):
         out_file = os.path.join(self.output_dir, 'out_test.h5')
         # run victree
         config = Config(n_nodes=4, n_cells=20, n_states=4,
-                        n_run_iter=3, sieving_size=2, n_sieving_iter=2)
+                        n_run_iter=3)
         adata = generate_2chr_adata(config)
         data_handler = DataHandler(adata=adata)
         obs = data_handler.norm_reads