model upload

branchsbm/branch_flow_net_train.py

@@ -0,0 +1,348 @@
+import os
+import sys
+sys.path.append("./BranchSBM")
+import torch
+import wandb
+import matplotlib.pyplot as plt
+import pytorch_lightning as pl
+from torch.optim import AdamW
+from torchmetrics.functional import mean_squared_error
+from torchdyn.core import NeuralODE
+from networks.utils import flow_model_torch_wrapper
+from utils import wasserstein_distance, plot_lidar
+from branchsbm.ema import EMA
+
+class BranchFlowNetTrainBase(pl.LightningModule):
+    def __init__(
+        self,
+        flow_matcher,
+        flow_nets,
+        skipped_time_points=None,
+        ot_sampler=None,
+        args=None,
+    ):
+        super().__init__()
+        self.args = args
+        
+        self.flow_matcher = flow_matcher
+        self.flow_nets = flow_nets # list of flow networks for each branch
+        self.ot_sampler = ot_sampler
+        self.skipped_time_points = skipped_time_points
+
+        self.optimizer_name = args.flow_optimizer
+        self.lr = args.flow_lr
+        self.weight_decay = args.flow_weight_decay
+        self.whiten = args.whiten
+        self.working_dir = args.working_dir
+        
+        #branching 
+        self.branches = len(flow_nets)
+
+    def forward(self, t, xt, branch_idx):
+        # output velocity given branch_idx
+        return self.flow_nets[branch_idx](t, xt)
+
+    def _compute_loss(self, main_batch):
+        
+        x0s = [main_batch["x0"][0]]
+        w0s = [main_batch["x0"][1]]
+        
+        x1s_list = []
+        w1s_list = [] 
+        
+        if self.branches > 1:
+            for i in range(self.branches):
+                x1s_list.append([main_batch[f"x1_{i+1}"][0]])
+                w1s_list.append([main_batch[f"x1_{i+1}"][1]])
+        else:
+            x1s_list.append([main_batch["x1"][0]])
+            w1s_list.append([main_batch["x1"][1]])
+        
+        assert len(x1s_list) == self.branches, "Mismatch between x1s_list and expected branches"
+        
+        loss = 0
+        for branch_idx in range(self.branches):
+            ts, xts, uts = self._process_flow(x0s, x1s_list[branch_idx], branch_idx)
+
+            t = torch.cat(ts)
+            xt = torch.cat(xts)
+            ut = torch.cat(uts)
+            vt = self(t[:, None], xt, branch_idx)
+
+            loss += mean_squared_error(vt, ut)
+
+        return loss
+
+    def _process_flow(self, x0s, x1s, branch_idx):
+        ts, xts, uts = [], [], []
+        t_start = self.timesteps[0]
+
+        for i, (x0, x1) in enumerate(zip(x0s, x1s)):
+            
+            x0, x1 = torch.squeeze(x0), torch.squeeze(x1)
+
+            if self.ot_sampler is not None:
+                x0, x1 = self.ot_sampler.sample_plan(
+                    x0,
+                    x1,
+                    replace=True,
+                )
+            if self.skipped_time_points and i + 1 >= self.skipped_time_points[0]:
+                t_start_next = self.timesteps[i + 2]
+            else:
+                t_start_next = self.timesteps[i + 1]
+            
+            # edit to sample from correct flow matcher
+            t, xt, ut = self.flow_matcher.sample_location_and_conditional_flow(
+                x0, x1, t_start, t_start_next, branch_idx
+            )
+
+            ts.append(t)
+
+            xts.append(xt)
+            uts.append(ut)
+            t_start = t_start_next
+        return ts, xts, uts
+
+    def training_step(self, batch, batch_idx):
+        if self.args.data_type in ["scrna", "tahoe"]:
+            main_batch = batch[0]["train_samples"][0]
+        else: 
+            main_batch = batch["train_samples"][0]
+            
+        print("Main batch length")
+        print(len(main_batch["x0"]))
+        self.timesteps = torch.linspace(0.0, 1.0, len(main_batch["x0"])).tolist()
+        loss = self._compute_loss(main_batch)
+        if self.flow_matcher.alpha != 0:
+            self.log(
+                "FlowNet/mean_geopath_cfm",
+                (self.flow_matcher.geopath_net_output.abs().mean()),
+                on_step=False,
+                on_epoch=True,
+                prog_bar=True,
+            )
+
+        self.log(
+            "FlowNet/train_loss_cfm",
+            loss,
+            on_step=False,
+            on_epoch=True,
+            prog_bar=True,
+            logger=True,
+        )
+        
+        
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+        if self.args.data_type in ["scrna", "tahoe"]:
+            main_batch = batch[0]["val_samples"][0]
+        else: 
+            main_batch = batch["val_samples"][0]
+            
+        self.timesteps = torch.linspace(0.0, 1.0, len(main_batch["x0"])).tolist()
+        val_loss = self._compute_loss(main_batch)
+        self.log(
+            "FlowNet/val_loss_cfm",
+            val_loss,
+            on_step=False,
+            on_epoch=True,
+            prog_bar=True,
+            logger=True,
+        )
+        return val_loss
+
+    def optimizer_step(self, *args, **kwargs):
+        super().optimizer_step(*args, **kwargs)
+        
+        for net in self.flow_nets:
+            if isinstance(net, EMA):
+                net.update_ema()
+
+    def configure_optimizers(self):
+        if self.optimizer_name == "adamw":
+            optimizer = AdamW(
+                self.parameters(),
+                lr=self.lr,
+                weight_decay=self.weight_decay,
+            )
+        elif self.optimizer_name == "adam":
+            optimizer = torch.optim.Adam(
+                self.parameters(),
+                lr=self.lr,
+            )
+
+        return optimizer
+
+
+class FlowNetTrainTrajectory(BranchFlowNetTrainBase):
+    def test_step(self, batch, batch_idx):
+        data_type = self.args.data_type
+        node = NeuralODE(
+            flow_model_torch_wrapper(self.flow_nets),
+            solver="euler",
+            sensitivity="adjoint",
+            atol=1e-5,
+            rtol=1e-5,
+        )
+
+        t_exclude = self.skipped_time_points[0] if self.skipped_time_points else None
+        if t_exclude is not None:
+            traj = node.trajectory(
+                batch[t_exclude - 1],
+                t_span=torch.linspace(
+                    self.timesteps[t_exclude - 1], self.timesteps[t_exclude], 101
+                ),
+            )
+            X_mid_pred = traj[-1]
+            traj = node.trajectory(
+                batch[t_exclude - 1],
+                t_span=torch.linspace(
+                    self.timesteps[t_exclude - 1],
+                    self.timesteps[t_exclude + 1],
+                    101,
+                ),
+            )
+            
+            EMD = wasserstein_distance(X_mid_pred, batch[t_exclude], p=1)
+            self.final_EMD = EMD
+
+            self.log("test_EMD", EMD, on_step=False, on_epoch=True, prog_bar=True)
+
+class FlowNetTrainCell(BranchFlowNetTrainBase):
+    def test_step(self, batch, batch_idx):
+        x0 = batch[0]["test_samples"][0]["x0"][0]  # [B, D]
+        dataset_points = batch[0]["test_samples"][0]["dataset"][0]  # full dataset, [N, D]
+        t_span = torch.linspace(0, 1, 101)
+
+        all_trajs = []
+
+        for i, flow_net in enumerate(self.flow_nets):
+            node = NeuralODE(
+                flow_model_torch_wrapper(flow_net),
+                solver="euler",
+                sensitivity="adjoint",
+            )
+
+            with torch.no_grad():
+                traj = node.trajectory(x0, t_span).cpu()  # [T, B, D]
+
+            if self.whiten:
+                traj_shape = traj.shape
+                traj = traj.reshape(-1, traj.shape[-1])
+                traj = self.trainer.datamodule.scaler.inverse_transform(
+                    traj.cpu().detach().numpy()
+                ).reshape(traj_shape)
+                dataset_points = self.trainer.datamodule.scaler.inverse_transform(
+                    dataset_points.cpu().detach().numpy()
+                )
+
+            traj = torch.tensor(traj)
+            traj = torch.transpose(traj, 0, 1)  # [B, T, D]
+            all_trajs.append(traj)
+
+        dataset_2d = dataset_points[:, :2] if isinstance(dataset_points, torch.Tensor) else dataset_points[:, :2]
+
+        # ===== Plot all 2D trajectories together with dataset and start/end points =====
+        fig, ax = plt.subplots(figsize=(6, 5))
+        dataset_2d = dataset_2d.cpu().numpy()
+        ax.scatter(dataset_2d[:, 0], dataset_2d[:, 1], c="gray", s=1, alpha=0.5, label="Dataset", zorder=1)
+        for traj in all_trajs:
+            traj_2d = traj[..., :2]  # [B, T, 2]
+            for i in range(traj_2d.shape[0]):
+                ax.plot(traj_2d[i, :, 0], traj_2d[i, :, 1], alpha=0.8, zorder=2)
+                ax.scatter(traj_2d[i, 0, 0], traj_2d[i, 0, 1], c='green', s=10, label="t=0" if i == 0 else "", zorder=3)
+                ax.scatter(traj_2d[i, -1, 0], traj_2d[i, -1, 1], c='red', s=10, label="t=1" if i == 0 else "", zorder=3)
+
+        ax.set_title("All Branch Trajectories (2D) with Dataset")
+        ax.set_xlabel("x")
+        ax.set_ylabel("y")
+        plt.axis("equal")
+        handles, labels = ax.get_legend_handles_labels()
+        if labels:
+            ax.legend()
+            
+        save_path = f'./figures/{self.args.data_name}'
+        
+        os.makedirs(save_path, exist_ok=True)
+        plt.savefig(f'{save_path}/{self.args.data_name}_all_branches.png', dpi=300)
+        plt.close()
+
+        # ===== Plot each 2D trajectory separately with dataset and endpoints =====
+        for i, traj in enumerate(all_trajs):
+            traj_2d = traj[..., :2]
+            fig, ax = plt.subplots(figsize=(6, 5))
+            ax.scatter(dataset_2d[:, 0], dataset_2d[:, 1], c="gray", s=1, alpha=0.5, label="Dataset", zorder=1)
+            for j in range(traj_2d.shape[0]):
+                ax.plot(traj_2d[j, :, 0], traj_2d[j, :, 1], alpha=0.9, zorder=2)
+                ax.scatter(traj_2d[j, 0, 0], traj_2d[j, 0, 1], c='green', s=12, label="t=0" if j == 0 else "", zorder=3)
+                ax.scatter(traj_2d[j, -1, 0], traj_2d[j, -1, 1], c='red', s=12, label="t=1" if j == 0 else "", zorder=3)
+
+            ax.set_title(f"Branch {i + 1} Trajectories (2D) with Dataset")
+            ax.set_xlabel("x")
+            ax.set_ylabel("y")
+            plt.axis("equal")
+            handles, labels = ax.get_legend_handles_labels()
+            if labels:
+                ax.legend()
+            plt.savefig(f'{save_path}/{self.args.data_name}_branch_{i + 1}.png', dpi=300)
+            plt.close()
+    
+class FlowNetTrainLidar(BranchFlowNetTrainBase):
+    def test_step(self, batch, batch_idx):
+        main_batch = batch["test_samples"][0]
+        metric_batch = batch["metric_samples"][0]
+                
+        x0 = main_batch["x0"][0] # [B, D]
+        cloud_points = main_batch["dataset"][0]  # full dataset, [N, D]
+        t_span = torch.linspace(0, 1, 101)
+
+        all_trajs = []
+
+        for i, flow_net in enumerate(self.flow_nets):
+            node = NeuralODE(
+                flow_model_torch_wrapper(flow_net),
+                solver="euler",
+                sensitivity="adjoint",
+            )
+
+            with torch.no_grad():
+                traj = node.trajectory(x0, t_span).cpu()  # [T, B, D]
+
+            if self.whiten:
+                traj_shape = traj.shape
+                traj = traj.reshape(-1, 3)
+                traj = self.trainer.datamodule.scaler.inverse_transform(
+                    traj.cpu().detach().numpy()
+                ).reshape(traj_shape)
+
+            traj = torch.tensor(traj)
+            traj = torch.transpose(traj, 0, 1)  # [B, T, D]
+            all_trajs.append(traj)
+
+        # Inverse-transform the point cloud once
+        if self.whiten:
+            cloud_points = torch.tensor(
+                self.trainer.datamodule.scaler.inverse_transform(
+                    cloud_points.cpu().detach().numpy()
+                )
+            )
+
+        # ===== Plot all trajectories together =====
+        fig = plt.figure(figsize=(6, 5))
+        ax = fig.add_subplot(111, projection="3d", computed_zorder=False)
+        ax.view_init(elev=30, azim=-115, roll=0)
+        for i, traj in enumerate(all_trajs):
+            plot_lidar(ax, cloud_points, xs=traj, branch_idx=i)
+        plt.savefig('./figures/lidar/lidar_all_branches.png', dpi=300)
+        plt.close()
+
+        # ===== Plot each trajectory separately =====
+        for i, traj in enumerate(all_trajs):
+            fig = plt.figure(figsize=(6, 5))
+            ax = fig.add_subplot(111, projection="3d", computed_zorder=False)
+            ax.view_init(elev=30, azim=-115, roll=0)
+            plot_lidar(ax, cloud_points, xs=traj, branch_idx=i)
+            plt.savefig(f'./figures/lidar/lidar_branch_{i + 1}.png', dpi=300)
+            plt.close()

branchsbm/branch_growth_net_train.py

@@ -0,0 +1,514 @@
+import os
+import sys
+sys.path.append("./BranchSBM")
+import torch
+import wandb
+import matplotlib.pyplot as plt
+import pytorch_lightning as pl
+from torch.optim import AdamW
+from torchmetrics.functional import mean_squared_error
+from torchdyn.core import NeuralODE
+import numpy as np
+import lpips
+from networks.utils import flow_model_torch_wrapper
+from utils import wasserstein_distance, plot_lidar
+from branchsbm.ema import EMA
+from torchdiffeq import odeint as odeint2
+from losses.energy_loss import EnergySolver, ReconsLoss
+
+class GrowthNetTrain(pl.LightningModule):
+    def __init__(
+        self,
+        flow_nets,
+        growth_nets,
+        skipped_time_points=None,
+        ot_sampler=None,
+        args=None,
+        
+        state_cost=None,
+        data_manifold_metric=None,
+        
+        joint = False
+    ):
+        super().__init__()
+        #self.save_hyperparameters()
+        self.flow_nets = flow_nets
+        
+        if not joint:
+            for param in self.flow_nets.parameters():
+                param.requires_grad = False
+        
+        self.growth_nets = growth_nets # list of growth networks for each branch
+        
+        self.ot_sampler = ot_sampler
+        self.skipped_time_points = skipped_time_points
+
+        self.optimizer_name = args.growth_optimizer
+        self.lr = args.growth_lr
+        self.weight_decay = args.growth_weight_decay
+        self.whiten = args.whiten
+        self.working_dir = args.working_dir
+        
+        self.args = args
+        
+        #branching 
+        self.state_cost = state_cost
+        self.data_manifold_metric = data_manifold_metric
+        self.branches = len(growth_nets)
+        self.metric_clusters = args.metric_clusters
+        
+        self.recons_loss = ReconsLoss()
+        
+        # loss weights
+        self.lambda_energy = args.lambda_energy
+        self.lambda_mass = args.lambda_mass
+        self.lambda_match = args.lambda_match
+        self.lambda_recons = args.lambda_recons
+        
+        self.joint = joint
+
+    def forward(self, t, xt, branch_idx):
+        # output growth rate given branch_idx
+        return self.growth_nets[branch_idx](t, xt)
+
+    def _compute_loss(self, main_batch,  metric_samples_batch=None, validation=False):
+        x0s = main_batch["x0"][0]
+        w0s = main_batch["x0"][1]
+        x1s_list = []
+        w1s_list = [] 
+        
+        if self.branches > 1:
+            for i in range(self.branches):
+                x1s_list.append([main_batch[f"x1_{i+1}"][0]])
+                w1s_list.append([main_batch[f"x1_{i+1}"][1]])
+        else:
+            x1s_list.append([main_batch["x1"][0]])
+            w1s_list.append([main_batch["x1"][1]])
+        
+        if self.args.manifold:
+            #changed
+            if self.metric_clusters == 4:
+                branch_sample_pairs = [
+                    (metric_samples_batch[0], metric_samples_batch[1]),  # x0 → x1_1 (branch 1)
+                    (metric_samples_batch[0], metric_samples_batch[2]),  # x0 → x1_2 (branch 2)
+                    (metric_samples_batch[0], metric_samples_batch[3]),
+                ]
+            elif self.metric_clusters == 3:
+                branch_sample_pairs = [
+                    (metric_samples_batch[0], metric_samples_batch[1]),  # x0 → x1_1 (branch 1)
+                    (metric_samples_batch[0], metric_samples_batch[2]),  # x0 → x1_2 (branch 2)
+                ]
+            elif self.metric_clusters == 2 and self.branches == 2:
+                branch_sample_pairs = [
+                    (metric_samples_batch[0], metric_samples_batch[1]),  # x0 → x1_1 (branch 1)
+                    (metric_samples_batch[0], metric_samples_batch[1]),  # x0 → x1_2 (branch 2)
+                ]
+            else:
+                branch_sample_pairs = [
+                    (metric_samples_batch[0], metric_samples_batch[1]),  # x0 → x1_1 (branch 1)
+                ]
+        
+        batch_size = x0s.shape[0]
+        
+        assert len(x1s_list) == self.branches, "Mismatch between x1s_list and expected branches"
+        
+        energy_loss = [0.] * self.branches
+        mass_loss = 0.
+        neg_weight_penalty = 0.
+        match_loss = [0.] * self.branches
+        recons_loss = [0.] * self.branches
+        
+        dtype = x0s[0].dtype
+        #w0s = torch.zeros((batch_size, 1), dtype=dtype) 
+        m0s = torch.zeros_like(w0s, dtype=dtype)
+        start_state = (x0s, w0s, m0s)
+        
+        xt = [x0s.clone() for _ in range(self.branches)]
+        w0_branch = torch.zeros_like(w0s, dtype=dtype)
+        w0_branches = []
+        w0_branches.append(w0s)
+        for _ in range(self.branches - 1):
+            w0_branches.append(w0_branch)
+        #w0_branches = [w0_branch.clone() for _ in range(self.branches - 1)]
+        wt = w0_branches
+        
+        mt = [m0s.clone() for _ in range(self.branches)]
+        
+        # loop through timesteps
+        for s, t in zip(self.timesteps[:-1], self.timesteps[1:]):
+            time = torch.Tensor([s, t])
+            
+            total_w_t = 0
+            # loop through branches
+            for i in range(self.branches):
+                
+                if self.args.manifold:
+                    start_samples, end_samples = branch_sample_pairs[i]
+                    samples = torch.cat([start_samples, end_samples], dim=0)
+                    
+                # initialize weight and energy
+                start_state = (xt[i], wt[i], mt[i])
+                
+                # loop over timesteps
+                xt_next, wt_next, mt_next = self.take_step(time, start_state, i, samples)
+                
+                # placeholders for next state
+                xt_last = xt_next[-1]
+                wt_last = wt_next[-1]
+                mt_last = mt_next[-1]
+                
+                total_w_t += wt_last
+
+                energy_loss[i] += (mt_last - mt[i])
+                neg_weight_penalty += torch.relu(-wt_last).sum()
+                
+                # update branch state
+                xt[i] = xt_last.clone().detach()
+                wt[i] = wt_last.clone().detach()
+                mt[i] = mt_last.clone().detach()
+
+            # calculate mass loss from all branches
+            target = torch.ones_like(total_w_t)
+            mass_loss += mean_squared_error(total_w_t, target)
+        
+        # calculate loss that matches final weights 
+        for i in range(self.branches):
+            match_loss[i] = mean_squared_error(wt[i], w1s_list[i][0])
+            # compute reconstruction loss
+            recons_loss[i] = self.recons_loss(xt[i], x1s_list[i][0])
+        
+        # average across times
+        mass_loss = mass_loss / len(self.timesteps)
+        
+        
+        # mean across branches
+        energy_loss = torch.mean(torch.stack(energy_loss))
+        match_loss = torch.mean(torch.stack(match_loss))
+        recons_loss = torch.mean(torch.stack(recons_loss))
+        
+        loss = (self.lambda_energy * energy_loss) + (self.lambda_mass * (mass_loss + neg_weight_penalty)) + (self.lambda_match * match_loss) \
+            + (self.lambda_recons * recons_loss)
+            
+        if self.joint:
+            if validation:
+                self.log("JointTrain/val_mass_loss", mass_loss, on_step=False, on_epoch=True, prog_bar=True)
+                self.log("JointTrain/val_neg_penalty_loss", neg_weight_penalty, on_step=False, on_epoch=True, prog_bar=True)
+                self.log("JointTrain/val_match_loss", match_loss, on_step=False, on_epoch=True, prog_bar=True)
+                self.log("JointTrain/val_energy_loss", energy_loss, on_step=False, on_epoch=True, prog_bar=True)
+                self.log("JointTrain/val_recons_loss", recons_loss, on_step=False, on_epoch=True, prog_bar=True)
+            else:
+                self.log("JointTrain/train_mass_loss", mass_loss, on_step=False, on_epoch=True, prog_bar=True)
+                self.log("JointTrain/train_neg_penalty_loss", neg_weight_penalty, on_step=False, on_epoch=True, prog_bar=True)
+                self.log("JointTrain/train_match_loss", match_loss, on_step=False, on_epoch=True, prog_bar=True)
+                self.log("JointTrain/train_energy_loss", energy_loss, on_step=False, on_epoch=True, prog_bar=True)
+                self.log("JointTrain/train_recons_loss", recons_loss, on_step=False, on_epoch=True, prog_bar=True)
+        else:
+            if validation:
+                self.log("GrowthNet/val_mass_loss", mass_loss, on_step=False, on_epoch=True, prog_bar=True)
+                self.log("GrowthNet/val_neg_penalty_loss", neg_weight_penalty, on_step=False, on_epoch=True, prog_bar=True)
+                self.log("GrowthNet/val_match_loss", match_loss, on_step=False, on_epoch=True, prog_bar=True)
+                self.log("GrowthNet/val_energy_loss", energy_loss, on_step=False, on_epoch=True, prog_bar=True)
+                self.log("GrowthNet/val_recons_loss", recons_loss, on_step=False, on_epoch=True, prog_bar=True)
+            else:
+                self.log("GrowthNet/train_mass_loss", mass_loss, on_step=False, on_epoch=True, prog_bar=True)
+                self.log("GrowthNet/train_neg_penalty_loss", neg_weight_penalty, on_step=False, on_epoch=True, prog_bar=True)
+                self.log("GrowthNet/train_match_loss", match_loss, on_step=False, on_epoch=True, prog_bar=True)
+                self.log("GrowthNet/train_energy_loss", energy_loss, on_step=False, on_epoch=True, prog_bar=True)
+                self.log("GrowthNet/train_recons_loss", recons_loss, on_step=False, on_epoch=True, prog_bar=True)
+        
+        return loss
+    
+    def take_step(self, t, start_state, branch_idx, samples=None):
+        
+        flow_net = self.flow_nets[branch_idx]
+        growth_net = self.growth_nets[branch_idx]
+        
+        
+        x_t, w_t, m_t = odeint2(EnergySolver(flow_net, growth_net, self.state_cost, self.data_manifold_metric, samples), start_state, t, options=dict(step_size=0.1),method='euler')
+        
+        return x_t, w_t, m_t
+    
+    def training_step(self, batch, batch_idx):
+        if self.args.data_type in ["scrna", "tahoe"]:
+            main_batch = batch[0]["train_samples"][0]
+            metric_batch = batch[0]["metric_samples"][0]
+        else: 
+            main_batch = batch["train_samples"][0]
+            metric_batch = batch["metric_samples"][0]
+        
+        self.timesteps = torch.linspace(0.0, 1.0, len(main_batch["x0"])).tolist()
+        loss = self._compute_loss(main_batch, metric_batch, validation=False)
+        
+        if self.joint:
+            self.log(
+                "JointTrain/train_loss",
+                loss,
+                on_step=False,
+                on_epoch=True,
+                prog_bar=True,
+                logger=True,
+            )
+        else:
+            self.log(
+                "GrowthNet/train_loss",
+                loss,
+                on_step=False,
+                on_epoch=True,
+                prog_bar=True,
+                logger=True,
+            )
+        
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+        if self.args.data_type in ["scrna", "tahoe"]:
+            main_batch = batch[0]["val_samples"][0]
+            metric_batch = batch[0]["metric_samples"][0]
+        else: 
+            main_batch = batch["val_samples"][0]
+            metric_batch = batch["metric_samples"][0]
+
+        self.timesteps = torch.linspace(0.0, 1.0, len(main_batch["x0"])).tolist()
+        val_loss = self._compute_loss(main_batch, metric_batch, validation=True)
+        
+        if self.joint:
+            self.log(
+                "JointTrain/val_loss",
+                val_loss,
+                on_step=False,
+                on_epoch=True,
+                prog_bar=True,
+                logger=True,
+            )
+        else:
+            self.log(
+                "GrowthNet/val_loss",
+                val_loss,
+                on_step=False,
+                on_epoch=True,
+                prog_bar=True,
+                logger=True,
+            )
+        return val_loss
+
+    def optimizer_step(self, *args, **kwargs):
+        super().optimizer_step(*args, **kwargs)
+        for net in self.growth_nets:
+            if isinstance(net, EMA):
+                net.update_ema()
+        if self.joint:
+            for net in self.flow_nets:
+                if isinstance(net, EMA):
+                    net.update_ema()
+
+    def configure_optimizers(self):
+        params = []
+        for net in self.growth_nets:
+            params += list(net.parameters())
+            
+        if self.joint:
+            for net in self.flow_nets:
+                params += list(net.parameters())
+        
+        if self.optimizer_name == "adamw":
+            optimizer = AdamW(
+                params,
+                lr=self.lr,
+                weight_decay=self.weight_decay,
+            )
+        elif self.optimizer_name == "adam":
+            optimizer = torch.optim.Adam(
+                params,
+                lr=self.lr,
+            )
+
+        return optimizer
+    
+    @torch.no_grad()
+    def _plot_mass_and_energy(self, main_batch, metric_samples_batch=None, save_dir="./figures"):
+        x0s = main_batch["x0"][0]
+        w0s = main_batch["x0"][1]
+
+        if self.args.manifold:
+            if self.metric_clusters == 4:
+                branch_sample_pairs = [
+                    (metric_samples_batch[0], metric_samples_batch[1]),  # x0 → x1_1 (branch 1)
+                    (metric_samples_batch[0], metric_samples_batch[2]),  # x0 → x1_2 (branch 2)
+                    (metric_samples_batch[0], metric_samples_batch[3]),
+                ]
+            elif self.metric_clusters == 3:
+                branch_sample_pairs = [
+                    (metric_samples_batch[0], metric_samples_batch[1]),  # x0 → x1_1 (branch 1)
+                    (metric_samples_batch[0], metric_samples_batch[2]),  # x0 → x1_2 (branch 2)
+                ]
+            elif self.metric_clusters == 2 and self.branches == 2:
+                branch_sample_pairs = [
+                    (metric_samples_batch[0], metric_samples_batch[1]),  # x0 → x1_1 (branch 1)
+                    (metric_samples_batch[0], metric_samples_batch[1]),  # x0 → x1_2 (branch 2)
+                ]
+            else:
+                branch_sample_pairs = [
+                    (metric_samples_batch[0], metric_samples_batch[1]),  # x0 → x1_1 (branch 1)
+                ]
+
+        batch_size = x0s.shape[0]
+        dtype = x0s[0].dtype
+
+        m0s = torch.zeros_like(w0s, dtype=dtype)
+        xt = [x0s.clone() for _ in range(self.branches)]
+
+        w0_branch = torch.zeros_like(w0s, dtype=dtype)
+        w0_branches = []
+        w0_branches.append(w0s)
+        for _ in range(self.branches - 1):
+            w0_branches.append(w0_branch)
+
+        wt = w0_branches
+        mt = [m0s.clone() for _ in range(self.branches)]
+
+        time_points = []
+        mass_over_time = [[] for _ in range(self.branches)]
+        energy_over_time = [[] for _ in range(self.branches)]
+
+        t_span = torch.linspace(0, 1, 101)
+        for s, t in zip(t_span[:-1], t_span[1:]):
+            time_points.append(t.item())
+            time = torch.Tensor([s, t])
+
+            for i in range(self.branches):
+                if self.args.manifold:
+                    start_samples, end_samples = branch_sample_pairs[i]
+                    samples = torch.cat([start_samples, end_samples], dim=0)
+                else:
+                    samples = None
+
+                start_state = (xt[i], wt[i], mt[i])
+                xt_next, wt_next, mt_next = self.take_step(time, start_state, i, samples)
+
+                xt[i] = xt_next[-1].clone().detach()
+                wt[i] = wt_next[-1].clone().detach()
+                mt[i] = mt_next[-1].clone().detach()
+
+                mass_over_time[i].append(wt[i].mean().item())
+                energy_over_time[i].append(mt[i].mean().item())
+
+        os.makedirs(os.path.join(save_dir, self.args.data_type), exist_ok=True)
+
+        # Use tab10 colormap to get visually distinct colors
+        if self.args.branches == 3:
+            branch_colors = ['#9793F8', '#50B2D7', '#D577FF']  # tuple of RGBs
+        else:
+            branch_colors = ['#50B2D7', '#D577FF']  # tuple of RGBs
+
+        # --- Plot Mass ---
+        plt.figure(figsize=(8, 5))
+        for i in range(self.branches):
+            color = branch_colors[i]
+            plt.plot(time_points, mass_over_time[i], color=color, linewidth=2.5, label=f"Mass Branch {i}")
+        plt.xlabel("Time")
+        plt.ylabel("Mass")
+        plt.title("Mass Evolution per Branch")
+        plt.legend()
+        plt.grid(True)
+        if self.joint:
+            mass_path = os.path.join(save_dir, f"{self.args.data_name}/{self.args.data_name}_joint_mass.png")
+        else:
+            mass_path = os.path.join(save_dir, f"{self.args.data_name}/{self.args.data_name}_growth_mass.png")
+        plt.savefig(mass_path, dpi=300, bbox_inches="tight")
+        plt.close()
+
+        # --- Plot Energy ---
+        plt.figure(figsize=(8, 5))
+        for i in range(self.branches):
+            color = branch_colors[i]
+            plt.plot(time_points, energy_over_time[i], color=color, linewidth=2.5, label=f"Energy Branch {i}")
+        plt.xlabel("Time")
+        plt.ylabel("Energy")
+        plt.title("Energy Evolution per Branch")
+        plt.legend()
+        plt.grid(True)
+        if self.joint:
+            energy_path = os.path.join(save_dir, f"{self.args.data_name}/{self.args.data_name}_joint_energy.png")
+        else: 
+            energy_path = os.path.join(save_dir, f"{self.args.data_name}/{self.args.data_name}_growth_energy.png")
+        plt.savefig(energy_path, dpi=300, bbox_inches="tight")
+        plt.close()
+        
+        
+class GrowthNetTrainLidar(GrowthNetTrain):
+    def test_step(self, batch, batch_idx):
+        main_batch = batch["test_samples"][0]
+        metric_batch = batch["metric_samples"][0]
+        
+        self._plot_mass_and_energy(main_batch, metric_batch)
+        
+        x0 = main_batch["x0"][0] # [B, D]
+        cloud_points = main_batch["dataset"][0]  # full dataset, [N, D]
+        t_span = torch.linspace(0, 1, 101)
+        
+
+        all_trajs = []
+
+        for i, flow_net in enumerate(self.flow_nets):
+            node = NeuralODE(
+                flow_model_torch_wrapper(flow_net),
+                solver="euler",
+                sensitivity="adjoint",
+            )
+
+            with torch.no_grad():
+                traj = node.trajectory(x0, t_span).cpu()  # [T, B, D]
+
+            if self.whiten:
+                traj_shape = traj.shape
+                traj = traj.reshape(-1, 3)
+                traj = self.trainer.datamodule.scaler.inverse_transform(
+                    traj.cpu().detach().numpy()
+                ).reshape(traj_shape)
+
+            traj = torch.tensor(traj)
+            traj = torch.transpose(traj, 0, 1)  # [B, T, D]
+            all_trajs.append(traj)
+
+        # Inverse-transform the point cloud once
+        if self.whiten:
+            cloud_points = torch.tensor(
+                self.trainer.datamodule.scaler.inverse_transform(
+                    cloud_points.cpu().detach().numpy()
+                )
+            )
+
+        # ===== Plot all trajectories together =====
+        fig = plt.figure(figsize=(6, 5))
+        ax = fig.add_subplot(111, projection="3d", computed_zorder=False)
+        ax.view_init(elev=30, azim=-115, roll=0)
+        for i, traj in enumerate(all_trajs):
+            plot_lidar(ax, cloud_points, xs=traj, branch_idx=i)
+        if self.joint:
+            plt.savefig('./figures/lidar/joint_lidar_all_branches.png', dpi=300)
+        else:
+            plt.savefig('./figures/lidar/growth_lidar_all_branches.png', dpi=300)
+        plt.close()
+
+        # ===== Plot each trajectory separately =====
+        for i, traj in enumerate(all_trajs):
+            fig = plt.figure(figsize=(6, 5))
+            ax = fig.add_subplot(111, projection="3d", computed_zorder=False)
+            ax.view_init(elev=30, azim=-115, roll=0)
+            plot_lidar(ax, cloud_points, xs=traj, branch_idx=i)
+            if self.joint:
+                plt.savefig(f'./figures/lidar/joint_lidar_branch_{i + 1}.png', dpi=300)
+            else:
+                plt.savefig(f'./figures/lidar/growth_lidar_branch_{i + 1}.png', dpi=300)
+            plt.close()
+        
+class GrowthNetTrainCell(GrowthNetTrain):
+    def test_step(self, batch, batch_idx):
+        if self.args.data_type in ["scrna", "tahoe"]:
+            main_batch = batch[0]["test_samples"][0]
+            metric_batch = batch[0]["metric_samples"][0]
+        else: 
+            main_batch = batch["test_samples"][0]
+            metric_batch = batch["metric_samples"][0]
+        
+        self._plot_mass_and_energy(main_batch, metric_batch)

