| | import sys |
| | 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 |
| | |
| | |
| | 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) |
| | |
| | |
| | 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 |
| | ) |
| | ) |
