torch-molecule 0.1.0.dev2__py3-none-any.whl

torch_molecule/generator/gdss/sde.py

@@ -0,0 +1,334 @@
+import abc
+import torch
+import numpy as np
+
+def get_score_fn(sde, model, train=True, continuous=True):
+    if not train:
+        model.eval()
+    model_fn = model
+    if isinstance(sde, VPSDE) or isinstance(sde, subVPSDE):
+        def score_fn(x, adj, flags, t):
+            # Scale neural network output by standard deviation and flip sign
+            if continuous:
+                score = model_fn(x, adj, flags)
+                std = sde.marginal_prob(torch.zeros_like(adj), t)[1]
+            else:
+                raise NotImplementedError(f"Discrete not supported")
+            score = -score / std[:, None, None]
+            return score
+
+    elif isinstance(sde, VESDE):
+        def score_fn(x, adj, flags, t):
+            if continuous:
+                score = model_fn(x, adj, flags)
+            else:
+                raise NotImplementedError(f"Discrete not supported")
+            return score
+    else:
+        raise NotImplementedError(f"SDE class {sde.__class__.__name__} not supported.")
+    return score_fn
+
+def load_sde(sde_type, beta_min, beta_max, num_scales):
+    if sde_type == "VP":
+        sde = VPSDE(beta_min=beta_min, beta_max=beta_max, N=num_scales)
+    elif sde_type == "VE":
+        sde = VESDE(sigma_min=beta_min, sigma_max=beta_max, N=num_scales)
+    elif sde_type == "subVP":
+        sde = subVPSDE(beta_min=beta_min, beta_max=beta_max, N=num_scales)
+    else:
+        raise NotImplementedError(f"SDE class {sde_type} not yet supported.")
+    return sde
+
+class SDE(abc.ABC):
+    """SDE abstract class. Functions are designed for a mini-batch of inputs."""
+
+    def __init__(self, N):
+        """Construct an SDE.
+        Args:
+          N: number of discretization time steps.
+        """
+        super().__init__()
+        self.N = N
+
+    @property
+    @abc.abstractmethod
+    def T(self):
+        """End time of the SDE."""
+        pass
+
+    @abc.abstractmethod
+    def sde(self, x, t):
+        pass
+
+    @abc.abstractmethod
+    def marginal_prob(self, x, t):
+        """Parameters to determine the marginal distribution of the SDE, $p_t(x)$."""
+        pass
+
+    @abc.abstractmethod
+    def prior_sampling(self, shape):
+        """Generate one sample from the prior distribution, $p_T(x)$."""
+        pass
+
+    @abc.abstractmethod
+    def prior_logp(self, z):
+        """Compute log-density of the prior distribution.
+        Useful for computing the log-likelihood via probability flow ODE.
+        Args:
+          z: latent code
+        Returns:
+          log probability density
+        """
+        pass
+
+    def discretize(self, x, t):
+        """Discretize the SDE in the form: x_{i+1} = x_i + f_i(x_i) + G_i z_i.
+        Useful for reverse diffusion sampling and probabiliy flow sampling.
+        Defaults to Euler-Maruyama discretization.
+        Args:
+          x: a torch tensor
+          t: a torch float representing the time step (from 0 to `self.T`)
+        Returns:
+          f, G
+        """
+        dt = 1 / self.N
+        drift, diffusion = self.sde(x, t)
+        f = drift * dt
+        G = diffusion * torch.sqrt(torch.tensor(dt, device=t.device))
+        return f, G
+
+    def reverse(self, score_fn, probability_flow=False):
+        """Create the reverse-time SDE/ODE.
+        Args:
+          score_fn: A time-dependent score-based model that takes x and t and returns the score.
+          probability_flow: If `True`, create the reverse-time ODE used for probability flow sampling.
+        """
+        N = self.N
+        T = self.T
+        sde_fn = self.sde
+        discretize_fn = self.discretize
+
+        # -------- Build the class for reverse-time SDE --------
+        class RSDE(self.__class__):
+            def __init__(self):
+                self.N = N
+                self.probability_flow = probability_flow
+
+            @property
+            def T(self):
+                return T
+
+            def sde(self, feature, x, flags, t, is_adj=True):
+                """Create the drift and diffusion functions for the reverse SDE/ODE."""
+                drift, diffusion = sde_fn(x, t) if is_adj else sde_fn(feature, t)
+                score = score_fn(feature, x, flags, t)
+                drift = drift - diffusion[:, None, None] ** 2 * score * (
+                    0.5 if self.probability_flow else 1.0
+                )
+                # -------- Set the diffusion function to zero for ODEs. --------
+                diffusion = 0.0 if self.probability_flow else diffusion
+                return drift, diffusion
+
+            def discretize(self, feature, x, flags, t, is_adj=True):
+                """Create discretized iteration rules for the reverse diffusion sampler."""
+                f, G = discretize_fn(x, t) if is_adj else discretize_fn(feature, t)
+                score = score_fn(feature, x, flags, t)
+                rev_f = f - G[:, None, None] ** 2 * score * (
+                    0.5 if self.probability_flow else 1.0
+                )
+                rev_G = torch.zeros_like(G) if self.probability_flow else G
+                return rev_f, rev_G
+
+        return RSDE()
+
+class VPSDE(SDE):
+    def __init__(self, beta_min=0.1, beta_max=20, N=1000):
+        """Construct a Variance Preserving SDE.
+        Args:
+          beta_min: value of beta(0)
+          beta_max: value of beta(1)
+          N: number of discretization steps
+        """
+        super().__init__(N)
+        self.beta_0 = beta_min
+        self.beta_1 = beta_max
+        self.N = N
+        self.discrete_betas = torch.linspace(beta_min / N, beta_max / N, N)
+        self.alphas = 1.0 - self.discrete_betas
+        self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
+        self.sqrt_alphas_cumprod = torch.sqrt(self.alphas_cumprod)
+        self.sqrt_1m_alphas_cumprod = torch.sqrt(1.0 - self.alphas_cumprod)
+
+    @property
+    def T(self):
+        return 1
+
+    def sde(self, x, t):
+        beta_t = self.beta_0 + t * (self.beta_1 - self.beta_0)
+        drift = -0.5 * beta_t[:, None, None] * x
+        diffusion = torch.sqrt(beta_t)
+        return drift, diffusion
+
+    # -------- mean, std of the perturbation kernel --------
+    def marginal_prob(self, x, t):
+        log_mean_coeff = (
+            -0.25 * t**2 * (self.beta_1 - self.beta_0) - 0.5 * t * self.beta_0
+        )
+        mean = torch.exp(log_mean_coeff[:, None, None]) * x
+        std = torch.sqrt(1.0 - torch.exp(2.0 * log_mean_coeff))
+        return mean, std
+
+    def prior_sampling(self, shape):
+        return torch.randn(*shape)
+
+    def prior_sampling_sym(self, shape):
+        x = torch.randn(*shape).triu(1)
+        return x + x.transpose(-1, -2)
+
+    def prior_logp(self, z):
+        shape = z.shape
+        N = np.prod(shape[1:])
+        logps = -N / 2.0 * np.log(2 * np.pi) - torch.sum(z**2, dim=(1, 2)) / 2.0
+        return logps
+
+    def discretize(self, x, t):
+        """DDPM discretization."""
+        timestep = (t * (self.N - 1) / self.T).long()
+        beta = self.discrete_betas.to(x.device)[timestep]
+        alpha = self.alphas.to(x.device)[timestep]
+        sqrt_beta = torch.sqrt(beta)
+        f = torch.sqrt(alpha)[:, None, None] * x - x
+        G = sqrt_beta
+        return f, G
+
+    def transition(self, x, t, dt):
+        # -------- negative timestep dt --------
+        log_mean_coeff = (
+            0.25 * dt * (2 * self.beta_0 + (2 * t + dt) * (self.beta_1 - self.beta_0))
+        )
+        mean = torch.exp(-log_mean_coeff[:, None, None]) * x
+        std = torch.sqrt(1.0 - torch.exp(2.0 * log_mean_coeff))
+        return mean, std
+
+
+class VESDE(SDE):
+    def __init__(self, sigma_min=0.01, sigma_max=50, N=1000):
+        """Construct a Variance Exploding SDE.
+        Args:
+          sigma_min: smallest sigma.
+          sigma_max: largest sigma.
+          N: number of discretization steps
+        """
+        super().__init__(N)
+        self.sigma_min = sigma_min
+        self.sigma_max = sigma_max
+        self.discrete_sigmas = torch.exp(
+            torch.linspace(np.log(self.sigma_min), np.log(self.sigma_max), N)
+        )
+        self.N = N
+
+    @property
+    def T(self):
+        return 1
+
+    def sde(self, x, t):
+        sigma = self.sigma_min * (self.sigma_max / self.sigma_min) ** t
+        drift = torch.zeros_like(x)
+        diffusion = sigma * torch.sqrt(
+            torch.tensor(
+                2 * (np.log(self.sigma_max) - np.log(self.sigma_min)), device=t.device
+            )
+        )
+        return drift, diffusion
+
+    def marginal_prob(self, x, t):
+        std = self.sigma_min * (self.sigma_max / self.sigma_min) ** t
+        mean = x
+        return mean, std
+
+    def prior_sampling(self, shape):
+        return torch.randn(*shape)
+
+    def prior_sampling_sym(self, shape):
+        x = torch.randn(*shape).triu(1)
+        x = x + x.transpose(-1, -2)
+        return x
+
+    def prior_logp(self, z):
+        shape = z.shape
+        N = np.prod(shape[1:])
+        return -N / 2.0 * np.log(2 * np.pi * self.sigma_max**2) - torch.sum(
+            z**2, dim=(1, 2, 3)
+        ) / (2 * self.sigma_max**2)
+
+    def discretize(self, x, t):
+        """SMLD(NCSN) discretization."""
+        timestep = (t * (self.N - 1) / self.T).long()
+        self.discrete_sigmas = self.discrete_sigmas.to(t.device)
+        sigma = self.discrete_sigmas.to(t.device)[timestep]
+        adjacent_sigma = torch.where(
+            timestep == 0,
+            torch.zeros_like(t),
+            self.discrete_sigmas[timestep - 1].to(t.device),
+        )
+        f = torch.zeros_like(x)
+        G = torch.sqrt(sigma**2 - adjacent_sigma**2)
+        return f, G
+
+    def transition(self, x, t, dt):
+        # -------- negative timestep dt --------
+        std = torch.square(
+            self.sigma_min * (self.sigma_max / self.sigma_min) ** t
+        ) - torch.square(self.sigma_min * (self.sigma_max / self.sigma_min) ** (t + dt))
+        std = torch.sqrt(std)
+        mean = x
+        return mean, std
+
+
+class subVPSDE(SDE):
+    def __init__(self, beta_min=0.1, beta_max=20, N=1000):
+        """Construct the sub-VP SDE that excels at likelihoods.
+        Args:
+          beta_min: value of beta(0)
+          beta_max: value of beta(1)
+          N: number of discretization steps
+        """
+        super().__init__(N)
+        self.beta_0 = beta_min
+        self.beta_1 = beta_max
+        self.N = N
+        self.discrete_betas = torch.linspace(beta_min / N, beta_max / N, N)
+        self.alphas = 1.0 - self.discrete_betas
+
+    @property
+    def T(self):
+        return 1
+
+    def sde(self, x, t):
+        beta_t = self.beta_0 + t * (self.beta_1 - self.beta_0)
+        drift = -0.5 * beta_t[:, None, None] * x
+        discount = 1.0 - torch.exp(
+            -2 * self.beta_0 * t - (self.beta_1 - self.beta_0) * t**2
+        )
+        diffusion = torch.sqrt(beta_t * discount)
+        return drift, diffusion
+
+    def marginal_prob(self, x, t):
+        log_mean_coeff = (
+            -0.25 * t**2 * (self.beta_1 - self.beta_0) - 0.5 * t * self.beta_0
+        )
+        mean = torch.exp(log_mean_coeff)[:, None, None] * x
+        std = 1 - torch.exp(2.0 * log_mean_coeff)
+        return mean, std
+
+    def prior_sampling(self, shape):
+        return torch.randn(*shape)
+
+    def prior_sampling_sym(self, shape):
+        x = torch.randn(*shape).triu(1)
+        return x + x.transpose(-1, -2)
+
+    def prior_logp(self, z):
+        shape = z.shape
+        N = np.prod(shape[1:])
+        return -N / 2.0 * np.log(2 * np.pi) - torch.sum(z**2, dim=(1, 2, 3)) / 2.0