branchsbm/branch_interpolant_train.py

@@ -0,0 +1,398 @@
+import sys
+sys.path.append("./BranchSBM")
+import torch
+import pytorch_lightning as pl
+from branchsbm.ema import EMA
+import itertools
+from utils import wasserstein_distance, plot_lidar
+import matplotlib.pyplot as plt
+
+class BranchInterpolantTrain(pl.LightningModule):
+    def __init__(
+        self,
+        flow_matcher,
+        args,
+        skipped_time_points: list = None,
+        ot_sampler=None,
+        
+        state_cost=None,
+        data_manifold_metric=None,
+    ):
+        super().__init__()
+        self.save_hyperparameters()
+        self.args = args
+        
+        self.flow_matcher = flow_matcher
+        
+        # list of geopath nets
+        self.geopath_nets = flow_matcher.geopath_nets
+        self.branches = len(self.geopath_nets)
+        self.metric_clusters = args.metric_clusters
+        
+        self.ot_sampler = ot_sampler
+        self.skipped_time_points = skipped_time_points if skipped_time_points else []
+        self.optimizer_name = args.geopath_optimizer
+        self.lr = args.geopath_lr
+        self.weight_decay = args.geopath_weight_decay
+        self.args = args
+        self.multiply_validation = 4
+
+        self.first_loss = None
+        self.timesteps = None
+        self.computing_reference_loss = False
+        
+        # updates
+        self.state_cost = state_cost
+        self.data_manifold_metric = data_manifold_metric
+        self.whiten = args.whiten
+
+    def forward(self, x0, x1, t, branch_idx):
+        # return specific branch interpolant
+        return self.geopath_nets[branch_idx](x0, x1, t)
+
+    def on_train_start(self):
+        self.first_loss = self.compute_initial_loss()
+        print("first loss")
+        print(self.first_loss)
+
+    # to edit
+    def compute_initial_loss(self):
+        # Set all GeoPath networks to eval mode
+        for net in self.geopath_nets:
+            net.train(mode=False)
+        
+        total_loss = 0
+        total_count = 0
+        with torch.enable_grad():
+            self.t_val = []
+            for i in range(
+                self.trainer.datamodule.num_timesteps - len(self.skipped_time_points)
+            ):
+                self.t_val.append(
+                    torch.rand(
+                        self.trainer.datamodule.batch_size * self.multiply_validation,
+                        requires_grad=True,
+                    )
+                )
+        self.computing_reference_loss = True
+        with torch.no_grad():
+            old_alpha = self.flow_matcher.alpha
+            self.flow_matcher.alpha = 0
+            for batch in self.trainer.datamodule.train_dataloader():
+                
+                self.timesteps = torch.linspace(
+                    0.0, 1.0, len(batch[0]["train_samples"][0])
+                )
+                
+                loss = self._compute_loss(
+                    batch[0]["train_samples"][0],
+                    batch[0]["metric_samples"][0],
+                )
+                print("initial loss")
+                print(loss)
+                total_loss += loss.item()
+                total_count += 1
+            self.flow_matcher.alpha = old_alpha
+            
+        self.computing_reference_loss = False
+        
+        # Set all GeoPath networks back to training mode
+        for net in self.geopath_nets:
+            net.train(mode=True)
+        return total_loss / total_count if total_count > 0 else 1.0
+
+    def _compute_loss(self, main_batch, metric_samples_batch=None):
+        
+        x0s = [main_batch["x0"][0]]
+        w0s = [main_batch["x0"][1]]
+        
+        x1s_list = []
+        w1s_list = [] 
+        
+        if self.branches > 1:
+            for i in range(self.branches):
+                x1s_list.append([main_batch[f"x1_{i+1}"][0]])
+                w1s_list.append([main_batch[f"x1_{i+1}"][1]])
+        else:
+            x1s_list.append([main_batch["x1"][0]])
+            w1s_list.append([main_batch["x1"][1]])
+        
+        if self.args.manifold:
+            #changed
+            if self.metric_clusters == 4:
+                branch_sample_pairs = [
+                    (metric_samples_batch[0], metric_samples_batch[1]),  # x0 → x1_1 (branch 1)
+                    (metric_samples_batch[0], metric_samples_batch[2]),  # x0 → x1_2 (branch 2)
+                    (metric_samples_batch[0], metric_samples_batch[3]),
+                ]
+            elif self.metric_clusters == 3:
+                branch_sample_pairs = [
+                    (metric_samples_batch[0], metric_samples_batch[1]),  # x0 → x1_1 (branch 1)
+                    (metric_samples_batch[0], metric_samples_batch[2]),  # x0 → x1_2 (branch 2)
+                ]
+            elif self.metric_clusters == 2 and self.branches == 2:
+                branch_sample_pairs = [
+                    (metric_samples_batch[0], metric_samples_batch[1]),  # x0 → x1_1 (branch 1)
+                    (metric_samples_batch[0], metric_samples_batch[1]),  # x0 → x1_2 (branch 2)
+                ]
+            else:
+                branch_sample_pairs = [
+                    (metric_samples_batch[0], metric_samples_batch[1]),  # x0 → x1_1 (branch 1)
+                ]
+            """samples0, samples1, samples2 = (
+                metric_samples_batch[0],
+                metric_samples_batch[1],
+                metric_samples_batch[2]
+            )"""
+                            
+        assert len(x1s_list) == self.branches, "Mismatch between x1s_list and expected branches"
+        
+        # compute sum of velocities for each branch
+        loss = 0
+        velocities = []
+        for branch_idx in range(self.branches):
+        
+            ts, xts, uts = self._process_flow(x0s, x1s_list[branch_idx], branch_idx)
+            
+            for i in range(len(ts)):
+                # calculate kinetic and potential energy of the predicted interpolant
+                if self.args.manifold:
+                    start_samples, end_samples = branch_sample_pairs[branch_idx]
+                    
+                    samples = torch.cat([start_samples, end_samples], dim=0)
+                    #print("metric sample shape")
+                    #print(samples.shape)
+                    vel, _, _ = self.data_manifold_metric.calculate_velocity(
+                        xts[i], uts[i], samples, i
+                    )
+                else:
+                    vel = torch.sqrt((uts[i]**2).sum(dim =-1) + self.state_cost(xts[i]))
+                    #vel = (uts[i]**2).sum(dim =-1)
+                
+                velocities.append(vel)
+            
+        loss = torch.mean(torch.cat(velocities) ** 2)
+        
+        self.log(
+            "BranchPathNet/mean_velocity_geopath",
+            loss,
+            on_step=False,
+            on_epoch=True,
+            prog_bar=True,
+        )
+
+        return loss
+
+    def _process_flow(self, x0s, x1s, branch_idx):
+        ts, xts, uts = [], [], []
+        t_start = self.timesteps[0]
+        i_start = 0
+
+        for i, (x0, x1) in enumerate(zip(x0s, x1s)):
+            x0, x1 = torch.squeeze(x0), torch.squeeze(x1)
+            if self.trainer.validating or self.computing_reference_loss:
+                repeat_tuple = (self.multiply_validation, 1) + (1,) * (
+                    len(x0.shape) - 2
+                )
+                x0 = x0.repeat(repeat_tuple)
+                x1 = x1.repeat(repeat_tuple)
+
+            if self.ot_sampler is not None:
+                x0, x1 = self.ot_sampler.sample_plan(
+                    x0,
+                    x1,
+                    replace=True,
+                )
+            if self.skipped_time_points and i + 1 >= self.skipped_time_points[0]:
+                t_start_next = self.timesteps[i + 2]
+            else:
+                t_start_next = self.timesteps[i + 1]
+
+            t = None
+            if self.trainer.validating or self.computing_reference_loss:
+                t = self.t_val[i]
+
+            t, xt, ut = self.flow_matcher.sample_location_and_conditional_flow(
+                x0, x1, t_start, t_start_next, branch_idx, training_geopath_net=True, t=t
+            )
+            ts.append(t)
+            xts.append(xt)
+            uts.append(ut)
+            t_start = t_start_next
+
+        return ts, xts, uts
+
+    def training_step(self, batch, batch_idx):
+        if self.args.data_type in ["scrna", "tahoe"]:
+            main_batch = batch[0]["train_samples"][0]
+            metric_batch = batch[0]["metric_samples"][0]
+        else: 
+            main_batch = batch["train_samples"][0]
+            metric_batch = batch["metric_samples"][0]
+        
+        tangential_velocity_loss = self._compute_loss(main_batch, metric_batch)
+        
+        if self.first_loss:
+            tangential_velocity_loss = tangential_velocity_loss / self.first_loss
+            
+        self.log(
+            "BranchPathNet/mean_geopath_geopath",
+            (self.flow_matcher.geopath_net_output.abs().mean()),
+            on_step=False,
+            on_epoch=True,
+            prog_bar=True,
+        )
+        
+        self.log(
+            "BranchPathNet/train_loss_geopath",
+            tangential_velocity_loss,
+            on_step=True,
+            on_epoch=True,
+            prog_bar=True,
+            logger=True,
+        )
+                
+        return tangential_velocity_loss
+
+    def validation_step(self, batch, batch_idx):
+        if self.args.data_type in ["scrna", "tahoe"]:
+            main_batch = batch[0]["val_samples"][0]
+            metric_batch = batch[0]["metric_samples"][0]
+        else: 
+            main_batch = batch["val_samples"][0]
+            metric_batch = batch["metric_samples"][0]
+
+        tangential_velocity_loss = self._compute_loss(main_batch, metric_batch)
+        if self.first_loss:
+            tangential_velocity_loss = tangential_velocity_loss / self.first_loss
+            
+        self.log(
+            "BranchPathNet/val_loss_geopath",
+            tangential_velocity_loss,
+            on_step=False,
+            on_epoch=True,
+            prog_bar=True,
+            logger=True,
+        )
+        return tangential_velocity_loss
+    
+    
+    def test_step(self, batch, batch_idx):
+        main_batch = batch["test_samples"][0]
+        metric_batch = batch["metric_samples"][0]
+                
+        x0 = main_batch["x0"][0]  # [B, D]
+        cloud_points = main_batch["dataset"][0]  # full dataset, [N, D]
+        
+        x0 = x0.to(self.device)
+        cloud_points = cloud_points.to(self.device)
+
+        t_vals = [0.25, 0.5, 0.75]
+        t_labels = ["t=1/4", "t=1/2", "t=3/4"]
+
+        colors = {
+            "x0": "#4D176C",
+            "t=1/4": "#5C3B9D",
+            "t=1/2": "#6172B9",
+            "t=3/4": "#AC4E51",
+            "x1": "#771F4F",
+        }
+
+        # Unwhiten cloud points if needed
+        if self.whiten:
+            cloud_points = torch.tensor(
+                self.trainer.datamodule.scaler.inverse_transform(cloud_points.cpu().numpy())
+            )
+
+        for i in range(self.branches):
+            geopath = self.geopath_nets[i]
+            x1_key = f"x1_{i + 1}"
+            if x1_key not in main_batch:
+                print(f"Skipping branch {i + 1}: no final distribution {x1_key}")
+                continue
+
+            x1 = main_batch[x1_key][0].to(self.device)
+            print(x1.shape)
+            print(x0.shape)
+            interpolated_points = []
+            with torch.no_grad():
+                for t_scalar in t_vals:
+                    t_tensor = torch.full((x0.shape[0], 1), t_scalar, device=self.device)  # [B, 1]
+                    xt = geopath(x0, x1, t_tensor).cpu()  # [B, D]
+                    if self.whiten:
+                        xt = torch.tensor(
+                            self.trainer.datamodule.scaler.inverse_transform(xt.numpy())
+                        )
+                    interpolated_points.append(xt)
+
+            if self.whiten:
+                x0_plot = torch.tensor(
+                    self.trainer.datamodule.scaler.inverse_transform(x0.cpu().numpy())
+                )
+                x1_plot = torch.tensor(
+                    self.trainer.datamodule.scaler.inverse_transform(x1.cpu().numpy())
+                )
+            else:
+                x0_plot = x0.cpu()
+                x1_plot = x1.cpu()
+
+            # Plot
+            fig = plt.figure(figsize=(6, 5))
+            ax = fig.add_subplot(111, projection="3d", computed_zorder=False)
+            ax.view_init(elev=30, azim=-115, roll=0)
+            plot_lidar(ax, cloud_points)
+
+            # Initial x₀
+            ax.scatter(
+                x0_plot[:, 0], x0_plot[:, 1], x0_plot[:, 2],
+                s=15, alpha=1.0, color=colors["x0"], label="x₀", depthshade=True,
+                edgecolors="white",
+                linewidths=0.3
+            )
+
+            # Interpolated points
+            for xt, t_label in zip(interpolated_points, t_labels):
+                ax.scatter(
+                    xt[:, 0], xt[:, 1], xt[:, 2],
+                    s=15, alpha=1.0, color=colors[t_label], label=t_label, depthshade=True,
+                    edgecolors="white",
+                    linewidths=0.3
+                )
+
+            # Final x₁
+            ax.scatter(
+                x1_plot[:, 0], x1_plot[:, 1], x1_plot[:, 2],
+                s=15, alpha=1.0, color=colors["x1"], label="x₁", depthshade=True,
+                edgecolors="white",
+                linewidths=0.3
+            )
+
+            ax.legend()
+            save_path = f"/raid/st512/branchsbm/figures/{self.args.data_type}/lidar_geopath_branch_{i+1}.png"
+            plt.savefig(save_path, dpi=300)
+            plt.close()
+
+    def optimizer_step(self, *args, **kwargs):
+        super().optimizer_step(*args, **kwargs)
+        if isinstance(self.geopath_nets, EMA):
+            self.geopath_nets.update_ema()
+
+    def configure_optimizers(self):
+        if self.optimizer_name == "adam":
+            """optimizer = torch.optim.Adam(
+                self.geopath_nets.parameters(),
+                lr=self.lr,
+            )"""
+            optimizer = torch.optim.Adam(
+                itertools.chain(*[net.parameters() for net in self.geopath_nets]), lr=self.lr
+            )
+        elif self.optimizer_name == "adamw":
+            """optimizer = torch.optim.AdamW(
+                self.geopath_nets.parameters(),
+                lr=self.lr,
+                weight_decay=self.weight_decay,
+            )"""
+            optimizer = torch.optim.AdamW(
+                itertools.chain(*[net.parameters() for net in self.geopath_nets]), lr=self.lr
+            )
+        return optimizer

branchsbm/branchsbm.py

@@ -0,0 +1,109 @@
+import sys
+sys.path.append("./BranchSBM")
+import torch
+from torchcfm.conditional_flow_matching import ConditionalFlowMatcher, pad_t_like_x
+import torch.nn as nn
+
+class BranchSBM(ConditionalFlowMatcher):
+    def __init__(
+        self, geopath_nets: nn.ModuleList = None, alpha: float = 1.0, *args, **kwargs
+    ):
+        super().__init__(*args, **kwargs)
+        self.alpha = alpha
+        self.geopath_nets = geopath_nets
+        if self.alpha != 0:
+            assert (
+                geopath_nets is not None
+            ), "GeoPath model must be provided if alpha != 0"
+            
+        self.branches = len(geopath_nets)
+
+    def gamma(self, t, t_min, t_max):
+        return (
+            1.0
+            - ((t - t_min) / (t_max - t_min)) ** 2
+            - ((t_max - t) / (t_max - t_min)) ** 2
+        )
+
+    def d_gamma(self, t, t_min, t_max):
+        return 2 * (-2 * t + t_max + t_min) / (t_max - t_min) ** 2
+
+    def compute_mu_t(self, x0, x1, t, t_min, t_max, branch_idx):
+        assert branch_idx < self.branches, "Index out of bounds"
+
+        with torch.enable_grad():
+            t = pad_t_like_x(t, x0)
+            if self.alpha == 0:
+                return (t_max - t) / (t_max - t_min) * x0 + (t - t_min) / (
+                    t_max - t_min
+                ) * x1
+                
+            # compute value for specific branch
+            self.geopath_net_output = self.geopath_nets[branch_idx](x0, x1, t)
+            if self.geopath_nets[branch_idx].time_geopath:
+                self.doutput_dt = torch.autograd.grad(
+                    self.geopath_net_output,
+                    t,
+                    grad_outputs=torch.ones_like(self.geopath_net_output),
+                    create_graph=False,
+                    retain_graph=True,
+                )[0]
+        return (
+            (t_max - t) / (t_max - t_min) * x0
+            + (t - t_min) / (t_max - t_min) * x1
+            + self.gamma(t, t_min, t_max) * self.geopath_net_output
+        )
+
+    def sample_xt(self, x0, x1, t, epsilon, t_min, t_max, branch_idx):
+        assert branch_idx < self.branches, "Index out of bounds"
+        mu_t = self.compute_mu_t(x0, x1, t, t_min, t_max, branch_idx)
+        sigma_t = self.compute_sigma_t(t)
+        sigma_t = pad_t_like_x(sigma_t, x0)
+        return mu_t + sigma_t * epsilon
+
+    def sample_location_and_conditional_flow(
+        self,
+        x0,
+        x1,
+        t_min,
+        t_max,
+        branch_idx,
+        training_geopath_net=False,
+        midpoint_only=False,
+        t=None,
+    ):
+
+        self.training_geopath_net = training_geopath_net
+        with torch.enable_grad():
+            if t is None:
+                t = torch.rand(x0.shape[0], requires_grad=True)
+            t = t.type_as(x0)
+            t = t * (t_max - t_min) + t_min
+            if midpoint_only:
+                t = (t_max + t_min) / 2 * torch.ones_like(t).type_as(x0)
+        
+        assert len(t) == x0.shape[0], "t has to have batch size dimension"
+
+        eps = self.sample_noise_like(x0)
+        
+        # compute xt and ut for branch_idx
+        xt = self.sample_xt(x0, x1, t, eps, t_min, t_max, branch_idx)
+        ut = self.compute_conditional_flow(x0, x1, t, xt, t_min, t_max, branch_idx)
+
+        return t, xt, ut
+
+    def compute_conditional_flow(self, x0, x1, t, xt, t_min, t_max, branch_idx):
+        del xt
+        t = pad_t_like_x(t, x0)
+        if self.alpha == 0:
+            return (x1 - x0) / (t_max - t_min)
+        
+        return (
+            (x1 - x0) / (t_max - t_min)
+            + self.d_gamma(t, t_min, t_max) * self.geopath_net_output
+            + (
+                self.gamma(t, t_min, t_max) * self.doutput_dt
+                if self.geopath_nets[branch_idx].time_geopath
+                else 0
+            )
+        )

branchsbm/ema.py

@@ -0,0 +1,64 @@
+import torch
+
+class EMA(torch.nn.Module):
+    def __init__(self, model: torch.nn.Module, decay: float = 0.999):
+        super().__init__()
+        self.model = model
+        self.decay = decay
+        if hasattr(self.model, "time_geopath"):
+            self.time_geopath = self.model.time_geopath
+
+        # Put this in a buffer so that it gets included in the state dict
+        self.register_buffer("num_updates", torch.tensor(0))
+
+        self.shadow_params = torch.nn.ParameterList(
+            [
+                torch.nn.Parameter(p.clone().detach(), requires_grad=False)
+                for p in model.parameters()
+                if p.requires_grad
+            ]
+        )
+        self.backup_params = []
+
+    def train(self, mode: bool):
+        if self.training and mode == False:
+            # Switching from train mode to eval mode.  Backup the model parameters and
+            # overwrite with shadow params
+            self.backup()
+            self.copy_to_model()
+        elif not self.training and mode == True:
+            # Switching from eval to train mode.  Restore the `backup_params`
+            self.restore_to_model()
+
+        super().train(mode)
+
+    def update_ema(self):
+        self.num_updates += 1
+        num_updates = self.num_updates.item()
+        decay = min(self.decay, (1 + num_updates) / (10 + num_updates))
+        with torch.no_grad():
+            params = [p for p in self.model.parameters() if p.requires_grad]
+            for shadow, param in zip(self.shadow_params, params):
+                shadow.sub_((1 - decay) * (shadow - param))
+
+    def forward(self, *args, **kwargs):
+        return self.model(*args, **kwargs)
+
+    def copy_to_model(self):
+        # copy the shadow (ema) parameters to the model
+        params = [p for p in self.model.parameters() if p.requires_grad]
+        for shaddow, param in zip(self.shadow_params, params):
+            param.data.copy_(shaddow.data)
+
+    def backup(self):
+        # Backup the current model parameters
+        if len(self.backup_params) > 0:
+            for p, b in zip(self.model.parameters(), self.backup_params):
+                b.data.copy_(p.data)
+        else:
+            self.backup_params = [param.clone() for param in self.model.parameters()]
+
+    def restore_to_model(self):
+        # Restores the backed up parameters to the model.
+        for param, backup in zip(self.model.parameters(), self.backup_params):
+            param.data.copy_(backup.data)

configs/experiment/cell_single_branch.yaml

@@ -0,0 +1,12 @@
+data_type: "scrna"
+data_name: "mouse"
+dim: 2
+whiten: false
+t_exclude: []
+velocity_metric: "land"
+gammas: [0.125]
+rho: 0.001
+branchsbm: true
+seeds: [42]
+patience_geopath: 50
+time_geopath: true

configs/experiment/clonidine_100D.yaml

@@ -0,0 +1,22 @@
+data_type: "tahoe"
+data_name: "clonidine100D"
+accelerator: "gpu"
+hidden_dims_geopath: [1024, 1024, 1024]
+hidden_dims_flow: [1024, 1024, 1024]
+hidden_dims_growth: [1024, 1024, 1024]
+dim: 100
+t_exclude: []
+time_geopath: true
+whiten: false
+velocity_metric: "rbf"
+metric_patience: 25
+patience: 25
+n_centers: 300
+kappa: 2
+rho: -2.75
+alpha_metric: 1
+metric_epochs: 200
+branchsbm: true
+seeds: [42]
+branches: 2
+metric_clusters: 3

configs/experiment/clonidine_150D.yaml

@@ -0,0 +1,22 @@
+data_type: "tahoe"
+data_name: "clonidine150D"
+accelerator: "gpu"
+hidden_dims_geopath: [1024, 1024, 1024]
+hidden_dims_flow: [1024, 1024, 1024]
+hidden_dims_growth: [1024, 1024, 1024]
+dim: 150
+t_exclude: []
+time_geopath: true
+whiten: false
+velocity_metric: "rbf"
+metric_patience: 25
+patience: 25
+n_centers: 300
+kappa: 3
+rho: -2.75
+alpha_metric: 1
+metric_epochs: 400
+branchsbm: true
+seeds: [42]
+branches: 2
+metric_clusters: 3

configs/experiment/clonidine_50D.yaml

@@ -0,0 +1,22 @@
+data_type: "tahoe"
+data_name: "clonidine50D"
+accelerator: "gpu"
+hidden_dims_geopath: [1024, 1024, 1024]
+hidden_dims_flow: [1024, 1024, 1024]
+hidden_dims_growth: [1024, 1024, 1024]
+dim: 50
+t_exclude: []
+time_geopath: true
+whiten: false
+velocity_metric: "rbf"
+metric_patience: 25
+patience: 25
+n_centers: 150
+kappa: 1.5
+rho: -2.75
+alpha_metric: 1
+metric_epochs: 200
+branchsbm: true
+seeds: [42]
+branches: 2
+metric_clusters: 3

