File size: 11,244 Bytes
f15d29e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 | # Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.
import torch
from omegaconf import DictConfig
from ...diffusion.corruption.corruption import B, BatchedData, maybe_expand
from ...diffusion.corruption.sde_lib import SDE as DiffSDE
from ...diffusion.corruption.sde_lib import VESDE as DiffVESDE
from ...diffusion.corruption.sde_lib import VPSDE
from ...diffusion.wrapped.wrapped_sde import WrappedVESDE
def expand(a, x_shape, left=False):
a_dim = len(a.shape)
if left:
return a.reshape(*(((1,) * (len(x_shape) - a_dim)) + a.shape))
else:
return a.reshape(*(a.shape + ((1,) * (len(x_shape) - a_dim))))
def make_noise_symmetric_preserve_variance(noise: torch.Tensor) -> torch.Tensor:
"""Makes the noise matrix symmetric, preserving the variance. Assumes i.i.d. noise for each dimension.
Args:
noise (torch.Tensor): Input noise matrix, must be a batched square matrix, i.e., have shape (batch_size, dim, dim).
Returns:
torch.Tensor: The symmetric noise matrix, with the same variance as the input.
"""
assert (
len(noise.shape) == 3 and noise.shape[1] == noise.shape[2]
), "Symmetric noise only works for square-matrix-shaped data."
# Var[1/sqrt(2) * (eps_i + eps_j)] = 0.5 Var[eps_i] + 0.5 Var[eps_j] = Var[noise]
# Special treatment of the diagonal elements, i.e., those we leave unchanged via masking.
return (1 / (2**0.5)) * (1 - torch.eye(3, device=noise.device)[None]) * (
noise + noise.transpose(1, 2)
) + torch.eye(3, device=noise.device)[None] * noise
class LatticeVPSDE(VPSDE):
@staticmethod
def from_vpsde_config(vpsde_config: DictConfig):
return LatticeVPSDE(
**vpsde_config,
)
def __init__(
self,
beta_min: float = 0.1,
beta_max: float = 20,
limit_density: float | None = 0.05,
limit_var_scaling_constant: float = 0.25,
**kwargs,
):
"""Variance-preserving SDE with drift coefficient changing linearly over time."""
super().__init__()
self.beta_0 = beta_min
self.beta_1 = beta_max
# each crystal is diffused to have expected lattice vectors
# based on the number of atoms per crystal and self.limit_density
# units=(atoms/Angstrom**3)
self.limit_density = limit_density
self.limit_var_scaling_constant = limit_var_scaling_constant
self._limit_info_key = "num_atoms"
@property
def limit_info_key(self) -> str:
return self._limit_info_key
def beta(self, t: torch.Tensor) -> torch.Tensor:
return self.beta_0 + t * (self.beta_1 - self.beta_0)
def _marginal_mean_coeff(self, t: torch.Tensor) -> torch.Tensor:
log_mean_coeff = -0.25 * t**2 * (self.beta_1 - self.beta_0) - 0.5 * t * self.beta_0
return torch.exp(log_mean_coeff)
def marginal_prob(
self,
x: torch.Tensor,
t: torch.Tensor,
batch_idx: B = None,
batch: BatchedData | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
assert batch is not None
mean_coeff = self._marginal_mean_coeff(t)
# x: shape [batch_size, *x.shape[1:]]
# t, limit_info: shape [batch_size,]
limit_mean = self.get_limit_mean(x=x, batch=batch)
limit_var = self.get_limit_var(x=x, batch=batch)
mean_coeff_expanded = maybe_expand(mean_coeff, batch_idx, x)
mean = mean_coeff_expanded * x + (1 - mean_coeff_expanded) * limit_mean
std = torch.sqrt((1.0 - mean_coeff_expanded**2) * limit_var)
return mean, std
def mean_coeff_and_std(
self,
x: torch.Tensor,
t: torch.Tensor,
batch_idx: B = None,
batch: BatchedData | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Returns mean coefficient and standard deviation of marginal distribution at time t."""
mean_coeff = self._marginal_mean_coeff(t)
std = self.marginal_prob(x, t, batch_idx, batch)[1]
return maybe_expand(mean_coeff, batch=None, like=x), std
def get_limit_mean(self, x: torch.Tensor, batch: BatchedData) -> torch.Tensor:
# x: shape [batch_size, *x.shape[1:]]
# limit_info: shape [batch_size,], a 1d tensor containing number of atoms per crystal
# self.limit_density = limit_info / mean lattice vector length**3
# shape=[Ncrystals,]
n_atoms = batch[self.limit_info_key]
# shape=[Ncrystals, 3, 3]
return torch.pow(
torch.eye(3, device=x.device).expand(len(n_atoms), 3, 3)
* n_atoms[:, None, None]
/ self.limit_density,
1.0 / 3,
).to(x.device)
def get_limit_var(self, x: torch.Tensor, batch: BatchedData) -> torch.Tensor:
"""
Returns the element-wise variance of the limit distribution.
NOTE: even though we have a different limit variance per data
dimension we still sample IID for each element per data point.
We do NOT do any correlated sampling over data dimensions per
data point.
Return shape=x.shape
"""
# x: shape [batch_size, *x.shape[1:]]
# limit_info: shape [batch_size,]
# necessary for mypy
n_atoms = batch[self.limit_info_key]
# expand to fit shape of data, shape = (n_crystals, 1, 1)
n_atoms_expanded = expand(n_atoms, x.shape)
# shape = (n_crystals, 3, 3)
n_atoms_expanded = torch.tile(n_atoms_expanded, (1, 3, 3))
# scale limit standard deviation to be proportional to number atoms = n_atoms**(1/3)
# per lattice vector. We hope that prod_i std_i scales as the standard deviation
# of the actual volume. NOTE: we return variance here, hence 2 in the power
# shape=(Ncrystals, 3, 3) for limit_info.shape=[Ncrystals,]
out = torch.pow(n_atoms_expanded, 2.0 / 3).to(x.device) * self.limit_var_scaling_constant
return out
def sample_marginal(
self,
x: torch.Tensor,
t: torch.Tensor,
batch_idx: B = None,
batch: BatchedData | None = None,
) -> torch.Tensor:
mean, std = self.marginal_prob(x=x, t=t, batch=batch)
z = torch.randn_like(x)
z = make_noise_symmetric_preserve_variance(z)
return mean + expand(std, z.shape) * z
def prior_sampling(
self,
shape: torch.Size | tuple,
conditioning_data: BatchedData | None = None,
batch_idx: B = None,
) -> torch.Tensor:
x_sample = torch.randn(*shape)
x_sample = make_noise_symmetric_preserve_variance(x_sample)
assert conditioning_data is not None
limit_info = conditioning_data[self.limit_info_key]
x_sample = x_sample.to(limit_info.device)
limit_mean = self.get_limit_mean(x=x_sample, batch=conditioning_data)
limit_var = self.get_limit_var(x=x_sample, batch=conditioning_data)
# shape=[Nbatch,...] for shape[0]=Nbatch
return x_sample * limit_var.sqrt() + limit_mean
def sde(
self,
x: torch.Tensor,
t: torch.Tensor,
batch_idx: B = None,
batch: BatchedData | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
assert batch is not None
# x: shape [batch_size, *x.shape[1:]]
# t, limit_info: shape [batch_size,]
# if a per data-point limit mean is supplied, expand to shape of data
# shape=x.shape
limit_mean = self.get_limit_mean(x=x, batch=batch)
# if a per data-point limit variance is supplied, shape=[x.shape[0], ]
limit_var = self.get_limit_var(x=x, batch=batch)
beta_t = self.beta(t)
drift = (
-0.5
* expand(
beta_t,
x.shape,
)
* (x - limit_mean)
)
diffusion = torch.sqrt(expand(beta_t, limit_var.shape) * limit_var)
# drift.shape=[Nbatch,...], diffusion.shape=[Nbatch,] for x.shape[0]=Nbatch
return maybe_expand(drift, batch_idx), maybe_expand(diffusion, batch_idx)
class NumAtomsVarianceAdjustedWrappedVESDE(WrappedVESDE):
"""Wrapped VESDE with variance adjusted by number of atoms. We divide the standard deviation by the cubic root of the number of atoms.
The goal is to reduce the influence by the cell size on the variance of the fractional coordinates.
"""
def __init__(
self,
wrapping_boundary: float | torch.Tensor = 1.0,
sigma_min: float = 0.01,
sigma_max: float = 5.0,
limit_info_key: str = "num_atoms",
):
super().__init__(
sigma_min=sigma_min, sigma_max=sigma_max, wrapping_boundary=wrapping_boundary
)
self.limit_info_key = limit_info_key
def std_scaling(self, batch: BatchedData) -> torch.Tensor:
return batch[self.limit_info_key] ** (-1 / 3)
def marginal_prob(
self,
x: torch.Tensor,
t: torch.Tensor,
batch_idx: B = None,
batch: BatchedData | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
mean, std = super().marginal_prob(x, t, batch_idx, batch)
assert (
batch is not None
), "batch must be provided when using NumAtomsVarianceAdjustedWrappedVESDEMixin"
std_scale = self.std_scaling(batch)
std = std * maybe_expand(std_scale, batch_idx, like=std)
return mean, std
def prior_sampling(
self,
shape: torch.Size | tuple,
conditioning_data: BatchedData | None = None,
batch_idx=None,
) -> torch.Tensor:
_super = super()
assert isinstance(self, DiffSDE) and hasattr(_super, "prior_sampling")
assert (
conditioning_data is not None
), "batch must be provided when using NumAtomsVarianceAdjustedWrappedVESDEMixin"
num_atoms = conditioning_data[self.limit_info_key]
batch_idx = torch.repeat_interleave(
torch.arange(num_atoms.shape[0], device=num_atoms.device), num_atoms, dim=0
)
std_scale = self.std_scaling(conditioning_data)
# prior sample is randn() * sigma_max, so we need additionally multiply by std_scale to get the correct variance.
# We call VESDE.prior_sampling (a "grandparent" function) because the super() prior_sampling already does the wrapping,
# which means we couldn't do the variance adjustment here anymore otherwise.
prior_sample = DiffVESDE.prior_sampling(self, shape=shape).to(num_atoms.device)
return self.wrap(prior_sample * maybe_expand(std_scale, batch_idx, like=prior_sample))
def sde(
self,
x: torch.Tensor,
t: torch.Tensor,
batch_idx: B = None,
batch: BatchedData | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
sigma = self.marginal_prob(x, t, batch_idx, batch)[1]
sigma_min = self.marginal_prob(x, torch.zeros_like(t), batch_idx, batch)[1]
sigma_max = self.marginal_prob(x, torch.ones_like(t), batch_idx, batch)[1]
drift = torch.zeros_like(x)
diffusion = sigma * torch.sqrt(2 * (sigma_max.log() - sigma_min.log()))
return drift, diffusion
|