torch_molecule/generator/gdss/solver.py

@@ -0,0 +1,227 @@
+import torch
+import numpy as np
+import abc
+
+from .utils import mask_adjs, mask_x, gen_noise
+from .sde import VPSDE, subVPSDE, get_score_fn
+
+### sampling
+# -------- PC sampler --------
+def get_pc_sampler(
+    sde_x,
+    sde_adj,
+    predictor="Euler",
+    corrector="None",
+    snr=0.1,
+    scale_eps=1.0,
+    n_steps=1,
+    probability_flow=False,
+    continuous=False,
+    denoise=True,
+    eps=1e-3,
+    device="cuda",
+):
+    def pc_sampler(model_x, model_adj, shape_x, shape_adj, init_flags):
+        score_fn_x = get_score_fn(sde_x, model_x, train=False, continuous=continuous)
+        score_fn_adj = get_score_fn(
+            sde_adj, model_adj, train=False, continuous=continuous
+        )
+
+        predictor_fn = (
+            ReverseDiffusionPredictor
+            if predictor == "Reverse"
+            else EulerMaruyamaPredictor
+        )
+        corrector_fn = LangevinCorrector if corrector == "Langevin" else NoneCorrector
+
+        predictor_obj_x = predictor_fn("x", sde_x, score_fn_x, probability_flow)
+        corrector_obj_x = corrector_fn(
+            "x", sde_x, score_fn_x, snr, scale_eps, n_steps
+        )
+
+        predictor_obj_adj = predictor_fn(
+            "adj", sde_adj, score_fn_adj, probability_flow
+        )
+        corrector_obj_adj = corrector_fn(
+            "adj", sde_adj, score_fn_adj, snr, scale_eps, n_steps
+        )
+
+        with torch.no_grad():
+            # -------- Initial sample --------
+            x = sde_x.prior_sampling(shape_x).to(device)
+            adj = sde_adj.prior_sampling_sym(shape_adj).to(device)
+            flags = init_flags
+            x = mask_x(x, flags)
+            adj = mask_adjs(adj, flags)
+            diff_steps = sde_adj.N
+            timesteps = torch.linspace(sde_adj.T, eps, diff_steps, device=device)
+
+            # -------- Reverse diffusion process --------
+            # for i in trange(
+            #     0, (diff_steps), desc="[Sampling]", position=1, leave=False
+            # ):
+            for i in range(0, diff_steps):
+                t = timesteps[i]
+                vec_t = torch.ones(shape_adj[0], device=t.device) * t
+
+                _x = x
+                x, x_mean = corrector_obj_x.update_fn(x, adj, flags, vec_t)
+                adj, adj_mean = corrector_obj_adj.update_fn(_x, adj, flags, vec_t)
+
+                _x = x
+                x, x_mean = predictor_obj_x.update_fn(x, adj, flags, vec_t)
+                adj, adj_mean = predictor_obj_adj.update_fn(_x, adj, flags, vec_t)
+            # print(" ")
+            return (
+                (x_mean if denoise else x),
+                (adj_mean if denoise else adj),
+                diff_steps * (n_steps + 1),
+            )
+
+    return pc_sampler
+
+
+class Predictor(abc.ABC):
+    """The abstract class for a predictor algorithm."""
+
+    def __init__(self, sde, score_fn, probability_flow=False):
+        super().__init__()
+        self.sde = sde
+        # Compute the reverse SDE/ODE
+        self.rsde = sde.reverse(score_fn, probability_flow)
+        self.score_fn = score_fn
+
+    @abc.abstractmethod
+    def update_fn(self, x, t, flags):
+        pass
+
+
+class Corrector(abc.ABC):
+    """The abstract class for a corrector algorithm."""
+
+    def __init__(self, sde, score_fn, snr, scale_eps, n_steps):
+        super().__init__()
+        self.sde = sde
+        self.score_fn = score_fn
+        self.snr = snr
+        self.scale_eps = scale_eps
+        self.n_steps = n_steps
+
+    @abc.abstractmethod
+    def update_fn(self, x, t, flags):
+        pass
+
+
+class EulerMaruyamaPredictor(Predictor):
+    def __init__(self, obj, sde, score_fn, probability_flow=False):
+        super().__init__(sde, score_fn, probability_flow)
+        self.obj = obj
+
+    def update_fn(self, x, adj, flags, t):
+        dt = -1.0 / self.rsde.N
+
+        if self.obj == "x":
+            z = gen_noise(x, flags, sym=False)
+            drift, diffusion = self.rsde.sde(x, adj, flags, t, is_adj=False)
+            x_mean = x + drift * dt
+            x = x_mean + diffusion[:, None, None] * np.sqrt(-dt) * z
+            return x, x_mean
+
+        elif self.obj == "adj":
+            z = gen_noise(adj, flags)
+            drift, diffusion = self.rsde.sde(x, adj, flags, t, is_adj=True)
+            adj_mean = adj + drift * dt
+            adj = adj_mean + diffusion[:, None, None] * np.sqrt(-dt) * z
+
+            return adj, adj_mean
+
+        else:
+            raise NotImplementedError(f"obj {self.obj} not yet supported.")
+
+
+class ReverseDiffusionPredictor(Predictor):
+    def __init__(self, obj, sde, score_fn, probability_flow=False):
+        super().__init__(sde, score_fn, probability_flow)
+        self.obj = obj
+
+    def update_fn(self, x, adj, flags, t):
+        if self.obj == "x":
+            f, G = self.rsde.discretize(x, adj, flags, t, is_adj=False)
+            z = gen_noise(x, flags, sym=False)
+            x_mean = x - f
+            x = x_mean + G[:, None, None] * z
+            return x, x_mean
+
+        elif self.obj == "adj":
+            f, G = self.rsde.discretize(x, adj, flags, t, is_adj=True)
+            z = gen_noise(adj, flags)
+            adj_mean = adj - f
+            adj = adj_mean + G[:, None, None] * z
+            return adj, adj_mean
+
+        else:
+            raise NotImplementedError(f"obj {self.obj} not yet supported.")
+
+
+class NoneCorrector(Corrector):
+    """An empty corrector that does nothing."""
+
+    def __init__(self, obj, sde, score_fn, snr, scale_eps, n_steps):
+        self.obj = obj
+        pass
+
+    def update_fn(self, x, adj, flags, t):
+        if self.obj == "x":
+            return x, x
+        elif self.obj == "adj":
+            return adj, adj
+        else:
+            raise NotImplementedError(f"obj {self.obj} not yet supported.")
+
+
+class LangevinCorrector(Corrector):
+    def __init__(self, obj, sde, score_fn, snr, scale_eps, n_steps):
+        super().__init__(sde, score_fn, snr, scale_eps, n_steps)
+        self.obj = obj
+
+    def update_fn(self, x, adj, flags, t):
+        sde = self.sde
+        score_fn = self.score_fn
+        n_steps = self.n_steps
+        target_snr = self.snr
+        seps = self.scale_eps
+
+        if isinstance(sde, VPSDE) or isinstance(sde, subVPSDE):
+            timestep = (t * (sde.N - 1) / sde.T).long()
+            alpha = sde.alphas.to(t.device)[timestep]
+        else:
+            alpha = torch.ones_like(t)
+
+        if self.obj == "x":
+            for i in range(n_steps):
+                grad = score_fn(x, adj, flags, t)
+                noise = gen_noise(x, flags, sym=False)
+                grad_norm = torch.norm(grad.reshape(grad.shape[0], -1), dim=-1).mean()
+                noise_norm = torch.norm(
+                    noise.reshape(noise.shape[0], -1), dim=-1
+                ).mean()
+                step_size = (target_snr * noise_norm / grad_norm) ** 2 * 2 * alpha
+                x_mean = x + step_size[:, None, None] * grad
+                x = x_mean + torch.sqrt(step_size * 2)[:, None, None] * noise * seps
+            return x, x_mean
+
+        elif self.obj == "adj":
+            for i in range(n_steps):
+                grad = score_fn(x, adj, flags, t)
+                noise = gen_noise(adj, flags)
+                grad_norm = torch.norm(grad.reshape(grad.shape[0], -1), dim=-1).mean()
+                noise_norm = torch.norm(
+                    noise.reshape(noise.shape[0], -1), dim=-1
+                ).mean()
+                step_size = (target_snr * noise_norm / grad_norm) ** 2 * 2 * alpha
+                adj_mean = adj + step_size[:, None, None] * grad
+                adj = adj_mean + torch.sqrt(step_size * 2)[:, None, None] * noise * seps
+            return adj, adj_mean
+
+        else:
+            raise NotImplementedError(f"obj {self.obj} not yet supported")