configs/experiment/clonidine_50Dsingle.yaml

@@ -0,0 +1,22 @@
+data_type: "tahoe"
+data_name: "clonidine50Dsingle"
+accelerator: "gpu"
+hidden_dims_geopath: [1024, 1024, 1024]
+hidden_dims_flow: [1024, 1024, 1024]
+hidden_dims_growth: [1024, 1024, 1024]
+dim: 50
+t_exclude: []
+time_geopath: true
+whiten: false
+velocity_metric: "rbf"
+metric_patience: 25
+patience: 25
+n_centers: 150
+kappa: 1.5
+rho: -2.75
+alpha_metric: 1
+metric_epochs: 200
+branchsbm: true
+seeds: [42]
+branches: 1
+metric_clusters: 2

configs/experiment/lidar.yaml

@@ -0,0 +1,14 @@
+data_type: "lidar"
+data_name: "lidar"
+dim: 3
+whiten: true
+t_exclude: []
+velocity_metric: "land"
+gammas: [0.125]
+rho: 0.001
+branchsbm: true
+seeds: [42]
+patience_geopath: 50
+time_geopath: true
+branches: 2
+metric_clusters: 3

configs/experiment/lidar_single.yaml

@@ -0,0 +1,14 @@
+data_type: "lidar"
+data_name: "lidarsingle"
+dim: 3
+whiten: true
+t_exclude: []
+velocity_metric: "land"
+gammas: [0.125]
+rho: 0.001
+branchsbm: true
+seeds: [42]
+patience_geopath: 50
+time_geopath: true
+branches: 1
+metric_clusters: 2

configs/experiment/mouse.yaml

@@ -0,0 +1,17 @@
+data_type: "scrna"
+data_name: "mouse"
+hidden_dims_geopath: [1024, 1024, 1024]
+hidden_dims_flow: [1024, 1024, 1024]
+hidden_dims_growth: [1024, 1024, 1024] 
+dim: 2
+whiten: false
+t_exclude: []
+velocity_metric: "land"
+gammas: [0.125] 
+rho: 0.001
+branchsbm: true
+seeds: [42]
+patience_geopath: 50
+time_geopath: true
+branches: 2
+metric_clusters: 3

configs/experiment/trametinib.yaml

@@ -0,0 +1,22 @@
+data_type: "tahoe"
+data_name: "trametinib"
+accelerator: "gpu"
+hidden_dims_geopath: [1024, 1024, 1024]
+hidden_dims_flow: [1024, 1024, 1024]
+hidden_dims_growth: [1024, 1024, 1024]
+dim: 50
+t_exclude: []
+time_geopath: true
+whiten: false
+velocity_metric: "rbf"
+metric_patience: 25
+patience: 25
+n_centers: 150
+kappa: 1.5
+rho: -2.75
+alpha_metric: 1
+metric_epochs: 200
+branchsbm: true
+seeds: [42]
+branches: 3
+metric_clusters: 4

configs/experiment/trametinib_single.yaml

@@ -0,0 +1,22 @@
+data_type: "tahoe"
+data_name: "trametinibsingle"
+accelerator: "gpu"
+hidden_dims_geopath: [1024, 1024, 1024]
+hidden_dims_flow: [1024, 1024, 1024]
+hidden_dims_growth: [1024, 1024, 1024]
+dim: 50
+t_exclude: []
+time_geopath: true
+whiten: false
+velocity_metric: "rbf"
+metric_patience: 25
+patience: 25
+n_centers: 150
+kappa: 1.5
+rho: -2.75
+alpha_metric: 1
+metric_epochs: 200
+branchsbm: true
+seeds: [42]
+branches: 1
+metric_clusters: 2

dataloaders/clonidine_data.py

@@ -0,0 +1,269 @@
+import torch
+import sys
+sys.argv = ['']
+from sklearn.preprocessing import StandardScaler
+import pytorch_lightning as pl
+from torch.utils.data import DataLoader
+from lightning.pytorch.utilities.combined_loader import CombinedLoader
+import pandas as pd
+import numpy as np
+from functools import partial
+from scipy.spatial import cKDTree
+from sklearn.cluster import KMeans
+from torch.utils.data import TensorDataset
+
+#from train.parsers_tahoe import parse_args
+#args = parse_args()
+
+class DrugResponseDataModule(pl.LightningDataModule):
+    def __init__(self, args):
+        super().__init__()
+        self.save_hyperparameters()
+
+        self.batch_size = args.batch_size
+        self.max_dim = args.dim
+        self.whiten = args.whiten
+        self.split_ratios = args.split_ratios
+        
+        # Path to your combined data
+        self.data_path = "/raid/st512/branchsbm/data/pca_and_leiden_labels.csv"
+        self.num_timesteps = 2
+        self.args = args
+        self._prepare_data()
+
+    def _prepare_data(self):
+        df = pd.read_csv(self.data_path, comment='#')
+        df = df.iloc[:, 1:]
+        df = df.replace('', np.nan)
+        pc_cols = df.columns[:50]
+        for col in pc_cols:
+            df[col] = pd.to_numeric(df[col], errors='coerce')
+        leiden_dmso_col = 'leiden_DMSO_TF_0.0uM'
+        leiden_clonidine_col = 'leiden_Clonidine (hydrochloride)_5.0uM'
+
+        dmso_mask = df[leiden_dmso_col].notna()        # Has leiden value in DMSO column
+        clonidine_mask = df[leiden_clonidine_col].notna()  # Has leiden value in Clonidine column
+        
+        dmso_data = df[dmso_mask].copy()
+        clonidine_data = df[clonidine_mask].copy()
+        
+        top_clonidine_clusters = ['0.0', '4.0']
+        
+        x1_1_data = clonidine_data[clonidine_data[leiden_clonidine_col].astype(str) == top_clonidine_clusters[0]]
+        x1_2_data = clonidine_data[clonidine_data[leiden_clonidine_col].astype(str) == top_clonidine_clusters[1]]
+        x1_1_coords = x1_1_data[pc_cols].values 
+        x1_2_coords = x1_2_data[pc_cols].values 
+        x1_1_coords = x1_1_coords.astype(float)
+        x1_2_coords = x1_2_coords.astype(float)
+
+        target_size = min(len(x1_1_coords), len(x1_2_coords))
+        
+        # Sample endpoint clusters to target size
+        np.random.seed(42)
+        if len(x1_1_coords) > target_size:
+            idx1 = np.random.choice(len(x1_1_coords), target_size, replace=False)
+            x1_1_coords = x1_1_coords[idx1]
+        
+        if len(x1_2_coords) > target_size:
+            idx2 = np.random.choice(len(x1_2_coords), target_size, replace=False)
+            x1_2_coords = x1_2_coords[idx2]
+        
+        dmso_cluster_counts = dmso_data[leiden_dmso_col].value_counts()
+        
+        # DMSO
+        largest_dmso_cluster = dmso_cluster_counts.index[0]
+        dmso_cluster_data = dmso_data[dmso_data[leiden_dmso_col] == largest_dmso_cluster]
+        
+        dmso_coords = dmso_cluster_data[pc_cols].values
+        
+        # Random sampling from largest DMSO cluster to match target size
+        np.random.seed(42)
+        if len(dmso_coords) >= target_size:
+            idx0 = np.random.choice(len(dmso_coords), target_size, replace=False)
+            x0_coords = dmso_coords[idx0]
+        else:
+            # If largest cluster is smaller than target, use all of it and pad with other DMSO cells
+            remaining_needed = target_size - len(dmso_coords)
+            other_dmso_data = dmso_data[dmso_data[leiden_dmso_col] != largest_dmso_cluster]
+            other_dmso_coords = other_dmso_data[pc_cols].values
+            
+            if len(other_dmso_coords) >= remaining_needed:
+                idx_other = np.random.choice(len(other_dmso_coords), remaining_needed, replace=False)
+                x0_coords = np.vstack([dmso_coords, other_dmso_coords[idx_other]])
+            else:
+                # Use all available DMSO cells and reduce target size
+                all_dmso_coords = dmso_data[pc_cols].values
+                target_size = min(target_size, len(all_dmso_coords))
+                idx0 = np.random.choice(len(all_dmso_coords), target_size, replace=False)
+                x0_coords = all_dmso_coords[idx0]
+        
+        # Also resample endpoint clusters to match final target size
+        if len(x1_1_coords) > target_size:
+            idx1 = np.random.choice(len(x1_1_coords), target_size, replace=False)
+            x1_1_coords = x1_1_coords[idx1]
+        
+        if len(x1_2_coords) > target_size:
+            idx2 = np.random.choice(len(x1_2_coords), target_size, replace=False)
+            x1_2_coords = x1_2_coords[idx2]
+        
+        self.n_samples = target_size
+        
+        x0 = torch.tensor(x0_coords, dtype=torch.float32)
+        x1_1 = torch.tensor(x1_1_coords, dtype=torch.float32)
+        x1_2 = torch.tensor(x1_2_coords, dtype=torch.float32)
+        
+        self.coords_t0 = x0
+        self.coords_t1 = torch.cat([x1_1, x1_2], dim=0)
+        self.time_labels = np.concatenate([
+            np.zeros(len(self.coords_t0)),    # t=0
+            np.ones(len(self.coords_t1)),     # t=1
+        ])
+        
+        split_index = int(target_size * self.split_ratios[0])
+        
+        if target_size - split_index < self.batch_size:
+            split_index = target_size - self.batch_size
+        
+        train_x0 = x0[:split_index]
+        val_x0 = x0[split_index:]
+        train_x1_1 = x1_1[:split_index]
+        val_x1_1 = x1_1[split_index:]
+        train_x1_2 = x1_2[:split_index]
+        val_x1_2 = x1_2[split_index:]
+        
+        self.val_x0 = val_x0
+        
+        train_x0_weights = torch.full((train_x0.shape[0], 1), fill_value=1.0)
+        train_x1_1_weights = torch.full((train_x1_1.shape[0], 1), fill_value=0.5)
+        train_x1_2_weights = torch.full((train_x1_2.shape[0], 1), fill_value=0.5)
+        
+        val_x0_weights = torch.full((val_x0.shape[0], 1), fill_value=1.0)
+        val_x1_1_weights = torch.full((val_x1_1.shape[0], 1), fill_value=0.5)
+        val_x1_2_weights = torch.full((val_x1_2.shape[0], 1), fill_value=0.5)
+        
+        self.train_dataloaders = {
+            "x0": DataLoader(TensorDataset(train_x0, train_x0_weights), batch_size=self.batch_size, shuffle=True, drop_last=True),
+            "x1_1": DataLoader(TensorDataset(train_x1_1, train_x1_1_weights), batch_size=self.batch_size, shuffle=True, drop_last=True),
+            "x1_2": DataLoader(TensorDataset(train_x1_2, train_x1_2_weights), batch_size=self.batch_size, shuffle=True, drop_last=True),
+        }
+        
+        self.val_dataloaders = {
+            "x0": DataLoader(TensorDataset(val_x0, val_x0_weights), batch_size=self.batch_size, shuffle=False, drop_last=True),
+            "x1_1": DataLoader(TensorDataset(val_x1_1, val_x1_1_weights), batch_size=self.batch_size, shuffle=True, drop_last=True),
+            "x1_2": DataLoader(TensorDataset(val_x1_2, val_x1_2_weights), batch_size=self.batch_size, shuffle=True, drop_last=True),
+        }
+        
+        all_coords = df[pc_cols].dropna().values.astype(float)
+        self.dataset = torch.tensor(all_coords, dtype=torch.float32)
+        self.tree = cKDTree(all_coords)
+
+        self.test_dataloaders = {
+            "x0": DataLoader(TensorDataset(val_x0, val_x0_weights), batch_size=self.val_x0.shape[0], shuffle=False, drop_last=False),
+            "dataset": DataLoader(TensorDataset(self.dataset), batch_size=self.dataset.shape[0], shuffle=False, drop_last=False),
+        }
+        
+        # Metric samples
+        km_all = KMeans(n_clusters=3, random_state=0).fit(self.dataset.numpy())
+        cluster_labels = km_all.labels_
+        
+        cluster_0_mask = cluster_labels == 0
+        cluster_1_mask = cluster_labels == 1
+        cluster_2_mask = cluster_labels == 2
+        
+        samples = self.dataset.cpu().numpy()
+            
+        cluster_0_data = samples[cluster_0_mask]
+        cluster_1_data = samples[cluster_1_mask]
+        cluster_2_data = samples[cluster_2_mask]
+        
+        self.metric_samples_dataloaders = [
+            DataLoader(
+                torch.tensor(cluster_2_data, dtype=torch.float32),
+                batch_size=cluster_2_data.shape[0],
+                shuffle=False,
+                drop_last=False,
+            ),
+            DataLoader(
+                torch.tensor(cluster_0_data, dtype=torch.float32),
+                batch_size=cluster_0_data.shape[0],
+                shuffle=False,
+                drop_last=False,
+            ),
+            
+            DataLoader(
+                torch.tensor(cluster_1_data, dtype=torch.float32),
+                batch_size=cluster_1_data.shape[0],
+                shuffle=False,
+                drop_last=False,
+            ),
+        ]
+        
+    def train_dataloader(self):
+        combined_loaders = {
+            "train_samples": CombinedLoader(self.train_dataloaders, mode="min_size"),
+            "metric_samples": CombinedLoader(
+                self.metric_samples_dataloaders, mode="min_size"
+            ),
+        }
+        return CombinedLoader(combined_loaders, mode="max_size_cycle")
+
+    def val_dataloader(self):
+        combined_loaders = {
+            "val_samples": CombinedLoader(self.val_dataloaders, mode="min_size"),
+            "metric_samples": CombinedLoader(
+                self.metric_samples_dataloaders, mode="min_size"
+            ),
+        }
+
+        return CombinedLoader(combined_loaders, mode="max_size_cycle")
+
+    
+    
+    def test_dataloader(self):
+        combined_loaders = {
+            "test_samples": CombinedLoader(self.test_dataloaders, mode="min_size"),
+            "metric_samples": CombinedLoader(
+                self.metric_samples_dataloaders, mode="min_size"
+            ),
+        }
+
+        return CombinedLoader(combined_loaders, mode="max_size_cycle")
+
+    def get_manifold_proj(self, points):
+        """Adapted for 2D cell data - uses local neighborhood averaging instead of plane fitting"""
+        return partial(self.local_smoothing_op, tree=self.tree, dataset=self.dataset)
+
+    @staticmethod
+    def local_smoothing_op(x, tree, dataset, k=10, temp=1e-3):
+        """
+        Apply local smoothing based on k-nearest neighbors in the full dataset
+        This replaces the plane projection for 2D manifold regularization
+        """
+        points_np = x.detach().cpu().numpy()
+        _, idx = tree.query(points_np, k=k)
+        nearest_pts = dataset[idx]  # Shape: (batch_size, k, 2)
+        
+        # Compute weighted average of neighbors
+        dists = (x.unsqueeze(1) - nearest_pts).pow(2).sum(-1, keepdim=True)
+        weights = torch.exp(-dists / temp)
+        weights = weights / weights.sum(dim=1, keepdim=True)
+        
+        # Weighted average of neighbors
+        smoothed = (weights * nearest_pts).sum(dim=1)
+        
+        # Blend original point with smoothed version
+        alpha = 0.3  # How much smoothing to apply
+        return (1 - alpha) * x + alpha * smoothed
+    
+    def get_timepoint_data(self):
+        """Return data organized by timepoints for visualization"""
+        return {
+            't0': self.coords_t0,
+            't1': self.coords_t1, 
+            'time_labels': self.time_labels
+        }
+    
+def get_datamodule():
+    datamodule = DrugResponseDataModule(args)
+    datamodule.setup(stage="fit")   
+    return datamodule    

dataloaders/clonidine_single_branch.py

@@ -0,0 +1,274 @@
+import torch
+import sys
+sys.argv = ['']
+from sklearn.preprocessing import StandardScaler
+import pytorch_lightning as pl
+from torch.utils.data import DataLoader
+from lightning.pytorch.utilities.combined_loader import CombinedLoader
+import pandas as pd
+import numpy as np
+from functools import partial
+from scipy.spatial import cKDTree
+from sklearn.cluster import KMeans
+from torch.utils.data import TensorDataset
+
+#uncomment for plotting
+#from train.parsers_tahoe import parse_args
+#args = parse_args()
+
+class ClonidineSingleBranchDataModule(pl.LightningDataModule):
+    def __init__(self, args):
+        super().__init__()
+        self.save_hyperparameters()
+
+        self.batch_size = args.batch_size
+        self.max_dim = args.dim
+        self.whiten = args.whiten
+        self.split_ratios = args.split_ratios
+        
+        self.dim = args.dim
+        print("dimension")
+        print(self.dim)
+        # Path to your combined data
+        self.data_path = "./data/pca_and_leiden_labels.csv"
+        self.num_timesteps = 2
+        self.args = args
+        self._prepare_data()
+
+    def _prepare_data(self):
+        df = pd.read_csv(self.data_path, comment='#')
+        df = df.iloc[:, 1:]
+        df = df.replace('', np.nan)
+        pc_cols = df.columns[:self.dim]
+        for col in pc_cols:
+            df[col] = pd.to_numeric(df[col], errors='coerce')
+        leiden_dmso_col = 'leiden_DMSO_TF_0.0uM'
+        leiden_clonidine_col = 'leiden_Clonidine (hydrochloride)_5.0uM'
+
+        dmso_mask = df[leiden_dmso_col].notna()        # Has leiden value in DMSO column
+        clonidine_mask = df[leiden_clonidine_col].notna()  # Has leiden value in Clonidine column
+        
+        dmso_data = df[dmso_mask].copy()
+        clonidine_data = df[clonidine_mask].copy()
+        
+        top_clonidine_clusters = ['0.0', '4.0']
+        
+        x1_1_data = clonidine_data[clonidine_data[leiden_clonidine_col].astype(str) == top_clonidine_clusters[0]]
+        x1_2_data = clonidine_data[clonidine_data[leiden_clonidine_col].astype(str) == top_clonidine_clusters[1]]
+        
+        x1_1_coords = x1_1_data[pc_cols].values 
+        x1_2_coords = x1_2_data[pc_cols].values 
+        
+        x1_1_coords = x1_1_coords.astype(float)
+        x1_2_coords = x1_2_coords.astype(float)
+
+        # Target size is now the minimum across all three endpoint clusters
+        target_size = min(len(x1_1_coords), len(x1_2_coords),)
+        
+        # Helper function to select points closest to centroid
+        def select_closest_to_centroid(coords, target_size):
+            if len(coords) <= target_size:
+                return coords
+            
+            # Calculate centroid
+            centroid = np.mean(coords, axis=0)
+            
+            # Calculate distances to centroid
+            distances = np.linalg.norm(coords - centroid, axis=1)
+            
+            # Get indices of closest points
+            closest_indices = np.argsort(distances)[:target_size]
+            
+            return coords[closest_indices]
+        
+        # Sample all endpoint clusters to target size using centroid-based selection
+        x1_1_coords = select_closest_to_centroid(x1_1_coords, target_size)
+        x1_2_coords = select_closest_to_centroid(x1_2_coords, target_size)
+        
+        dmso_cluster_counts = dmso_data[leiden_dmso_col].value_counts()
+        
+        # DMSO (unchanged)
+        largest_dmso_cluster = dmso_cluster_counts.index[0]
+        dmso_cluster_data = dmso_data[dmso_data[leiden_dmso_col] == largest_dmso_cluster]
+        
+        dmso_coords = dmso_cluster_data[pc_cols].values
+        
+        # Random sampling from largest DMSO cluster to match target size
+        # For DMSO, we'll also use centroid-based selection for consistency
+        if len(dmso_coords) >= target_size:
+            x0_coords = select_closest_to_centroid(dmso_coords, target_size)
+        else:
+            # If largest cluster is smaller than target, use all of it and pad with other DMSO cells
+            remaining_needed = target_size - len(dmso_coords)
+            other_dmso_data = dmso_data[dmso_data[leiden_dmso_col] != largest_dmso_cluster]
+            other_dmso_coords = other_dmso_data[pc_cols].values
+            
+            if len(other_dmso_coords) >= remaining_needed:
+                # Select closest to centroid from other DMSO cells
+                other_selected = select_closest_to_centroid(other_dmso_coords, remaining_needed)
+                x0_coords = np.vstack([dmso_coords, other_selected])
+            else:
+                # Use all available DMSO cells and reduce target size
+                all_dmso_coords = dmso_data[pc_cols].values
+                target_size = min(target_size, len(all_dmso_coords))
+                x0_coords = select_closest_to_centroid(all_dmso_coords, target_size)
+                
+                # Re-select endpoint clusters with updated target size
+                x1_1_coords = select_closest_to_centroid(x1_1_data[pc_cols].values.astype(float), target_size)
+                x1_2_coords = select_closest_to_centroid(x1_2_data[pc_cols].values.astype(float), target_size)
+        
+        # No need to resample since we already selected the right number
+        # The endpoint clusters are already at target_size from centroid-based selection
+        
+        self.n_samples = target_size
+        
+        x0 = torch.tensor(x0_coords, dtype=torch.float32)
+        x1_1 = torch.tensor(x1_1_coords, dtype=torch.float32)
+        x1_2 = torch.tensor(x1_2_coords, dtype=torch.float32)
+        x1 = torch.cat([x1_1, x1_2], dim=0)
+        
+        self.coords_t0 = x0
+        self.coords_t1 = x1
+        self.time_labels = np.concatenate([
+            np.zeros(len(self.coords_t0)),    # t=0
+            np.ones(len(self.coords_t1)),     # t=1
+        ])
+        
+        split_index = int(target_size * self.split_ratios[0])
+        
+        if target_size - split_index < self.batch_size:
+            split_index = target_size - self.batch_size
+        print('total count is:', target_size)
+        
+        train_x0 = x0[:split_index]
+        val_x0 = x0[split_index:]
+        train_x1 = x1[:split_index]
+        val_x1 = x1[split_index:]
+
+        
+        self.val_x0 = val_x0
+        
+        train_x0_weights = torch.full((train_x0.shape[0], 1), fill_value=1.0)
+        train_x1_weights = torch.full((train_x1.shape[0], 1), fill_value=1.0)
+        
+        val_x0_weights = torch.full((val_x0.shape[0], 1), fill_value=1.0)
+        val_x1_weights = torch.full((val_x1.shape[0], 1), fill_value=1.0)
+        
+        # Updated train dataloaders to include x1_3
+        self.train_dataloaders = {
+            "x0": DataLoader(TensorDataset(train_x0, train_x0_weights), batch_size=self.batch_size, shuffle=True, drop_last=True),
+            "x1": DataLoader(TensorDataset(train_x1, train_x1_weights), batch_size=self.batch_size, shuffle=True, drop_last=True),
+        }
+        
+        self.val_dataloaders = {
+            "x0": DataLoader(TensorDataset(val_x0, val_x0_weights), batch_size=self.batch_size, shuffle=False, drop_last=True),
+            "x1": DataLoader(TensorDataset(val_x1, val_x1_weights), batch_size=self.batch_size, shuffle=True, drop_last=True),
+        }
+        
+        all_coords = df[pc_cols].dropna().values.astype(float)
+        self.dataset = torch.tensor(all_coords, dtype=torch.float32)
+        self.tree = cKDTree(all_coords)
+
+        self.test_dataloaders = {
+            "x0": DataLoader(TensorDataset(val_x0, val_x0_weights), batch_size=self.val_x0.shape[0], shuffle=False, drop_last=False),
+            "dataset": DataLoader(TensorDataset(self.dataset), batch_size=self.dataset.shape[0], shuffle=False, drop_last=False),
+        }
+
+        # Updated metric samples - now using 4 clusters instead of 3
+        #km_all = KMeans(n_clusters=4, random_state=42).fit(self.dataset.numpy())
+        km_all = KMeans(n_clusters=2, random_state=0).fit(self.dataset.numpy())
+
+        cluster_labels = km_all.labels_
+        
+        cluster_0_mask = cluster_labels == 0
+        cluster_1_mask = cluster_labels == 1
+        
+        samples = self.dataset.cpu().numpy()
+            
+        cluster_0_data = samples[cluster_0_mask]
+        cluster_1_data = samples[cluster_1_mask]
+                
+        self.metric_samples_dataloaders = [
+            DataLoader(
+                torch.tensor(cluster_1_data, dtype=torch.float32),
+                batch_size=cluster_1_data.shape[0],
+                shuffle=False,
+                drop_last=False,
+            ),
+            DataLoader(
+                torch.tensor(cluster_0_data, dtype=torch.float32),
+                batch_size=cluster_0_data.shape[0],
+                shuffle=False,
+                drop_last=False,
+            ),
+        ]
+        
+    def train_dataloader(self):
+        combined_loaders = {
+            "train_samples": CombinedLoader(self.train_dataloaders, mode="min_size"),
+            "metric_samples": CombinedLoader(
+                self.metric_samples_dataloaders, mode="min_size"
+            ),
+        }
+        return CombinedLoader(combined_loaders, mode="max_size_cycle")
+
+    def val_dataloader(self):
+        combined_loaders = {
+            "val_samples": CombinedLoader(self.val_dataloaders, mode="min_size"),
+            "metric_samples": CombinedLoader(
+                self.metric_samples_dataloaders, mode="min_size"
+            ),
+        }
+
+        return CombinedLoader(combined_loaders, mode="max_size_cycle")
+
+    
+    
+    def test_dataloader(self):
+        combined_loaders = {
+            "test_samples": CombinedLoader(self.test_dataloaders, mode="min_size"),
+            "metric_samples": CombinedLoader(
+                self.metric_samples_dataloaders, mode="min_size"
+            ),
+        }
+
+        return CombinedLoader(combined_loaders, mode="max_size_cycle")
+
+    def get_manifold_proj(self, points):
+        """Adapted for 2D cell data - uses local neighborhood averaging instead of plane fitting"""
+        return partial(self.local_smoothing_op, tree=self.tree, dataset=self.dataset)
+
+    @staticmethod
+    def local_smoothing_op(x, tree, dataset, k=10, temp=1e-3):
+        """
+        Apply local smoothing based on k-nearest neighbors in the full dataset
+        This replaces the plane projection for 2D manifold regularization
+        """
+        points_np = x.detach().cpu().numpy()
+        _, idx = tree.query(points_np, k=k)
+        nearest_pts = dataset[idx]  # Shape: (batch_size, k, 2)
+        
+        # Compute weighted average of neighbors
+        dists = (x.unsqueeze(1) - nearest_pts).pow(2).sum(-1, keepdim=True)
+        weights = torch.exp(-dists / temp)
+        weights = weights / weights.sum(dim=1, keepdim=True)
+        
+        # Weighted average of neighbors
+        smoothed = (weights * nearest_pts).sum(dim=1)
+        
+        # Blend original point with smoothed version
+        alpha = 0.3  # How much smoothing to apply
+        return (1 - alpha) * x + alpha * smoothed
+    
+    def get_timepoint_data(self):
+        """Return data organized by timepoints for visualization"""
+        return {
+            't0': self.coords_t0,
+            't1': self.coords_t1, 
+            'time_labels': self.time_labels
+        }
+    
+def get_datamodule():
+    datamodule = ClonidineSingleBranchDataModule(args)
+    datamodule.setup(stage="fit")   
+    return datamodule

dataloaders/clonidine_v2_data.py

