Added --save-latent flag to train-vae for latent mu/sigma and loss hi…

docs/cli.md

@@ -48,6 +48,7 @@ embkit model train-vae INPUT_PATH [OPTIONS]
 | `--seed` | `42` | Random seed |
 | `--bfloat16` | false | Use bfloat16 dtype for reduced memory usage |
 | `--save-stats` | false | Save training statistics alongside the model |
+| `--save-latent` | false | Save latent statistics alongside the model (mu, sigma, losses) |
 | `--sampling/--no-sampling` | `true` | Enable reparameterization sampling (VAE). Use `--no-sampling` for a standard autoencoder. |
 
 For HDF5 input (`--group`), normalization must be `none`.

src/embkit/commands/model.py

@@ -37,6 +37,7 @@
 @click.option("--schedule", "-s", type=str, default=None, help="20:0,20:0.1,40:.3,40:.4")
 @click.option("--loss", type=click.Choice(["mse", "bce", "bce-logit"]), default="bce-logit")
 @click.option("--save-stats", is_flag=True)
+@click.option("--save-latent", is_flag=True)
 @click.option("--zero-mask", default=None, type=float)
 @click.option("--seed", default=42, type=int)
 @click.option("--bfloat16", is_flag=True)
@@ -54,6 +55,7 @@ def train_vae(input_path: str,
               schedule:str,
               zero_mask: float,
               save_stats: bool,
+              save_latent:bool,
               seed: int,
               bfloat16: bool,
               sampling: bool
@@ -143,6 +145,46 @@ def train_vae(input_path: str,
         stats.to_csv(stats_path, sep="\t")
         click.echo(f"Stats saved, to {stats_path}")
 
+    if save_latent:
+        vae.eval() # Makes output deterministic. Otherwise, model in training.
+        losses_df = pd.DataFrame({
+            "epoch": range(1, len(vae.history["loss"]) + 1),
+            "loss": vae.history["loss"],
+            "recon": vae.history["recon"],
+            "kl": vae.history["kl"]})
+
+        losses_df.to_csv(f"{out}.losses_stats.tsv", sep="\t", index=False)
+        click.echo(f"Training losses saved, to {out}.losses_stats.tsv")
+
+        exportloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
+        all_mu = []
+        all_logvar = []
+
+        with torch.no_grad():
+            for batch in exportloader:
+                    x_tensor = batch[0] if isinstance(batch, (tuple, list)) else batch
+                    x_tensor = x_tensor.to(device=device, dtype=dtype)
+                    _, mu, logvar, _ = vae(x_tensor)
+                    all_mu.append(mu.cpu())
+                    all_logvar.append(logvar.cpu())
+
+        mu = torch.cat(all_mu, dim=0)
+        logvar = torch.cat(all_logvar, dim=0)
+        std = torch.sqrt(torch.exp(logvar)) # sigma
+
+        index = df.index if df is not None else None
+
+        mu_np = mu.to(device="cpu", dtype=torch.float32).detach().numpy()
+        std_np = std.to(device="cpu", dtype=torch.float32).detach().numpy()
+        mu_df = pd.DataFrame(mu_np, index=index, columns=[f"mu_{i}" for i in range(mu.shape[1])])
+        std_df = pd.DataFrame(std_np, index=index, columns=[f"std_{i}" for i in range(std.shape[1])])
+        mu_path = f"{out}.latent_mu.tsv"
+        std_path = f"{out}.latent_std.tsv"
+        mu_df.to_csv(mu_path, sep="\t")
+        std_df.to_csv(std_path, sep="\t")
+        click.echo(f"Latent mu saved, to {mu_path}")
+        click.echo(f"Latent std (sigma) saved, to {std_path}")
+
 
 @model.command()
 @click.argument("input_path", type=click.Path(exists=True, dir_okay=False, readable=True, path_type=str))

tests/commands/test_model.py

@@ -128,11 +128,88 @@ def test_train_vae_tsv_branches(self, vae_cls, fit_mock, save_mock):
         self.assertEqual(result.exit_code, 0, msg=result.output)
         self.assertIn("No output path provided, using default naming.", result.output)
         self.assertIn("Stats saved, to vae_latent256_epochs20.model.stats.tsv", result.output)
+
         fit_mock.assert_called_once()
         self.assertEqual(fit_mock.call_args.kwargs["loss"], model_cmd.mse)
         self.assertEqual(fit_mock.call_args.kwargs["beta_schedule"], [(0.2, 1), (0.4, 1)])
         save_mock.assert_called_once()
 
+    def test_train_vae_save_latent(self):
+        with self.runner.isolated_filesystem():
+            with open("rna.tsv", "w", encoding="utf-8") as f:
+                f.write(
+                    "sample\tG1\tG2\tG3\tG4\n"
+                    "s1\t1\t2\t3\t4\n"
+                    "s2\t4\t3\t2\t1\n"
+                    "s3\t2\t2\t2\t2\n"
+                )
+            # not bflot16... possible training bug?  
+            result = self.runner.invoke(
+                cli_main,
+                [
+                    "model",
+                    "train-vae",
+                    "rna.tsv",
+                    "--epochs",
+                    "2",
+                    "--latent",
+                    "2",
+                    "--encode-layers",
+                    "4",
+                    "--decode-layers",
+                    "4",
+                    "--save-stats",
+                    "--save-latent",
+                    "--out",
+                    "vae.model",
+                ],
+            )
+
+            self.assertEqual(result.exit_code, 0, msg=result.output)
+            self.assertTrue(Path("vae.model.stats.tsv").exists())
+            losses_lines = Path("vae.model.losses_stats.tsv").read_text(encoding="utf-8").splitlines()
+            self.assertEqual(losses_lines[0].split("\t"), ["epoch", "loss", "recon", "kl"])
+            self.assertEqual(len(losses_lines) - 1, 2)
+            mu_lines = Path("vae.model.latent_mu.tsv").read_text(encoding="utf-8").splitlines()
+            self.assertEqual(mu_lines[0].split("\t"), ["sample", "mu_0", "mu_1"])
+            self.assertEqual([line.split("\t")[0] for line in mu_lines[1:]], ["s1", "s2", "s3"])
+            std_lines = Path("vae.model.latent_std.tsv").read_text(encoding="utf-8").splitlines()
+            std_values = [float(v) for line in std_lines[1:] for v in line.split("\t")[1:]]
+            self.assertTrue(all(v > 0 for v in std_values))
+
+        with self.runner.isolated_filesystem():
+            writer = H5Writer("matrix.h5", "rna", index=["s1", "s2"], columns=["G1", "G2"])
+            writer.set_irow(0, [1.0, 2.0])
+            writer.set_irow(1, [3.0, 4.0])
+            writer.close()
+
+            result = self.runner.invoke(
+                cli_main,
+                [
+                    "model",
+                    "train-vae",
+                    "matrix.h5",
+                    "--group",
+                    "rna",
+                    "--epochs",
+                    "1",
+                    "--latent",
+                    "2",
+                    "--encode-layers",
+                    "4",
+                    "--decode-layers",
+                    "4",
+                    "--save-latent",
+                    "--out",
+                    "vae_h5.model",
+                ],
+            )
+
+            self.assertEqual(result.exit_code, 0, msg=result.output)
+            mu_lines = Path("vae_h5.model.latent_mu.tsv").read_text(encoding="utf-8").splitlines()
+            self.assertEqual(mu_lines[0].split("\t"), ["", "mu_0", "mu_1"])
+            self.assertEqual([line.split("\t")[0] for line in mu_lines[1:]], ["0", "1"])
+
     @patch.object(model_cmd, "save")
     @patch.object(model_cmd, "fit_vae")
     @patch.object(model_cmd, "dataframe_loader", return_value="loader")