pytorch 2.6 weights_only compatibility

Author: jykr

src/simba_plus/model_prox.py

@@ -622,7 +622,7 @@ def training_step(self, batch, batch_idx):
             t0 = time.time()
             if "peak" in batch.node_types:
                 pid = batch["peak"].n_id.cpu()
-                herit_loss_value = self.herit_loss_lam * self.herit_loss(
+                herit_loss_value = self.herit_loss(
                     mu_dict["peak"],
                     pid,
                 )

src/simba_plus/post_training/factors.py

@@ -144,12 +144,10 @@ def summarize_enrichments(gene_enrichment, gene_bot_enrichment):
 
 def main(args, logger=None):
     if not logger:
-        logger = setup_logging(
-            "simba+factor", log_dir=os.path.dirname(args.adata_prefix)
-        )
+        logger = setup_logging("simba+factor", log_dir=args.adata_prefix)
     # Loading pretrained results as `simba+ heritability` takes long time to run
     if args.output_dir is None:
-        args.output_dir = f"{os.path.dirname(args.adata_prefix)}/factors/"
+        args.output_dir = f"{args.adata_prefix}/factors/"
         os.makedirs(args.output_dir, exist_ok=True)
 
     adata_C = sc.read_h5ad(f"{args.adata_prefix}/adata_C{args.version_suffix}.h5ad")
@@ -174,6 +172,7 @@ def main(args, logger=None):
             plt.close(fig_bot)
         sc.set_figure_params(vector_friendly=True)
         fig = sc.pl.umap(adata_C, color=args.cell_type_label, return_fig=True)
+        fig.axes[0].set_box_aspect(1)
         pdf.savefig(fig, bbox_inches="tight")
         plt.close(fig)
         factor_labels = adata_G.uns["factor_enrichments_summary"]
@@ -182,6 +181,9 @@ def main(args, logger=None):
             return_fig=True,
             factor_labels=factor_labels,
         )
+        for i, ax in enumerate(factor_umap.axes):
+            if i % 2 == 0:
+                ax.set_box_aspect(1)
         pdf.savefig(factor_umap, bbox_inches="tight")
         plt.close(factor_umap)
     logger.info(f"Generated factor report in {output_filename}.")

src/simba_plus/train.py

@@ -54,7 +54,7 @@ def setup_logging(checkpoint_dir):
 
     # Create timestamped log file
     timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-    log_file = os.path.join(checkpoint_dir, f"train_{timestamp}.log")
+    log_file = os.path.join(checkpoint_dir, f"simba+train_{timestamp}.log")
 
     # Configure root logger
     logger = logging.getLogger()
@@ -329,6 +329,8 @@ def train(
                 if load_checkpoint
                 else None
             ),
+            weights_only=False,
+
         )
         return checkpoint_callback.last_model_path