@@ -0,0 +1,287 @@
+import torch
+import sys
+sys.argv = ['']
+from sklearn.preprocessing import StandardScaler
+import pytorch_lightning as pl
+from torch.utils.data import DataLoader
+from lightning.pytorch.utilities.combined_loader import CombinedLoader
+import pandas as pd
+import numpy as np
+from functools import partial
+from scipy.spatial import cKDTree
+from sklearn.cluster import KMeans
+from torch.utils.data import TensorDataset
+
+from train.parsers_tahoe import parse_args
+args = parse_args()
+
+class ClonidineV2DataModule(pl.LightningDataModule):
+    def __init__(self, args):
+        super().__init__()
+        self.save_hyperparameters()
+
+        self.batch_size = args.batch_size
+        self.max_dim = args.dim
+        self.whiten = args.whiten
+        self.split_ratios = args.split_ratios
+        
+        self.dim = args.dim
+        print("dimension")
+        print(self.dim)
+        # Path to your combined data
+        self.data_path = "./data/pca_and_leiden_labels.csv"
+        self.num_timesteps = 2
+        self.args = args
+        self._prepare_data()
+
+    def _prepare_data(self):
+        df = pd.read_csv(self.data_path, comment='#')
+        df = df.iloc[:, 1:]
+        df = df.replace('', np.nan)
+        pc_cols = df.columns[:150]
+        for col in pc_cols:
+            df[col] = pd.to_numeric(df[col], errors='coerce')
+        leiden_dmso_col = 'leiden_DMSO_TF_0.0uM'
+        leiden_clonidine_col = 'leiden_Clonidine (hydrochloride)_5.0uM'
+
+        dmso_mask = df[leiden_dmso_col].notna()        # Has leiden value in DMSO column
+        clonidine_mask = df[leiden_clonidine_col].notna()  # Has leiden value in Clonidine column
+        
+        dmso_data = df[dmso_mask].copy()
+        clonidine_data = df[clonidine_mask].copy()
+        
+        top_clonidine_clusters = ['0.0', '4.0']
+        
+        x1_1_data = clonidine_data[clonidine_data[leiden_clonidine_col].astype(str) == top_clonidine_clusters[0]]
+        x1_2_data = clonidine_data[clonidine_data[leiden_clonidine_col].astype(str) == top_clonidine_clusters[1]]
+        
+        x1_1_coords = x1_1_data[pc_cols].values 
+        x1_2_coords = x1_2_data[pc_cols].values 
+        
+        x1_1_coords = x1_1_coords.astype(float)
+        x1_2_coords = x1_2_coords.astype(float)
+
+        # Target size is now the minimum across all three endpoint clusters
+        target_size = min(len(x1_1_coords), len(x1_2_coords),)
+        
+        # Helper function to select points closest to centroid
+        def select_closest_to_centroid(coords, target_size):
+            if len(coords) <= target_size:
+                return coords
+            
+            # Calculate centroid
+            centroid = np.mean(coords, axis=0)
+            
+            # Calculate distances to centroid
+            distances = np.linalg.norm(coords - centroid, axis=1)
+            
+            # Get indices of closest points
+            closest_indices = np.argsort(distances)[:target_size]
+            
+            return coords[closest_indices]
+        
+        # Sample all endpoint clusters to target size using centroid-based selection
+        x1_1_coords = select_closest_to_centroid(x1_1_coords, target_size)
+        x1_2_coords = select_closest_to_centroid(x1_2_coords, target_size)
+        
+        dmso_cluster_counts = dmso_data[leiden_dmso_col].value_counts()
+        
+        # DMSO (unchanged)
+        largest_dmso_cluster = dmso_cluster_counts.index[0]
+        dmso_cluster_data = dmso_data[dmso_data[leiden_dmso_col] == largest_dmso_cluster]
+        
+        dmso_coords = dmso_cluster_data[pc_cols].values
+        
+        # Random sampling from largest DMSO cluster to match target size
+        # For DMSO, we'll also use centroid-based selection for consistency
+        if len(dmso_coords) >= target_size:
+            x0_coords = select_closest_to_centroid(dmso_coords, target_size)
+        else:
+            # If largest cluster is smaller than target, use all of it and pad with other DMSO cells
+            remaining_needed = target_size - len(dmso_coords)
+            other_dmso_data = dmso_data[dmso_data[leiden_dmso_col] != largest_dmso_cluster]
+            other_dmso_coords = other_dmso_data[pc_cols].values
+            
+            if len(other_dmso_coords) >= remaining_needed:
+                # Select closest to centroid from other DMSO cells
+                other_selected = select_closest_to_centroid(other_dmso_coords, remaining_needed)
+                x0_coords = np.vstack([dmso_coords, other_selected])
+            else:
+                # Use all available DMSO cells and reduce target size
+                all_dmso_coords = dmso_data[pc_cols].values
+                target_size = min(target_size, len(all_dmso_coords))
+                x0_coords = select_closest_to_centroid(all_dmso_coords, target_size)
+                
+                # Re-select endpoint clusters with updated target size
+                x1_1_coords = select_closest_to_centroid(x1_1_data[pc_cols].values.astype(float), target_size)
+                x1_2_coords = select_closest_to_centroid(x1_2_data[pc_cols].values.astype(float), target_size)
+        
+        # No need to resample since we already selected the right number
+        # The endpoint clusters are already at target_size from centroid-based selection
+        
+        self.n_samples = target_size
+        
+        x0 = torch.tensor(x0_coords, dtype=torch.float32)
+        x1_1 = torch.tensor(x1_1_coords, dtype=torch.float32)
+        x1_2 = torch.tensor(x1_2_coords, dtype=torch.float32)
+        
+        self.coords_t0 = x0
+        self.coords_t1_1 = x1_1
+        self.coords_t1_2 = x1_2
+        self.time_labels = np.concatenate([
+            np.zeros(len(self.coords_t0)),    # t=0
+            np.ones(len(self.coords_t1_1)),     # t=1
+            np.ones(len(self.coords_t1_2)),
+        ])
+        
+        split_index = int(target_size * self.split_ratios[0])
+        
+        if target_size - split_index < self.batch_size:
+            split_index = target_size - self.batch_size
+        print('total count is:', target_size)
+        
+        train_x0 = x0[:split_index]
+        val_x0 = x0[split_index:]
+        train_x1_1 = x1_1[:split_index]
+        val_x1_1 = x1_1[split_index:]
+        train_x1_2 = x1_2[:split_index]
+        val_x1_2 = x1_2[split_index:]
+
+        
+        self.val_x0 = val_x0
+        
+        train_x0_weights = torch.full((train_x0.shape[0], 1), fill_value=1.0)
+        train_x1_1_weights = torch.full((train_x1_1.shape[0], 1), fill_value=0.5)
+        train_x1_2_weights = torch.full((train_x1_2.shape[0], 1), fill_value=0.5)
+        
+        val_x0_weights = torch.full((val_x0.shape[0], 1), fill_value=1.0)
+        val_x1_1_weights = torch.full((val_x1_1.shape[0], 1), fill_value=0.5)
+        val_x1_2_weights = torch.full((val_x1_2.shape[0], 1), fill_value=0.5)
+        
+        # Updated train dataloaders to include x1_3
+        self.train_dataloaders = {
+            "x0": DataLoader(TensorDataset(train_x0, train_x0_weights), batch_size=self.batch_size, shuffle=True, drop_last=True),
+            "x1_1": DataLoader(TensorDataset(train_x1_1, train_x1_1_weights), batch_size=self.batch_size, shuffle=True, drop_last=True),
+            "x1_2": DataLoader(TensorDataset(train_x1_2, train_x1_2_weights), batch_size=self.batch_size, shuffle=True, drop_last=True),
+        }
+        
+        self.val_dataloaders = {
+            "x0": DataLoader(TensorDataset(val_x0, val_x0_weights), batch_size=self.batch_size, shuffle=False, drop_last=True),
+            "x1_1": DataLoader(TensorDataset(val_x1_1, val_x1_1_weights), batch_size=self.batch_size, shuffle=True, drop_last=True),
+            "x1_2": DataLoader(TensorDataset(val_x1_2, val_x1_2_weights), batch_size=self.batch_size, shuffle=True, drop_last=True),
+        }
+        
+        all_coords = df[pc_cols].dropna().values.astype(float)
+        self.dataset = torch.tensor(all_coords, dtype=torch.float32)
+        self.tree = cKDTree(all_coords)
+
+        self.test_dataloaders = {
+            "x0": DataLoader(TensorDataset(val_x0, val_x0_weights), batch_size=self.val_x0.shape[0], shuffle=False, drop_last=False),
+            "dataset": DataLoader(TensorDataset(self.dataset), batch_size=self.dataset.shape[0], shuffle=False, drop_last=False),
+        }
+
+        km_all = KMeans(n_clusters=3, random_state=0).fit(self.dataset.numpy())
+
+        cluster_labels = km_all.labels_
+        
+        cluster_0_mask = cluster_labels == 0
+        cluster_1_mask = cluster_labels == 1
+        cluster_2_mask = cluster_labels == 2
+        
+        samples = self.dataset.cpu().numpy()
+            
+        cluster_0_data = samples[cluster_0_mask]
+        cluster_1_data = samples[cluster_1_mask]
+        cluster_2_data = samples[cluster_2_mask]
+                
+        self.metric_samples_dataloaders = [
+            DataLoader(
+                torch.tensor(cluster_2_data, dtype=torch.float32),
+                batch_size=cluster_2_data.shape[0],
+                shuffle=False,
+                drop_last=False,
+            ),
+            DataLoader(
+                torch.tensor(cluster_0_data, dtype=torch.float32),
+                batch_size=cluster_0_data.shape[0],
+                shuffle=False,
+                drop_last=False,
+            ),
+            
+            DataLoader(
+                torch.tensor(cluster_1_data, dtype=torch.float32),
+                batch_size=cluster_1_data.shape[0],
+                shuffle=False,
+                drop_last=False,
+            ),
+        ]
+        
+    def train_dataloader(self):
+        combined_loaders = {
+            "train_samples": CombinedLoader(self.train_dataloaders, mode="min_size"),
+            "metric_samples": CombinedLoader(
+                self.metric_samples_dataloaders, mode="min_size"
+            ),
+        }
+        return CombinedLoader(combined_loaders, mode="max_size_cycle")
+
+    def val_dataloader(self):
+        combined_loaders = {
+            "val_samples": CombinedLoader(self.val_dataloaders, mode="min_size"),
+            "metric_samples": CombinedLoader(
+                self.metric_samples_dataloaders, mode="min_size"
+            ),
+        }
+
+        return CombinedLoader(combined_loaders, mode="max_size_cycle")
+    
+    
+    def test_dataloader(self):
+        combined_loaders = {
+            "test_samples": CombinedLoader(self.test_dataloaders, mode="min_size"),
+            "metric_samples": CombinedLoader(
+                self.metric_samples_dataloaders, mode="min_size"
+            ),
+        }
+
+        return CombinedLoader(combined_loaders, mode="max_size_cycle")
+
+    def get_manifold_proj(self, points):
+        """Adapted for 2D cell data - uses local neighborhood averaging instead of plane fitting"""
+        return partial(self.local_smoothing_op, tree=self.tree, dataset=self.dataset)
+
+    @staticmethod
+    def local_smoothing_op(x, tree, dataset, k=10, temp=1e-3):
+        """
+        Apply local smoothing based on k-nearest neighbors in the full dataset
+        This replaces the plane projection for 2D manifold regularization
+        """
+        points_np = x.detach().cpu().numpy()
+        _, idx = tree.query(points_np, k=k)
+        nearest_pts = dataset[idx]  # Shape: (batch_size, k, 2)
+        
+        # Compute weighted average of neighbors
+        dists = (x.unsqueeze(1) - nearest_pts).pow(2).sum(-1, keepdim=True)
+        weights = torch.exp(-dists / temp)
+        weights = weights / weights.sum(dim=1, keepdim=True)
+        
+        # Weighted average of neighbors
+        smoothed = (weights * nearest_pts).sum(dim=1)
+        
+        # Blend original point with smoothed version
+        alpha = 0.3  # How much smoothing to apply
+        return (1 - alpha) * x + alpha * smoothed
+    
+    def get_timepoint_data(self):
+        """Return data organized by timepoints for visualization"""
+        return {
+            't0': self.coords_t0,
+            't1_1': self.coords_t1_1, 
+            't1_2': self.coords_t1_2, 
+            'time_labels': self.time_labels
+        }
+    
+def get_datamodule():
+    datamodule = ClonidineV2DataModule(args)
+    datamodule.setup(stage="fit")   
+    return datamodule

dataloaders/lidar_data.py

@@ -0,0 +1,532 @@
+import torch
+import sys
+sys.argv = ['']
+from sklearn.preprocessing import StandardScaler
+import pytorch_lightning as pl
+from torch.utils.data import DataLoader
+from pytorch_lightning.utilities.combined_loader import CombinedLoader
+import laspy
+import numpy as np
+from scipy.spatial import cKDTree
+import math
+from functools import partial
+from torch.utils.data import TensorDataset
+
+#from train.parsers import parse_args
+#args = parse_args()
+
+class GaussianMM:
+    def __init__(self, mu, var):
+        super().__init__()
+        self.centers = torch.tensor(mu)
+        self.logstd = torch.tensor(var).log() / 2.0
+        self.K = self.centers.shape[0]
+
+    def logprob(self, x):
+        logprobs = self.normal_logprob(
+            x.unsqueeze(1), self.centers.unsqueeze(0), self.logstd
+        )
+        logprobs = torch.sum(logprobs, dim=2)
+        return torch.logsumexp(logprobs, dim=1) - math.log(self.K)
+
+    def normal_logprob(self, z, mean, log_std):
+        mean = mean + torch.tensor(0.0)
+        log_std = log_std + torch.tensor(0.0)
+        c = torch.tensor([math.log(2 * math.pi)]).to(z)
+        inv_sigma = torch.exp(-log_std)
+        tmp = (z - mean) * inv_sigma
+        return -0.5 * (tmp * tmp + 2 * log_std + c)
+
+    def __call__(self, n_samples):
+        idx = torch.randint(self.K, (n_samples,)).to(self.centers.device)
+        mean = self.centers[idx]
+        return torch.randn(*mean.shape).to(mean) * torch.exp(self.logstd) + mean
+    
+class BranchedLidarDataModule(pl.LightningDataModule):
+    def __init__(self, args):
+        super().__init__()
+        self.save_hyperparameters()
+
+        self.data_path = args.data_path
+        self.batch_size = args.batch_size
+        self.max_dim = args.dim
+        self.whiten = args.whiten
+        self.p0_mu = [
+            [-4.5, -4.0, 0.5],
+            [-4.2, -3.5, 0.5],
+            [-4.0, -3.0, 0.5],
+            [-3.75, -2.5, 0.5],
+        ]
+        self.p0_var = 0.02
+        
+        self.p1_1_mu = [
+            [-2.5, -0.25, 0.5],
+            [-2.25, 0.675, 0.5],
+            [-2, 1.5, 0.5],
+        ]
+        self.p1_2_mu = [
+            [2, -2, 0.5], 
+            [2.6, -1.25, 0.5], 
+            [3.2, -0.5, 0.5]
+        ]
+        
+        self.p1_var = 0.03
+        self.k = 20
+        self.n_samples = 5000
+        self.num_timesteps = 2
+        self.split_ratios = args.split_ratios
+        self._prepare_data()
+        
+    def assign_region(self):
+        all_centers = {
+            0: torch.tensor(self.p0_mu),     # Region 0: p0
+            1: torch.tensor(self.p1_1_mu),   # Region 1: p1_1
+            2: torch.tensor(self.p1_2_mu),   # Region 2: p1_2
+        }
+
+        dataset = self.dataset.to(torch.float32)
+        N = dataset.shape[0]
+        assignments = torch.zeros(N, dtype=torch.long)
+
+        # For each point, compute min distance to each region's centers
+        for i in range(N):
+            point = dataset[i]
+            min_dist = float("inf")
+            best_region = 0
+            for region, centers in all_centers.items():
+                dists = ((centers - point)**2).sum(dim=1)
+                region_min = dists.min()
+                if region_min < min_dist:
+                    min_dist = region_min
+                    best_region = region
+            assignments[i] = best_region
+        return assignments
+
+    def _prepare_data(self):
+        las = laspy.read(self.data_path)
+        # Extract only "ground" points.
+        self.mask = las.classification == 2
+        # Original Preprocessing
+        x_offset, x_scale = las.header.offsets[0], las.header.scales[0]
+        y_offset, y_scale = las.header.offsets[1], las.header.scales[1]
+        z_offset, z_scale = las.header.offsets[2], las.header.scales[2]
+        dataset = np.vstack(
+            (
+                las.X[self.mask] * x_scale + x_offset,
+                las.Y[self.mask] * y_scale + y_offset,
+                las.Z[self.mask] * z_scale + z_offset,
+            )
+        ).transpose()
+        mi = dataset.min(axis=0, keepdims=True)
+        ma = dataset.max(axis=0, keepdims=True)
+        dataset = (dataset - mi) / (ma - mi) * [10.0, 10.0, 2.0] + [-5.0, -5.0, 0.0]
+
+        self.dataset = torch.tensor(dataset, dtype=torch.float32)
+        self.tree = cKDTree(dataset)
+
+        x0_gaussian = GaussianMM(self.p0_mu, self.p0_var)(self.n_samples)
+        x1_1_gaussian = GaussianMM(self.p1_1_mu, self.p1_var)(self.n_samples)
+        x1_2_gaussian = GaussianMM(self.p1_2_mu, self.p1_var)(self.n_samples)
+
+        x0 = self.get_tangent_proj(x0_gaussian)(x0_gaussian)
+        x1_1 = self.get_tangent_proj(x1_1_gaussian)(x1_1_gaussian)
+        x1_2 = self.get_tangent_proj(x1_2_gaussian)(x1_2_gaussian)
+
+        split_index = int(self.n_samples * self.split_ratios[0])
+
+        self.scaler = StandardScaler()
+        if self.whiten:
+            self.dataset = torch.tensor(
+                self.scaler.fit_transform(dataset), dtype=torch.float32
+            )
+            x0 = torch.tensor(self.scaler.transform(x0), dtype=torch.float32)
+            x1_1 = torch.tensor(self.scaler.transform(x1_1), dtype=torch.float32)
+            x1_2 = torch.tensor(self.scaler.transform(x1_2), dtype=torch.float32)
+
+        train_x0 = x0[:split_index]
+        val_x0 = x0[split_index:]
+        
+        # branches
+        train_x1_1 = x1_1[:split_index]
+        print("train_x1_1")
+        print(train_x1_1.shape)
+        val_x1_1 = x1_1[split_index:]
+        train_x1_2 = x1_2[:split_index]
+        val_x1_2 = x1_2[split_index:]
+        
+        self.val_x0 = val_x0
+
+        # Adjust split_index to ensure minimum validation samples
+        if self.n_samples - split_index < self.batch_size:
+            split_index = self.n_samples - self.batch_size
+
+        self.train_dataloaders = {
+            "x0": DataLoader(train_x0, batch_size=self.batch_size, shuffle=True, drop_last=True),
+            "x1_1": DataLoader(train_x1_1, batch_size=self.batch_size, shuffle=True, drop_last=True),
+            "x1_2": DataLoader(train_x1_2, batch_size=self.batch_size, shuffle=True, drop_last=True),
+        }
+        self.val_dataloaders = {
+            "x0": DataLoader(val_x0, batch_size=self.batch_size, shuffle=False, drop_last=True),
+            "x1_1": DataLoader(val_x1_1, batch_size=self.batch_size, shuffle=True, drop_last=True),
+            "x1_2": DataLoader(val_x1_2, batch_size=self.batch_size, shuffle=True, drop_last=True),
+        }
+        # to edit?
+        self.test_dataloaders = [
+            DataLoader(
+                self.val_x0,
+                batch_size=self.val_x0.shape[0],
+                shuffle=False,
+                drop_last=False,
+            ),
+            DataLoader(
+                self.dataset,
+                batch_size=self.dataset.shape[0],
+                shuffle=False,
+                drop_last=False,
+            ),
+        ]
+        
+        points = self.dataset.cpu().numpy()
+        x, y = points[:, 0], points[:, 1]
+        # Diagonal-based coordinates (rotated 45°)
+        u = (x + y) / np.sqrt(2)  # along x=y
+        # start region (A) using u
+        u_thresh = np.percentile(u, 30)  # tweak this threshold to control size
+        mask_A = u <= u_thresh
+
+        # among the rest, split by x=y diagonal
+        remaining = ~mask_A
+        mask_B = remaining & (x < y)  # left of diagonal
+        mask_C = remaining & (x >= y)  # right of diagonal
+
+        # Assign dataloaders
+        self.metric_samples_dataloaders = [
+            DataLoader(torch.tensor(points[mask_A], dtype=torch.float32), batch_size=points[mask_A].shape[0], shuffle=False),
+            DataLoader(torch.tensor(points[mask_B], dtype=torch.float32), batch_size=points[mask_B].shape[0], shuffle=False),
+            DataLoader(torch.tensor(points[mask_C], dtype=torch.float32), batch_size=points[mask_C].shape[0], shuffle=False),
+        ]
+
+    def train_dataloader(self):
+        combined_loaders = {
+            "train_samples": CombinedLoader(self.train_dataloaders, mode="min_size"),
+            "metric_samples": CombinedLoader(
+                self.metric_samples_dataloaders, mode="min_size"
+            ),
+        }
+        return CombinedLoader(combined_loaders, mode="max_size_cycle")
+
+    def val_dataloader(self):
+        combined_loaders = {
+            "val_samples": CombinedLoader(self.val_dataloaders, mode="min_size"),
+            "metric_samples": CombinedLoader(
+                self.metric_samples_dataloaders, mode="min_size"
+            ),
+        }
+
+        return CombinedLoader(combined_loaders, mode="max_size_cycle")
+
+    def test_dataloader(self):
+        return CombinedLoader(self.test_dataloaders)
+
+    def get_tangent_proj(self, points):
+        w = self.get_tangent_plane(points)
+        return partial(BranchedLidarDataModule.projection_op, w=w)
+
+    def get_tangent_plane(self, points, temp=1e-3):
+        points_np = points.detach().cpu().numpy()
+        _, idx = self.tree.query(points_np, k=self.k)
+        nearest_pts = self.dataset[idx]
+        nearest_pts = torch.tensor(nearest_pts).to(points)
+
+        dists = (points.unsqueeze(1) - nearest_pts).pow(2).sum(-1, keepdim=True)
+        weights = torch.exp(-dists / temp)
+
+        # Fits plane with least vertical distance.
+        w = BranchedLidarDataModule.fit_plane(nearest_pts, weights)
+        return w
+
+    @staticmethod
+    def fit_plane(points, weights=None):
+        """Expects points to be of shape (..., 3).
+        Returns [a, b, c] such that the plane is defined as
+            ax + by + c = z
+        """
+        D = torch.cat([points[..., :2], torch.ones_like(points[..., 2:3])], dim=-1)
+        z = points[..., 2]
+        if weights is not None:
+            Dtrans = D.transpose(-1, -2)
+        else:
+            DW = D * weights
+            Dtrans = DW.transpose(-1, -2)
+        w = torch.linalg.solve(
+            torch.matmul(Dtrans, D), torch.matmul(Dtrans, z.unsqueeze(-1))
+        ).squeeze(-1)
+        return w
+
+    @staticmethod
+    def projection_op(x, w):
+        """Projects points to a plane defined by w."""
+        # Normal vector to the tangent plane.
+        n = torch.cat([w[..., :2], -torch.ones_like(w[..., 2:3])], dim=1)
+
+        pn = torch.sum(x * n, dim=-1, keepdim=True)
+        nn = torch.sum(n * n, dim=-1, keepdim=True)
+
+        # Offset.
+        d = w[..., 2:3]
+
+        # Projection of x onto n.
+        projn_x = ((pn + d) / nn) * n
+
+        # Remove component in the normal direction.
+        return x - projn_x
+
+class WeightedBranchedLidarDataModule(pl.LightningDataModule):
+    def __init__(self, args):
+        super().__init__()
+        self.save_hyperparameters()
+
+        self.data_path = args.data_path
+        self.batch_size = args.batch_size
+        self.max_dim = args.dim
+        self.whiten = args.whiten
+        self.p0_mu = [
+            [-4.5, -4.0, 0.5],
+            [-4.2, -3.5, 0.5],
+            [-4.0, -3.0, 0.5],
+            [-3.75, -2.5, 0.5],
+        ]
+        self.p0_var = 0.02
+        # multiple p1 for each branch
+        #changed
+        self.p1_1_mu = [
+            [-2.5, -0.25, 0.5],
+            [-2.25, 0.675, 0.5],
+            [-2, 1.5, 0.5],
+        ]
+        self.p1_2_mu = [
+            [2, -2, 0.5], 
+            [2.6, -1.25, 0.5], 
+            [3.2, -0.5, 0.5]
+        ]
+        
+        self.p1_var = 0.03
+        self.k = 20
+        self.n_samples = 5000
+        self.num_timesteps = 2
+        self.split_ratios = args.split_ratios
+        
+        self.num_timesteps = 2
+        self.metric_clusters = 3
+        self.args = args
+        self._prepare_data()
+
+    def _prepare_data(self):
+        las = laspy.read(self.data_path)
+        # Extract only "ground" points.
+        self.mask = las.classification == 2
+        # Original Preprocessing
+        x_offset, x_scale = las.header.offsets[0], las.header.scales[0]
+        y_offset, y_scale = las.header.offsets[1], las.header.scales[1]
+        z_offset, z_scale = las.header.offsets[2], las.header.scales[2]
+        dataset = np.vstack(
+            (
+                las.X[self.mask] * x_scale + x_offset,
+                las.Y[self.mask] * y_scale + y_offset,
+                las.Z[self.mask] * z_scale + z_offset,
+            )
+        ).transpose()
+        mi = dataset.min(axis=0, keepdims=True)
+        ma = dataset.max(axis=0, keepdims=True)
+        dataset = (dataset - mi) / (ma - mi) * [10.0, 10.0, 2.0] + [-5.0, -5.0, 0.0]
+
+        self.dataset = torch.tensor(dataset, dtype=torch.float32)
+        self.tree = cKDTree(dataset)
+
+        x0_gaussian = GaussianMM(self.p0_mu, self.p0_var)(self.n_samples)
+        x1_1_gaussian = GaussianMM(self.p1_1_mu, self.p1_var)(self.n_samples)
+        x1_2_gaussian = GaussianMM(self.p1_2_mu, self.p1_var)(self.n_samples)
+
+        x0 = self.get_tangent_proj(x0_gaussian)(x0_gaussian)
+        x1_1 = self.get_tangent_proj(x1_1_gaussian)(x1_1_gaussian)
+        x1_2 = self.get_tangent_proj(x1_2_gaussian)(x1_2_gaussian)
+
+        split_index = int(self.n_samples * self.split_ratios[0])
+
+        self.scaler = StandardScaler()
+        if self.whiten:
+            self.dataset = torch.tensor(
+                self.scaler.fit_transform(dataset), dtype=torch.float32
+            )
+            x0 = torch.tensor(self.scaler.transform(x0), dtype=torch.float32)
+            x1_1 = torch.tensor(self.scaler.transform(x1_1), dtype=torch.float32)
+            x1_2 = torch.tensor(self.scaler.transform(x1_2), dtype=torch.float32)
+
+        self.coords_t0 = x0
+        self.coords_t1_1 = x1_1
+        self.coords_t1_2 = x1_2
+        self.time_labels = np.concatenate([
+            np.zeros(len(self.coords_t0)),    # t=0
+            np.ones(len(self.coords_t1_1)),     # t=1
+            np.ones(len(self.coords_t1_2)),     # t=1
+        ])
+        
+        train_x0 = x0[:split_index]
+        val_x0 = x0[split_index:]
+        
+        # branches
+        train_x1_1 = x1_1[:split_index]
+        
+        val_x1_1 = x1_1[split_index:]
+        train_x1_2 = x1_2[:split_index]
+        val_x1_2 = x1_2[split_index:]
+        
+        self.val_x0 = val_x0
+        
+        train_x0_weights = torch.full((train_x0.shape[0], 1), fill_value=1.0)
+        train_x1_1_weights = torch.full((train_x1_1.shape[0], 1), fill_value=0.5)
+        train_x1_2_weights = torch.full((train_x1_2.shape[0], 1), fill_value=0.5)
+        
+        val_x0_weights = torch.full((val_x0.shape[0], 1), fill_value=1.0)
+        val_x1_1_weights = torch.full((val_x1_1.shape[0], 1), fill_value=0.5)
+        val_x1_2_weights = torch.full((val_x1_2.shape[0], 1), fill_value=0.5)
+
+        # Adjust split_index to ensure minimum validation samples
+        if self.n_samples - split_index < self.batch_size:
+            split_index = self.n_samples - self.batch_size
+            
+        self.train_dataloaders = {
+            "x0": DataLoader(TensorDataset(train_x0, train_x0_weights), batch_size=self.batch_size, shuffle=True, drop_last=True),
+            "x1_1": DataLoader(TensorDataset(train_x1_1, train_x1_1_weights), batch_size=self.batch_size, shuffle=True, drop_last=True),
+            "x1_2": DataLoader(TensorDataset(train_x1_2, train_x1_2_weights), batch_size=self.batch_size, shuffle=True, drop_last=True),
+        }
+        
+        self.val_dataloaders = {
+            "x0": DataLoader(TensorDataset(val_x0, val_x0_weights), batch_size=self.batch_size, shuffle=False, drop_last=True),
+            "x1_1": DataLoader(TensorDataset(val_x1_1, val_x1_1_weights), batch_size=self.batch_size, shuffle=True, drop_last=True),
+            "x1_2": DataLoader(TensorDataset(val_x1_2, val_x1_2_weights), batch_size=self.batch_size, shuffle=True, drop_last=True),
+        }
+        
+        # to edit?
+        self.test_dataloaders = {
+            "x0": DataLoader(TensorDataset(val_x0, val_x0_weights), batch_size=self.val_x0.shape[0], shuffle=False, drop_last=False),
+            "x1_1": DataLoader(TensorDataset(val_x1_1, val_x1_1_weights), batch_size=self.val_x0.shape[0], shuffle=True, drop_last=True),
+            "x1_2": DataLoader(TensorDataset(val_x1_2, val_x1_2_weights), batch_size=self.val_x0.shape[0], shuffle=True, drop_last=True),
+            "dataset": DataLoader(TensorDataset(self.dataset), batch_size=self.dataset.shape[0], shuffle=False, drop_last=False),
+        }
+        
+        points = self.dataset.cpu().numpy()
+        x, y = points[:, 0], points[:, 1]
+        # Diagonal-based coordinates (rotated 45°)
+        u = (x + y) / np.sqrt(2)  # along x=y
+        # start region (A) using u
+        u_thresh = np.percentile(u, 30)  # tweak this threshold to control size
+        mask_A = u <= u_thresh
+
+        # among the rest, split by x=y diagonal
+        remaining = ~mask_A
+        mask_B = remaining & (x < y)  # left of diagonal
+        mask_C = remaining & (x >= y)  # right of diagonal
+
+        # Assign dataloaders
+        self.metric_samples_dataloaders = [
+            DataLoader(torch.tensor(points[mask_A], dtype=torch.float32), batch_size=points[mask_A].shape[0], shuffle=False),
+            DataLoader(torch.tensor(points[mask_B], dtype=torch.float32), batch_size=points[mask_B].shape[0], shuffle=False),
+            DataLoader(torch.tensor(points[mask_C], dtype=torch.float32), batch_size=points[mask_C].shape[0], shuffle=False),
+        ]
+
+    def train_dataloader(self):
+        combined_loaders = {
+            "train_samples": CombinedLoader(self.train_dataloaders, mode="min_size"),
+            "metric_samples": CombinedLoader(
+                self.metric_samples_dataloaders, mode="min_size"
+            ),
+        }
+        return CombinedLoader(combined_loaders, mode="max_size_cycle")
+
+    def val_dataloader(self):
+        combined_loaders = {
+            "val_samples": CombinedLoader(self.val_dataloaders, mode="min_size"),
+            "metric_samples": CombinedLoader(
+                self.metric_samples_dataloaders, mode="min_size"
+            ),
+        }
+
+        return CombinedLoader(combined_loaders, mode="max_size_cycle")
+
+    def test_dataloader(self):
+        combined_loaders = {
+            "test_samples": CombinedLoader(self.test_dataloaders, mode="min_size"),
+            "metric_samples": CombinedLoader(
+                self.metric_samples_dataloaders, mode="min_size"
+            ),
+        }
+
+        return CombinedLoader(combined_loaders, mode="max_size_cycle")
+
+    def get_tangent_proj(self, points):
+        w = self.get_tangent_plane(points)
+        return partial(BranchedLidarDataModule.projection_op, w=w)
+
+    def get_tangent_plane(self, points, temp=1e-3):
+        points_np = points.detach().cpu().numpy()
+        _, idx = self.tree.query(points_np, k=self.k)
+        nearest_pts = self.dataset[idx]
+        nearest_pts = torch.tensor(nearest_pts).to(points)
+
+        dists = (points.unsqueeze(1) - nearest_pts).pow(2).sum(-1, keepdim=True)
+        weights = torch.exp(-dists / temp)
+
+        # Fits plane with least vertical distance.
+        w = BranchedLidarDataModule.fit_plane(nearest_pts, weights)
+        return w
+
+    @staticmethod
+    def fit_plane(points, weights=None):
+        """Expects points to be of shape (..., 3).
+        Returns [a, b, c] such that the plane is defined as
+            ax + by + c = z
+        """
+        D = torch.cat([points[..., :2], torch.ones_like(points[..., 2:3])], dim=-1)
+        z = points[..., 2]
+        if weights is not None:
+            Dtrans = D.transpose(-1, -2)
+        else:
+            DW = D * weights
+            Dtrans = DW.transpose(-1, -2)
+        w = torch.linalg.solve(
+            torch.matmul(Dtrans, D), torch.matmul(Dtrans, z.unsqueeze(-1))
+        ).squeeze(-1)
+        return w
+
+    @staticmethod
+    def projection_op(x, w):
+        """Projects points to a plane defined by w."""
+        # Normal vector to the tangent plane.
+        n = torch.cat([w[..., :2], -torch.ones_like(w[..., 2:3])], dim=1)
+
+        pn = torch.sum(x * n, dim=-1, keepdim=True)
+        nn = torch.sum(n * n, dim=-1, keepdim=True)
+
+        # Offset.
+        d = w[..., 2:3]
+
+        # Projection of x onto n.
+        projn_x = ((pn + d) / nn) * n
+
+        # Remove component in the normal direction.
+        return x - projn_x
+    
+    def get_timepoint_data(self):
+        """Return data organized by timepoints for visualization"""
+        return {
+            't0': self.coords_t0,
+            't1_1': self.coords_t1_1, 
+            't1_2': self.coords_t1_2, 
+            'time_labels': self.time_labels
+        }
+
+def get_datamodule():
+    datamodule = WeightedBranchedLidarDataModule(args)
+    datamodule.setup(stage="fit")   
+    return datamodule

dataloaders/lidar_data_single.py

@@ -0,0 +1,282 @@
+import torch
+import sys
+sys.argv = ['']
+from sklearn.preprocessing import StandardScaler
+import pytorch_lightning as pl
+from torch.utils.data import DataLoader
+from pytorch_lightning.utilities.combined_loader import CombinedLoader
+import laspy
+import numpy as np
+from scipy.spatial import cKDTree
+import math
+from functools import partial
+from torch.utils.data import TensorDataset
+
+from train.parsers import parse_args
+args = parse_args()
+
+class GaussianMM:
+    def __init__(self, mu, var):
+        super().__init__()
+        self.centers = torch.tensor(mu)
+        self.logstd = torch.tensor(var).log() / 2.0
+        self.K = self.centers.shape[0]
+
+    def logprob(self, x):
+        logprobs = self.normal_logprob(
+            x.unsqueeze(1), self.centers.unsqueeze(0), self.logstd
+        )
+        logprobs = torch.sum(logprobs, dim=2)
+        return torch.logsumexp(logprobs, dim=1) - math.log(self.K)
+
+    def normal_logprob(self, z, mean, log_std):
+        mean = mean + torch.tensor(0.0)
+        log_std = log_std + torch.tensor(0.0)
+        c = torch.tensor([math.log(2 * math.pi)]).to(z)
+        inv_sigma = torch.exp(-log_std)
+        tmp = (z - mean) * inv_sigma
+        return -0.5 * (tmp * tmp + 2 * log_std + c)
+
+    def __call__(self, n_samples):
+        idx = torch.randint(self.K, (n_samples,)).to(self.centers.device)
+        mean = self.centers[idx]
+        return torch.randn(*mean.shape).to(mean) * torch.exp(self.logstd) + mean
+    
+class LidarSingleDataModule(pl.LightningDataModule):
+    def __init__(self, args):
+        super().__init__()
+        self.save_hyperparameters()
+
+        self.data_path = args.data_path
+        self.batch_size = args.batch_size
+        self.max_dim = args.dim
+        self.whiten = args.whiten
+        self.p0_mu = [
+            [-4.5, -4.0, 0.5],
+            [-4.2, -3.5, 0.5],
+            [-4.0, -3.0, 0.5],
+            [-3.75, -2.5, 0.5],
+        ]
+        self.p0_var = 0.02
+        # multiple p1 for each branch
+        #changed
+        self.p1_1_mu = [
+            [-2.5, -0.25, 0.5],
+            [-2.25, 0.675, 0.5],
+            [-2, 1.5, 0.5],
+        ]
+        self.p1_2_mu = [
+            [2, -2, 0.5], 
+            [2.6, -1.25, 0.5], 
+            [3.2, -0.5, 0.5]
+        ]
+        
+        self.p1_var = 0.03
+        self.k = 20
+        self.n_samples = 5000
+        self.num_timesteps = 2
+        self.split_ratios = args.split_ratios
+        
+        self.num_timesteps = 2
+        self.metric_clusters = 3
+        self.args = args
+        self._prepare_data()
+
+    def _prepare_data(self):
+        las = laspy.read(self.data_path)
+        # Extract only "ground" points.
+        self.mask = las.classification == 2
+        # Original Preprocessing
+        x_offset, x_scale = las.header.offsets[0], las.header.scales[0]
+        y_offset, y_scale = las.header.offsets[1], las.header.scales[1]
+        z_offset, z_scale = las.header.offsets[2], las.header.scales[2]
+        dataset = np.vstack(
+            (
+                las.X[self.mask] * x_scale + x_offset,
+                las.Y[self.mask] * y_scale + y_offset,
+                las.Z[self.mask] * z_scale + z_offset,
+            )
+        ).transpose()
+        mi = dataset.min(axis=0, keepdims=True)
+        ma = dataset.max(axis=0, keepdims=True)
+        dataset = (dataset - mi) / (ma - mi) * [10.0, 10.0, 2.0] + [-5.0, -5.0, 0.0]
+
+        self.dataset = torch.tensor(dataset, dtype=torch.float32)
+        self.tree = cKDTree(dataset)
+
+        x0_gaussian = GaussianMM(self.p0_mu, self.p0_var)(self.n_samples)
+        x1_1_gaussian = GaussianMM(self.p1_1_mu, self.p1_var)(self.n_samples)
+        x1_2_gaussian = GaussianMM(self.p1_2_mu, self.p1_var)(self.n_samples)
+
+        x0 = self.get_tangent_proj(x0_gaussian)(x0_gaussian)
+        x1_1 = self.get_tangent_proj(x1_1_gaussian)(x1_1_gaussian)
+        x1_2 = self.get_tangent_proj(x1_2_gaussian)(x1_2_gaussian)
+
+        split_index = int(self.n_samples * self.split_ratios[0])
+
+        self.scaler = StandardScaler()
+        if self.whiten:
+            self.dataset = torch.tensor(
+                self.scaler.fit_transform(dataset), dtype=torch.float32
+            )
+            x0 = torch.tensor(self.scaler.transform(x0), dtype=torch.float32)
+            x1_1 = torch.tensor(self.scaler.transform(x1_1), dtype=torch.float32)
+            x1_2 = torch.tensor(self.scaler.transform(x1_2), dtype=torch.float32)
+            x1 = torch.cat([x1_1, x1_2], dim=0)
+
+        self.coords_t0 = x0
+        self.coords_t1 = x1
+        self.time_labels = np.concatenate([
+            np.zeros(len(self.coords_t0)),    # t=0
+            np.ones(len(self.coords_t1)),     # t=1
+        ])
+        
+        train_x0 = x0[:split_index]
+        val_x0 = x0[split_index:]
+        
+        # branches
+        train_x1 = x1[:split_index]
+        val_x1 = x1[split_index:]
+        
+        self.val_x0 = val_x0
+        
+        train_x0_weights = torch.full((train_x0.shape[0], 1), fill_value=1.0)
+        train_x1_weights = torch.full((train_x1.shape[0], 1), fill_value=1.0)
+        
+        val_x0_weights = torch.full((val_x0.shape[0], 1), fill_value=1.0)
+        val_x1_weights = torch.full((val_x1.shape[0], 1), fill_value=1.0)
+
+        # Adjust split_index to ensure minimum validation samples
+        if self.n_samples - split_index < self.batch_size:
+            split_index = self.n_samples - self.batch_size
+            
+        self.train_dataloaders = {
+            "x0": DataLoader(TensorDataset(train_x0, train_x0_weights), batch_size=self.batch_size, shuffle=True, drop_last=True),
+            "x1": DataLoader(TensorDataset(train_x1, train_x1_weights), batch_size=self.batch_size, shuffle=True, drop_last=True),
+        }
+        
+        self.val_dataloaders = {
+            "x0": DataLoader(TensorDataset(val_x0, val_x0_weights), batch_size=self.batch_size, shuffle=True, drop_last=True),
+            "x1": DataLoader(TensorDataset(val_x1, val_x1_weights), batch_size=self.batch_size, shuffle=True, drop_last=True),
+        }
+        
+        # to edit?
+        self.test_dataloaders = {
+            "x0": DataLoader(TensorDataset(val_x0, val_x0_weights), batch_size=self.val_x0.shape[0], shuffle=True, drop_last=False),
+            "x1": DataLoader(TensorDataset(val_x1, val_x1_weights), batch_size=self.val_x0.shape[0], shuffle=True, drop_last=True),
+            "dataset": DataLoader(TensorDataset(self.dataset), batch_size=self.dataset.shape[0], shuffle=True, drop_last=False),
+        }
+        
+        points = self.dataset.cpu().numpy()
+        x, y = points[:, 0], points[:, 1]
+        # Diagonal-based coordinates (rotated 45°)
+        u = (x + y) / np.sqrt(2)  # along x=y
+        # start region (A) using u
+        u_thresh = np.percentile(u, 30)  # tweak this threshold to control size
+        mask_A = u <= u_thresh
+
+        # among the rest, split by x=y diagonal
+        remaining = ~mask_A
+        mask_B = remaining & (x < y)  # left of diagonal
+        mask_C = remaining & (x >= y)  # right of diagonal
+
+        # Assign dataloaders
+        self.metric_samples_dataloaders = [
+            DataLoader(torch.tensor(points[mask_A], dtype=torch.float32), batch_size=points[mask_A].shape[0], shuffle=False),
+            DataLoader(torch.tensor(points[remaining], dtype=torch.float32), batch_size=points[remaining].shape[0], shuffle=False),
+        ]
+
+    def train_dataloader(self):
+        combined_loaders = {
+            "train_samples": CombinedLoader(self.train_dataloaders, mode="min_size"),
+            "metric_samples": CombinedLoader(
+                self.metric_samples_dataloaders, mode="min_size"
+            ),
+        }
+        return CombinedLoader(combined_loaders, mode="max_size_cycle")
+
+    def val_dataloader(self):
+        combined_loaders = {
+            "val_samples": CombinedLoader(self.val_dataloaders, mode="min_size"),
+            "metric_samples": CombinedLoader(
+                self.metric_samples_dataloaders, mode="min_size"
+            ),
+        }
+
+        return CombinedLoader(combined_loaders, mode="max_size_cycle")
+
+    def test_dataloader(self):
+        combined_loaders = {
+            "test_samples": CombinedLoader(self.test_dataloaders, mode="min_size"),
+            "metric_samples": CombinedLoader(
+                self.metric_samples_dataloaders, mode="min_size"
+            ),
+        }
+
+        return CombinedLoader(combined_loaders, mode="max_size_cycle")
+
+    def get_tangent_proj(self, points):
+        w = self.get_tangent_plane(points)
+        return partial(LidarSingleDataModule.projection_op, w=w)
+
+    def get_tangent_plane(self, points, temp=1e-3):
+        points_np = points.detach().cpu().numpy()
+        _, idx = self.tree.query(points_np, k=self.k)
+        nearest_pts = self.dataset[idx]
+        nearest_pts = torch.tensor(nearest_pts).to(points)
+
+        dists = (points.unsqueeze(1) - nearest_pts).pow(2).sum(-1, keepdim=True)
+        weights = torch.exp(-dists / temp)
+
+        # Fits plane with least vertical distance.
+        w = LidarSingleDataModule.fit_plane(nearest_pts, weights)
+        return w
+
+    @staticmethod
+    def fit_plane(points, weights=None):
+        """Expects points to be of shape (..., 3).
+        Returns [a, b, c] such that the plane is defined as
+            ax + by + c = z
+        """
+        D = torch.cat([points[..., :2], torch.ones_like(points[..., 2:3])], dim=-1)
+        z = points[..., 2]
+        if weights is not None:
+            Dtrans = D.transpose(-1, -2)
+        else:
+            DW = D * weights
+            Dtrans = DW.transpose(-1, -2)
+        w = torch.linalg.solve(
+            torch.matmul(Dtrans, D), torch.matmul(Dtrans, z.unsqueeze(-1))
+        ).squeeze(-1)
+        return w
+
+    @staticmethod
+    def projection_op(x, w):
+        """Projects points to a plane defined by w."""
+        # Normal vector to the tangent plane.
+        n = torch.cat([w[..., :2], -torch.ones_like(w[..., 2:3])], dim=1)
+
+        pn = torch.sum(x * n, dim=-1, keepdim=True)
+        nn = torch.sum(n * n, dim=-1, keepdim=True)
+
+        # Offset.
+        d = w[..., 2:3]
+
+        # Projection of x onto n.
+        projn_x = ((pn + d) / nn) * n
+
+        # Remove component in the normal direction.
+        return x - projn_x
+
+    def get_timepoint_data(self):
+        """Return data organized by timepoints for visualization"""
+        return {
+            't0': self.coords_t0,
+            't1': self.coords_t1, 
+            'time_labels': self.time_labels
+        }
+    
+def get_datamodule():
+    datamodule = LidarSingleDataModule(args)
+    datamodule.setup(stage="fit")   
+    return datamodule

dataloaders/mouse_data.py

@@ -0,0 +1,438 @@
+import torch
+import sys
+sys.argv = ['']
+from sklearn.preprocessing import StandardScaler
+import pytorch_lightning as pl
+from torch.utils.data import DataLoader
+from lightning.pytorch.utilities.combined_loader import CombinedLoader
+import numpy as np
+from scipy.spatial import cKDTree
+import math
+from functools import partial
+from sklearn.cluster import KMeans, DBSCAN
+import matplotlib.pyplot as plt
+import pandas as pd
+from torch.utils.data import TensorDataset
+
+from train.parsers_sc import parse_args
+args = parse_args()
+
+class WeightedBranchedCellDataModule(pl.LightningDataModule):
+    def __init__(self, args):
+        super().__init__()
+        self.save_hyperparameters()
+
+        self.data_path = "./data/mouse_hematopoiesis.csv"
+        self.batch_size = args.batch_size
+        self.max_dim = args.dim
+        self.whiten = args.whiten
+        self.k = 20
+        self.n_samples = 1429 
+        self.num_timesteps = 3  # t=0, t=1, t=2
+        self.split_ratios = args.split_ratios
+        self.metric_clusters = args.metric_clusters
+        self.args = args
+        self._prepare_data()
+        
+
+    def _prepare_data(self):
+        print("Preparing cell data in BranchedCellDataModule")
+        
+        df = pd.read_csv(self.data_path)
+        
+        # Build dictionary of coordinates by time
+        coords_by_t = {
+            t: df[df["samples"] == t][["x1","x2"]].values
+            for t in sorted(df["samples"].unique())
+        }
+        n0 = coords_by_t[0].shape[0]  # Number of T=0 points
+        self.n_samples = n0  # Update n_samples to match actual data if changes
+
+        # Cluster the t=2 cells into two branches
+        km = KMeans(n_clusters=2, random_state=42).fit(coords_by_t[2])
+        df2 = df[df["samples"] == 2].copy()
+        df2["branch"] = km.labels_
+        
+        cluster_counts = df2["branch"].value_counts().sort_index()
+        print(cluster_counts)
+
+        # Sample n0 points from each branch
+        endpoints = {}
+        for b in (0, 1):
+            endpoints[b] = (
+                df2[df2["branch"] == b]
+                .sample(n=n0, random_state=42)[["x1","x2"]]
+                .values
+            )
+            
+        x0 = torch.tensor(coords_by_t[0], dtype=torch.float32) # T=0 coordinates index
+        x_inter = torch.tensor(coords_by_t[1], dtype=torch.float32)
+        x1_1 = torch.tensor(endpoints[0], dtype=torch.float32) # Branch index
+        x1_2 = torch.tensor(endpoints[1], dtype=torch.float32) # Branch index
+
+        self.coords_t0 = x0
+        self.coords_t1 = x_inter
+        self.coords_t2_1 = x1_1
+        self.coords_t2_2 = x1_2
+        self.time_labels = np.concatenate([
+            np.zeros(len(self.coords_t0)),    # t=0
+            np.ones(len(self.coords_t1)),     # t=1
+            np.ones(len(self.coords_t2_1)) * 2,     # t=1
+            np.ones(len(self.coords_t2_2)) * 2,
+        ])
+        
+        split_index = int(n0 * self.split_ratios[0])
+        
+        if n0 - split_index < self.batch_size:
+            split_index = n0 - self.batch_size
+
+        train_x0 = x0[:split_index]
+        val_x0 = x0[split_index:]
+        train_x1_1 = x1_1[:split_index]
+        val_x1_1 = x1_1[split_index:]
+        train_x1_2 = x1_2[:split_index]
+        val_x1_2 = x1_2[split_index:]
+        
+        self.val_x0 = val_x0
+        
+        train_x0_weights = torch.full((train_x0.shape[0], 1), fill_value=1.0)
+        train_x1_1_weights = torch.full((train_x1_1.shape[0], 1), fill_value=0.5)
+        train_x1_2_weights = torch.full((train_x1_2.shape[0], 1), fill_value=0.5)
+        
+        val_x0_weights = torch.full((val_x0.shape[0], 1), fill_value=1.0)
+        val_x1_1_weights = torch.full((val_x1_1.shape[0], 1), fill_value=0.5)
+        val_x1_2_weights = torch.full((val_x1_2.shape[0], 1), fill_value=0.5)
+
+        if self.n_samples - split_index < self.batch_size:
+            split_index = self.n_samples - self.batch_size
+            
+        self.train_dataloaders = {
+            "x0": DataLoader(TensorDataset(train_x0, train_x0_weights), batch_size=self.batch_size, shuffle=True, drop_last=True),
+            "x1_1": DataLoader(TensorDataset(train_x1_1, train_x1_1_weights), batch_size=self.batch_size, shuffle=True, drop_last=True),
+            "x1_2": DataLoader(TensorDataset(train_x1_2, train_x1_2_weights), batch_size=self.batch_size, shuffle=True, drop_last=True),
+        }
+        
+        self.val_dataloaders = {
+            "x0": DataLoader(TensorDataset(val_x0, val_x0_weights), batch_size=self.batch_size, shuffle=False, drop_last=True),
+            "x1_1": DataLoader(TensorDataset(val_x1_1, val_x1_1_weights), batch_size=self.batch_size, shuffle=True, drop_last=True),
+            "x1_2": DataLoader(TensorDataset(val_x1_2, val_x1_2_weights), batch_size=self.batch_size, shuffle=True, drop_last=True),
+        }
+
+        all_data = np.vstack([coords_by_t[t] for t in sorted(coords_by_t.keys())])
+        self.dataset = torch.tensor(all_data, dtype=torch.float32)
+        self.tree = cKDTree(all_data)
+        
+        # if whitening is enabled, need to apply this to the full dataset
+        #if self.whiten:
+            #self.scaler = StandardScaler()
+            #self.dataset = torch.tensor(
+                #self.scaler.fit_transform(all_data), dtype=torch.float32
+            #)
+            
+        self.test_dataloaders = {
+            "x0": DataLoader(TensorDataset(val_x0, val_x0_weights), batch_size=self.val_x0.shape[0], shuffle=False, drop_last=False),
+            "dataset": DataLoader(TensorDataset(self.dataset), batch_size=self.dataset.shape[0], shuffle=False, drop_last=False),
+        }
+        
+        # Metric Dataloader
+        # K-means clustering of ALL points into 2 groups
+        if self.metric_clusters == 3:
+            km_all = KMeans(n_clusters=3, random_state=45).fit(self.dataset.numpy())
+            cluster_labels = km_all.labels_
+            
+            cluster_0_mask = cluster_labels == 0
+            cluster_1_mask = cluster_labels == 1
+            cluster_2_mask = cluster_labels == 2
+            
+            samples = self.dataset.cpu().numpy()
+            
+            cluster_0_data = samples[cluster_0_mask]
+            cluster_1_data = samples[cluster_1_mask]
+            cluster_2_data = samples[cluster_2_mask]
+            
+            self.metric_samples_dataloaders = [
+                DataLoader(
+                    torch.tensor(cluster_1_data, dtype=torch.float32),
+                    batch_size=cluster_1_data.shape[0],
+                    shuffle=False,
+                    drop_last=False,
+                ),
+                DataLoader(
+                    torch.tensor(cluster_2_data, dtype=torch.float32),
+                    batch_size=cluster_2_data.shape[0],
+                    shuffle=False,
+                    drop_last=False,
+                ),
+                
+                DataLoader(
+                    torch.tensor(cluster_0_data, dtype=torch.float32),
+                    batch_size=cluster_0_data.shape[0],
+                    shuffle=False,
+                    drop_last=False,
+                ),
+            ]
+        else:
+            km_all = KMeans(n_clusters=2, random_state=45).fit(self.dataset.numpy())
+            cluster_labels = km_all.labels_
+            
+            cluster_0_mask = cluster_labels == 0
+            cluster_1_mask = cluster_labels == 1
+            
+            samples = self.dataset.cpu().numpy()
+            
+            cluster_0_data = samples[cluster_0_mask]
+            cluster_1_data = samples[cluster_1_mask]
+            
+            self.metric_samples_dataloaders = [
+                DataLoader(
+                    torch.tensor(cluster_1_data, dtype=torch.float32),
+                    batch_size=cluster_1_data.shape[0],
+                    shuffle=False,
+                    drop_last=False,
+                ),
+                DataLoader(
+                    torch.tensor(cluster_0_data, dtype=torch.float32),
+                    batch_size=cluster_0_data.shape[0],
+                    shuffle=False,
+                    drop_last=False,
+                ),
+            ]
+
+
+    def train_dataloader(self):
+        combined_loaders = {
+            "train_samples": CombinedLoader(self.train_dataloaders, mode="min_size"),
+            "metric_samples": CombinedLoader(
+                self.metric_samples_dataloaders, mode="min_size"
+            ),
+        }
+        return CombinedLoader(combined_loaders, mode="max_size_cycle")
+
+    def val_dataloader(self):
+        combined_loaders = {
+            "val_samples": CombinedLoader(self.val_dataloaders, mode="min_size"),
+            "metric_samples": CombinedLoader(
+                self.metric_samples_dataloaders, mode="min_size"
+            ),
+        }
+
+        return CombinedLoader(combined_loaders, mode="max_size_cycle")
+
+    def test_dataloader(self):
+        combined_loaders = {
+            "test_samples": CombinedLoader(self.test_dataloaders, mode="min_size"),
+            "metric_samples": CombinedLoader(
+                self.metric_samples_dataloaders, mode="min_size"
+            ),
+        }
+
+        return CombinedLoader(combined_loaders, mode="max_size_cycle")
+
+    def get_manifold_proj(self, points):
+        """Adapted for 2D cell data - uses local neighborhood averaging instead of plane fitting"""
+        return partial(self.local_smoothing_op, tree=self.tree, dataset=self.dataset)
+
+    @staticmethod
+    def local_smoothing_op(x, tree, dataset, k=10, temp=1e-3):
+        """
+        Apply local smoothing based on k-nearest neighbors in the full dataset
+        This replaces the plane projection for 2D manifold regularization
+        """
+        points_np = x.detach().cpu().numpy()
+        _, idx = tree.query(points_np, k=k)
+        nearest_pts = dataset[idx]  # Shape: (batch_size, k, 2)
+        
+        # Compute weighted average of neighbors
+        dists = (x.unsqueeze(1) - nearest_pts).pow(2).sum(-1, keepdim=True)
+        weights = torch.exp(-dists / temp)
+        weights = weights / weights.sum(dim=1, keepdim=True)
+        
+        # Weighted average of neighbors
+        smoothed = (weights * nearest_pts).sum(dim=1)
+        
+        # Blend original point with smoothed version
+        alpha = 0.3  # How much smoothing to apply
+        return (1 - alpha) * x + alpha * smoothed
+    
+    def get_timepoint_data(self):
+        """Return data organized by timepoints for visualization"""
+        return {
+            't0': self.coords_t0,
+            't1': self.coords_t1, 
+            't2_1': self.coords_t2_1, 
+            't2_2': self.coords_t2_2, 
+            'time_labels': self.time_labels
+        }
+
+
+
+class SingleBranchCellDataModule(pl.LightningDataModule):
+    def __init__(self, args):
+        super().__init__()
+        self.save_hyperparameters()
+
+        self.data_path = "./data/mouse_hematopoiesis.csv"
+        self.batch_size = args.batch_size
+        self.max_dim = args.dim
+        self.whiten = args.whiten
+        self.k = 20
+        self.n_samples = 1429 
+        self.num_timesteps = 3  # t=0, t=1, t=2
+        self.split_ratios = args.split_ratios
+        self.metric_clusters = 3
+        self.args = args
+        self._prepare_data()
+        
+
+    def _prepare_data(self):
+        print("Preparing cell data in BranchedCellDataModule")
+        
+        df = pd.read_csv(self.data_path)
+        
+        # Build dictionary of coordinates by time
+        coords_by_t = {
+            t: df[df["samples"] == t][["x1","x2"]].values
+            for t in sorted(df["samples"].unique())
+        }
+        n0 = coords_by_t[0].shape[0]  # Number of T=0 points
+        self.n_samples = n0  # Update n_samples to match actual data if changes
+
+        x0 = torch.tensor(coords_by_t[0], dtype=torch.float32) # T=0 coordinates index
+        x_inter = torch.tensor(coords_by_t[1], dtype=torch.float32)
+        x1 = torch.tensor(coords_by_t[2], dtype=torch.float32) # Branch index
+
+        split_index = int(n0 * self.split_ratios[0])
+        
+        if n0 - split_index < self.batch_size:
+            split_index = n0 - self.batch_size
+
+        train_x0 = x0[:split_index]
+        val_x0 = x0[split_index:]
+        train_x1 = x1[:split_index]
+        val_x1 = x1[split_index:]
+        
+        self.val_x0 = val_x0
+        
+        train_x0_weights = torch.full((train_x0.shape[0], 1), fill_value=1.0)
+        train_x1_weights = torch.full((train_x1.shape[0], 1), fill_value=0.5)
+        
+        val_x0_weights = torch.full((val_x0.shape[0], 1), fill_value=1.0)
+        val_x1_weights = torch.full((val_x1.shape[0], 1), fill_value=0.5)
+
+        if self.n_samples - split_index < self.batch_size:
+            split_index = self.n_samples - self.batch_size
+            
+        self.train_dataloaders = {
+            "x0": DataLoader(TensorDataset(train_x0, train_x0_weights), batch_size=self.batch_size, shuffle=True, drop_last=True),
+            "x1": DataLoader(TensorDataset(train_x1, train_x1_weights), batch_size=self.batch_size, shuffle=True, drop_last=True),
+        }
+        
+        self.val_dataloaders = {
+            "x0": DataLoader(TensorDataset(val_x0, val_x0_weights), batch_size=self.batch_size, shuffle=False, drop_last=True),
+            "x1": DataLoader(TensorDataset(val_x1, val_x1_weights), batch_size=self.batch_size, shuffle=True, drop_last=True),
+        }
+
+        all_data = np.vstack([coords_by_t[t] for t in sorted(coords_by_t.keys())])
+        self.dataset = torch.tensor(all_data, dtype=torch.float32)
+        self.tree = cKDTree(all_data)
+        
+        # if whitening is enabled, need to apply this to the full dataset
+        if self.whiten:
+            self.scaler = StandardScaler()
+            self.dataset = torch.tensor(
+                self.scaler.fit_transform(all_data), dtype=torch.float32
+            )
+            
+        self.test_dataloaders = {
+            "x0": DataLoader(TensorDataset(val_x0, val_x0_weights), batch_size=self.val_x0.shape[0], shuffle=False, drop_last=False),
+            "dataset": DataLoader(TensorDataset(self.dataset), batch_size=self.dataset.shape[0], shuffle=False, drop_last=False),
+        }
+        
+        # Metric Dataloader
+        # K-means clustering of ALL points into 2 groups
+        km_all = KMeans(n_clusters=2, random_state=45).fit(self.dataset.numpy())
+        cluster_labels = km_all.labels_
+        
+        cluster_0_mask = cluster_labels == 0
+        cluster_1_mask = cluster_labels == 1
+        
+        samples = self.dataset.cpu().numpy()
+        
+        cluster_0_data = samples[cluster_0_mask]
+        cluster_1_data = samples[cluster_1_mask]
+        
+        self.metric_samples_dataloaders = [
+            DataLoader(
+                torch.tensor(cluster_1_data, dtype=torch.float32),
+                batch_size=cluster_1_data.shape[0],
+                shuffle=False,
+                drop_last=False,
+            ),
+            DataLoader(
+                torch.tensor(cluster_0_data, dtype=torch.float32),
+                batch_size=cluster_0_data.shape[0],
+                shuffle=False,
+                drop_last=False,
+            ),
+        ]
+
+
+    def train_dataloader(self):
+        combined_loaders = {
+            "train_samples": CombinedLoader(self.train_dataloaders, mode="min_size"),
+            "metric_samples": CombinedLoader(
+                self.metric_samples_dataloaders, mode="min_size"
+            ),
+        }
+        return CombinedLoader(combined_loaders, mode="max_size_cycle")
+
+    def val_dataloader(self):
+        combined_loaders = {
+            "val_samples": CombinedLoader(self.val_dataloaders, mode="min_size"),
+            "metric_samples": CombinedLoader(
+                self.metric_samples_dataloaders, mode="min_size"
+            ),
+        }
+
+        return CombinedLoader(combined_loaders, mode="max_size_cycle")
+
+    def test_dataloader(self):
+        combined_loaders = {
+            "test_samples": CombinedLoader(self.test_dataloaders, mode="min_size"),
+            "metric_samples": CombinedLoader(
+                self.metric_samples_dataloaders, mode="min_size"
+            ),
+        }
+
+        return CombinedLoader(combined_loaders, mode="max_size_cycle")
+
+    def get_manifold_proj(self, points):
+        """Adapted for 2D cell data - uses local neighborhood averaging instead of plane fitting"""
+        return partial(self.local_smoothing_op, tree=self.tree, dataset=self.dataset)
+
+    @staticmethod
+    def local_smoothing_op(x, tree, dataset, k=10, temp=1e-3):
+        """
+        Apply local smoothing based on k-nearest neighbors in the full dataset
+        This replaces the plane projection for 2D manifold regularization
+        """
+        points_np = x.detach().cpu().numpy()
+        _, idx = tree.query(points_np, k=k)
+        nearest_pts = dataset[idx]  # Shape: (batch_size, k, 2)
+        
+        # Compute weighted average of neighbors
+        dists = (x.unsqueeze(1) - nearest_pts).pow(2).sum(-1, keepdim=True)
+        weights = torch.exp(-dists / temp)
+        weights = weights / weights.sum(dim=1, keepdim=True)
+        
+        # Weighted average of neighbors
+        smoothed = (weights * nearest_pts).sum(dim=1)
+        
+        # Blend original point with smoothed version
+        alpha = 0.3  # How much smoothing to apply
+        return (1 - alpha) * x + alpha * smoothed
+
+def get_datamodule():
+    datamodule = WeightedBranchedCellDataModule(args)
+    datamodule.setup(stage="fit")   
+    return datamodule

dataloaders/three_branch_data.py

@@ -0,0 +1,310 @@
+import torch
+import sys
+sys.argv = ['']
+from sklearn.preprocessing import StandardScaler
+import pytorch_lightning as pl
+from torch.utils.data import DataLoader
+from lightning.pytorch.utilities.combined_loader import CombinedLoader
+import pandas as pd
+import numpy as np
+from functools import partial
+from scipy.spatial import cKDTree
+from sklearn.cluster import KMeans
+from torch.utils.data import TensorDataset
+
+from train.parsers_tahoe import parse_args
+args = parse_args()
+
+class ThreeBranchTahoeDataModule(pl.LightningDataModule):
+    def __init__(self, args):
+        super().__init__()
+        self.save_hyperparameters()
+
+        self.batch_size = args.batch_size
+        self.max_dim = args.dim
+        self.whiten = args.whiten
+        self.split_ratios = args.split_ratios
+        self.num_timesteps = 2
+        self.data_path = "./data/Trametinib_5.0uM_pca_and_leidenumap_labels.csv"
+        self.args = args
+
+        self._prepare_data()
+
+    def _prepare_data(self):
+        df = pd.read_csv(self.data_path, comment='#')
+        df = df.iloc[:, 1:]
+        df = df.replace('', np.nan)
+        pc_cols = df.columns[:50]
+        for col in pc_cols:
+            df[col] = pd.to_numeric(df[col], errors='coerce')
+        leiden_dmso_col = 'leiden_DMSO_TF_0.0uM'
+        leiden_clonidine_col = 'leiden_Trametinib_5.0uM'
+
+        dmso_mask = df[leiden_dmso_col].notna()        # Has leiden value in DMSO column
+        clonidine_mask = df[leiden_clonidine_col].notna()  # Has leiden value in Clonidine column
+        
+        dmso_data = df[dmso_mask].copy()
+        clonidine_data = df[clonidine_mask].copy()
+        
+        # Updated to include all three clusters: 0, 4, and 6
+        top_clonidine_clusters = ['1.0', '3.0', '5.0']
+        
+        x1_1_data = clonidine_data[clonidine_data[leiden_clonidine_col].astype(str) == top_clonidine_clusters[0]]
+        x1_2_data = clonidine_data[clonidine_data[leiden_clonidine_col].astype(str) == top_clonidine_clusters[1]]
+        x1_3_data = clonidine_data[clonidine_data[leiden_clonidine_col].astype(str) == top_clonidine_clusters[2]]
+        
+        x1_1_coords = x1_1_data[pc_cols].values 
+        x1_2_coords = x1_2_data[pc_cols].values 
+        x1_3_coords = x1_3_data[pc_cols].values
+        
+        x1_1_coords = x1_1_coords.astype(float)
+        x1_2_coords = x1_2_coords.astype(float)
+        x1_3_coords = x1_3_coords.astype(float)
+
+        # Target size is now the minimum across all three endpoint clusters
+        target_size = min(len(x1_1_coords), len(x1_2_coords), len(x1_3_coords))
+        
+        # Helper function to select points closest to centroid
+        def select_closest_to_centroid(coords, target_size):
+            if len(coords) <= target_size:
+                return coords
+            
+            # Calculate centroid
+            centroid = np.mean(coords, axis=0)
+            
+            # Calculate distances to centroid
+            distances = np.linalg.norm(coords - centroid, axis=1)
+            
+            # Get indices of closest points
+            closest_indices = np.argsort(distances)[:target_size]
+            
+            return coords[closest_indices]
+        
+        # Sample all endpoint clusters to target size using centroid-based selection
+        x1_1_coords = select_closest_to_centroid(x1_1_coords, target_size)
+        x1_2_coords = select_closest_to_centroid(x1_2_coords, target_size)
+        x1_3_coords = select_closest_to_centroid(x1_3_coords, target_size)
+        
+        dmso_cluster_counts = dmso_data[leiden_dmso_col].value_counts()
+        
+        # DMSO (unchanged)
+        largest_dmso_cluster = dmso_cluster_counts.index[0]
+        dmso_cluster_data = dmso_data[dmso_data[leiden_dmso_col] == largest_dmso_cluster]
+        
+        dmso_coords = dmso_cluster_data[pc_cols].values
+        
+        # Random sampling from largest DMSO cluster to match target size
+        # For DMSO, we'll also use centroid-based selection for consistency
+        if len(dmso_coords) >= target_size:
+            x0_coords = select_closest_to_centroid(dmso_coords, target_size)
+        else:
+            # If largest cluster is smaller than target, use all of it and pad with other DMSO cells
+            remaining_needed = target_size - len(dmso_coords)
+            other_dmso_data = dmso_data[dmso_data[leiden_dmso_col] != largest_dmso_cluster]
+            other_dmso_coords = other_dmso_data[pc_cols].values
+            
+            if len(other_dmso_coords) >= remaining_needed:
+                # Select closest to centroid from other DMSO cells
+                other_selected = select_closest_to_centroid(other_dmso_coords, remaining_needed)
+                x0_coords = np.vstack([dmso_coords, other_selected])
+            else:
+                # Use all available DMSO cells and reduce target size
+                all_dmso_coords = dmso_data[pc_cols].values
+                target_size = min(target_size, len(all_dmso_coords))
+                x0_coords = select_closest_to_centroid(all_dmso_coords, target_size)
+                
+                # Re-select endpoint clusters with updated target size
+                x1_1_coords = select_closest_to_centroid(x1_1_data[pc_cols].values.astype(float), target_size)
+                x1_2_coords = select_closest_to_centroid(x1_2_data[pc_cols].values.astype(float), target_size)
+                x1_3_coords = select_closest_to_centroid(x1_3_data[pc_cols].values.astype(float), target_size)
+        
+        # No need to resample since we already selected the right number
+        # The endpoint clusters are already at target_size from centroid-based selection
+        
+        self.n_samples = target_size
+        
+        # for plotting
+        self.coords_t0 = torch.tensor(x0_coords, dtype=torch.float32)
+        self.coords_t1_1 = torch.tensor(x1_1_coords, dtype=torch.float32)
+        self.coords_t1_2 = torch.tensor(x1_2_coords, dtype=torch.float32)
+        self.coords_t1_3 = torch.tensor(x1_3_coords, dtype=torch.float32)
+        
+        self.time_labels = np.concatenate([
+            np.zeros(len(self.coords_t0)),    # t=0
+            np.ones(len(self.coords_t1_1)),     # t=1
+            np.ones(len(self.coords_t1_2)),     # t=1
+            np.ones(len(self.coords_t1_3)),     # t=1
+        ])
+        
+        x0 = torch.tensor(x0_coords, dtype=torch.float32)
+        x1_1 = torch.tensor(x1_1_coords, dtype=torch.float32)
+        x1_2 = torch.tensor(x1_2_coords, dtype=torch.float32)
+        x1_3 = torch.tensor(x1_3_coords, dtype=torch.float32)
+        
+        split_index = int(target_size * self.split_ratios[0])
+        
+        if target_size - split_index < self.batch_size:
+            split_index = target_size - self.batch_size
+        
+        train_x0 = x0[:split_index]
+        val_x0 = x0[split_index:]
+        train_x1_1 = x1_1[:split_index]
+        val_x1_1 = x1_1[split_index:]
+        train_x1_2 = x1_2[:split_index]
+        val_x1_2 = x1_2[split_index:]
+        train_x1_3 = x1_3[:split_index]
+        val_x1_3 = x1_3[split_index:]
+        
+        self.val_x0 = val_x0
+        
+        train_x0_weights = torch.full((train_x0.shape[0], 1), fill_value=1.0)
+        train_x1_1_weights = torch.full((train_x1_1.shape[0], 1), fill_value=0.603)
+        train_x1_2_weights = torch.full((train_x1_2.shape[0], 1), fill_value=0.255)
+        train_x1_3_weights = torch.full((train_x1_3.shape[0], 1), fill_value=0.142)
+        
+        val_x0_weights = torch.full((val_x0.shape[0], 1), fill_value=1.0)
+        val_x1_1_weights = torch.full((val_x1_1.shape[0], 1), fill_value=0.603)
+        val_x1_2_weights = torch.full((val_x1_2.shape[0], 1), fill_value=0.255)
+        val_x1_3_weights = torch.full((val_x1_3.shape[0], 1), fill_value=0.142)
+        
+        # Updated train dataloaders to include x1_3
+        self.train_dataloaders = {
+            "x0": DataLoader(TensorDataset(train_x0, train_x0_weights), batch_size=self.batch_size, shuffle=True, drop_last=True),
+            "x1_1": DataLoader(TensorDataset(train_x1_1, train_x1_1_weights), batch_size=self.batch_size, shuffle=True, drop_last=True),
+            "x1_2": DataLoader(TensorDataset(train_x1_2, train_x1_2_weights), batch_size=self.batch_size, shuffle=True, drop_last=True),
+            "x1_3": DataLoader(TensorDataset(train_x1_3, train_x1_3_weights), batch_size=self.batch_size, shuffle=True, drop_last=True),
+        }
+        
+        # Updated val dataloaders to include x1_3
+        self.val_dataloaders = {
+            "x0": DataLoader(TensorDataset(val_x0, val_x0_weights), batch_size=self.batch_size, shuffle=False, drop_last=True),
+            "x1_1": DataLoader(TensorDataset(val_x1_1, val_x1_1_weights), batch_size=self.batch_size, shuffle=True, drop_last=True),
+            "x1_2": DataLoader(TensorDataset(val_x1_2, val_x1_2_weights), batch_size=self.batch_size, shuffle=True, drop_last=True),
+            "x1_3": DataLoader(TensorDataset(val_x1_3, val_x1_3_weights), batch_size=self.batch_size, shuffle=True, drop_last=True),
+        }
+        
+        all_coords = df[pc_cols].dropna().values.astype(float)
+        self.dataset = torch.tensor(all_coords, dtype=torch.float32)
+        self.tree = cKDTree(all_coords)
+
+        self.test_dataloaders = {
+            "x0": DataLoader(TensorDataset(val_x0, val_x0_weights), batch_size=self.val_x0.shape[0], shuffle=False, drop_last=False),
+            "dataset": DataLoader(TensorDataset(self.dataset), batch_size=self.dataset.shape[0], shuffle=False, drop_last=False),
+        }
+
+        # Updated metric samples - now using 4 clusters instead of 3
+        #km_all = KMeans(n_clusters=4, random_state=42).fit(self.dataset.numpy())
+        km_all = KMeans(n_clusters=4, random_state=0).fit(self.dataset[:, :3].numpy())
+
+        cluster_labels = km_all.labels_
+        
+        cluster_0_mask = cluster_labels == 0
+        cluster_1_mask = cluster_labels == 1
+        cluster_2_mask = cluster_labels == 2
+        cluster_3_mask = cluster_labels == 3
+        
+        samples = self.dataset.cpu().numpy()
+        
+        cluster_0_data = samples[cluster_0_mask]
+        cluster_1_data = samples[cluster_1_mask]
+        cluster_2_data = samples[cluster_2_mask]
+        cluster_3_data = samples[cluster_3_mask]
+                
+        self.metric_samples_dataloaders = [
+            DataLoader(
+                torch.tensor(cluster_1_data, dtype=torch.float32),
+                batch_size=cluster_1_data.shape[0],
+                shuffle=False,
+                drop_last=False,
+            ),
+            DataLoader(
+                torch.tensor(cluster_3_data, dtype=torch.float32),
+                batch_size=cluster_3_data.shape[0],
+                shuffle=False,
+                drop_last=False,
+            ),
+            
+            
+            DataLoader(
+                torch.tensor(cluster_2_data, dtype=torch.float32),
+                batch_size=cluster_2_data.shape[0],
+                shuffle=False,
+                drop_last=False,
+            ),
+            DataLoader(
+                torch.tensor(cluster_0_data, dtype=torch.float32),
+                batch_size=cluster_0_data.shape[0],
+                shuffle=False,
+                drop_last=False,
+            ),
+        ]
+        
+    def train_dataloader(self):
+        combined_loaders = {
+            "train_samples": CombinedLoader(self.train_dataloaders, mode="min_size"),
+            "metric_samples": CombinedLoader(
+                self.metric_samples_dataloaders, mode="min_size"
+            ),
+        }
+        return CombinedLoader(combined_loaders, mode="max_size_cycle")
+
+    def val_dataloader(self):
+        combined_loaders = {
+            "val_samples": CombinedLoader(self.val_dataloaders, mode="min_size"),
+            "metric_samples": CombinedLoader(
+                self.metric_samples_dataloaders, mode="min_size"
+            ),
+        }
+
+        return CombinedLoader(combined_loaders, mode="max_size_cycle")
+    
+    def test_dataloader(self):
+        combined_loaders = {
+            "test_samples": CombinedLoader(self.test_dataloaders, mode="min_size"),
+            "metric_samples": CombinedLoader(
+                self.metric_samples_dataloaders, mode="min_size"
+            ),
+        }
+
+        return CombinedLoader(combined_loaders, mode="max_size_cycle")
+
+    def get_manifold_proj(self, points):
+        """Adapted for 2D cell data - uses local neighborhood averaging instead of plane fitting"""
+        return partial(self.local_smoothing_op, tree=self.tree, dataset=self.dataset)
+
+    @staticmethod
+    def local_smoothing_op(x, tree, dataset, k=10, temp=1e-3):
+        """
+        Apply local smoothing based on k-nearest neighbors in the full dataset
+        This replaces the plane projection for 2D manifold regularization
+        """
+        points_np = x.detach().cpu().numpy()
+        _, idx = tree.query(points_np, k=k)
+        nearest_pts = dataset[idx]  # Shape: (batch_size, k, 2)
+        
+        # Compute weighted average of neighbors
+        dists = (x.unsqueeze(1) - nearest_pts).pow(2).sum(-1, keepdim=True)
+        weights = torch.exp(-dists / temp)
+        weights = weights / weights.sum(dim=1, keepdim=True)
+        
+        # Weighted average of neighbors
+        smoothed = (weights * nearest_pts).sum(dim=1)
+        
+        # Blend original point with smoothed version
+        alpha = 0.3  # How much smoothing to apply
+        return (1 - alpha) * x + alpha * smoothed
+    
+    def get_timepoint_data(self):
+        """Return data organized by timepoints for visualization"""
+        return {
+            't0': self.coords_t0,
+            't1_1': self.coords_t1_1, 
+            't1_2': self.coords_t1_2, 
+            't1_3': self.coords_t1_3, 
+            'time_labels': self.time_labels
+        }
+    
+def get_datamodule():
+    datamodule = ThreeBranchTahoeDataModule(args)
+    datamodule.setup(stage="fit")   
+    return datamodule    

dataloaders/trametinib_single.py

@@ -0,0 +1,279 @@
+import torch
+import sys
+sys.argv = ['']
+from sklearn.preprocessing import StandardScaler
+import pytorch_lightning as pl
+from torch.utils.data import DataLoader
+from lightning.pytorch.utilities.combined_loader import CombinedLoader
+import pandas as pd
+import numpy as np
+from functools import partial
+from scipy.spatial import cKDTree
+from sklearn.cluster import KMeans
+from torch.utils.data import TensorDataset
+
+from train.parsers_tahoe import parse_args
+args = parse_args()
+
+class TrametinibSingleBranchDataModule(pl.LightningDataModule):
+    def __init__(self, args):
+        super().__init__()
+        self.save_hyperparameters()
+
+        self.batch_size = args.batch_size
+        self.max_dim = args.dim
+        self.whiten = args.whiten
+        self.split_ratios = args.split_ratios
+        self.num_timesteps = 2
+        self.data_path = "./data/Trametinib_5.0uM_pca_and_leidenumap_labels.csv"
+        self.args = args
+
+        self._prepare_data()
+
+    def _prepare_data(self):
+        df = pd.read_csv(self.data_path, comment='#')
+        df = df.iloc[:, 1:]
+        df = df.replace('', np.nan)
+        pc_cols = df.columns[:50]
+        for col in pc_cols:
+            df[col] = pd.to_numeric(df[col], errors='coerce')
+        leiden_dmso_col = 'leiden_DMSO_TF_0.0uM'
+        leiden_clonidine_col = 'leiden_Trametinib_5.0uM'
+
+        dmso_mask = df[leiden_dmso_col].notna()        # Has leiden value in DMSO column
+        clonidine_mask = df[leiden_clonidine_col].notna()  # Has leiden value in Clonidine column
+        
+        dmso_data = df[dmso_mask].copy()
+        clonidine_data = df[clonidine_mask].copy()
+        
+        # Updated to include all three clusters: 0, 4, and 6
+        top_clonidine_clusters = ['1.0', '3.0', '5.0']
+        
+        x1_1_data = clonidine_data[clonidine_data[leiden_clonidine_col].astype(str) == top_clonidine_clusters[0]]
+        x1_2_data = clonidine_data[clonidine_data[leiden_clonidine_col].astype(str) == top_clonidine_clusters[1]]
+        x1_3_data = clonidine_data[clonidine_data[leiden_clonidine_col].astype(str) == top_clonidine_clusters[2]]
+        
+        x1_1_coords = x1_1_data[pc_cols].values 
+        x1_2_coords = x1_2_data[pc_cols].values 
+        x1_3_coords = x1_3_data[pc_cols].values
+        
+        x1_1_coords = x1_1_coords.astype(float)
+        x1_2_coords = x1_2_coords.astype(float)
+        x1_3_coords = x1_3_coords.astype(float)
+
+        # Target size is now the minimum across all three endpoint clusters
+        target_size = min(len(x1_1_coords), len(x1_2_coords), len(x1_3_coords))
+        
+        # Helper function to select points closest to centroid
+        def select_closest_to_centroid(coords, target_size):
+            if len(coords) <= target_size:
+                return coords
+            
+            # Calculate centroid
+            centroid = np.mean(coords, axis=0)
+            
+            # Calculate distances to centroid
+            distances = np.linalg.norm(coords - centroid, axis=1)
+            
+            # Get indices of closest points
+            closest_indices = np.argsort(distances)[:target_size]
+            
+            return coords[closest_indices]
+        
+        # Sample all endpoint clusters to target size using centroid-based selection
+        x1_1_coords = select_closest_to_centroid(x1_1_coords, target_size)
+        x1_2_coords = select_closest_to_centroid(x1_2_coords, target_size)
+        x1_3_coords = select_closest_to_centroid(x1_3_coords, target_size)
+        
+        dmso_cluster_counts = dmso_data[leiden_dmso_col].value_counts()
+        
+        # DMSO (unchanged)
+        largest_dmso_cluster = dmso_cluster_counts.index[0]
+        dmso_cluster_data = dmso_data[dmso_data[leiden_dmso_col] == largest_dmso_cluster]
+        
+        dmso_coords = dmso_cluster_data[pc_cols].values
+        
+        # Random sampling from largest DMSO cluster to match target size
+        # For DMSO, we'll also use centroid-based selection for consistency
+        if len(dmso_coords) >= target_size:
+            x0_coords = select_closest_to_centroid(dmso_coords, target_size)
+        else:
+            # If largest cluster is smaller than target, use all of it and pad with other DMSO cells
+            remaining_needed = target_size - len(dmso_coords)
+            other_dmso_data = dmso_data[dmso_data[leiden_dmso_col] != largest_dmso_cluster]
+            other_dmso_coords = other_dmso_data[pc_cols].values
+            
+            if len(other_dmso_coords) >= remaining_needed:
+                # Select closest to centroid from other DMSO cells
+                other_selected = select_closest_to_centroid(other_dmso_coords, remaining_needed)
+                x0_coords = np.vstack([dmso_coords, other_selected])
+            else:
+                # Use all available DMSO cells and reduce target size
+                all_dmso_coords = dmso_data[pc_cols].values
+                target_size = min(target_size, len(all_dmso_coords))
+                x0_coords = select_closest_to_centroid(all_dmso_coords, target_size)
+                
+                # Re-select endpoint clusters with updated target size
+                x1_1_coords = select_closest_to_centroid(x1_1_data[pc_cols].values.astype(float), target_size)
+                x1_2_coords = select_closest_to_centroid(x1_2_data[pc_cols].values.astype(float), target_size)
+                x1_3_coords = select_closest_to_centroid(x1_3_data[pc_cols].values.astype(float), target_size)
+        
+        # No need to resample since we already selected the right number
+        # The endpoint clusters are already at target_size from centroid-based selection
+        
+        self.n_samples = target_size
+        
+        # for plotting
+
+        
+        x0 = torch.tensor(x0_coords, dtype=torch.float32)
+        x1_1 = torch.tensor(x1_1_coords, dtype=torch.float32)
+        x1_2 = torch.tensor(x1_2_coords, dtype=torch.float32)
+        x1_3 = torch.tensor(x1_3_coords, dtype=torch.float32)
+        x1 = torch.cat([x1_1, x1_2, x1_3], dim=0)
+        
+        self.coords_t0 = x0
+        self.coords_t1 = x1
+        
+        self.time_labels = np.concatenate([
+            np.zeros(len(self.coords_t0)),    # t=0
+            np.ones(len(self.coords_t1)),     # t=1
+        ])
+        
+        split_index = int(target_size * self.split_ratios[0])
+        
+        if target_size - split_index < self.batch_size:
+            split_index = target_size - self.batch_size
+        
+        train_x0 = x0[:split_index]
+        val_x0 = x0[split_index:]
+        train_x1 = x1_1[:split_index]
+        val_x1 = x1_1[split_index:]
+        
+        self.val_x0 = val_x0
+        
+        train_x0_weights = torch.full((train_x0.shape[0], 1), fill_value=1.0)
+        train_x1_weights = torch.full((train_x1.shape[0], 1), fill_value=1.0)
+        
+        val_x0_weights = torch.full((val_x0.shape[0], 1), fill_value=1.0)
+        val_x1_weights = torch.full((val_x1.shape[0], 1), fill_value=1.0)
+        
+        # Updated train dataloaders to include x1_3
+        self.train_dataloaders = {
+            "x0": DataLoader(TensorDataset(train_x0, train_x0_weights), batch_size=self.batch_size, shuffle=True, drop_last=True),
+            "x1": DataLoader(TensorDataset(train_x1, train_x1_weights), batch_size=self.batch_size, shuffle=True, drop_last=True),
+        }
+        
+        # Updated val dataloaders to include x1_3
+        self.val_dataloaders = {
+            "x0": DataLoader(TensorDataset(val_x0, val_x0_weights), batch_size=self.batch_size, shuffle=False, drop_last=True),
+            "x1": DataLoader(TensorDataset(val_x1, val_x1_weights), batch_size=self.batch_size, shuffle=True, drop_last=True),
+        }
+        
+        all_coords = df[pc_cols].dropna().values.astype(float)
+        self.dataset = torch.tensor(all_coords, dtype=torch.float32)
+        self.tree = cKDTree(all_coords)
+
+        self.test_dataloaders = {
+            "x0": DataLoader(TensorDataset(val_x0, val_x0_weights), batch_size=self.val_x0.shape[0], shuffle=False, drop_last=False),
+            "dataset": DataLoader(TensorDataset(self.dataset), batch_size=self.dataset.shape[0], shuffle=False, drop_last=False),
+        }
+
+        # Updated metric samples - now using 4 clusters instead of 3
+        #km_all = KMeans(n_clusters=4, random_state=42).fit(self.dataset.numpy())
+        km_all = KMeans(n_clusters=2, random_state=0).fit(self.dataset[:, :3].numpy())
+
+        cluster_labels = km_all.labels_
+        
+        cluster_0_mask = cluster_labels == 0
+        cluster_1_mask = cluster_labels == 1
+        
+        samples = self.dataset.cpu().numpy()
+        
+        cluster_0_data = samples[cluster_0_mask]
+        cluster_1_data = samples[cluster_1_mask]
+                
+        self.metric_samples_dataloaders = [
+            DataLoader(
+                torch.tensor(cluster_1_data, dtype=torch.float32),
+                batch_size=cluster_1_data.shape[0],
+                shuffle=False,
+                drop_last=False,
+            ),
+            DataLoader(
+                torch.tensor(cluster_0_data, dtype=torch.float32),
+                batch_size=cluster_0_data.shape[0],
+                shuffle=False,
+                drop_last=False,
+            ),
+        ]
+        
+    def train_dataloader(self):
+        combined_loaders = {
+            "train_samples": CombinedLoader(self.train_dataloaders, mode="min_size"),
+            "metric_samples": CombinedLoader(
+                self.metric_samples_dataloaders, mode="min_size"
+            ),
+        }
+        return CombinedLoader(combined_loaders, mode="max_size_cycle")
+
+    def val_dataloader(self):
+        combined_loaders = {
+            "val_samples": CombinedLoader(self.val_dataloaders, mode="min_size"),
+            "metric_samples": CombinedLoader(
+                self.metric_samples_dataloaders, mode="min_size"
+            ),
+        }
+
+        return CombinedLoader(combined_loaders, mode="max_size_cycle")
+
+    
+    
+    def test_dataloader(self):
+        combined_loaders = {
+            "test_samples": CombinedLoader(self.test_dataloaders, mode="min_size"),
+            "metric_samples": CombinedLoader(
+                self.metric_samples_dataloaders, mode="min_size"
+            ),
+        }
+
+        return CombinedLoader(combined_loaders, mode="max_size_cycle")
+
+    def get_manifold_proj(self, points):
+        """Adapted for 2D cell data - uses local neighborhood averaging instead of plane fitting"""
+        return partial(self.local_smoothing_op, tree=self.tree, dataset=self.dataset)
+
+    @staticmethod
+    def local_smoothing_op(x, tree, dataset, k=10, temp=1e-3):
+        """
+        Apply local smoothing based on k-nearest neighbors in the full dataset
+        This replaces the plane projection for 2D manifold regularization
+        """
+        points_np = x.detach().cpu().numpy()
+        _, idx = tree.query(points_np, k=k)
+        nearest_pts = dataset[idx]  # Shape: (batch_size, k, 2)
+        
+        # Compute weighted average of neighbors
+        dists = (x.unsqueeze(1) - nearest_pts).pow(2).sum(-1, keepdim=True)
+        weights = torch.exp(-dists / temp)
+        weights = weights / weights.sum(dim=1, keepdim=True)
+        
+        # Weighted average of neighbors
+        smoothed = (weights * nearest_pts).sum(dim=1)
+        
+        # Blend original point with smoothed version
+        alpha = 0.3  # How much smoothing to apply
+        return (1 - alpha) * x + alpha * smoothed
+    
+    def get_timepoint_data(self):
+        """Return data organized by timepoints for visualization"""
+        return {
+            't0': self.coords_t0,
+            't1': self.coords_t1,
+            'time_labels': self.time_labels
+        }
+    
+def get_datamodule():
+    datamodule = TrametinibSingleBranchDataModule(args)
+    datamodule.setup(stage="fit")   
+    return datamodule    

losses/energy_loss.py

@@ -0,0 +1,73 @@
+import os, math, numpy as np
+import torch
+import torch.nn as nn
+from torchdiffeq import odeint as odeint2
+from torchmetrics.functional import mean_squared_error
+import ot
+    
+class EnergySolver(nn.Module):
+    def __init__(self, flow_net, growth_net, state_cost, data_manifold_metric=None, samples=None):
+        super(EnergySolver, self).__init__()
+        self.flow_net = flow_net
+        self.growth_net = growth_net
+        self.state_cost = state_cost
+        
+        self.data_manifold_metric = data_manifold_metric
+        self.samples = samples
+
+    def forward(self, t, state):
+        xt, wt, mt = state
+        
+        xt.requires_grad_(True)
+        wt.requires_grad_(True)
+        mt.requires_grad_(True)
+        
+        t.requires_grad_(True)
+
+        ut = self.flow_net(t, xt)
+        gt = self.growth_net(t, xt)
+
+        time=t.expand(xt.shape[0], 1)
+        time.requires_grad_(True) 
+        
+        dx_dt = ut
+        dw_dt = gt
+                        
+        if self.data_manifold_metric is not None:
+            vel, _, _ = self.data_manifold_metric.calculate_velocity(
+                xt, ut, self.samples, 0
+            )
+                
+            dm_dt = torch.mean(vel ** 2) * wt
+        else:
+            dm_dt = ((ut**2).sum(dim =-1) + self.state_cost(xt)) * wt
+        
+        assert xt.shape == dx_dt.shape, f"dx mismatch: expected {xt.shape}, got {dx_dt.shape}"
+        assert wt.shape == dw_dt.shape, f"dw mismatch: expected {wt.shape}, got {dw_dt.shape}"
+        assert mt.shape == dm_dt.shape, f"dm mismatch: expected {mt.shape}, got {dm_dt.shape}"
+        return dx_dt, dw_dt, dm_dt
+    
+class ReconsLoss(nn.Module):
+    def __init__(self, hinge_value=0.01):
+        super(ReconsLoss, self).__init__()
+        self.hinge_value = hinge_value
+
+    def __call__(self, source, target, groups = None, to_ignore = None, top_k = 5):
+        if groups is not None:
+            # for global loss
+            c_dist = torch.stack([
+                torch.cdist(source[i], target[i]) 
+
+                for i in range(1,len(groups))
+                if groups[i] != to_ignore
+            ])
+        else:
+            # for local loss
+             c_dist = torch.stack([
+                torch.cdist(source, target)                 
+            ])
+        values, _ = torch.topk(c_dist, top_k, dim=2, largest=False, sorted=False)
+        values -= self.hinge_value
+        values[values<0] = 0
+        loss = torch.mean(values)
+        return loss

networks/flow_mlp.py

@@ -0,0 +1,18 @@
+import sys
+sys.path.append("./BranchSBM")
+import torch
+from networks.mlp_base import SimpleDenseNet
+
+
+class VelocityNet(SimpleDenseNet):
+    def __init__(self, dim: int, *args, **kwargs):
+        super().__init__(input_size=dim + 1, target_size=dim, *args, **kwargs)
+
+    def forward(self, t, x):
+
+        if t.dim() < 1 or t.shape[0] != x.shape[0]:
+            t = t.repeat(x.shape[0])[:, None]
+        if t.dim() < 2:
+            t = t[:, None]
+        x = torch.cat([t, x], dim=-1)
+        return self.model(x)

networks/growth_mlp.py

@@ -0,0 +1,37 @@
+import sys
+sys.path.append("./BranchSBM")
+import torch
+import torch.nn as nn
+from typing import List, Optional
+
+from networks.mlp_base import SimpleDenseNet
+
+class GrowthNet(SimpleDenseNet):
+    def __init__(
+        self,
+        dim: int,
+        activation: str,
+        hidden_dims: List[int] = None,
+        batch_norm: bool = False,
+        negative: bool = False
+    ):
+        super().__init__(input_size=dim + 1, target_size=1,
+                         activation=activation, 
+                         batch_norm=batch_norm, 
+                         hidden_dims=hidden_dims)
+        
+        self.softplus = nn.Softplus()
+        self.negative = negative
+        
+    def forward(self, t, x):
+
+        if t.dim() < 1 or t.shape[0] != x.shape[0]:
+            t = t.repeat(x.shape[0])[:, None]
+        if t.dim() < 2:
+            t = t[:, None]
+        x = torch.cat([t, x], dim=-1)
+        x = self.model(x)
+        x = self.softplus(self.model(x))
+        if self.negative:
+            x = -x
+        return x

networks/interpolant_mlp.py

@@ -0,0 +1,35 @@
+import sys
+sys.path.append("./BranchSBM")
+import torch
+import torch.nn as nn
+from typing import List, Optional
+
+from networks.mlp_base import SimpleDenseNet
+
+class GeoPathMLP(nn.Module):
+    def __init__(
+        self,
+        input_dim: int,
+        activation: str,
+        batch_norm: bool = True,
+        hidden_dims: Optional[List[int]] = None,
+        time_geopath: bool = False,
+    ):
+        super().__init__()
+        self.input_dim = input_dim
+        self.time_geopath = time_geopath
+        self.mainnet = SimpleDenseNet(
+            input_size=2 * input_dim + (1 if time_geopath else 0),
+            target_size=input_dim,
+            activation=activation,
+            batch_norm=batch_norm,
+            hidden_dims=hidden_dims,
+        )
+
+    def forward(
+        self, x0: torch.Tensor, x1: torch.Tensor, t: torch.Tensor
+    ) -> torch.Tensor:
+        x = torch.cat([x0, x1], dim=1)
+        if self.time_geopath:
+            x = torch.cat([x, t], dim=1)
+        return self.mainnet(x)

networks/mlp_base.py

@@ -0,0 +1,46 @@
+import sys
+sys.path.append("./BranchSBM")
+import torch.nn as nn
+import torch
+from typing import List, Optional
+
+class swish(nn.Module):
+    def forward(self, x):
+        return x * torch.sigmoid(x)
+
+
+ACTIVATION_MAP = {
+    "relu": nn.ReLU,
+    "sigmoid": nn.Sigmoid,
+    "tanh": nn.Tanh,
+    "selu": nn.SELU,
+    "elu": nn.ELU,
+    "lrelu": nn.LeakyReLU,
+    "softplus": nn.Softplus,
+    "silu": nn.SiLU,
+    "swish": swish,
+}
+
+
+class SimpleDenseNet(nn.Module):
+    def __init__(
+        self,
+        input_size: int,
+        target_size: int,
+        activation: str,
+        batch_norm: bool = False,
+        hidden_dims: List[int] = None,
+    ):
+        super().__init__()
+        dims = [input_size, *hidden_dims, target_size]
+        layers = []
+        for i in range(len(dims) - 2):
+            layers.append(nn.Linear(dims[i], dims[i + 1]))
+            if batch_norm:
+                layers.append(nn.BatchNorm1d(dims[i + 1]))
+            layers.append(ACTIVATION_MAP[activation]())
+        layers.append(nn.Linear(dims[-2], dims[-1]))
+        self.model = nn.Sequential(*layers)
+
+    def forward(self, x):
+        return self.model(x)

networks/utils.py

@@ -0,0 +1,13 @@
+import sys
+sys.path.append("./BranchSBM")
+import torch
+
+class flow_model_torch_wrapper(torch.nn.Module):
+    """Wraps model to torchdyn compatible format."""
+
+    def __init__(self, model):
+        super().__init__()
+        self.model = model
+
+    def forward(self, t, x, *args, **kwargs):
+        return self.model(t, x)

state_costs/land.py

@@ -0,0 +1,26 @@
+import torch
+
+
+def weighting_function(x, samples, gamma):
+    pairwise_sq_diff = (x[:, None, :] - samples[None, :, :]) ** 2
+    pairwise_sq_dist = pairwise_sq_diff.sum(-1)
+    weights = torch.exp(-pairwise_sq_dist / (2 * gamma**2))
+    return weights
+
+
+def land_metric_tensor(x, samples, gamma, rho):
+    weights = weighting_function(x, samples, gamma)  # Shape [B, N]
+    differences = samples[None, :, :] - x[:, None, :]  # Shape [B, N, D]
+    squared_differences = differences**2  # Shape [B, N, D]
+
+    # Compute the sum of weighted squared differences for each dimension
+    M_dd_diag = torch.einsum("bn,bnd->bd", weights, squared_differences) + rho
+
+    # Invert the metric tensor diagonal for each x_t
+    M_dd_inv_diag = 1.0 / M_dd_diag  # Shape [B, D] since it's diagonal
+    return M_dd_inv_diag
+
+
+def weighting_function_dt(x, dx_dt, samples, gamma, weights):
+    pairwise_sq_diff_dt = (x[:, None, :] - samples[None, :, :]) * dx_dt[:, None, :]
+    return -pairwise_sq_diff_dt.sum(-1) * weights / (gamma**2)

state_costs/metric_factory.py

@@ -0,0 +1,105 @@
+import sys
+sys.path.append("./BranchSBM")
+import torch
+import pytorch_lightning as pl
+from pytorch_lightning.loggers import WandbLogger
+from torch.utils.data import Dataset, DataLoader
+
+from state_costs.land import land_metric_tensor
+from state_costs.rbf import RBFNetwork
+
+class DataManifoldMetric:
+    def __init__(
+        self,
+        args,
+        skipped_time_points=None,
+        datamodule=None,
+    ):
+        self.skipped_time_points = skipped_time_points
+        self.datamodule = datamodule
+
+        self.gamma = args.gamma_current
+        self.rho = args.rho
+        self.metric = args.velocity_metric
+        self.n_centers = args.n_centers
+        self.kappa = args.kappa
+        self.metric_epochs = args.metric_epochs
+        self.metric_patience = args.metric_patience
+        self.lr = args.metric_lr
+        self.alpha_metric = args.alpha_metric
+        self.image_data = args.data_type == "image"
+        self.accelerator = args.accelerator
+
+        self.called_first_time = True
+        self.args = args
+
+    def calculate_metric(self, x_t, samples, current_timestep):
+        if self.metric == "land":
+            M_dd_x_t = (
+                land_metric_tensor(x_t, samples, self.gamma, self.rho)
+                ** self.alpha_metric
+            )
+            
+        elif self.metric == "rbf":
+            if self.called_first_time:
+                self.rbf_networks = []
+                for timestep in range(self.datamodule.num_timesteps - 1):
+                    if timestep in self.skipped_time_points:
+                        continue
+                    print("Learning RBF networks, timestep: ", timestep)
+                    rbf_network = RBFNetwork(
+                        current_timestep=timestep,
+                        next_timestep=timestep
+                        + 1
+                        + (1 if timestep + 1 in self.skipped_time_points else 0),
+                        n_centers=self.n_centers,
+                        kappa=self.kappa,
+                        lr=self.lr,
+                        datamodule=self.datamodule,
+                        args=self.args
+                    )
+                    early_stop_callback = pl.callbacks.EarlyStopping(
+                        monitor="MetricModel/val_loss_learn_metric",
+                        patience=self.metric_patience,
+                        mode="min",
+                    )
+                    trainer = pl.Trainer(
+                        max_epochs=self.metric_epochs,
+                        accelerator=self.accelerator,
+                        logger=WandbLogger(),
+                        num_sanity_val_steps=0,
+                        callbacks=(
+                            [early_stop_callback] if not self.image_data else None
+                        ),
+                    )
+                    if self.image_data:
+                        self.dataloader = DataLoader(
+                            self.datamodule.all_data,
+                            batch_size=128,
+                            shuffle=True,
+                        )
+                        trainer.fit(rbf_network, self.dataloader)
+                    else:
+                        trainer.fit(rbf_network, self.datamodule)
+                    self.rbf_networks.append(rbf_network)
+                self.called_first_time = False
+                print("Learning RBF networksss... Done")
+            M_dd_x_t = self.rbf_networks[current_timestep].compute_metric(
+                x_t,
+                epsilon=self.rho,
+                alpha=self.alpha_metric,
+                image_hx=self.image_data,
+            )
+        return M_dd_x_t
+
+    def calculate_velocity(self, x_t, u_t, samples, timestep):
+
+        if len(u_t.shape) > 2:
+            u_t = u_t.reshape(u_t.shape[0], -1)
+            x_t = x_t.reshape(x_t.shape[0], -1)
+        M_dd_x_t = self.calculate_metric(x_t, samples, timestep).to(u_t.device)
+
+        velocity = torch.sqrt(((u_t**2) * M_dd_x_t).sum(dim=-1))
+        ut_sum = (u_t**2).sum(dim=-1)
+        metric_sum = M_dd_x_t.sum(dim=-1)
+        return velocity, ut_sum, metric_sum

state_costs/rbf.py

@@ -0,0 +1,156 @@
+import pytorch_lightning as pl
+import torch
+from sklearn.cluster import KMeans
+import numpy as np
+
+class RBFNetwork(pl.LightningModule):
+    def __init__(
+        self,
+        current_timestep,
+        next_timestep,
+        n_centers: int = 100,
+        kappa: float = 1.0,
+        lr=1e-2,
+        datamodule=None,
+        image_data=False,
+        args=None
+    ):
+        super().__init__()
+        self.K = n_centers
+        self.current_timestep = current_timestep
+        self.next_timestep = next_timestep
+        self.clustering_model = KMeans(n_clusters=self.K)
+        self.kappa = kappa
+        self.last_val_loss = 1
+        self.lr = lr
+        self.W = torch.nn.Parameter(torch.rand(self.K, 1))
+        self.datamodule = datamodule
+        self.image_data = image_data
+        self.args = args
+
+    def on_before_zero_grad(self, *args, **kwargs):
+        self.W.data = torch.clamp(self.W.data, min=0.0001)
+
+    def on_train_start(self):
+        with torch.no_grad():
+            
+            batch = next(iter(self.trainer.datamodule.train_dataloader()))
+            
+            metric_samples = batch[0]["metric_samples"][0]
+            all_data = torch.cat(metric_samples)
+            data_to_fit = all_data
+
+            print("Fitting Clustering model...")
+            self.clustering_model.fit(data_to_fit)
+
+            clusters = (
+                self.calculate_centroids(all_data, self.clustering_model.labels_)
+                if self.image_data
+                else self.clustering_model.cluster_centers_
+            )
+
+            self.C = torch.tensor(clusters, dtype=torch.float32).to(self.device)
+            labels = self.clustering_model.labels_
+            sigmas = np.zeros((self.K, 1))
+
+            for k in range(self.K):
+                points = all_data[labels == k, :]
+                variance = ((points - clusters[k]) ** 2).mean(axis=0)
+                sigmas[k, :] = np.sqrt(
+                    variance.sum() if self.image_data else variance.mean()
+                )
+
+            self.lamda = torch.tensor(
+                0.5 / (self.kappa * sigmas) ** 2, dtype=torch.float32
+            ).to(self.device)
+
+    def forward(self, x):
+        if len(x.shape) > 2:
+            x = x.reshape(x.shape[0], -1).to(self.C.device)
+        
+        x = x.to(self.C.device)
+        dist2 = torch.cdist(x, self.C) ** 2
+        self.phi_x = torch.exp(-0.5 * self.lamda[None, :, :] * dist2[:, :, None])
+        
+        h_x = (self.W.to(x.device) * self.phi_x).sum(dim=1)
+        
+        return h_x
+
+    def training_step(self, batch, batch_idx):
+        if self.args.data_type == "scrna" or self.args.data_type == "tahoe":
+            main_batch = batch[0]["train_samples"][0]
+        else: 
+            main_batch = batch["train_samples"][0]
+
+        x0 = main_batch["x0"][0]
+        if self.args.branches == 1:
+            x1 = main_batch["x1"][0]
+            inputs = torch.cat([x0, x1], dim=0).to(self.device)
+        else:
+            x1_1 = main_batch["x1_1"][0]
+            x1_2 = main_batch["x1_2"][0]
+            
+            inputs = torch.cat([x0, x1_1, x1_2], dim=0).to(self.device)
+        print("inputs shape")
+        print(inputs.shape)
+        
+        loss = ((1 - self.forward(inputs)) ** 2).mean()
+        self.log(
+            "MetricModel/train_loss_learn_metric",
+            loss,
+            on_step=True,
+            on_epoch=True,
+            prog_bar=True,
+        )
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+        if self.args.data_type == "scrna" or self.args.data_type == "tahoe":
+            main_batch = batch[0]["val_samples"][0]
+        else: 
+            main_batch = batch["val_samples"][0]
+
+        x0 = main_batch["x0"][0]
+        if self.args.branches == 1:
+            x1 = main_batch["x1"][0]
+            inputs = torch.cat([x0, x1], dim=0).to(self.device)
+        else:
+            x1_1 = main_batch["x1_1"][0]
+            x1_2 = main_batch["x1_2"][0]
+            
+            inputs = torch.cat([x0, x1_1, x1_2], dim=0).to(self.device)
+
+        h = self.forward(inputs)
+        
+        loss = ((1 - h) ** 2).mean()
+        self.log(
+            "MetricModel/val_loss_learn_metric",
+            loss,
+            on_step=True,
+            on_epoch=True,
+            prog_bar=True,
+        )
+        self.last_val_loss = loss.detach()
+        return loss
+
+    def calculate_centroids(self, all_data, labels):
+        unique_labels = np.unique(labels)
+        centroids = np.zeros((len(unique_labels), all_data.shape[1]))
+        for i, label in enumerate(unique_labels):
+            centroids[i] = all_data[labels == label].mean(axis=0)
+        return centroids
+
+    def configure_optimizers(self):
+        optimizer = torch.optim.Adam(self.parameters(), lr=self.lr)
+        return optimizer
+
+    def compute_metric(self, x, alpha=1, epsilon=1e-2, image_hx=False):
+        if epsilon < 0:
+            epsilon = (1 - self.last_val_loss.item()) / abs(epsilon)
+        h_x = self.forward(x)
+        if image_hx:
+            h_x = 1 - torch.abs(1 - h_x)
+            M_x = 1 / (h_x**alpha + epsilon)
+        else:
+            M_x = 1 / (h_x + epsilon) ** alpha
+        return M_x

train/main_branches.py

@@ -0,0 +1,342 @@
+import sys
+sys.path.append("./BranchSBM")
+import os
+import sys
+import argparse
+import copy
+from pytorch_lightning import Trainer
+from pytorch_lightning.loggers import WandbLogger
+import wandb
+import hydra
+from omegaconf import DictConfig, OmegaConf
+from torchcfm.optimal_transport import OTPlanSampler
+
+from branchsbm.branchsbm import BranchSBM
+from branchsbm.branch_flow_net_train import FlowNetTrainCell, FlowNetTrainLidar
+from branchsbm.branch_interpolant_train import BranchInterpolantTrain
+
+from dataloaders.trajectory_data import TemporalDataModule
+from dataloaders.mouse_data import WeightedBranchedCellDataModule
+from dataloaders.three_branch_data import ThreeBranchTahoeDataModule
+from dataloaders.clonidine_v2_data import ClonidineV2DataModule
+from dataloaders.clonidine_single_branch import ClonidineSingleBranchDataModule
+from dataloaders.trametinib_single import TrametinibSingleBranchDataModule
+from dataloaders.lidar_data import WeightedBranchedLidarDataModule
+from dataloaders.lidar_data_single import LidarSingleDataModule
+
+from networks.flow_networks.mlp import VelocityNet
+from networks.growth_networks.mlp import GrowthNet
+from networks.geopath_networks.mlp import GeoPathMLP 
+from networks.unet_base import UNetModelWrapper as UNetModel
+from networks.geopath_networks.unet import GeoPathUNet
+from utils import set_seed
+
+from train.parsers import parse_args
+from flow_matchers.ema import EMA
+from train.train_utils import (
+    load_config,
+    merge_config,
+    generate_group_string,
+    dataset_name2datapath,
+    create_callbacks,
+)
+from geo_metrics.metric_factory import DataManifoldMetric
+import torch.nn as nn
+from flow_matchers.branch_growth_net_train import GrowthNetTrain, GrowthNetTrainCell, GrowthNetTrainLidar
+
+def main(args: argparse.Namespace, seed: int, t_exclude: int) -> None:
+    set_seed(seed)
+    branches = args.branches
+
+    skipped_time_points = [t_exclude] if t_exclude else []
+    
+    ### DATAMODULES ###
+    if args.data_name == "lidar":
+        datamodule = WeightedBranchedLidarDataModule(args=args)
+    elif args.data_name == "lidarsingle":
+        datamodule = LidarSingleDataModule(args=args)
+    elif args.data_name == "mouse":
+        datamodule = WeightedBranchedCellDataModule(args=args)
+    elif args.data_name in ["clonidine50D", "clonidine100D", "clonidine150D"]:
+        datamodule = ClonidineV2DataModule(args=args)  
+    elif args.data_name == "clonidine50Dsingle":
+        datamodule = ClonidineSingleBranchDataModule(args=args)
+    elif args.data_name == "trametinib":
+        datamodule = ThreeBranchTahoeDataModule(args=args)
+    elif args.data_name == "trametinibsingle":
+        datamodule = TrametinibSingleBranchDataModule(args=args)
+    
+    flow_nets = nn.ModuleList()
+    geopath_nets = nn.ModuleList()
+    growth_nets = nn.ModuleList()
+        
+    ##### initialize branched flow and growth networks #####
+    for i in range(branches):
+        flow_net = VelocityNet(
+            dim=args.dim,
+            hidden_dims=args.hidden_dims_flow,
+            activation=args.activation_flow,
+            batch_norm=False,
+        )            
+        geopath_net = GeoPathMLP(
+            input_dim=args.dim,
+            hidden_dims=args.hidden_dims_geopath,
+            time_geopath=args.time_geopath,
+            activation=args.activation_geopath,
+            batch_norm=False,
+        )
+        
+        if i == 0:
+            growth_net = GrowthNet(
+                dim=args.dim,
+                hidden_dims=args.hidden_dims_growth,
+                activation=args.activation_growth,
+                batch_norm=False,
+                negative=True
+            )
+        else: 
+            growth_net = GrowthNet(
+                dim=args.dim,
+                hidden_dims=args.hidden_dims_growth,
+                activation=args.activation_growth,
+                batch_norm=False,
+                negative=False
+            )
+        
+        if args.ema_decay is not None:
+            flow_net = EMA(model=flow_net, decay=args.ema_decay)
+            geopath_net = EMA(model=geopath_net, decay=args.ema_decay)
+            growth_net = EMA(model=growth_net, decay=args.ema_decay)
+        
+        flow_nets.append(flow_net)
+        geopath_nets.append(geopath_net)
+        growth_nets.append(growth_net)
+        
+    
+    ot_sampler = (
+        OTPlanSampler(method=args.optimal_transport_method)
+        if args.optimal_transport_method != "None"
+        else None
+    )
+
+    wandb.init(
+        project=f"branchsbm-{args.data_name}-{branches}-branches",
+        group=args.group_name,
+        config=vars(args),
+        dir=args.working_dir,
+    )
+
+    flow_matcher_base = BranchSBM(
+        geopath_nets=geopath_nets,
+        sigma=args.sigma,
+        alpha=int(args.branchsbm),
+    )
+
+    ##### STAGE 1: Training of Geodesic Interpolants Beginning #####
+
+    geopath_callbacks = create_callbacks(
+        args, phase="geopath", data_type=args.data_type, run_id=wandb.run.id
+    )
+    
+    # define state cost
+    data_manifold_metric = DataManifoldMetric(
+        args=args,
+        skipped_time_points=skipped_time_points,
+        datamodule=datamodule,
+    )
+    geopath_model = BranchInterpolantTrain(
+        flow_matcher=flow_matcher_base,
+        skipped_time_points=skipped_time_points,
+        ot_sampler=ot_sampler,
+        args=args,
+        data_manifold_metric=data_manifold_metric
+    )
+        
+    wandb_logger = WandbLogger()
+
+    trainer = Trainer(
+        max_epochs=args.epochs,
+        callbacks=geopath_callbacks,
+        accelerator=args.accelerator,
+        logger=wandb_logger,
+        num_sanity_val_steps=0,
+        default_root_dir=args.working_dir,
+        gradient_clip_val=(1.0 if args.data_type == "image" else None),
+    )
+    
+    if args.load_geopath_model_ckpt:
+        best_model_path = args.load_geopath_model_ckpt
+    else:
+        trainer.fit(
+            geopath_model,
+            datamodule=datamodule,
+        )
+        
+        best_model_path = geopath_callbacks[0].best_model_path
+        
+    geopath_model = BranchInterpolantTrain.load_from_checkpoint(best_model_path)
+
+    flow_matcher_base.geopath_nets = geopath_model.geopath_nets
+
+    ##### STAGE 1: Training of Geodesic Interpolants End #####
+    
+
+    ##### STAGE 2: Flow Matching Beginning #####
+    flow_callbacks = create_callbacks(
+        args,
+        phase="flow",
+        data_type=args.data_type,
+        run_id=wandb.run.id,
+        datamodule=datamodule,
+    )
+    if args.data_type == "lidar":
+        FlowNetTrain = FlowNetTrainLidar
+    else:
+        FlowNetTrain = FlowNetTrainCell
+
+    flow_train = FlowNetTrain(
+        flow_matcher=flow_matcher_base,
+        flow_nets=flow_nets,
+        ot_sampler=ot_sampler,
+        skipped_time_points=skipped_time_points,
+        args=args,
+    )
+
+    wandb_logger = WandbLogger()
+
+    trainer = Trainer(
+        max_epochs=args.epochs,
+        callbacks=flow_callbacks,
+        check_val_every_n_epoch=args.check_val_every_n_epoch,
+        accelerator=args.accelerator,
+        logger=wandb_logger,
+        default_root_dir=args.working_dir,
+        gradient_clip_val=(1.0 if args.data_type == "image" else None),
+        num_sanity_val_steps=(0 if args.data_type == "image" else None),
+    )
+
+    trainer.fit(
+        flow_train, datamodule=datamodule, ckpt_path=args.resume_flow_model_ckpt
+    )
+    if args.data_type == "lidar":
+        trainer.test(flow_train, datamodule=datamodule)
+    ##### STAGE 2: Flow Matching End #####
+    
+    
+    ##### STAGE 3: Training Growth Networks Beginning ####
+    flow_nets = flow_train.flow_nets
+    
+    growth_callbacks = create_callbacks(
+        args,
+        phase="growth",
+        data_type=args.data_type,
+        run_id=wandb.run.id,
+        datamodule=datamodule,
+    )
+    
+    if args.data_type == "lidar":
+        GrowthNetTrain = GrowthNetTrainLidar
+    else:
+        GrowthNetTrain = GrowthNetTrainCell
+    
+    growth_train = GrowthNetTrain(
+        flow_nets = flow_nets,
+        growth_nets = growth_nets,
+        ot_sampler=ot_sampler,
+        skipped_time_points=skipped_time_points,
+        args=args,
+        data_manifold_metric=data_manifold_metric,
+        joint = False
+    )
+    
+    wandb_logger = WandbLogger()
+
+    trainer = Trainer(
+        max_epochs=args.epochs,
+        callbacks=growth_callbacks,
+        check_val_every_n_epoch=args.check_val_every_n_epoch,
+        accelerator=args.accelerator,
+        logger=wandb_logger,
+        default_root_dir=args.working_dir,
+        gradient_clip_val=(1.0 if args.data_type == "image" else None),
+        num_sanity_val_steps=(0 if args.data_type == "image" else None),
+    )
+    
+    trainer.fit(
+        growth_train, datamodule=datamodule, ckpt_path=None
+    )
+    trainer.test(growth_train, datamodule=datamodule)
+    
+    ##### STAGE 3: Training Growth Networks End ####
+    
+    ##### STAGE 4: Joint Training Beginning ####
+    
+    growth_nets = growth_train.growth_nets
+    
+    joint_callbacks = create_callbacks(
+        args,
+        phase="joint",
+        data_type=args.data_type,
+        run_id=wandb.run.id,
+        datamodule=datamodule,
+    )
+    
+    if args.data_type == "lidar":
+        GrowthNetTrain = GrowthNetTrainLidar
+    else:
+        GrowthNetTrain = GrowthNetTrainCell
+        
+    joint_train = GrowthNetTrain(
+        flow_nets = flow_nets,
+        growth_nets = growth_nets,
+        ot_sampler=ot_sampler,
+        skipped_time_points=skipped_time_points,
+        args=args,
+        data_manifold_metric=data_manifold_metric,
+        joint = True
+    )
+    
+    wandb_logger = WandbLogger()
+
+    trainer = Trainer(
+        max_epochs=args.epochs,
+        callbacks=joint_callbacks,
+        check_val_every_n_epoch=args.check_val_every_n_epoch,
+        accelerator=args.accelerator,
+        logger=wandb_logger,
+        default_root_dir=args.working_dir,
+        gradient_clip_val=(1.0 if args.data_type == "image" else None),
+        num_sanity_val_steps=(0 if args.data_type == "image" else None),
+    )
+    
+    trainer.fit(
+        joint_train, datamodule=datamodule, ckpt_path=None
+    )
+    trainer.test(joint_train, datamodule=datamodule)
+    
+    ##### STAGE 4: Joint Training End ####
+    
+    wandb.finish()
+    
+if __name__ == "__main__":
+    args = parse_args()
+    updated_args = copy.deepcopy(args)
+    if args.config_path:
+        config = load_config(args.config_path)
+        updated_args = merge_config(updated_args, config)
+
+    updated_args.group_name = generate_group_string()
+    updated_args.data_path = dataset_name2datapath(
+        updated_args.data_name, updated_args.working_dir
+    )
+    for seed in updated_args.seeds:
+        if updated_args.t_exclude:
+            for i, t_exclude in enumerate(updated_args.t_exclude):
+                updated_args.t_exclude_current = t_exclude
+                updated_args.seed_current = seed
+                updated_args.gamma_current = updated_args.gammas[i]
+                main(updated_args, seed=seed, t_exclude=t_exclude)
+        else:
+            updated_args.seed_current = seed
+            updated_args.gamma_current = updated_args.gammas[0]
+            main(updated_args, seed=seed, t_exclude=None)

train/parsers.py

@@ -0,0 +1,419 @@
+import argparse
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Train BranchSBM")
+
+    parser.add_argument(
+        "--config_path", type=str, 
+        default='./configs/experiment/lidar.yaml', 
+        help="Path to config file"
+    )
+
+    ####### ITERATES IN THE CODE #######
+    parser.add_argument(
+        "--seeds",
+        nargs="+",
+        type=int,
+        default=[42, 43, 44, 45, 46],
+        help="Random seeds to iterate over",
+    )
+    parser.add_argument(
+        "--t_exclude",
+        nargs="+",
+        type=int,
+        default=[1, 2],
+        help="Time points to exclude (iterating over)",
+    )
+    ####################################
+
+    parser.add_argument(
+        "--working_dir",
+        type=str,
+        default="./",
+        help="Working directory",
+    )
+    parser.add_argument(
+        "--resume_flow_model_ckpt",
+        type=str,
+        default=None,
+        help="Path to the flow model to resume training",
+    )
+    parser.add_argument(
+        "--resume_growth_model_ckpt",
+        type=str,
+        default=None,
+        help="Path to the flow model to resume training",
+    )
+    parser.add_argument(
+        "--load_geopath_model_ckpt",
+        type=str,
+        default=None,
+        help="Path to the geopath model to resume training",
+    )
+    parser.add_argument(
+        "--branches",
+        type=int,
+        default=2,
+        help="Number of branches",
+    )
+    parser.add_argument(
+        "--metric_clusters",
+        type=int,
+        default=3,
+        help="Number of metric clusters",
+    )
+
+    ######### DATASETS #################
+    parser = datasets_parser(parser)
+    ####################################
+
+    ######### IMAGE DATASETS ###########
+    parser = image_datasets_parser(parser)
+    ####################################
+
+    ######### METRICS ##################
+    parser = metric_parser(parser)
+    ####################################
+
+    ######### General Training #########
+    parser = general_training_parser(parser)
+    ####################################
+
+    ######### Training GeoPath Network ####
+    parser = geopath_network_parser(parser)
+    ####################################
+
+    ######### Training Flow Network ####
+    parser = flow_network_parser(parser)
+    ####################################
+    
+    parser = growth_network_parser(parser)
+
+    return parser.parse_args()
+
+
+def datasets_parser(parser):
+    parser.add_argument("--dim", type=int, default=3, help="Dimension of data")
+
+    parser.add_argument(
+        "--data_type",
+        type=str,
+        default="lidar",
+        help="Type of data, now wither scrna or one of toys",
+    )
+    parser.add_argument(
+        "--data_path",
+        type=str,
+        default="./data/rainier2-thin.las",
+        help="lidar data path",
+    )
+    parser.add_argument(
+        "--data_name",
+        type=str,
+        default="lidar",
+        help="Path to the dataset",
+    )
+    parser.add_argument(
+        "--whiten",
+        action=argparse.BooleanOptionalAction,
+        default=True,
+        help="Whiten the data",
+    )
+    return parser
+
+
+def image_datasets_parser(parser):
+    parser.add_argument(
+        "--image_size",
+        type=int,
+        default=128,
+        help="Size of the image",
+    )
+    parser.add_argument(
+        "--x0_label",
+        type=str,
+        default="dog",
+        help="Label for x0",
+    )
+    parser.add_argument(
+        "--x1_label",
+        type=str,
+        default="cat",
+        help="Label for x1",
+    )
+    return parser
+
+
+def metric_parser(parser):
+    parser.add_argument(
+        "--branchsbm",
+        action=argparse.BooleanOptionalAction,
+        default=True,
+        help="If branched SBM",
+    )
+    parser.add_argument(
+        "--n_centers",
+        type=int,
+        default=100,
+        help="Number of centers for RBF network",
+    )
+    parser.add_argument(
+        "--kappa",
+        type=float,
+        default=1.0,
+        help="Kappa parameter for RBF network",
+    )
+    parser.add_argument(
+        "--rho",
+        type=float,
+        default=0.001,
+        help="Rho parameter in Riemanian Velocity Calculation",
+    )
+    parser.add_argument(
+        "--velocity_metric",
+        type=str,
+        default="rbf",
+        help="Metric for velocity calculation",
+    )
+    parser.add_argument(
+        "--gammas",
+        nargs="+",
+        type=float,
+        default=[0.2, 0.2],
+        help="Gamma parameter in Riemanian Velocity Calculation",
+    )
+    parser.add_argument(
+        "--metric_epochs",
+        type=int,
+        default=50,
+        help="Number of epochs for metric learning",
+    )
+    parser.add_argument(
+        "--metric_patience",
+        type=int,
+        default=5,
+        help="Patience for metric learning",
+    )
+    parser.add_argument(
+        "--metric_lr",
+        type=float,
+        default=1e-2,
+        help="Learning rate for metric learning",
+    )
+    parser.add_argument(
+        "--alpha_metric",
+        type=float,
+        default=1.0,
+        help="Alpha parameter for metric learning",
+    )
+
+    return parser
+
+
+def general_training_parser(parser):
+    parser.add_argument(
+        "--batch_size", type=int, default=128, help="Batch size for training"
+    )
+    parser.add_argument(
+        "--optimal_transport_method",
+        type=str,
+        default="exact",
+        help="Use optimal transport in CFM training",
+    )
+    parser.add_argument(
+        "--ema_decay",
+        type=float,
+        default=None,
+        help="Decay for EMA",
+    )
+    parser.add_argument(
+        "--split_ratios",
+        nargs=2,
+        type=float,
+        default=[0.9, 0.1],
+        help="Split ratios for training/validation data in CFM training",
+    )
+    parser.add_argument("--epochs", type=int, default=100, help="Number of epochs")
+    parser.add_argument(
+        "--accelerator", type=str, default="cpu", help="Training accelerator"
+    )
+    parser.add_argument(
+        "--sim_num_steps",
+        type=int,
+        default=1000,
+        help="Number of steps in simulation",
+    )
+    return parser
+
+
+def geopath_network_parser(parser):
+    parser.add_argument(
+        "--manifold",
+        action=argparse.BooleanOptionalAction,
+        default=True,
+        help="If use data manifold metric",
+    )
+    parser.add_argument(
+        "--patience_geopath",
+        type=int,
+        default=50,
+        help="Patience for training geopath model",
+    )
+    parser.add_argument(
+        "--hidden_dims_geopath",
+        nargs="+",
+        type=int,
+        default=[64, 64, 64],
+        help="Dimensions of hidden layers for GeoPath model training",
+    )
+    parser.add_argument(
+        "--time_geopath",
+        action=argparse.BooleanOptionalAction,
+        default=False,
+        help="Use time in GeoPath model",
+    )
+    parser.add_argument(
+        "--activation_geopath",
+        type=str,
+        default="selu",
+        help="Activation function for GeoPath",
+    )
+    parser.add_argument(
+        "--geopath_optimizer",
+        type=str,
+        default="adam",
+        help="Optimizer for GeoPath training",
+    )
+    parser.add_argument(
+        "--geopath_lr",
+        type=float,
+        default=1e-4,
+        help="Learning rate for GeoPath training",
+    )
+    parser.add_argument(
+        "--geopath_weight_decay",
+        type=float,
+        default=1e-5,
+        help="Weight decay for GeoPath training",
+    )
+    return parser
+
+
+def flow_network_parser(parser):
+    parser.add_argument(
+        "--sigma", type=float, default=0.1, help="Sigma parameter for CFM (variance)"
+    )
+    parser.add_argument(
+        "--patience",
+        type=int,
+        default=5,
+        help="Patience for early stopping in CFM training",
+    )
+    parser.add_argument(
+        "--hidden_dims_flow",
+        nargs="+",
+        type=int,
+        default=[64, 64, 64],
+        help="Dimensions of hidden layers for CFM training",
+    )
+    parser.add_argument(
+        "--check_val_every_n_epoch",
+        type=int,
+        default=10,
+        help="Check validation every N epochs during CFM training",
+    )
+    parser.add_argument(
+        "--activation_flow",
+        type=str,
+        default="selu",
+        help="Activation function for CFM",
+    )
+    parser.add_argument(
+        "--flow_optimizer",
+        type=str,
+        default="adamw",
+        help="Optimizer for GeoPath training",
+    )
+    parser.add_argument(
+        "--flow_lr",
+        type=float,
+        default=1e-3,
+        help="Learning rate for GeoPath training",
+    )
+    parser.add_argument(
+        "--flow_weight_decay",
+        type=float,
+        default=1e-5,
+        help="Weight decay for GeoPath training",
+    )
+    return parser
+
+def growth_network_parser(parser):
+    parser.add_argument(
+        "--patience_growth",
+        type=int,
+        default=5,
+        help="Patience for early stopping in CFM training",
+    )
+    parser.add_argument(
+        "--time_growth",
+        action=argparse.BooleanOptionalAction,
+        default=False,
+        help="Use time in GeoPath model",
+    )
+    parser.add_argument(
+        "--hidden_dims_growth",
+        nargs="+",
+        type=int,
+        default=[64, 64, 64],
+        help="Dimensions of hidden layers for growth net training",
+    )
+    parser.add_argument(
+        "--activation_growth",
+        type=str,
+        default="tanh",
+        help="Activation function for CFM",
+    )
+    parser.add_argument(
+        "--growth_optimizer",
+        type=str,
+        default="adamw",
+        help="Optimizer for GeoPath training",
+    )
+    parser.add_argument(
+        "--growth_lr",
+        type=float,
+        default=1e-3,
+        help="Learning rate for GeoPath training",
+    )
+    parser.add_argument(
+        "--growth_weight_decay",
+        type=float,
+        default=1e-5,
+        help="Weight decay for GeoPath training",
+    )
+    parser.add_argument(
+        "--lambda_energy",
+        type=float,
+        default=1.0,
+        help="Weight for energy loss",
+    )
+    parser.add_argument(
+        "--lambda_mass",
+        type=float,
+        default=100.0,
+        help="Weight for mass loss",
+    )
+    parser.add_argument(
+        "--lambda_match",
+        type=float,
+        default=1000.0,
+        help="Weight for matching loss",
+    )
+    parser.add_argument(
+        "--lambda_recons",
+        type=float,
+        default=1.0,
+        help="Weight for reconstruction loss",
+    )
+    return parser

train/train_utils.py

@@ -0,0 +1,154 @@
+import sys
+sys.path.append("./BranchSBM")
+import yaml
+import string
+import secrets
+import os
+import torch
+import wandb
+from pytorch_lightning.callbacks import Callback, EarlyStopping, ModelCheckpoint
+from torchdyn.core import NeuralODE
+from utils import plot_images_trajectory
+from networks.utils import flow_model_torch_wrapper
+
+
+def load_config(path):
+    with open(path, "r") as file:
+        config = yaml.safe_load(file)
+    return config
+
+
+def merge_config(args, config_updates):
+    for key, value in config_updates.items():
+        if not hasattr(args, key):
+            raise ValueError(
+                f"Unknown configuration parameter '{key}' found in the config file."
+            )
+        setattr(args, key, value)
+    return args
+
+
+def generate_group_string(length=16):
+    alphabet = string.ascii_letters + string.digits
+    return "".join(secrets.choice(alphabet) for _ in range(length))
+
+
+def dataset_name2datapath(dataset_name, working_dir):
+    if dataset_name in ["lidar", "lidarsingle"]:
+        return os.path.join(working_dir, "/raid/st512/branchsbm/data", "rainier2-thin.las")
+    elif dataset_name == "mouse":
+        return os.path.join(working_dir, "/raid/st512/branchsbm/data", "mouse_hematopoiesis.csv")
+    elif dataset_name in ["clonidine50D", "clonidine100D", "clonidine150D", "clonidine50Dsingle", "clonidine100Dsingle", "clonidine150Dsingle"]:
+        return os.path.join(working_dir, "/raid/st512/branchsbm/data", "pca_and_leiden_labels.csv")
+    elif dataset_name in ["trametinib", "trametinibsingle"]:
+        return os.path.join(working_dir, "/raid/st512/branchsbm/data", "Trametinib_5.0uM_pca_and_leidenumap_labels.csv")
+    else:
+        raise ValueError("Dataset not recognized")
+
+
+def create_callbacks(args, phase, data_type, run_id, datamodule=None):
+
+    dirpath = os.path.join(
+        args.working_dir,
+        "checkpoints",
+        data_type,
+        str(run_id),
+        f"{phase}_model",
+    )
+
+    if phase == "geopath":
+        early_stop_callback = EarlyStopping(
+            monitor="BranchPathNet/val_loss_geopath",
+            patience=args.patience_geopath,
+            mode="min",
+        )
+        checkpoint_callback = ModelCheckpoint(
+            dirpath=dirpath,
+            monitor="BranchPathNet/val_loss_geopath",
+            mode="min",
+            save_top_k=1,
+        )
+        callbacks = [checkpoint_callback, early_stop_callback]
+    elif phase == "flow":
+        early_stop_callback = EarlyStopping(
+            monitor="FlowNet/val_loss_cfm",
+            patience=args.patience,
+            mode="min",
+        )
+        checkpoint_callback = ModelCheckpoint(
+            dirpath=dirpath,
+            mode="min",
+            save_top_k=1,
+        )
+        callbacks = [checkpoint_callback, early_stop_callback]
+    elif phase == "growth":
+        early_stop_callback = EarlyStopping(
+            monitor="GrowthNet/val_loss",
+            patience=args.patience,
+            mode="min",
+        )
+        checkpoint_callback = ModelCheckpoint(
+            dirpath=dirpath,
+            mode="min",
+            save_top_k=1,
+        )
+        callbacks = [checkpoint_callback, early_stop_callback]
+    elif phase == "joint":
+        early_stop_callback = EarlyStopping(
+            monitor="JointTrain/val_loss",
+            patience=args.patience,
+            mode="min",
+        )
+        checkpoint_callback = ModelCheckpoint(
+            dirpath=dirpath,
+            mode="min",
+            save_top_k=1,
+        )
+        callbacks = [checkpoint_callback, early_stop_callback]
+    else:
+        raise ValueError("Unknown phase")
+    return callbacks
+
+
+class PlottingCallback(Callback):
+    def __init__(self, plot_interval, datamodule):
+        self.plot_interval = plot_interval
+        self.datamodule = datamodule
+
+    def on_train_epoch_end(self, trainer, pl_module):
+        epoch = trainer.current_epoch
+        pl_module.flow_net.train(mode=False)
+        if epoch % self.plot_interval == 0 and epoch != 0:
+            node = NeuralODE(
+                flow_model_torch_wrapper(pl_module.flow_net).to(self.datamodule.device),
+                solver="tsit5",
+                sensitivity="adjoint",
+                atol=1e-5,
+                rtol=1e-5,
+            )
+
+            for mode in ["train", "val"]:
+                x0 = getattr(self.datamodule, f"{mode}_x0")
+                x0 = x0[0:15]
+                fig = self.trajectory_and_plot(x0, node, self.datamodule)
+                wandb.log({f"Trajectories {mode.capitalize()}": wandb.Image(fig)})
+        pl_module.flow_net.train(mode=True)
+
+    def trajectory_and_plot(self, x0, node, datamodule):
+        selected_images = x0[0:15]
+        with torch.no_grad():
+            traj = node.trajectory(
+                selected_images.to(datamodule.device),
+                t_span=torch.linspace(0, 1, 100).to(datamodule.device),
+            )
+
+        traj = traj.transpose(0, 1)
+        traj = traj.reshape(*traj.shape[0:2], *datamodule.dim)
+
+        fig = plot_images_trajectory(
+            traj.to(datamodule.device),
+            datamodule.vae.to(datamodule.device),
+            datamodule.process,
+            num_steps=5,
+        )
+        return fig

utils.py

@@ -0,0 +1,198 @@
+import numpy as np
+import torch
+import random
+import matplotlib
+import matplotlib.pyplot as plt
+import math
+import umap
+import scanpy as sc
+from sklearn.decomposition import PCA
+
+import ot as pot
+from tqdm import tqdm
+from functools import partial
+from typing import Optional
+
+from matplotlib.colors import LinearSegmentedColormap
+
+
+def set_seed(seed):
+    """
+    Sets the seed for reproducibility in PyTorch, Numpy, and Python's Random.
+
+    Parameters:
+    seed (int): The seed for the random number generators.
+    """
+    random.seed(seed)  # Python random module
+    np.random.seed(seed)  # Numpy
+    torch.manual_seed(seed)  # CPU and GPU (deterministic)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(seed)  # CUDA
+        torch.cuda.manual_seed_all(seed)  # all GPU devices
+        torch.backends.cudnn.deterministic = True  # CuDNN behavior
+        torch.backends.cudnn.benchmark = False
+
+
+def wasserstein_distance(
+    x0: torch.Tensor,
+    x1: torch.Tensor,
+    method: Optional[str] = None,
+    reg: float = 0.05,
+    power: int = 1,
+    **kwargs,
+) -> float:
+    assert power == 1 or power == 2
+    if method == "exact" or method is None:
+        ot_fn = pot.emd2
+    elif method == "sinkhorn":
+        ot_fn = partial(pot.sinkhorn2, reg=reg)
+    else:
+        raise ValueError(f"Unknown method: {method}")
+
+    a, b = pot.unif(x0.shape[0]), pot.unif(x1.shape[0])
+    if x0.dim() > 2:
+        x0 = x0.reshape(x0.shape[0], -1)
+    if x1.dim() > 2:
+        x1 = x1.reshape(x1.shape[0], -1)
+    M = torch.cdist(x0, x1)
+    if power == 2:
+        M = M**2
+    ret = ot_fn(a, b, M.detach().cpu().numpy(), numItermax=1e7)
+    if power == 2:
+        ret = math.sqrt(ret)
+    return ret
+
+
+def plot_lidar(ax, dataset, xs=None, S=25, branch_idx=None):
+    # Combine the dataset and trajectory points for sorting
+    combined_points = []
+    combined_colors = []
+    combined_sizes = []
+    
+    
+    custom_colors_1 = ["#05009E", "#A19EFF", "#50B2D7"]
+    custom_colors_2 = ["#05009E", "#A19EFF", "#D577FF"]
+    
+    custom_cmap_1 = LinearSegmentedColormap.from_list("my_cmap", custom_colors_1)
+    custom_cmap_2 = LinearSegmentedColormap.from_list("my_cmap", custom_colors_2)
+
+    # Normalize the z-coordinates for alpha scaling
+    z_coords = (
+        dataset[:, 2].numpy() if torch.is_tensor(dataset[:, 2]) else dataset[:, 2]
+    )
+    z_min, z_max = z_coords.min(), z_coords.max()
+    z_norm = (z_coords - z_min) / (z_max - z_min)
+
+    # Add surface points with a lower z-order
+    for i, point in enumerate(dataset):
+        grey_value = 0.95 - 0.7 * z_norm[i]
+        combined_points.append(point.numpy())
+        combined_colors.append(
+            (
+                grey_value,
+                grey_value,
+                grey_value,
+                1.0
+            )
+        )  # Grey color with transparency
+        combined_sizes.append(0.1)
+
+    # Add trajectory points with a higher z-order
+    if xs is not None:
+        if branch_idx == 0:
+            cmap = custom_cmap_1
+        else:
+            cmap = custom_cmap_2
+            
+        B, T, D = xs.shape
+        steps_to_log = np.linspace(0, T - 1, S).astype(int)
+        xs = xs.cpu().detach().clone()
+        for idx, step in enumerate(steps_to_log):
+            for point in xs[:512, step]:
+                combined_points.append(
+                    point.numpy() if torch.is_tensor(point) else point
+                )
+                combined_colors.append(cmap(idx / (len(steps_to_log) - 1)))
+                combined_sizes.append(0.8)
+
+    # Convert to numpy array for easier manipulation
+    combined_points = np.array(combined_points)
+    combined_colors = np.array(combined_colors)
+    combined_sizes = np.array(combined_sizes)
+
+    # Sort by z-coordinate (depth)
+    sorted_indices = np.argsort(combined_points[:, 2])
+    combined_points = combined_points[sorted_indices]
+    combined_colors = combined_colors[sorted_indices]
+    combined_sizes = combined_sizes[sorted_indices]
+
+    # Plot the sorted points
+    ax.scatter(
+        combined_points[:, 0],
+        combined_points[:, 1],
+        combined_points[:, 2],
+        s=combined_sizes,
+        c=combined_colors,
+        depthshade=True,
+    )
+
+    ax.set_xlim3d(left=-4.8, right=4.8)
+    ax.set_ylim3d(bottom=-4.8, top=4.8)
+    ax.set_zlim3d(bottom=0.0, top=2.0)
+    ax.set_zticks([0, 1.0, 2.0])
+    ax.grid(False)
+    plt.axis("off")
+
+    return ax
+
+
+def plot_images_trajectory(trajectories, vae, processor, num_steps):
+
+    # Compute trajectories for each image
+    t_span = torch.linspace(0, trajectories.shape[1] - 1, num_steps)
+    t_span = [int(t) for t in t_span]
+    num_images = trajectories.shape[0]
+
+    # Decode images at each step in each trajectory
+    decoded_images = [
+        [
+            processor.postprocess(
+                vae.decode(
+                    trajectories[i_image, traj_step].unsqueeze(0)
+                ).sample.detach()
+            )[0]
+            for traj_step in t_span
+        ]
+        for i_image in range(num_images)
+    ]
+
+    # Plotting
+    fig, axes = plt.subplots(
+        num_images, num_steps, figsize=(num_steps * 2, num_images * 2)
+    )
+    if num_images == 1:
+        axes = [axes]  # Ensure axes is iterable
+    for img_idx, img_traj in enumerate(decoded_images):
+        for step_idx, img in enumerate(img_traj):
+            ax = axes[img_idx][step_idx] if num_images > 1 else axes[step_idx]
+            if (
+                isinstance(img, np.ndarray) and img.shape[0] == 3
+            ):  # Assuming 3 channels (RGB)
+                img = img.transpose(1, 2, 0)
+            ax.imshow(img)
+            ax.axis("off")
+            if img_idx == 0:
+                ax.set_title(f"t={t_span[step_idx]/t_span[-1]:.2f}")
+    plt.tight_layout()
+    return fig
+
+
+def plot_growth(dataset, growth_nets, xs, output_file='plot.pdf'):
+    x0s = [dataset["x0"][0]]
+    w0s = [dataset["x0"][1]]
+    x1s_list = [[dataset["x1_1"][0]], [dataset["x1_2"][0]]]
+    w1s_list = [[dataset["x1_1"][1]], [dataset["x1_2"][1]]] 
+    
+    
+    
+    plt.show()