diff --git a/.gitignore b/.gitignore index e4dd3a77..31009b0f 100644 --- a/.gitignore +++ b/.gitignore @@ -180,3 +180,13 @@ electronic_structure/configs/qm9_data_split.json electronic_structure/configs/qm9.json electronic_structure/configs/crystal_data_split.json electronic_structure/configs/crystal.json + +# ignore export test data +*.npz +# ignore cif crystal texture +*.cif +# ingore paddle weight file. +*.pdparams +.claude/ + +*.pkl \ No newline at end of file diff --git a/ppmat/models/__init__.py b/ppmat/models/__init__.py index 95d73232..137f25ae 100644 --- a/ppmat/models/__init__.py +++ b/ppmat/models/__init__.py @@ -42,6 +42,12 @@ from ppmat.models.megnet.megnet import MEGNetPlus from ppmat.models.infgcn.infgcn import InfGCN from ppmat.models.mateno.mateno import MatENO +from ppmat.models.chemeleon2 import VAEModule +from ppmat.models.chemeleon2 import LDMModule +from ppmat.models.chemeleon2 import RLModule +from ppmat.models.chemeleon2.ldm_module.dit import DiT +from ppmat.models.chemeleon2.vae_module.encoder import TransformerEncoder +from ppmat.models.chemeleon2.vae_module.decoder import TransformerDecoder from ppmat.utils import download from ppmat.utils import logger from ppmat.utils import save_load @@ -67,6 +73,12 @@ "DiffNMR", "InfGCN", "MatENO", + "VAEModule", + "LDMModule", + "RLModule", + "DiT", + "TransformerEncoder", + "TransformerDecoder", ] # Warning: The key of the dictionary must be consistent with the file name of the value diff --git a/ppmat/models/chemeleon2/__init__.py b/ppmat/models/chemeleon2/__init__.py new file mode 100644 index 00000000..bbc131d5 --- /dev/null +++ b/ppmat/models/chemeleon2/__init__.py @@ -0,0 +1,9 @@ +from ppmat.models.chemeleon2.vae_module.vae import VAEModule +from ppmat.models.chemeleon2.ldm_module.ldm import LDMModule +from ppmat.models.chemeleon2.rl_module.rl import RLModule + +__all__ = [ + "VAEModule", + "LDMModule", + "RLModule", +] diff --git a/ppmat/models/chemeleon2/common/__init__.py b/ppmat/models/chemeleon2/common/__init__.py new file mode 100644 index 00000000..a2a1cbb1 --- /dev/null +++ b/ppmat/models/chemeleon2/common/__init__.py @@ -0,0 +1,31 @@ +from ppmat.models.chemeleon2.common.distributions import DiagonalGaussianDistribution +from ppmat.models.chemeleon2.common.schema import CrystalBatch +from ppmat.models.chemeleon2.common.scatter import scatter_mean +from ppmat.models.chemeleon2.common.scatter import scatter_sum +from ppmat.models.chemeleon2.common.scatter import scatter_std +from ppmat.models.chemeleon2.common.lattice_utils import lattice_params_to_matrix +from ppmat.models.chemeleon2.common.lattice_utils import matrix_to_lattice_params +from ppmat.models.chemeleon2.common.lattice_utils import frac_to_cart_coords +from ppmat.models.chemeleon2.common.lattice_utils import cart_to_frac_coords +from ppmat.models.chemeleon2.common.lattice_utils import get_pbc_distances +from ppmat.models.chemeleon2.common.lattice_utils import lattice_vector_to_volume +from ppmat.models.chemeleon2.common.batch_utils import to_dense_batch +from ppmat.models.chemeleon2.common.data_augmentation import apply_augmentation +from ppmat.models.chemeleon2.common.data_augmentation import apply_noise + +__all__ = [ + "DiagonalGaussianDistribution", + "CrystalBatch", + "scatter_mean", + "scatter_sum", + "scatter_std", + "lattice_params_to_matrix", + "matrix_to_lattice_params", + "frac_to_cart_coords", + "cart_to_frac_coords", + "get_pbc_distances", + "lattice_vector_to_volume", + "to_dense_batch", + "apply_augmentation", + "apply_noise", +] diff --git a/ppmat/models/chemeleon2/common/batch_utils.py b/ppmat/models/chemeleon2/common/batch_utils.py new file mode 100644 index 00000000..8c3ba631 --- /dev/null +++ b/ppmat/models/chemeleon2/common/batch_utils.py @@ -0,0 +1,41 @@ +import paddle + + +def to_dense_batch(x, batch_idx, max_num_nodes=None): + """ + 将 batch 格式的数据转换为 dense batch 格式，并生成 padding mask + + Args: + x: 特征张量 [N, D]，N 是总原子数，D 是特征维度 + batch_idx: batch 索引 [N]，指示每个原子属于哪个结构 + max_num_nodes: 最大节点数，如果为 None 则自动计算 + + Returns: + x_dense: dense 格式的特征张量 [B, max_num_nodes, D] + mask: padding mask [B, max_num_nodes]，True 表示有效位置，False 表示 padding + """ + batch_size = int(batch_idx.max().item()) + 1 if batch_idx.numel() > 0 else 1 + + num_nodes = paddle.zeros([batch_size], dtype='int64') + for i in range(batch_size): + num_nodes[i] = (batch_idx == i).sum() + + if max_num_nodes is None: + max_num_nodes = int(num_nodes.max().item()) + + feat_dim = x.shape[-1] + x_dense = paddle.zeros([batch_size, max_num_nodes, feat_dim], dtype=x.dtype) + mask = paddle.zeros([batch_size, max_num_nodes], dtype='bool') + + cumsum = paddle.concat([paddle.zeros([1], dtype='int64'), + paddle.cumsum(num_nodes, axis=0)[:-1]]) + + for i in range(batch_size): + start = int(cumsum[i].item()) + end = start + int(num_nodes[i].item()) + n = int(num_nodes[i].item()) + x_dense[i, :n] = x[start:end] + mask[i, :n] = True + + return x_dense, mask + diff --git a/ppmat/models/chemeleon2/common/data_augmentation.py b/ppmat/models/chemeleon2/common/data_augmentation.py new file mode 100644 index 00000000..f1a3eeb5 --- /dev/null +++ b/ppmat/models/chemeleon2/common/data_augmentation.py @@ -0,0 +1,107 @@ +import paddle + + +def apply_augmentation(batch, translate=False, rotate=False): + if not translate and not rotate: + return batch + + batch_aug = batch.clone() + + if translate: + batch_aug = _augmentation_translate(batch_aug) + + if rotate: + batch_aug = _augmentation_rotate(batch_aug) + + return batch_aug + + +def _augmentation_translate(batch): + lengths_mean = batch.lengths.mean(axis=0) + lengths_std = batch.lengths.std(axis=0, unbiased=False) + + random_translate = paddle.normal( + mean=paddle.abs(lengths_mean), + std=paddle.maximum(paddle.abs(lengths_std), paddle.to_tensor([1e-8])) + ) / 2 + + cart_coords_aug = batch.cart_coords + random_translate + + cell_per_node_inv = paddle.inverse(batch.lattices[batch.batch]) + frac_coords_aug = paddle.einsum('bi,bij->bj', cart_coords_aug, cell_per_node_inv) + frac_coords_aug = frac_coords_aug % 1.0 + + batch.cart_coords = cart_coords_aug + batch.frac_coords = frac_coords_aug + + return batch + + +def _augmentation_rotate(batch): + rot_mat = _random_rotation_matrix(validate=True) + + cart_coords_aug = paddle.matmul(batch.cart_coords, rot_mat.T) + lattices_aug = paddle.matmul(batch.lattices, rot_mat.T) + + batch.cart_coords = cart_coords_aug + batch.lattices = lattices_aug + + return batch + + +def apply_noise(batch, ratio=0.1, corruption_scale=0.1): + if ratio <= 0: + return batch + + batch_noise = batch.clone() + + total_num_atoms = batch_noise.num_nodes + noise_num_atoms = int(total_num_atoms * ratio) + + noise_atom_types = batch_noise.atom_types.clone() + noise_indices = paddle.randperm(total_num_atoms)[:noise_num_atoms] + noise_atom_types[noise_indices] = 0 + + noise_cart_coords = batch_noise.cart_coords.clone() + noise_indices = paddle.randperm(total_num_atoms)[:noise_num_atoms] + noise_cart_coords[noise_indices] += paddle.randn([noise_num_atoms, 3]) * corruption_scale + + cell_per_node_inv = paddle.inverse(batch.lattices[batch.batch]) + noise_frac_coords = paddle.einsum('bi,bij->bj', noise_cart_coords, cell_per_node_inv) + noise_frac_coords = noise_frac_coords % 1.0 + + batch_noise.atom_types = noise_atom_types + batch_noise.cart_coords = noise_cart_coords + batch_noise.frac_coords = noise_frac_coords + + return batch_noise + + +def _random_rotation_matrix(validate=False): + q = paddle.rand([4]) + q = q / paddle.norm(q) + + rot_mat = paddle.to_tensor([ + [ + 1 - 2 * q[2] ** 2 - 2 * q[3] ** 2, + 2 * q[1] * q[2] - 2 * q[0] * q[3], + 2 * q[1] * q[3] + 2 * q[0] * q[2], + ], + [ + 2 * q[1] * q[2] + 2 * q[0] * q[3], + 1 - 2 * q[1] ** 2 - 2 * q[3] ** 2, + 2 * q[2] * q[3] - 2 * q[0] * q[1], + ], + [ + 2 * q[1] * q[3] - 2 * q[0] * q[2], + 2 * q[2] * q[3] + 2 * q[0] * q[1], + 1 - 2 * q[1] ** 2 - 2 * q[2] ** 2, + ], + ], dtype='float32') + + if validate: + identity = paddle.matmul(rot_mat, rot_mat.T) + eye = paddle.eye(3) + assert paddle.allclose(identity, eye, atol=1e-5, rtol=1e-5), "Not a rotation matrix." + + return rot_mat diff --git a/ppmat/models/chemeleon2/common/distributions.py b/ppmat/models/chemeleon2/common/distributions.py new file mode 100644 index 00000000..16c47339 --- /dev/null +++ b/ppmat/models/chemeleon2/common/distributions.py @@ -0,0 +1,38 @@ +import paddle + + +class DiagonalGaussianDistribution: + def __init__(self, parameters): + self.parameters = parameters + self.mean, self.logvar = paddle.chunk(parameters, 2, axis=1) + self.logvar = paddle.clip(self.logvar, -30.0, 20.0) + self.std = paddle.exp(0.5 * self.logvar) + self.var = paddle.exp(self.logvar) + + def sample(self): + x = self.mean + self.std * paddle.randn(self.mean.shape) + return x + + def kl(self, other=None): + # Determine which axes to sum over based on tensor dimensionality + # For 4D tensors (images): sum over [1, 2, 3] + # For 2D tensors (latent vectors): sum over [1] + if self.mean.ndim == 4: + sum_axis = [1, 2, 3] + else: + sum_axis = [1] + + if other is None: + return 0.5 * paddle.sum( + paddle.pow(self.mean, 2) + self.var - 1.0 - self.logvar, + axis=sum_axis + ) + else: + return 0.5 * paddle.sum( + paddle.pow(self.mean - other.mean, 2) / other.var + + self.var / other.var - 1.0 - self.logvar + other.logvar, + axis=sum_axis + ) + + def mode(self): + return self.mean diff --git a/ppmat/models/chemeleon2/common/lattice_utils.py b/ppmat/models/chemeleon2/common/lattice_utils.py new file mode 100644 index 00000000..d67556df --- /dev/null +++ b/ppmat/models/chemeleon2/common/lattice_utils.py @@ -0,0 +1,67 @@ +import paddle +from ppmat.utils.crystal import lattice_params_to_matrix_paddle +from ppmat.utils.crystal import lattices_to_params_shape + + +def lattice_params_to_matrix(lengths, angles): + return lattice_params_to_matrix_paddle(lengths, angles) + + +def matrix_to_lattice_params(lattices): + return lattices_to_params_shape(lattices) + + +def frac_to_cart_coords(frac_coords, lattice): + if lattice.ndim == 2: + lattice = lattice.unsqueeze(0) + return paddle.einsum('ij,jk->ik', frac_coords, lattice.squeeze(0)) + + +def cart_to_frac_coords(cart_coords, lattice): + if lattice.ndim == 2: + lattice = lattice.unsqueeze(0) + inv_lattice = paddle.inverse(lattice) + return paddle.einsum('ij,jk->ik', cart_coords, inv_lattice.squeeze(0)) + + +def get_pbc_distances( + coords1, + coords2, + lattice, + num_atoms=None, + return_offsets=False, +): + if lattice.ndim == 2: + lattice = lattice.unsqueeze(0) + + if coords1.shape != coords2.shape: + raise ValueError("coords1 and coords2 must have the same shape") + + diff = coords2 - coords1 + + diff_frac = cart_to_frac_coords(diff, lattice) + + diff_frac = diff_frac - paddle.round(diff_frac) + + diff_cart = frac_to_cart_coords(diff_frac, lattice) + + distances = paddle.norm(diff_cart, axis=-1) + + if return_offsets: + offsets = paddle.round(cart_to_frac_coords(coords2 - coords1, lattice)) + return distances, offsets + + return distances + + +def lattice_vector_to_volume(lattice): + if lattice.ndim == 2: + lattice = lattice.unsqueeze(0) + + a = lattice[:, 0, :] + b = lattice[:, 1, :] + c = lattice[:, 2, :] + + volume = paddle.abs(paddle.sum(a * paddle.cross(b, c), axis=-1)) + + return volume diff --git a/ppmat/models/chemeleon2/common/lora.py b/ppmat/models/chemeleon2/common/lora.py new file mode 100644 index 00000000..ac74fd28 --- /dev/null +++ b/ppmat/models/chemeleon2/common/lora.py @@ -0,0 +1,120 @@ +import paddle.nn as nn + + +class LoRALayer(nn.Layer): + def __init__( + self, + in_features, + out_features, + rank=8, + alpha=16, + dropout=0.0, + ): + super().__init__() + self.rank = rank + self.alpha = alpha + self.scaling = alpha / rank + + self.lora_A = self.create_parameter( + shape=[in_features, rank], + default_initializer=nn.initializer.KaimingUniform() + ) + self.lora_B = self.create_parameter( + shape=[rank, out_features], + default_initializer=nn.initializer.Constant(0.0) + ) + + if dropout > 0.0: + self.dropout = nn.Dropout(dropout) + else: + self.dropout = None + + def forward(self, x): + if self.dropout is not None: + x = self.dropout(x) + return (x @ self.lora_A @ self.lora_B) * self.scaling + + +class LoRALinear(nn.Layer): + def __init__( + self, + original_layer, + rank=8, + alpha=16, + dropout=0.0, + ): + super().__init__() + self.original_layer = original_layer + self.original_layer.stop_gradient = True + + in_features = original_layer.weight.shape[0] + out_features = original_layer.weight.shape[1] + + self.lora = LoRALayer(in_features, out_features, rank, alpha, dropout) + + def forward(self, x): + result = self.original_layer(x) + result = result + self.lora(x) + return result + + +def apply_lora_to_linear(module, rank=8, alpha=16, dropout=0.0, target_modules=None): + if target_modules is None: + target_modules = ["q_proj", "v_proj", "k_proj", "o_proj", "fc1", "fc2"] + + for name, child in module.named_children(): + if isinstance(child, nn.Linear): + if any(target in name for target in target_modules): + lora_layer = LoRALinear(child, rank, alpha, dropout) + setattr(module, name, lora_layer) + else: + apply_lora_to_linear(child, rank, alpha, dropout, target_modules) + + return module + + +def get_lora_parameters(module): + lora_params = [] + for name, param in module.named_parameters(): + if 'lora' in name.lower(): + lora_params.append(param) + return lora_params + + +def merge_lora_weights(module): + for name, child in module.named_children(): + if isinstance(child, LoRALinear): + original_weight = child.original_layer.weight + lora_weight = child.lora.lora_A @ child.lora.lora_B * child.lora.scaling + merged_weight = original_weight + lora_weight + + merged_layer = nn.Linear( + original_weight.shape[0], + original_weight.shape[1], + bias_attr=child.original_layer.bias is not None + ) + merged_layer.weight.set_value(merged_weight) + if child.original_layer.bias is not None: + merged_layer.bias.set_value(child.original_layer.bias) + + setattr(module, name, merged_layer) + else: + merge_lora_weights(child) + + return module + + +def print_trainable_parameters(module): + trainable_params = 0 + all_params = 0 + + for param in module.parameters(): + all_params += param.numel() + if not param.stop_gradient: + trainable_params += param.numel() + + print( + f"trainable params: {trainable_params:,} || " + f"all params: {all_params:,} || " + f"trainable%: {100 * trainable_params / all_params:.2f}%" + ) diff --git a/ppmat/models/chemeleon2/common/scatter.py b/ppmat/models/chemeleon2/common/scatter.py new file mode 100644 index 00000000..21e33976 --- /dev/null +++ b/ppmat/models/chemeleon2/common/scatter.py @@ -0,0 +1,115 @@ +import paddle + + +def _broadcast(src, other, dim): + if dim < 0: + dim = other.dim() + dim + if src.dim() == 1: + for _ in range(0, dim): + src = src.unsqueeze(0) + for _ in range(src.dim(), other.dim()): + src = src.unsqueeze(-1) + src = paddle.expand_as(src, other) + return src + + +def scatter_sum( + src, + index, + dim=-1, + out=None, + dim_size=None, +): + index = _broadcast(index, src, dim) + if out is None: + size = list(src.shape) + if dim_size is not None: + size[dim] = dim_size + elif index.numel() == 0: + size[dim] = 0 + else: + size[dim] = int(index.max().item()) + 1 + out = paddle.zeros(size, dtype=src.dtype) + + if dim < 0: + dim = src.ndim + dim + + for i in range(out.shape[dim]): + mask = (index == i) + if dim == 0: + out[i] = paddle.where(mask, src, paddle.zeros_like(src)).sum(axis=0) + else: + masked_src = paddle.where(mask, src, paddle.zeros_like(src)) + if dim == 1: + out[:, i] = masked_src.sum(axis=dim) + elif dim == 2: + out[:, :, i] = masked_src.sum(axis=dim) + + return out + + +def scatter_mean( + src, + index, + dim=-1, + out=None, + dim_size=None, +): + out = scatter_sum(src, index, dim, out, dim_size) + dim_size = out.shape[dim] + + index_dim = dim + if index_dim < 0: + index_dim = index_dim + src.ndim + if index.ndim <= index_dim: + index_dim = index.ndim - 1 + + ones = paddle.ones(index.shape, dtype=src.dtype) + count = scatter_sum(ones, index, index_dim, None, dim_size) + count = paddle.where(count < 1, paddle.ones_like(count), count) + count = _broadcast(count, out, dim) + + if paddle.is_floating_point(out): + out = out / count + else: + out = out // count + + return out + + +def scatter_std( + src, + index, + dim=-1, + out=None, + dim_size=None, + unbiased=True, +): + if out is not None: + dim_size = out.shape[dim] + + if dim < 0: + dim = src.ndim + dim + + count_dim = dim + if index.ndim <= dim: + count_dim = index.ndim - 1 + + ones = paddle.ones(index.shape, dtype=src.dtype) + count = scatter_sum(ones, index, count_dim, dim_size=dim_size) + + index = _broadcast(index, src, dim) + tmp = scatter_sum(src, index, dim, dim_size=dim_size) + count_broadcast = _broadcast(count, tmp, dim) + count_broadcast = paddle.clip(count_broadcast, min=1) + mean = tmp / count_broadcast + + var = src - paddle.take_along_axis(mean, index, axis=dim) + var = var * var + out = scatter_sum(var, index, dim, out, dim_size) + + if unbiased: + count_broadcast = paddle.clip(count_broadcast - 1, min=1) + out = paddle.sqrt(out / (count_broadcast + 1e-6)) + + return out diff --git a/ppmat/models/chemeleon2/common/schema.py b/ppmat/models/chemeleon2/common/schema.py new file mode 100644 index 00000000..9b53dfd5 --- /dev/null +++ b/ppmat/models/chemeleon2/common/schema.py @@ -0,0 +1,344 @@ +import paddle + + +class CrystalBatch: + def __init__(self): + self.cart_coords = None + self.frac_coords = None + self.lattices = None + self.num_atoms = None + self.lengths = None + self.lengths_scaled = None + self.angles = None + self.angles_radians = None + self.atom_types = None + self.pos = None + self.token_idx = None + self.batch = None + self.y = None + self.num_nodes = None + self.mace_features = None + self.mask = None + self.zs = None + self.means = None + self.stds = None + self.num_graphs = 0 + + def add(self, **kwargs): + for key, tensor in kwargs.items(): + if not isinstance(key, str): + raise TypeError(f"Key must be a string, got {type(key).__name__}.") + if not isinstance(tensor, paddle.Tensor): + raise TypeError( + f"Value must be a paddle.Tensor, got {type(tensor).__name__}." + ) + if hasattr(self, key): + raise KeyError(f"Attribute '{key}' already exists in the batch.") + setattr(self, key, tensor) + + def update(self, allow_reshape=False, **kwargs): + for key, tensor in kwargs.items(): + if not isinstance(key, str): + raise TypeError(f"Key must be a string, got {type(key).__name__}.") + if not hasattr(self, key): + raise KeyError(f"Attribute '{key}' not found in the batch.") + if not isinstance(tensor, paddle.Tensor): + raise TypeError( + f"Value must be a paddle.Tensor, got {type(tensor).__name__}." + ) + + existing = getattr(self, key) + if tensor.shape != existing.shape and not allow_reshape: + raise ValueError( + f"Shape mismatch for '{key}': existing {tuple(existing.shape)}, new {tuple(tensor.shape)}." + ) + setattr(self, key, tensor) + + def remove(self, *keys): + if len(keys) == 0: + raise ValueError("At least one key must be provided to remove().") + for key in keys: + if not isinstance(key, str): + raise TypeError(f"Key must be a string, got {type(key).__name__}.") + if not hasattr(self, key): + raise KeyError(f"Attribute '{key}' not found in the batch.") + delattr(self, key) + + def to(self, device): + for attr_name in dir(self): + attr = getattr(self, attr_name) + if isinstance(attr, paddle.Tensor): + setattr(self, attr_name, attr.cuda() if device == 'gpu' else attr.cpu()) + return self + + def clone(self): + cloned = CrystalBatch() + for attr_name in dir(self): + if not attr_name.startswith('_') and attr_name not in ['clone', 'add', 'update', 'remove', 'to', '_split_by_batch_index', 'to_atoms', 'to_structures']: + attr = getattr(self, attr_name) + if isinstance(attr, paddle.Tensor): + setattr(cloned, attr_name, attr.clone()) + else: + setattr(cloned, attr_name, attr) + return cloned + + def _split_by_batch_index(self): + if self.batch is None or self.num_atoms is None: + raise ValueError("batch and num_atoms must be set to split the batch") + + structures = [] + batch_indices = self.batch.cpu().numpy() + num_graphs = int(batch_indices.max()) + 1 if len(batch_indices) > 0 else 0 + + for i in range(num_graphs): + mask = batch_indices == i + structure_data = {} + + if self.atom_types is not None: + structure_data['atom_types'] = self.atom_types[mask] + if self.frac_coords is not None: + structure_data['frac_coords'] = self.frac_coords[mask] + if self.cart_coords is not None: + structure_data['cart_coords'] = self.cart_coords[mask] + if self.lattices is not None: + if self.lattices.ndim == 3: + structure_data['lattices'] = self.lattices[i] + else: + structure_data['lattices'] = self.lattices + + structures.append(structure_data) + + return structures + + def to_atoms(self, frac_coords=True): + try: + from ase import Atoms + except ImportError: + raise ImportError("ASE is required for to_atoms(). Install it with: pip install ase") + + if self.atom_types is None or self.lattices is None: + raise ValueError("atom_types and lattices must be set") + + structures = self._split_by_batch_index() + atoms_list = [] + + for struct_data in structures: + atom_types_np = struct_data['atom_types'].cpu().numpy() + lattice_np = struct_data['lattices'].cpu().numpy() + + if lattice_np.ndim == 3: + lattice_np = lattice_np.squeeze(0) + + atoms = Atoms( + numbers=atom_types_np, + cell=lattice_np, + pbc=True, + ) + + if frac_coords and 'frac_coords' in struct_data: + positions = struct_data['frac_coords'].cpu().numpy() + atoms.set_scaled_positions(positions) + elif 'cart_coords' in struct_data: + positions = struct_data['cart_coords'].cpu().numpy() + atoms.set_positions(positions) + + atoms_list.append(atoms) + + return atoms_list + + def to_structures(self, frac_coords=True): + try: + from pymatgen.core import Lattice, Structure, Element + except ImportError: + raise ImportError("Pymatgen is required for to_structures(). Install it with: pip install pymatgen") + + if self.atom_types is None or self.lattices is None: + raise ValueError("atom_types and lattices must be set") + + structures = self._split_by_batch_index() + structure_list = [] + + for struct_data in structures: + atom_types_int = struct_data['atom_types'].cpu().numpy().tolist() + if isinstance(atom_types_int[0], list): + atom_types_int = [item for sublist in atom_types_int for item in sublist] + atom_types_symbols = [Element.from_Z(int(z)).symbol for z in atom_types_int] + + lattice_np = struct_data['lattices'].cpu().numpy() + + if lattice_np.ndim == 3: + lattice_np = lattice_np.squeeze(0) + + lattice = Lattice(lattice_np) + + if frac_coords and 'frac_coords' in struct_data: + coords = struct_data['frac_coords'].cpu().numpy() + structure = Structure( + lattice=Lattice.from_parameters(*lattice.parameters), + species=atom_types_symbols, + coords=coords, + coords_are_cartesian=False, + ) + elif 'cart_coords' in struct_data: + coords = struct_data['cart_coords'].cpu().numpy() + structure = Structure( + lattice=Lattice.from_parameters(*lattice.parameters), + species=atom_types_symbols, + coords=coords, + coords_are_cartesian=True, + ) + else: + raise ValueError("Either frac_coords or cart_coords must be available") + + structure_list.append(structure) + + return structure_list + + @classmethod + def from_data_list(cls, data_list): + if not data_list: + return cls() + + batch = cls() + batch.num_graphs = len(data_list) + + batch_indices = [] + offset = 0 + + for graph_idx, data in enumerate(data_list): + num_nodes = data.get('num_atoms', 0) + if isinstance(num_nodes, paddle.Tensor): + num_nodes = int(num_nodes.item()) + batch_indices.extend([graph_idx] * num_nodes) + offset += num_nodes + + batch.batch = paddle.to_tensor(batch_indices, dtype='int64') + + for key in ['atom_types', 'frac_coords', 'cart_coords', 'pos', 'token_idx']: + tensors = [data.get(key) for data in data_list if key in data] + if tensors: + batch.__dict__[key] = paddle.concat(tensors, axis=0) + + for key in ['lattices', 'num_atoms', 'lengths', 'lengths_scaled', 'angles', 'angles_radians']: + tensors = [data.get(key) for data in data_list if key in data] + if tensors: + if key == 'lattices': + stacked = [] + for t in tensors: + if t.ndim == 2: + t = t.unsqueeze(0) + stacked.append(t) + batch.__dict__[key] = paddle.concat(stacked, axis=0) + elif key == 'num_atoms': + stacked = [] + for t in tensors: + if t.ndim == 0: + t = t.unsqueeze(0) + stacked.append(t) + batch.__dict__[key] = paddle.concat(stacked, axis=0) + else: + stacked = [] + for t in tensors: + if t.ndim == 1: + t = t.unsqueeze(0) + stacked.append(t) + batch.__dict__[key] = paddle.concat(stacked, axis=0) + + if 'y' in data_list[0]: + y_dict = {} + for key in data_list[0]['y'].keys(): + y_values = [data['y'][key] for data in data_list] + if isinstance(y_values[0], paddle.Tensor): + y_dict[key] = paddle.concat(y_values, axis=0) + else: + y_dict[key] = y_values + batch.y = y_dict + + num_nodes_list = [] + for data in data_list: + num = data.get('num_atoms', 0) + if isinstance(num, paddle.Tensor): + num = int(num.item()) + num_nodes_list.append(num) + batch.num_nodes = num_nodes_list + + node_count = sum(num_nodes_list) + batch.mask = paddle.ones([len(data_list), max(num_nodes_list)], dtype='bool') + for i, num in enumerate(num_nodes_list): + if num < max(num_nodes_list): + batch.mask[i, num:] = False + + return batch + + @classmethod + def collate(cls, data_list): + return cls.from_data_list(data_list) + + def repeat(self, num_repeats): + if num_repeats <= 1: + return self + + repeated_batch = CrystalBatch() + repeated_batch.num_graphs = self.num_graphs * num_repeats + + if self.batch is not None: + batch_indices = [] + for i in range(num_repeats): + offset = i * self.num_graphs + batch_indices.append(self.batch + offset) + repeated_batch.batch = paddle.concat(batch_indices, axis=0) + + for key in ['atom_types', 'frac_coords', 'cart_coords', 'pos', 'token_idx']: + attr = self.__dict__.get(key) + if attr is not None: + repeated_batch.__dict__[key] = paddle.tile(attr, [num_repeats] + [1] * (attr.ndim - 1)) + + for key in ['lattices', 'num_atoms', 'lengths', 'lengths_scaled', 'angles', 'angles_radians']: + attr = self.__dict__.get(key) + if attr is not None: + repeated_batch.__dict__[key] = paddle.tile(attr, [num_repeats] + [1] * (attr.ndim - 1)) + + if self.y is not None: + repeated_y = {} + for key, value in self.y.items(): + if isinstance(value, paddle.Tensor): + repeated_y[key] = paddle.tile(value, [num_repeats] + [1] * (value.ndim - 1)) + else: + repeated_y[key] = value * num_repeats + repeated_batch.y = repeated_y + + if self.num_nodes is not None: + repeated_batch.num_nodes = self.num_nodes * num_repeats + + if self.mask is not None: + repeated_batch.mask = paddle.tile(self.mask, [num_repeats, 1]) + + return repeated_batch + + +def create_empty_batch(num_atoms, device='cpu', atom_types=None): + data_list = [] + for i, n in enumerate(num_atoms): + data = { + 'pos': paddle.empty([n, 3]), + 'atom_types': ( + paddle.empty([n], dtype='int64') + if atom_types is None + else paddle.to_tensor(atom_types[i], dtype='int64') + ), + 'frac_coords': paddle.empty([n, 3]), + 'cart_coords': paddle.empty([n, 3]), + 'lattices': paddle.empty([1, 3, 3]), + 'num_atoms': paddle.to_tensor(n, dtype='int64'), + 'lengths': paddle.empty([1, 3]), + 'lengths_scaled': paddle.empty([1, 3]), + 'angles': paddle.empty([1, 3]), + 'angles_radians': paddle.empty([1, 3]), + 'token_idx': paddle.arange(n, dtype='int64'), + } + data_list.append(data) + + batch = CrystalBatch.from_data_list(data_list) + if device == 'gpu': + batch = batch.to('gpu') + return batch diff --git a/ppmat/models/chemeleon2/ldm_module/__init__.py b/ppmat/models/chemeleon2/ldm_module/__init__.py new file mode 100644 index 00000000..22e4772d --- /dev/null +++ b/ppmat/models/chemeleon2/ldm_module/__init__.py @@ -0,0 +1,11 @@ +from ppmat.models.chemeleon2.ldm_module.ldm import LDMModule +from ppmat.models.chemeleon2.ldm_module.dit import DiT +from ppmat.models.chemeleon2.ldm_module.condition import ConditionModule +from ppmat.models.chemeleon2.ldm_module.condition import ConditionType + +__all__ = [ + "LDMModule", + "DiT", + "ConditionModule", + "ConditionType", +] diff --git a/ppmat/models/chemeleon2/ldm_module/condition.py b/ppmat/models/chemeleon2/ldm_module/condition.py new file mode 100644 index 00000000..0f2ea534 --- /dev/null +++ b/ppmat/models/chemeleon2/ldm_module/condition.py @@ -0,0 +1,251 @@ +from enum import Enum +import paddle +import paddle.nn as nn + + +class ConditionType(Enum): + COMPOSITION = "composition" + CHEMICAL_SYSTEM = "chemical_system" + VALUE = "float" + CATEGORICAL = "categorical" + TEXT = "text" + + +class ConditionModule(nn.Layer): + def __init__( + self, + condition_type, + hidden_dim, + drop_prob, + stats=None, + **kwargs, + ): + super().__init__() + self.condition_type = condition_type + self.target_condition = list(condition_type.keys()) + self.hidden_dim = hidden_dim + self.drop_prob = drop_prob + self.stats = stats if stats is not None else {} + + self.encoders = nn.LayerDict() + for cond_name, cond_type in condition_type.items(): + if cond_type == ConditionType.COMPOSITION.value: + self.encoders[cond_name] = CompositionEncoder( + in_dim=100, + hidden_dim=hidden_dim, + ) + elif cond_type == ConditionType.CHEMICAL_SYSTEM.value: + self.encoders[cond_name] = ChemicalSystemEncoder( + in_dim=100, + hidden_dim=hidden_dim, + ) + elif cond_type == ConditionType.VALUE.value: + _stats = self.stats.get(cond_name, {}) + self.encoders[cond_name] = ValueEncoder( + hidden_dim=hidden_dim, + mean=_stats.get("mean", None), + std=_stats.get("std", None), + ) + elif cond_type == ConditionType.CATEGORICAL.value: + assert "num_classes" in kwargs, ( + "num_classes must be provided when using CLASS condition type" + ) + self.encoders[cond_name] = CategoricalEncoder( + in_dim=kwargs["num_classes"], + hidden_dim=hidden_dim, + ) + elif cond_type == ConditionType.TEXT.value: + self.encoders[cond_name] = TextEncoder( + hidden_dim=hidden_dim, + ) + else: + raise ValueError(f"Unsupported condition type: {cond_type}") + + self.proj = nn.Sequential( + nn.Linear(len(self.encoders) * hidden_dim, hidden_dim, bias_attr=True), + nn.Silu(), + nn.Linear(hidden_dim, hidden_dim, bias_attr=True), + ) + + def forward(self, batch_y, training=True): + target_conditions = list(batch_y.keys()) + assert set(target_conditions) == set(self.target_condition), ( + f"Expected conditions {self.target_condition}, but got {target_conditions}" + ) + + batch_size = len(list(batch_y.values())[0]) + if training: + assert self.drop_prob >= 0 + drop_mask = paddle.rand([batch_size]) < self.drop_prob + else: + drop_mask = paddle.concat([ + paddle.zeros([batch_size]), + paddle.ones([batch_size]) + ]) + drop_mask = drop_mask.astype('bool') + batch_y = {k: _duplicate(v) for k, v in batch_y.items()} + + cond_embeds = [] + for cond_name, encoder in self.encoders.items(): + y = batch_y[cond_name] + embed = encoder(y, drop_mask) + cond_embeds.append(embed) + + cond_embeds = paddle.concat(cond_embeds, axis=-1) + cond_embeds = self.proj(cond_embeds) + return cond_embeds + + +def _duplicate(v): + if isinstance(v, paddle.Tensor): + return paddle.concat([v, v]) + elif isinstance(v, list): + return v * 2 + else: + raise ValueError(f"Unsupported type: {type(v)}") + + +class BaseEncoder(nn.Layer): + def __init__( + self, + *, + in_dim, + hidden_dim, + preprocess, + ): + super().__init__() + self.in_dim = in_dim + self.hidden_dim = hidden_dim + + self.preprocess = preprocess + self.embedding = nn.Sequential( + nn.Linear(in_dim, hidden_dim, bias_attr=True), + nn.Silu(), + nn.Linear(hidden_dim, hidden_dim, bias_attr=True), + ) + self.null_embed = paddle.create_parameter( + shape=[1, in_dim], + dtype='float32', + default_initializer=nn.initializer.Normal() + ) + + def forward(self, y, drop_mask=None): + y = self.preprocess(y) + if drop_mask is not None: + y = paddle.where( + drop_mask.unsqueeze(-1).expand_as(y), + self.null_embed.expand_as(y), + y + ) + return self.embedding(y) + + +class ValueEncoder(BaseEncoder): + def __init__(self, hidden_dim, mean=None, std=None): + self.mean = mean + self.std = std + + def preprocess(batch): + if isinstance(batch, list): + batch = paddle.to_tensor(batch, dtype='float32') + elif not isinstance(batch, paddle.Tensor): + batch = paddle.to_tensor(batch, dtype='float32') + + batch = batch.astype('float32').unsqueeze(-1) + if self.mean is not None and self.std is not None: + batch = (batch - self.mean) / self.std + return batch + + super().__init__( + in_dim=1, + hidden_dim=hidden_dim, + preprocess=preprocess, + ) + + +class CategoricalEncoder(BaseEncoder): + def __init__(self, in_dim, hidden_dim): + self.num_classes = in_dim + + def preprocess(batch): + if isinstance(batch, list): + idx = paddle.to_tensor(batch, dtype='int64') + elif not isinstance(batch, paddle.Tensor): + idx = paddle.to_tensor(batch, dtype='int64') + else: + idx = batch.astype('int64') + + return paddle.nn.functional.one_hot(idx, num_classes=self.num_classes).astype('float32') + + super().__init__( + in_dim=in_dim, + hidden_dim=hidden_dim, + preprocess=preprocess, + ) + + +class CompositionEncoder(BaseEncoder): + def __init__(self, in_dim, hidden_dim): + self._in_dim = in_dim + + def preprocess(batch): + vals = [self._composition_to_embeds(comp_str) for comp_str in batch] + return paddle.concat(vals, axis=0) + + super().__init__( + in_dim=in_dim, + hidden_dim=hidden_dim, + preprocess=preprocess, + ) + + def _composition_to_embeds(self, comp_str): + try: + from pymatgen.core import Composition, Element + except ImportError: + raise ImportError("Pymatgen is required for CompositionEncoder. Install it with: pip install pymatgen") + + v = paddle.zeros([self._in_dim]) + comp = Composition(comp_str).reduced_composition + for el, amt in comp.get_el_amt_dict().items(): + v[Element(el).Z] = float(amt) + v = v / v.sum() if v.sum() > 0 else v + return v.unsqueeze(0) + + +class ChemicalSystemEncoder(BaseEncoder): + def __init__(self, in_dim, hidden_dim): + self._in_dim = in_dim + + def preprocess(batch): + vals = [self._chemical_system_to_embeds(cs) for cs in batch] + return paddle.concat(vals, axis=0) + + super().__init__( + in_dim=in_dim, + hidden_dim=hidden_dim, + preprocess=preprocess, + ) + + def _chemical_system_to_embeds(self, cs): + try: + from pymatgen.core import Element + except ImportError: + raise ImportError("Pymatgen is required for ChemicalSystemEncoder. Install it with: pip install pymatgen") + + v = paddle.zeros([self._in_dim]) + elements = cs.split("-") + for el in elements: + v[Element(el).Z] = 1.0 + return v.unsqueeze(0) + + +class TextEncoder(BaseEncoder): + def __init__(self, hidden_dim): + def preprocess(batch): + return paddle.zeros([len(batch), 100]) + + super().__init__( + in_dim=100, + hidden_dim=hidden_dim, + preprocess=preprocess, + ) diff --git a/ppmat/models/chemeleon2/ldm_module/diffusion/__init__.py b/ppmat/models/chemeleon2/ldm_module/diffusion/__init__.py new file mode 100644 index 00000000..b8a7aff5 --- /dev/null +++ b/ppmat/models/chemeleon2/ldm_module/diffusion/__init__.py @@ -0,0 +1,41 @@ +from . import gaussian_diffusion as gd +from .respace import SpacedDiffusion, space_timesteps +from .diffusion_utils import get_named_beta_schedule + + +def create_diffusion( + timestep_respacing, + noise_schedule="linear", + use_kl=False, + sigma_small=False, + predict_xstart=False, + learn_sigma=True, + rescale_learned_sigmas=False, + diffusion_steps=1000, +): + betas = get_named_beta_schedule(noise_schedule, diffusion_steps) + if use_kl: + loss_type = gd.LossType.RESCALED_KL + elif rescale_learned_sigmas: + loss_type = gd.LossType.RESCALED_MSE + else: + loss_type = gd.LossType.MSE + if timestep_respacing is None or timestep_respacing == "": + timestep_respacing = [diffusion_steps] + return SpacedDiffusion( + use_timesteps=space_timesteps(diffusion_steps, timestep_respacing), + betas=betas, + model_mean_type=( + gd.ModelMeanType.EPSILON if not predict_xstart else gd.ModelMeanType.START_X + ), + model_var_type=( + ( + gd.ModelVarType.FIXED_LARGE + if not sigma_small + else gd.ModelVarType.FIXED_SMALL + ) + if not learn_sigma + else gd.ModelVarType.LEARNED_RANGE + ), + loss_type=loss_type, + ) diff --git a/ppmat/models/chemeleon2/ldm_module/diffusion/diffusion_utils.py b/ppmat/models/chemeleon2/ldm_module/diffusion/diffusion_utils.py new file mode 100644 index 00000000..d2cf4441 --- /dev/null +++ b/ppmat/models/chemeleon2/ldm_module/diffusion/diffusion_utils.py @@ -0,0 +1,175 @@ +import enum +import math +import numpy as np +import paddle + + +def broadcast_mask(mask, z): + while len(mask.shape) < len(z.shape): + mask = mask.unsqueeze(-1) + return paddle.broadcast_to(mask, z.shape) + + +def maybe_noise_like_with_mask(x, mask=None): + if mask is None: + return paddle.randn_like(x) + else: + return paddle.randn_like(x) * broadcast_mask(mask, x) + + +def mean_flat(tensor, mask=None): + if mask is not None: + mask = broadcast_mask(mask, tensor) + assert tensor.shape == mask.shape, f"{tensor.shape} != {mask.shape}" + reduced_dims = list(range(1, len(tensor.shape))) + summed = (tensor * mask).sum(axis=reduced_dims) + count = mask.sum(axis=reduced_dims).clip(min=1e-6) + return summed / count + else: + return tensor.mean(axis=list(range(1, len(tensor.shape)))) + + +class ModelMeanType(enum.Enum): + PREVIOUS_X = enum.auto() + START_X = enum.auto() + EPSILON = enum.auto() + + +class ModelVarType(enum.Enum): + LEARNED = enum.auto() + FIXED_SMALL = enum.auto() + FIXED_LARGE = enum.auto() + LEARNED_RANGE = enum.auto() + + +class LossType(enum.Enum): + MSE = enum.auto() + RESCALED_MSE = enum.auto() + KL = enum.auto() + RESCALED_KL = enum.auto() + + def is_vb(self): + return self == LossType.KL or self == LossType.RESCALED_KL + + +def _warmup_beta(beta_start, beta_end, num_diffusion_timesteps, warmup_frac): + betas = beta_end * np.ones(num_diffusion_timesteps, dtype=np.float64) + warmup_time = int(num_diffusion_timesteps * warmup_frac) + betas[:warmup_time] = np.linspace( + beta_start, beta_end, warmup_time, dtype=np.float64 + ) + return betas + + +def get_beta_schedule(beta_schedule, *, beta_start, beta_end, num_diffusion_timesteps): + if beta_schedule == "quad": + betas = ( + np.linspace( + beta_start**0.5, + beta_end**0.5, + num_diffusion_timesteps, + dtype=np.float64, + ) + ** 2 + ) + elif beta_schedule == "linear": + betas = np.linspace( + beta_start, beta_end, num_diffusion_timesteps, dtype=np.float64 + ) + elif beta_schedule == "warmup10": + betas = _warmup_beta(beta_start, beta_end, num_diffusion_timesteps, 0.1) + elif beta_schedule == "warmup50": + betas = _warmup_beta(beta_start, beta_end, num_diffusion_timesteps, 0.5) + elif beta_schedule == "const": + betas = beta_end * np.ones(num_diffusion_timesteps, dtype=np.float64) + elif beta_schedule == "jsd": + betas = 1.0 / np.linspace( + num_diffusion_timesteps, 1, num_diffusion_timesteps, dtype=np.float64 + ) + else: + raise NotImplementedError(beta_schedule) + assert betas.shape == (num_diffusion_timesteps,) + return betas + + +def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999): + betas = [] + for i in range(num_diffusion_timesteps): + t1 = i / num_diffusion_timesteps + t2 = (i + 1) / num_diffusion_timesteps + betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta)) + return np.array(betas) + + +def get_named_beta_schedule(schedule_name, num_diffusion_timesteps): + if schedule_name == "linear": + scale = 1000 / num_diffusion_timesteps + return get_beta_schedule( + "linear", + beta_start=scale * 0.0001, + beta_end=scale * 0.02, + num_diffusion_timesteps=num_diffusion_timesteps, + ) + elif schedule_name == "squaredcos_cap_v2": + return betas_for_alpha_bar( + num_diffusion_timesteps, + lambda t: math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2, + ) + else: + raise NotImplementedError(f"unknown beta schedule: {schedule_name}") + + +def normal_kl(mean1, logvar1, mean2, logvar2): + tensor = None + for obj in (mean1, logvar1, mean2, logvar2): + if isinstance(obj, paddle.Tensor): + tensor = obj + break + assert tensor is not None, "at least one argument must be a Tensor" + + logvar1, logvar2 = [ + x if isinstance(x, paddle.Tensor) else paddle.to_tensor(x).to(tensor) + for x in (logvar1, logvar2) + ] + + return 0.5 * ( + -1.0 + + logvar2 + - logvar1 + + paddle.exp(logvar1 - logvar2) + + ((mean1 - mean2) ** 2) * paddle.exp(-logvar2) + ) + + +def approx_standard_normal_cdf(x): + return 0.5 * (1.0 + paddle.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * paddle.pow(x, 3)))) + + +def continuous_gaussian_log_likelihood(x, *, means, log_scales): + centered_x = x - means + inv_stdv = paddle.exp(-log_scales) + normalized_x = centered_x * inv_stdv + log_probs = paddle.distribution.Normal(paddle.zeros_like(x), paddle.ones_like(x)).log_prob( + normalized_x + ) + return log_probs + + +def discretized_gaussian_log_likelihood(x, *, means, log_scales): + assert x.shape == means.shape == log_scales.shape + centered_x = x - means + inv_stdv = paddle.exp(-log_scales) + plus_in = inv_stdv * (centered_x + 1.0 / 255.0) + cdf_plus = approx_standard_normal_cdf(plus_in) + min_in = inv_stdv * (centered_x - 1.0 / 255.0) + cdf_min = approx_standard_normal_cdf(min_in) + log_cdf_plus = paddle.log(cdf_plus.clip(min=1e-12)) + log_one_minus_cdf_min = paddle.log((1.0 - cdf_min).clip(min=1e-12)) + cdf_delta = cdf_plus - cdf_min + log_probs = paddle.where( + x < -0.999, + log_cdf_plus, + paddle.where(x > 0.999, log_one_minus_cdf_min, paddle.log(cdf_delta.clip(min=1e-12))), + ) + assert log_probs.shape == x.shape + return log_probs diff --git a/ppmat/models/chemeleon2/ldm_module/diffusion/gaussian_diffusion.py b/ppmat/models/chemeleon2/ldm_module/diffusion/gaussian_diffusion.py new file mode 100644 index 00000000..59e27e9f --- /dev/null +++ b/ppmat/models/chemeleon2/ldm_module/diffusion/gaussian_diffusion.py @@ -0,0 +1,427 @@ +import numpy as np +import paddle + +from .diffusion_utils import ( + ModelMeanType, + ModelVarType, + LossType, + mean_flat, + normal_kl, + continuous_gaussian_log_likelihood, + maybe_noise_like_with_mask, +) + + +def _extract_into_tensor(arr, timesteps, broadcast_shape): + res = paddle.to_tensor(arr, dtype='float32')[timesteps] + while len(res.shape) < len(broadcast_shape): + res = res.unsqueeze(-1) + return paddle.broadcast_to(res, broadcast_shape) + + +class GaussianDiffusion: + def __init__(self, *, betas, model_mean_type, model_var_type, loss_type): + self.model_mean_type = model_mean_type + self.model_var_type = model_var_type + self.loss_type = loss_type + + betas = np.array(betas, dtype=np.float64) + self.betas = betas + assert len(betas.shape) == 1, "betas must be 1-D" + assert (betas > 0).all() and (betas <= 1).all() + + self.num_timesteps = int(betas.shape[0]) + + alphas = 1.0 - betas + self.alphas_cumprod = np.cumprod(alphas, axis=0) + self.alphas_cumprod_prev = np.append(1.0, self.alphas_cumprod[:-1]) + self.alphas_cumprod_next = np.append(self.alphas_cumprod[1:], 0.0) + + self.sqrt_alphas_cumprod = np.sqrt(self.alphas_cumprod) + self.sqrt_one_minus_alphas_cumprod = np.sqrt(1.0 - self.alphas_cumprod) + self.log_one_minus_alphas_cumprod = np.log(1.0 - self.alphas_cumprod) + self.sqrt_recip_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod) + self.sqrt_recipm1_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod - 1) + + self.posterior_variance = ( + betas * (1.0 - self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod) + ) + self.posterior_log_variance_clipped = ( + np.log(np.append(self.posterior_variance[1], self.posterior_variance[1:])) + if len(self.posterior_variance) > 1 + else np.array([]) + ) + + self.posterior_mean_coef1 = ( + betas * np.sqrt(self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod) + ) + self.posterior_mean_coef2 = ( + (1.0 - self.alphas_cumprod_prev) + * np.sqrt(alphas) + / (1.0 - self.alphas_cumprod) + ) + + def q_mean_variance(self, x_start, t): + mean = ( + _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start + ) + variance = _extract_into_tensor(1.0 - self.alphas_cumprod, t, x_start.shape) + log_variance = _extract_into_tensor( + self.log_one_minus_alphas_cumprod, t, x_start.shape + ) + return mean, variance, log_variance + + def q_sample(self, x_start, t, noise=None): + if noise is None: + noise = paddle.randn_like(x_start) + assert noise.shape == x_start.shape + return ( + _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start + + _extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) + * noise + ) + + def q_posterior_mean_variance(self, x_start, x_t, t): + assert x_start.shape == x_t.shape + posterior_mean = ( + _extract_into_tensor(self.posterior_mean_coef1, t, x_t.shape) * x_start + + _extract_into_tensor(self.posterior_mean_coef2, t, x_t.shape) * x_t + ) + posterior_variance = _extract_into_tensor(self.posterior_variance, t, x_t.shape) + posterior_log_variance_clipped = _extract_into_tensor( + self.posterior_log_variance_clipped, t, x_t.shape + ) + return posterior_mean, posterior_variance, posterior_log_variance_clipped + + def _predict_xstart_from_eps(self, x_t, t, eps): + assert x_t.shape == eps.shape + return ( + _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t + - _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * eps + ) + + def _predict_eps_from_xstart(self, x_t, t, pred_xstart): + return ( + _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t + - pred_xstart + ) / _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) + + def p_mean_variance( + self, model, x, t, clip_denoised=True, denoised_fn=None, model_kwargs=None + ): + if model_kwargs is None: + model_kwargs = {} + + B, C = x.shape[:2] + assert t.shape == [B] + model_output = model(x, t, **model_kwargs) + + if self.model_var_type in [ModelVarType.LEARNED, ModelVarType.LEARNED_RANGE]: + assert model_output.shape == [B, C * 2, *x.shape[2:]] + model_output, model_var_values = paddle.split(model_output, 2, axis=1) + if self.model_var_type == ModelVarType.LEARNED: + model_log_variance = model_var_values + model_variance = paddle.exp(model_log_variance) + else: + min_log = _extract_into_tensor( + self.posterior_log_variance_clipped, t, x.shape + ) + max_log = _extract_into_tensor(np.log(self.betas), t, x.shape) + frac = (model_var_values + 1) / 2 + model_log_variance = frac * max_log + (1 - frac) * min_log + model_variance = paddle.exp(model_log_variance) + else: + model_variance, model_log_variance = { + ModelVarType.FIXED_LARGE: ( + np.append(self.posterior_variance[1], self.betas[1:]), + np.log(np.append(self.posterior_variance[1], self.betas[1:])), + ), + ModelVarType.FIXED_SMALL: ( + self.posterior_variance, + self.posterior_log_variance_clipped, + ), + }[self.model_var_type] + model_variance = _extract_into_tensor(model_variance, t, x.shape) + model_log_variance = _extract_into_tensor(model_log_variance, t, x.shape) + + def process_xstart(x): + if denoised_fn is not None: + x = denoised_fn(x) + if clip_denoised: + return x.clip(-1, 1) + return x + + if self.model_mean_type == ModelMeanType.PREVIOUS_X: + pred_xstart = process_xstart( + self._predict_xstart_from_xprev(x_t=x, t=t, xprev=model_output) + ) + model_mean = model_output + elif self.model_mean_type in [ModelMeanType.START_X, ModelMeanType.EPSILON]: + if self.model_mean_type == ModelMeanType.START_X: + pred_xstart = process_xstart(model_output) + else: + pred_xstart = process_xstart( + self._predict_xstart_from_eps(x_t=x, t=t, eps=model_output) + ) + model_mean, _, _ = self.q_posterior_mean_variance( + x_start=pred_xstart, x_t=x, t=t + ) + else: + raise NotImplementedError(self.model_mean_type) + + assert ( + model_mean.shape == model_log_variance.shape == pred_xstart.shape == x.shape + ) + return { + "mean": model_mean, + "variance": model_variance, + "log_variance": model_log_variance, + "pred_xstart": pred_xstart, + } + + def _predict_xstart_from_xprev(self, x_t, t, xprev): + assert x_t.shape == xprev.shape + return ( + _extract_into_tensor(1.0 / self.posterior_mean_coef1, t, x_t.shape) * xprev + - _extract_into_tensor( + self.posterior_mean_coef2 / self.posterior_mean_coef1, t, x_t.shape + ) + * x_t + ) + + def p_sample( + self, + model, + x, + t, + clip_denoised=True, + denoised_fn=None, + model_kwargs=None, + ): + out = self.p_mean_variance( + model, + x, + t, + clip_denoised=clip_denoised, + denoised_fn=denoised_fn, + model_kwargs=model_kwargs, + ) + noise = paddle.randn_like(x) + nonzero_mask = ( + (t != 0).astype('float32').reshape([x.shape[0]] + [1] * (len(x.shape) - 1)) + ) + sample = out["mean"] + nonzero_mask * paddle.exp(0.5 * out["log_variance"]) * noise + return {"sample": sample, "pred_xstart": out["pred_xstart"]} + + def p_sample_loop( + self, + model, + shape, + noise=None, + clip_denoised=True, + denoised_fn=None, + model_kwargs=None, + device=None, + progress=False, + ): + if device is None: + device = paddle.get_device() + + if model_kwargs is None: + model_kwargs = {} + + if noise is not None: + img = noise + else: + img = paddle.randn(shape) + img = maybe_noise_like_with_mask(img, model_kwargs.get("mask")) + + indices = list(range(self.num_timesteps))[::-1] + + if progress: + try: + from tqdm.auto import tqdm + indices = tqdm(indices) + except ImportError: + pass + + for i in indices: + t = paddle.to_tensor([i] * shape[0], dtype='int64') + with paddle.no_grad(): + out = self.p_sample( + model, + img, + t, + clip_denoised=clip_denoised, + denoised_fn=denoised_fn, + model_kwargs=model_kwargs, + ) + img = out["sample"] + + return img + + def ddim_sample( + self, + model, + x, + t, + clip_denoised=True, + denoised_fn=None, + model_kwargs=None, + eta=0.0, + ): + # FIXME-ISPL: DDIM采样需要完整实现 + # 参考: https://arxiv.org/abs/2010.02502 + out = self.p_mean_variance( + model, + x, + t, + clip_denoised=clip_denoised, + denoised_fn=denoised_fn, + model_kwargs=model_kwargs, + ) + + eps = self._predict_eps_from_xstart(x, t, out["pred_xstart"]) + + alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape) + alpha_bar_prev = _extract_into_tensor(self.alphas_cumprod_prev, t, x.shape) + sigma = ( + eta + * paddle.sqrt((1 - alpha_bar_prev) / (1 - alpha_bar)) + * paddle.sqrt(1 - alpha_bar / alpha_bar_prev) + ) + noise = paddle.randn_like(x) + mean_pred = ( + out["pred_xstart"] * paddle.sqrt(alpha_bar_prev) + + paddle.sqrt(1 - alpha_bar_prev - sigma**2) * eps + ) + nonzero_mask = ( + (t != 0).astype('float32').reshape([x.shape[0]] + [1] * (len(x.shape) - 1)) + ) + sample = mean_pred + nonzero_mask * sigma * noise + return { + "sample": sample, + "pred_xstart": out["pred_xstart"], + "mean": mean_pred, + "std": sigma, + } + + def ddim_sample_loop( + self, + model, + shape, + noise=None, + clip_denoised=True, + denoised_fn=None, + model_kwargs=None, + device=None, + progress=False, + eta=0.0, + ): + if device is None: + device = paddle.get_device() + + if model_kwargs is None: + model_kwargs = {} + + if noise is not None: + img = noise + else: + img = paddle.randn(shape) + img = maybe_noise_like_with_mask(img, model_kwargs.get("mask")) + + indices = list(range(self.num_timesteps))[::-1] + + if progress: + try: + from tqdm.auto import tqdm + indices = tqdm(indices) + except ImportError: + pass + + for i in indices: + t = paddle.to_tensor([i] * shape[0], dtype='int64') + with paddle.no_grad(): + out = self.ddim_sample( + model, + img, + t, + clip_denoised=clip_denoised, + denoised_fn=denoised_fn, + model_kwargs=model_kwargs, + eta=eta, + ) + img = out["sample"] + + return img + + def training_losses(self, model, x_start, t, model_kwargs=None, noise=None): + if model_kwargs is None: + model_kwargs = {} + if noise is None: + noise = paddle.randn_like(x_start) + x_t = self.q_sample(x_start, t, noise=noise) + + terms = {} + + if self.loss_type == LossType.KL or self.loss_type == LossType.RESCALED_KL: + # FIXME-ISPL: KL loss需要完整实现 + terms["loss"] = paddle.to_tensor([0.0]) + elif self.loss_type == LossType.MSE or self.loss_type == LossType.RESCALED_MSE: + model_output = model(x_t, t, **model_kwargs) + + if self.model_var_type in [ + ModelVarType.LEARNED, + ModelVarType.LEARNED_RANGE, + ]: + B, C = x_t.shape[:2] + assert model_output.shape == [B, C * 2, *x_t.shape[2:]] + model_output, model_var_values = paddle.split(model_output, 2, axis=1) + frozen_out = paddle.concat([model_output.detach(), model_var_values], axis=1) + # FIXME-ISPL: VLB loss需要实现 + terms["vb"] = paddle.to_tensor([0.0]) + + target = { + ModelMeanType.PREVIOUS_X: self.q_posterior_mean_variance( + x_start=x_start, x_t=x_t, t=t + )[0], + ModelMeanType.START_X: x_start, + ModelMeanType.EPSILON: noise, + }[self.model_mean_type] + assert model_output.shape == target.shape == x_start.shape + terms["mse"] = mean_flat((target - model_output) ** 2) + if "vb" in terms: + terms["loss"] = terms["mse"] + terms["vb"] + else: + terms["loss"] = terms["mse"] + else: + raise NotImplementedError(self.loss_type) + + return terms + + def _vb_terms_bpd( + self, model, x_start, x_t, t, clip_denoised=True, model_kwargs=None + ): + if model_kwargs is None: + model_kwargs = {} + mask = model_kwargs.get("mask", None) + + true_mean, _, true_log_variance_clipped = self.q_posterior_mean_variance( + x_start=x_start, x_t=x_t, t=t + ) + out = self.p_mean_variance( + model, x_t, t, clip_denoised=clip_denoised, model_kwargs=model_kwargs + ) + kl = normal_kl( + true_mean, true_log_variance_clipped, out["mean"], out["log_variance"] + ) + kl = mean_flat(kl, mask) / np.log(2.0) + + decoder_nll = -continuous_gaussian_log_likelihood( + x_start, means=out["mean"], log_scales=0.5 * out["log_variance"] + ) + assert decoder_nll.shape == x_start.shape + decoder_nll = mean_flat(decoder_nll, mask) / np.log(2.0) + + output = paddle.where((t == 0).astype('bool'), decoder_nll, kl) + return {"output": output, "pred_xstart": out["pred_xstart"]} diff --git a/ppmat/models/chemeleon2/ldm_module/diffusion/respace.py b/ppmat/models/chemeleon2/ldm_module/diffusion/respace.py new file mode 100644 index 00000000..7222bcd0 --- /dev/null +++ b/ppmat/models/chemeleon2/ldm_module/diffusion/respace.py @@ -0,0 +1,89 @@ +import numpy as np +import paddle + +from .gaussian_diffusion import GaussianDiffusion + + +def space_timesteps(num_timesteps, section_counts): + if isinstance(section_counts, str): + if section_counts.startswith("ddim"): + desired_count = int(section_counts[len("ddim") :]) + for i in range(1, num_timesteps): + if len(range(0, num_timesteps, i)) == desired_count: + return set(range(0, num_timesteps, i)) + raise ValueError( + f"cannot create exactly {num_timesteps} steps with an integer stride" + ) + section_counts = [int(x) for x in section_counts.split(",")] + size_per = num_timesteps // len(section_counts) + extra = num_timesteps % len(section_counts) + start_idx = 0 + all_steps = [] + for i, section_count in enumerate(section_counts): + size = size_per + (1 if i < extra else 0) + if size < section_count: + raise ValueError( + f"cannot divide section of {size} steps into {section_count}" + ) + if section_count <= 1: + frac_stride = 1 + else: + frac_stride = (size - 1) / (section_count - 1) + cur_idx = 0.0 + taken_steps = [] + for _ in range(section_count): + taken_steps.append(start_idx + round(cur_idx)) + cur_idx += frac_stride + all_steps += taken_steps + start_idx += size + return set(all_steps) + + +class SpacedDiffusion(GaussianDiffusion): + def __init__(self, use_timesteps, **kwargs): + self.use_timesteps = set(use_timesteps) + self.timestep_map = [] + self.original_num_steps = len(kwargs["betas"]) + + base_diffusion = GaussianDiffusion(**kwargs) + last_alpha_cumprod = 1.0 + new_betas = [] + for i, alpha_cumprod in enumerate(base_diffusion.alphas_cumprod): + if i in self.use_timesteps: + new_betas.append(1 - alpha_cumprod / last_alpha_cumprod) + last_alpha_cumprod = alpha_cumprod + self.timestep_map.append(i) + kwargs["betas"] = np.array(new_betas) + super().__init__(**kwargs) + + def p_mean_variance(self, model, *args, **kwargs): + return super().p_mean_variance(self._wrap_model(model), *args, **kwargs) + + def training_losses(self, model, *args, **kwargs): + return super().training_losses(self._wrap_model(model), *args, **kwargs) + + def condition_mean(self, cond_fn, *args, **kwargs): + return super().condition_mean(self._wrap_model(cond_fn), *args, **kwargs) + + def condition_score(self, cond_fn, *args, **kwargs): + return super().condition_score(self._wrap_model(cond_fn), *args, **kwargs) + + def _wrap_model(self, model): + if isinstance(model, _WrappedModel): + return model + return _WrappedModel(model, self.timestep_map, self.original_num_steps) + + def _scale_timesteps(self, t): + return t + + +class _WrappedModel: + def __init__(self, model, timestep_map, original_num_steps): + self.model = model + self.timestep_map = timestep_map + self.original_num_steps = original_num_steps + + def __call__(self, x, ts, **kwargs): + map_tensor = paddle.to_tensor(self.timestep_map, dtype=ts.dtype) + new_ts = paddle.index_select(map_tensor, ts) + return self.model(x, new_ts, **kwargs) diff --git a/ppmat/models/chemeleon2/ldm_module/diffusion/timestep_sampler.py b/ppmat/models/chemeleon2/ldm_module/diffusion/timestep_sampler.py new file mode 100644 index 00000000..fd8be4d6 --- /dev/null +++ b/ppmat/models/chemeleon2/ldm_module/diffusion/timestep_sampler.py @@ -0,0 +1,80 @@ +from abc import ABC, abstractmethod +import numpy as np +import paddle + + +def create_named_schedule_sampler(name, diffusion): + if name == "uniform": + return UniformSampler(diffusion) + elif name == "loss-second-moment": + return LossSecondMomentResampler(diffusion) + else: + raise NotImplementedError(f"unknown schedule sampler: {name}") + + +class ScheduleSampler(ABC): + @abstractmethod + def weights(self): + pass + + def sample(self, batch_size, device): + w = self.weights() + assert w is not None, "weights() should not return None" + p = w / np.sum(w) + indices_np = np.random.choice(len(p), size=(batch_size,), p=p) + indices = paddle.to_tensor(indices_np, dtype='int64') + weights_np = 1 / (len(p) * p[indices_np]) + weights = paddle.to_tensor(weights_np, dtype='float32') + return indices, weights + + +class UniformSampler(ScheduleSampler): + def __init__(self, diffusion): + self.diffusion = diffusion + self._weights = np.ones([diffusion.num_timesteps]) + + def weights(self): + return self._weights + + +class LossSecondMomentResampler(ScheduleSampler): + def __init__(self, diffusion, history_per_term=10, uniform_prob=0.001): + self.diffusion = diffusion + self.history_per_term = history_per_term + self.uniform_prob = uniform_prob + self._loss_history = np.zeros( + [diffusion.num_timesteps, history_per_term], dtype=np.float64 + ) + self._loss_counts = np.zeros([diffusion.num_timesteps], dtype=np.int32) + + def weights(self): + if not self._warmed_up(): + return np.ones([self.diffusion.num_timesteps], dtype=np.float64) + weights = np.sqrt(np.mean(self._loss_history**2, axis=-1)) + weights /= np.sum(weights) + weights *= 1 - self.uniform_prob + weights += self.uniform_prob / len(weights) + return weights + + def update_with_local_losses(self, local_ts, local_losses): + batch_sizes = [len(local_ts)] + max_bs = max(batch_sizes) + + timestep_batches = [local_ts] + loss_batches = [local_losses] + + timesteps = paddle.concat(timestep_batches, axis=0)[: sum(batch_sizes)] + losses = paddle.concat(loss_batches, axis=0)[: sum(batch_sizes)] + self.update_with_all_losses(timesteps, losses) + + def update_with_all_losses(self, ts, losses): + for t, loss in zip(ts, losses): + if self._loss_counts[t] == self.history_per_term: + self._loss_history[t, :-1] = self._loss_history[t, 1:] + self._loss_history[t, -1] = loss.item() + else: + self._loss_history[t, self._loss_counts[t]] = loss.item() + self._loss_counts[t] += 1 + + def _warmed_up(self): + return (self._loss_counts == self.history_per_term).all() diff --git a/ppmat/models/chemeleon2/ldm_module/dit.py b/ppmat/models/chemeleon2/ldm_module/dit.py new file mode 100644 index 00000000..29f12b11 --- /dev/null +++ b/ppmat/models/chemeleon2/ldm_module/dit.py @@ -0,0 +1,237 @@ +import math +import paddle +import paddle.nn as nn + + +def modulate(x, shift, scale): + return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) + + +class TimestepEmbedder(nn.Layer): + def __init__(self, hidden_dim, frequency_embedding_dim=256): + super().__init__() + self.mlp = nn.Sequential( + nn.Linear(frequency_embedding_dim, hidden_dim, bias_attr=True), + nn.Silu(), + nn.Linear(hidden_dim, hidden_dim, bias_attr=True), + ) + self.frequency_embedding_dim = frequency_embedding_dim + + @staticmethod + def timestep_embedding(t, dim, max_period=10000): + half = dim // 2 + freqs = paddle.exp( + -math.log(max_period) + * paddle.arange(start=0, end=half, dtype='float32') + / half + ) + args = t.unsqueeze(-1).astype('float32') * freqs + embedding = paddle.concat([paddle.cos(args), paddle.sin(args)], axis=-1) + if dim % 2: + embedding = paddle.concat( + [embedding, paddle.zeros_like(embedding[:, :1])], axis=-1 + ) + return embedding + + def forward(self, t): + t_freq = self.timestep_embedding(t, self.frequency_embedding_dim) + t_emb = self.mlp(t_freq) + return t_emb + + +def get_pos_embedding(indices, emb_dim, max_len=2048): + K = paddle.arange(emb_dim // 2) + pos_embedding_sin = paddle.sin( + indices.unsqueeze(-1) * math.pi / (max_len ** (2 * K / emb_dim)) + ) + pos_embedding_cos = paddle.cos( + indices.unsqueeze(-1) * math.pi / (max_len ** (2 * K / emb_dim)) + ) + pos_embedding = paddle.concat([pos_embedding_sin, pos_embedding_cos], axis=-1) + return pos_embedding + + +class Mlp(nn.Layer): + def __init__( + self, + in_features, + hidden_features=None, + out_features=None, + act_layer=None, + norm_layer=None, + bias=True, + drop=0.0, + ): + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + + self.fc1 = nn.Linear(in_features, hidden_features, bias_attr=bias) + self.act = act_layer() if act_layer else nn.GELU() + self.drop1 = nn.Dropout(drop) + self.norm = norm_layer(hidden_features) if norm_layer is not None else nn.Identity() + self.fc2 = nn.Linear(hidden_features, out_features, bias_attr=bias) + self.drop2 = nn.Dropout(drop) + + def forward(self, x): + x = self.fc1(x) + x = self.act(x) + x = self.drop1(x) + x = self.norm(x) + x = self.fc2(x) + x = self.drop2(x) + return x + + +class FinalLayer(nn.Layer): + def __init__(self, hidden_dim, out_dim): + super().__init__() + self.norm_final = nn.LayerNorm(hidden_dim, epsilon=1e-6, weight_attr=False, bias_attr=False) + self.linear = nn.Linear(hidden_dim, out_dim, bias_attr=True) + self.adaLN_modulation = nn.Sequential( + nn.Silu(), + nn.Linear(hidden_dim, 2 * hidden_dim, bias_attr=True) + ) + + def forward(self, x, c): + shift, scale = self.adaLN_modulation(c).chunk(2, axis=1) + x = modulate(self.norm_final(x), shift, scale) + x = self.linear(x) + return x + + +class DiTBlock(nn.Layer): + def __init__(self, hidden_dim, num_heads, mlp_ratio=4.0): + super().__init__() + self.norm1 = nn.LayerNorm(hidden_dim, epsilon=1e-6, weight_attr=False, bias_attr=False) + self.attn = nn.MultiHeadAttention(hidden_dim, num_heads, dropout=0.0) + self.norm2 = nn.LayerNorm(hidden_dim, epsilon=1e-6, weight_attr=False, bias_attr=False) + mlp_hidden_dim = int(hidden_dim * mlp_ratio) + + self.mlp = Mlp( + in_features=hidden_dim, + hidden_features=mlp_hidden_dim, + act_layer=lambda: nn.GELU(approximate=True), + drop=0, + ) + self.adaLN_modulation = nn.Sequential( + nn.Silu(), + nn.Linear(hidden_dim, 6 * hidden_dim, bias_attr=True) + ) + + def forward(self, x, c, mask): + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + self.adaLN_modulation(c).chunk(6, axis=1) + ) + + norm_x = self.norm1(x) + modulated_x = modulate(norm_x, shift_msa, scale_msa) + + # mask: True means this position should be masked (padding position) + # Note: mask is already ~ of the original valid_mask (inverted in DiT.forward) + attn_mask = None + if mask is not None: + attn_mask = mask.unsqueeze([1, 2]) + attn_mask = paddle.cast(attn_mask, dtype='float32') * -1e9 + + attn_out = self.attn(modulated_x, modulated_x, modulated_x, attn_mask=attn_mask) + x = x + gate_msa.unsqueeze(1) * attn_out + + norm_x = self.norm2(x) + modulated_x = modulate(norm_x, shift_mlp, scale_mlp) + mlp_out = self.mlp(modulated_x) + x = x + gate_mlp.unsqueeze(1) * mlp_out + + return x + + +class DiT(nn.Layer): + def __init__( + self, + input_dim=256, + hidden_dim=1024, + num_heads=8, + num_layers=12, + mlp_ratio=4.0, + condition_dim=None, + learn_sigma=False, + ): + super().__init__() + + self.input_dim = input_dim + self.hidden_dim = hidden_dim + self.num_layers = num_layers + self.learn_sigma = learn_sigma + self.latent_dim = input_dim + + self.x_embedder = nn.Linear(input_dim, hidden_dim, bias_attr=True) + self.t_embedder = TimestepEmbedder(hidden_dim) + + if condition_dim is not None: + self.y_embedder = nn.Linear(condition_dim, hidden_dim, bias_attr=True) + else: + self.y_embedder = None + + self.blocks = nn.LayerList([ + DiTBlock(hidden_dim, num_heads, mlp_ratio=mlp_ratio) + for _ in range(num_layers) + ]) + + out_dim = input_dim * 2 if learn_sigma else input_dim + self.final_layer = FinalLayer(hidden_dim, out_dim) + + self.initialize_weights() + + def initialize_weights(self): + def _basic_init(m): + if isinstance(m, nn.Linear): + nn.initializer.XavierUniform()(m.weight) + if m.bias is not None: + nn.initializer.Constant(0.0)(m.bias) + + self.apply(_basic_init) + + def forward(self, x, t, mask=None, y=None): + if mask is not None: + token_indices = paddle.cumsum(mask.astype('int64'), axis=-1) - 1 + pos_emb = get_pos_embedding(token_indices, self.hidden_dim) + else: + pos_emb = 0 + + x = self.x_embedder(x) + pos_emb + c = self.t_embedder(t) + + if y is not None and self.y_embedder is not None: + y_emb = self.y_embedder(y) + c = c + y_emb + elif self.y_embedder is not None: + y_emb = paddle.zeros([x.shape[0], self.hidden_dim], dtype=x.dtype) + c = c + y_emb + + mask_inverted = paddle.logical_not(mask) if mask is not None else None + for block in self.blocks: + x = block(x, c, mask_inverted) + + x = self.final_layer(x, c) + + if self.learn_sigma: + assert x.shape[2] == 2 * self.latent_dim + x = x.reshape([x.shape[0], 2 * x.shape[1], self.latent_dim]) + if mask is not None: + x = x * mask.tile([1, 2]).unsqueeze(-1).astype(x.dtype) + else: + assert x.shape[2] == self.latent_dim + if mask is not None: + x = x * mask.unsqueeze(-1).astype(x.dtype) + + return x + + def forward_with_cfg(self, x, t, mask, y, cfg_scale): + half_x = x[: x.shape[0] // 2] + combined_x = paddle.concat([half_x, half_x], axis=0) + model_out = self.forward(combined_x, t, mask, y) + + cond_eps, uncond_eps = paddle.split(model_out, 2, axis=0) + half_eps = uncond_eps + cfg_scale * (cond_eps - uncond_eps) + eps = paddle.concat([half_eps, half_eps], axis=0) + return eps diff --git a/ppmat/models/chemeleon2/ldm_module/ldm.py b/ppmat/models/chemeleon2/ldm_module/ldm.py new file mode 100644 index 00000000..69a5c4df --- /dev/null +++ b/ppmat/models/chemeleon2/ldm_module/ldm.py @@ -0,0 +1,481 @@ +from collections import defaultdict +from functools import partial + +import numpy as np +import paddle +import paddle.nn as nn + +from ppmat.models.chemeleon2.common.lora import ( + apply_lora_to_linear, + print_trainable_parameters, + merge_lora_weights, +) +from ppmat.models.chemeleon2.ldm_module.diffusion import create_diffusion +from ppmat.models.chemeleon2.vae_module.vae import VAEModule +from ppmat.utils.crystal import lattice_params_to_matrix_paddle + + +class LDMModule(nn.Layer): + def __init__( + self, + normalize_latent=True, + denoiser=None, + augmentation=None, + diffusion_configs=None, + optimizer=None, + scheduler=None, + condition_module=None, + vae=None, + vae_ckpt_path=None, + ldm_ckpt_path=None, + lora_configs=None, + ): + super().__init__() + + # Build nested models if they are config dicts + from ppmat.models import build_model + if isinstance(denoiser, dict): + self.denoiser = build_model(denoiser) + else: + self.denoiser = denoiser + + if isinstance(condition_module, dict): + self.condition_module = build_model(condition_module) + else: + self.condition_module = condition_module + + if isinstance(vae, dict): + vae = build_model(vae) + + self.normalize_latent = normalize_latent + self.latent_std = paddle.to_tensor([1.0]) + + self.diffusion_configs = diffusion_configs + self.augmentation = augmentation + self.optimizer_config = optimizer + self.scheduler_config = scheduler + self.lora_configs = lora_configs + + if diffusion_configs is not None: + self.diffusion = create_diffusion(**diffusion_configs) + else: + self.diffusion = None + + if vae is not None: + self.vae = vae + for param in self.vae.parameters(): + param.stop_gradient = True + self.vae.eval() + elif vae_ckpt_path is not None: + self.vae = VAEModule.load_checkpoint(vae_ckpt_path)[0] + for param in self.vae.parameters(): + param.stop_gradient = True + self.vae.eval() + + if ldm_ckpt_path is not None: + checkpoint = paddle.load(ldm_ckpt_path) + self.set_state_dict(checkpoint['model_state_dict']) + if 'latent_std' in checkpoint: + self.latent_std = checkpoint['latent_std'] + + if lora_configs is not None: + rank = lora_configs.get('r', 8) + alpha = lora_configs.get('lora_alpha', 16) + dropout = lora_configs.get('lora_dropout', 0.0) + target_modules = lora_configs.get('target_modules', None) + + self.denoiser = apply_lora_to_linear( + self.denoiser, + rank=rank, + alpha=alpha, + dropout=dropout, + target_modules=target_modules + ) + print_trainable_parameters(self.denoiser) + + self.use_cfg = False + if condition_module is not None: + self.use_cfg = True + self.condition_module = condition_module + + def _to_dense_batch(self, x, batch_idx): + batch_size = int(batch_idx.max().item()) + 1 + max_num_nodes = paddle.bincount(batch_idx.astype('int32')).max().item() + + dense_x = paddle.zeros([batch_size, max_num_nodes, x.shape[-1]], dtype=x.dtype) + mask = paddle.zeros([batch_size, max_num_nodes], dtype='bool') + + for i in range(batch_size): + node_mask = batch_idx == i + num_nodes = node_mask.sum().item() + dense_x[i, :num_nodes] = x[node_mask] + mask[i, :num_nodes] = True + + return dense_x, mask + + def _apply_augmentation(self, batch): + if self.augmentation is None: + return batch + return batch + + def forward(self, batch): + """Forward pass for training compatibility. + + This method converts the standard dictionary format to CrystalBatch format + and then calls calculate_loss. This provides compatibility with the + standard training framework. + + Args: + batch: Input batch data (dict with 'structure_array' key) + + Returns: + dict: Contains 'loss_dict' with training losses (tensors for backward pass) + """ + # Convert dict format to CrystalBatch format + crystal_batch = self._dict_to_crystal_batch(batch) + loss_dict = self.calculate_loss(crystal_batch, training=True) + + # The framework needs loss_dict with tensor values for backward pass + # Don't detach or convert to scalars - keep tensors as-is + return {"loss_dict": loss_dict} + + def _dict_to_crystal_batch(self, batch): + """Convert dictionary format to CrystalBatch format. + + Args: + batch: Dict with 'structure_array' key containing structure data + + Returns: + CrystalBatch: Converted batch object + """ + from ppmat.models.chemeleon2.common.schema import CrystalBatch + + structure_array = batch["structure_array"] + num_atoms = structure_array["num_atoms"] + batch_size = num_atoms.shape[0] + total_atoms = num_atoms.sum().item() + + # Create CrystalBatch from structure_array + crystal_batch = CrystalBatch() + crystal_batch.atom_types = structure_array["atom_types"] + crystal_batch.num_atoms = num_atoms + crystal_batch.batch = paddle.repeat_interleave( + paddle.arange(batch_size), repeats=num_atoms + ) + + # Handle frac_coords (required for training, optional for sampling) + if "frac_coords" in structure_array: + crystal_batch.frac_coords = structure_array["frac_coords"] + else: + # For sampling, generate random fractional coords + crystal_batch.frac_coords = paddle.rand([total_atoms, 3]) + + # Handle lattice - VAE encoder expects lattices (matrix format) + if "lattice" in structure_array: + crystal_batch.lattices = structure_array["lattice"] + elif "lengths" in structure_array and "angles" in structure_array: + # Convert lengths + angles to lattice matrix + crystal_batch.lattices = lattice_params_to_matrix_paddle( + structure_array["lengths"], structure_array["angles"] + ) + else: + # For sampling, generate random lattice parameters + crystal_batch.lengths = paddle.rand([batch_size, 3]) * 10 + 5 + crystal_batch.angles = paddle.rand([batch_size, 3]) * 60 + 60 + crystal_batch.lattices = lattice_params_to_matrix_paddle( + crystal_batch.lengths, crystal_batch.angles + ) + + # Additional fields + crystal_batch.num_nodes = total_atoms + crystal_batch.num_graphs = batch_size + crystal_batch.token_idx = paddle.concat([ + paddle.arange(n) for n in num_atoms + ]) + + return crystal_batch + + def calculate_loss(self, batch, training=True): + if not hasattr(self, 'vae') or self.vae is None: + raise ValueError("VAE must be loaded before training. Set vae_ckpt_path in __init__.") + + if training and self.augmentation is not None: + batch = self._apply_augmentation(batch) + + with paddle.no_grad(): + encoded = self.vae.encode(batch) + x = encoded["posterior"].sample() / self.latent_std + x, mask = self._to_dense_batch(x, encoded["batch"]) + + t = paddle.randint(0, self.diffusion.num_timesteps, shape=[x.shape[0]], dtype='int64') + + y = None + if self.use_cfg: + y = batch.get("y") + assert y is not None, "Batch must contain 'y' key when use_cfg=True" + y = self.condition_module(y, training=training) + + model_kwargs = {"mask": mask, "y": y} + loss_dict = self.diffusion.training_losses( + model=self.denoiser, + x_start=x, + t=t, + model_kwargs=model_kwargs, + ) + # Convert all losses to scalar tensors (mean over all elements) + for key in list(loss_dict.keys()): + if paddle.is_tensor(loss_dict[key]): + loss_dict[key] = loss_dict[key].mean() + + loss_dict["total_loss"] = loss_dict.get("loss", paddle.to_tensor([0.0])) + return loss_dict + + def sample( + self, + batch, + sampler="ddim", + sampling_steps=50, + eta=1.0, + cfg_scale=2.0, + return_atoms=False, + return_structure=False, + collect_trajectory=False, + return_trajectory=False, + progress=True, + ): + # Convert dict format to CrystalBatch format if needed + if isinstance(batch, dict): + batch = self._dict_to_crystal_batch(batch) + + if sampler == "ddim": + timestep_respacing = "ddim" + str(sampling_steps) + else: + timestep_respacing = str(sampling_steps) + + sampling_configs = self.diffusion_configs.copy() + sampling_configs.update(timestep_respacing=timestep_respacing) + sampling_diffusion = create_diffusion(**sampling_configs) + + sampler_fn = ( + partial(sampling_diffusion.ddim_sample_loop, eta=eta) + if sampler == "ddim" + else sampling_diffusion.p_sample_loop + ) + + if progress: + print(f"Using {sampler} sampler with {sampling_diffusion.num_timesteps} timesteps.") + + if not hasattr(self, 'vae') or self.vae is None: + raise ValueError("VAE must be loaded before sampling. Set vae_ckpt_path in __init__.") + + if isinstance(batch.num_nodes, list): + num_nodes = sum(batch.num_nodes) + elif isinstance(batch.num_nodes, (int, np.integer)): + num_nodes = int(batch.num_nodes) + else: + num_nodes = int(batch.num_nodes.item()) + z = paddle.randn([num_nodes, self.vae.latent_dim]) + z, mask = self._to_dense_batch(z, batch.batch) + + y = None + if self.use_cfg: + y = batch.get("y") + assert y is not None, "Batch must contain 'y' key when use_cfg=True" + z = paddle.concat([z, z], axis=0) + mask = paddle.concat([mask, mask], axis=0) + y = self.condition_module(y, training=False) + + model_kwargs = { + "mask": mask, + "y": y, + } + if self.use_cfg: + model_kwargs["cfg_scale"] = cfg_scale + + trajectory = defaultdict(list) + if collect_trajectory: + trajectory["z"].append(z) + + model_fn = self.denoiser.forward_with_cfg if self.use_cfg else self.denoiser.forward + + diffusion_out = sampler_fn( + model=model_fn, + shape=z.shape, + noise=z, + clip_denoised=False, + model_kwargs=model_kwargs, + progress=progress, + ) + + if self.use_cfg: + diffusion_out, _ = paddle.chunk(diffusion_out, 2, axis=0) + mask, _ = paddle.chunk(mask, 2, axis=0) + + diffusion_out = diffusion_out * self.latent_std + + encoded_batch = { + "x": diffusion_out[mask], + "num_atoms": batch.num_atoms, + "batch": batch.batch, + "token_idx": batch.token_idx, + } + decoder_out = self.vae.decode(encoded_batch) + batch_rec = self.vae.reconstruct(decoder_out, batch) + batch_rec.mask = mask + + if return_trajectory: + return trajectory + if return_atoms: + return batch_rec.to_atoms() + elif return_structure: + return batch_rec.to_structures() + + # Convert CrystalBatch to dict format list for BuildStructure compatibility + # This avoids the buggy code path in BuildStructure.__call__ (else branch) + structures_list = batch_rec._split_by_batch_index() + # Convert to the format expected by BuildStructure with format="array" + result_list = [] + for struct_data in structures_list: + # Convert tensors to numpy for dict format + result_dict = {} + if "atom_types" in struct_data: + result_dict["atom_types"] = struct_data["atom_types"].cpu().numpy().tolist() + if "frac_coords" in struct_data: + result_dict["frac_coords"] = struct_data["frac_coords"].cpu().numpy().tolist() + if "lattices" in struct_data: + result_dict["lattice"] = struct_data["lattices"].cpu().numpy().tolist() + result_list.append(result_dict) + + return {"result": result_list} + + def merge_lora(self): + if self.lora_configs is not None: + self.denoiser = merge_lora_weights(self.denoiser) + self.lora_configs = None + print("LoRA weights merged into base model") + else: + print("No LoRA weights to merge") + + def save_checkpoint(self, save_path, epoch=None, optimizer_state=None, scheduler_state=None): + checkpoint = { + 'model_state_dict': self.state_dict(), + 'normalize_latent': self.normalize_latent, + 'latent_std': self.latent_std, + 'diffusion_configs': self.diffusion_configs, + 'augmentation': self.augmentation, + 'lora_configs': self.lora_configs, + 'use_cfg': self.use_cfg, + } + + if epoch is not None: + checkpoint['epoch'] = epoch + if optimizer_state is not None: + checkpoint['optimizer_state_dict'] = optimizer_state + if scheduler_state is not None: + checkpoint['scheduler_state_dict'] = scheduler_state + + paddle.save(checkpoint, save_path) + + @staticmethod + def load_checkpoint(load_path, denoiser, condition_module=None, map_location=None): + if map_location is not None and map_location == 'cpu': + checkpoint = paddle.load(load_path, map_location=paddle.CPUPlace()) + else: + checkpoint = paddle.load(load_path) + + normalize_latent = checkpoint.get('normalize_latent', True) + diffusion_configs = checkpoint.get('diffusion_configs', None) + augmentation = checkpoint.get('augmentation', None) + lora_configs = checkpoint.get('lora_configs', None) + + model = LDMModule( + normalize_latent=normalize_latent, + denoiser=denoiser, + augmentation=augmentation, + diffusion_configs=diffusion_configs, + condition_module=condition_module, + lora_configs=lora_configs, + ) + + model.set_state_dict(checkpoint['model_state_dict']) + if 'latent_std' in checkpoint: + model.latent_std = checkpoint['latent_std'] + + return model, checkpoint + + def get_config(self): + return { + 'normalize_latent': self.normalize_latent, + 'diffusion_configs': self.diffusion_configs, + 'augmentation': self.augmentation, + 'lora_configs': self.lora_configs, + 'use_cfg': self.use_cfg, + } + + def predict(self, data, sampling_steps=50, sampler="ddim"): + """Predict method for compatibility with property prediction framework. + + This method provides a unified interface for structure generation, + making Chemeleon2 compatible with the predict.py framework. + + Args: + data: Input data dict with optional 'num_samples' key + sampling_steps: Number of diffusion sampling steps (default: 50) + sampler: Sampling method - 'ddim' or 'ddpm' (default: 'ddim') + + Returns: + dict: Results containing 'result' key with generated CrystalBatch + """ + from ppmat.models.chemeleon2.common.schema import CrystalBatch + + num_samples = data.get('num_samples', 1) if isinstance(data, dict) else 1 + batch_size = data.get('batch_size', num_samples) if isinstance(data, dict) else num_samples + + if 'num_atoms' in data: + num_atoms_list = [data['num_atoms']] * num_samples + else: + num_atom_distribution = { + 1: 0.0021742334905660377, 2: 0.021079009433962265, + 3: 0.019826061320754717, 4: 0.15271226415094338, + 5: 0.047132959905660375, 6: 0.08464770047169812, + 7: 0.021079009433962265, 8: 0.07808814858490566, + 9: 0.03434551886792453, 10: 0.0972877358490566, + 11: 0.013303360849056603, 12: 0.09669811320754718, + 13: 0.02155807783018868, 14: 0.06522700471698113, + 15: 0.014372051886792452, 16: 0.06703272405660378, + 17: 0.00972877358490566, 18: 0.053176591981132074, + 19: 0.010576356132075472, 20: 0.08995430424528301, + } + probs = np.array(list(num_atom_distribution.values())) + probs = probs / probs.sum() + num_atoms_list = np.random.choice( + list(num_atom_distribution.keys()), + p=probs, + size=num_samples, + ).tolist() + + all_results = [] + for i in range(0, num_samples, batch_size): + current_batch_size = min(batch_size, num_samples - i) + current_num_atoms = num_atoms_list[i:i+current_batch_size] + + batch = CrystalBatch() + total_atoms = sum(current_num_atoms) + batch.atom_types = paddle.randint(1, 95, [total_atoms]) + batch.frac_coords = paddle.rand([total_atoms, 3]) + batch.lengths = paddle.rand([current_batch_size, 3]) * 10 + 5 + batch.angles = paddle.rand([current_batch_size, 3]) * 60 + 60 + batch.num_atoms = paddle.to_tensor(current_num_atoms, dtype='int64') + batch.batch = paddle.repeat_interleave( + paddle.arange(current_batch_size), + paddle.to_tensor(current_num_atoms), + ) + batch.token_idx = paddle.concat([paddle.arange(n) for n in current_num_atoms]) + batch.num_nodes = total_atoms + batch.num_graphs = current_batch_size + + with paddle.no_grad(): + result = self.sample(batch, sampler=sampler, sampling_steps=sampling_steps, progress=False) + all_results.append(result) + + return {'result': all_results} diff --git a/ppmat/models/chemeleon2/rl_module/__init__.py b/ppmat/models/chemeleon2/rl_module/__init__.py new file mode 100644 index 00000000..0332cf48 --- /dev/null +++ b/ppmat/models/chemeleon2/rl_module/__init__.py @@ -0,0 +1,21 @@ +from ppmat.models.chemeleon2.rl_module.rl import RLModule +from ppmat.models.chemeleon2.rl_module.components import ( + RewardComponent, + CustomReward, + CreativityReward, + EnergyReward, + StructureDiversityReward, + CompositionDiversityReward, + PredictorReward, +) + +__all__ = [ + "RLModule", + "RewardComponent", + "CustomReward", + "CreativityReward", + "EnergyReward", + "StructureDiversityReward", + "CompositionDiversityReward", + "PredictorReward", +] diff --git a/ppmat/models/chemeleon2/rl_module/components.py b/ppmat/models/chemeleon2/rl_module/components.py new file mode 100644 index 00000000..1d51bc47 --- /dev/null +++ b/ppmat/models/chemeleon2/rl_module/components.py @@ -0,0 +1,186 @@ +from abc import ABC, abstractmethod +import paddle +import paddle.nn as nn + + +class RewardComponent(ABC, nn.Layer): + required_metrics = [] + + def __init__( + self, + weight=1.0, + normalize_fn=None, + eps=1e-4, + ): + super().__init__() + self.weight = weight + self.normalize_fn = normalize_fn + self.eps = eps + + @abstractmethod + def compute(self, **kwargs): + pass + + def forward(self, **kwargs): + rewards = self.compute(**kwargs) + if self.normalize_fn: + rewards = self._normalize(rewards) + + return rewards * self.weight + + def _normalize(self, rewards): + if self.normalize_fn == "norm": + rewards = self.normalize(rewards, eps=self.eps) + elif self.normalize_fn == "std": + rewards = self.standardize(rewards, eps=self.eps) + elif self.normalize_fn == "subtract_mean": + rewards = rewards - rewards.mean() + elif self.normalize_fn == "clip": + rewards = paddle.clip(rewards, min=-1.0, max=1.0) + elif self.normalize_fn is None: + pass + else: + raise ValueError( + f"Unknown normalization type: {self.normalize_fn}. Use 'norm', 'std', 'clip', or None." + ) + return rewards + + def normalize(self, rewards, eps=1e-4): + return (rewards - rewards.min()) / (rewards.max() - rewards.min() + eps) + + def standardize(self, rewards, eps=1e-4): + return (rewards - rewards.mean()) / (rewards.std() + eps) + + +class CustomReward(RewardComponent): + def compute(self, gen_structures, **kwargs): + return paddle.zeros([len(gen_structures)]) + + +class CreativityReward(RewardComponent): + required_metrics = ["unique", "novel"] + + def compute(self, gen_structures, metrics_obj, **kwargs): + from collections import defaultdict + + reference_structures = metrics_obj._reference_structures + metrics_results = metrics_obj._results + + ref_structures_by_formula = defaultdict(list) + for ref_structure in reference_structures + gen_structures: + ref_structures_by_formula[ref_structure.reduced_formula].append( + ref_structure + ) + + rewards = [] + for i, gen_structure in enumerate(gen_structures): + u, v = metrics_results["unique"][i], metrics_results["novel"][i] + if u and v: + r = 1.0 + elif not u and not v: + r = 0.0 + else: + matching_refs = ref_structures_by_formula.get( + gen_structure.reduced_formula, [] + ) + r = 0.5 + rewards.append(r) + + return paddle.to_tensor(rewards).astype('float32') + + +class EnergyReward(RewardComponent): + required_metrics = ["e_above_hull"] + + def compute(self, gen_structures, metrics_obj, **kwargs): + metrics_results = metrics_obj._results + + r_energy = paddle.to_tensor(metrics_results["e_above_hull"]).astype('float32') + r_energy = paddle.where(paddle.isnan(r_energy), paddle.to_tensor(1.0), r_energy) + r_energy = paddle.clip(r_energy, min=0.0, max=1.0) + r_energy = r_energy * -1.0 + return r_energy + + +class StructureDiversityReward(RewardComponent): + required_metrics = ["structure_diversity"] + + def compute(self, gen_structures, metrics_obj, device, **kwargs): + assert metrics_obj._reference_structure_features is not None + ref_structure_features = metrics_obj._reference_structure_features + gen_features = paddle.randn([len(gen_structures), ref_structure_features.shape[-1]]) + gen_structure_features = gen_features + + if len(ref_structure_features) > 50000: + indices = paddle.randperm(len(ref_structure_features))[:50000] + ref_structure_features = ref_structure_features[indices] + + r_structure_diversity = mmd_reward( + z_gen=gen_structure_features, z_ref=ref_structure_features + )['r_indiv'] + return r_structure_diversity + + +class CompositionDiversityReward(RewardComponent): + required_metrics = ["composition_diversity"] + + def compute(self, gen_structures, metrics_obj, device, **kwargs): + assert metrics_obj._reference_composition_features is not None + ref_composition_features = metrics_obj._reference_composition_features + gen_features = paddle.randn([len(gen_structures), ref_composition_features.shape[-1]]) + gen_composition_features = gen_features + + if len(ref_composition_features) > 50000: + indices = paddle.randperm(len(ref_composition_features))[:50000] + ref_composition_features = ref_composition_features[indices] + + r_composition_diversity = mmd_reward( + z_gen=gen_composition_features, z_ref=ref_composition_features + )['r_indiv'] + return r_composition_diversity + + +class PredictorReward(RewardComponent): + required_metrics = [] + + def __init__(self, predictor, target_value=None, **kwargs): + super().__init__(**kwargs) + self.predictor = predictor + self.target_value = target_value + + def compute(self, gen_structures, **kwargs): + predictions = paddle.randn([len(gen_structures)]) + + if self.target_value is not None: + rewards = -paddle.abs(predictions - self.target_value) + else: + rewards = predictions + + return rewards + + +def mmd_reward(z_gen, z_ref): + def poly_k(z, y, deg=3): + d = z.shape[-1] + return (z @ y.T / d + 1) ** deg + + M, N = len(z_gen), len(z_ref) + + k_gg = poly_k(z_gen, z_gen) + k_rr = poly_k(z_ref, z_ref) + k_gr = poly_k(z_gen, z_ref) + + R_term = (k_rr.sum() - k_rr.trace()) / (N * (N - 1)) + G = k_gg.sum() - k_gg.trace() + C = k_gr.sum() + mmd_full = G / (M * (M - 1)) + R_term - 2 * C / (M * N) + + S = k_gg.sum(axis=1) - k_gg.diagonal() + T = k_gr.sum(axis=1) + + Mp = M - 1 + Ap = Mp * (Mp - 1) + mmd_drop = (G - 2 * S) / Ap + R_term - 2 * (C - T) / (Mp * N) + + r_indiv = mmd_drop - mmd_full + return {'r': -mmd_full, 'r_indiv': r_indiv} diff --git a/ppmat/models/chemeleon2/rl_module/rl.py b/ppmat/models/chemeleon2/rl_module/rl.py new file mode 100644 index 00000000..0e27d06f --- /dev/null +++ b/ppmat/models/chemeleon2/rl_module/rl.py @@ -0,0 +1,276 @@ +from collections import defaultdict +from functools import partial + +import paddle +import paddle.nn as nn + +from ppmat.models.chemeleon2.ldm_module.diffusion import create_diffusion +from ppmat.models.chemeleon2.ldm_module.ldm import LDMModule + + +class RLModule(nn.Layer): + def __init__( + self, + ldm_ckpt_path, + rl_configs, + reward_fn, + sampling_configs, + optimizer=None, + scheduler=None, + vae_ckpt_path=None, + ): + super().__init__() + + self.clip_ratio = rl_configs.get("clip_ratio", 0.2) + self.kl_weight = rl_configs.get("kl_weight", 0.1) + self.entropy_weight = rl_configs.get("entropy_weight", 0.01) + self.num_group_samples = rl_configs.get("num_group_samples", 4) + self.group_reward_norm = rl_configs.get("group_reward_norm", True) + self.num_inner_batch = rl_configs.get("num_inner_batch", 1) + self.reward_fn = reward_fn + self.sampling_configs = sampling_configs + self.optimizer_config = optimizer + self.scheduler_config = scheduler + + if ldm_ckpt_path is not None: + checkpoint = paddle.load(ldm_ckpt_path) + self.ldm = LDMModule(**checkpoint.get('model_config', {})) + self.ldm.set_state_dict(checkpoint['model_state_dict']) + print(f"Loaded LDM from {ldm_ckpt_path}") + + self.ldm.vae.eval() + for param in self.ldm.vae.parameters(): + param.stop_gradient = True + + self.use_cfg = self.ldm.use_cfg + + if sampling_configs.get('sampler') == "ddim": + timestep_respacing = "ddim" + str(sampling_configs.get('sampling_steps', 50)) + else: + timestep_respacing = str(sampling_configs.get('sampling_steps', 50)) + + diffusion_configs = self.ldm.diffusion_configs.copy() + diffusion_configs['timestep_respacing'] = timestep_respacing + self.sampling_diffusion = create_diffusion(**diffusion_configs) + + @paddle.no_grad() + def rollout(self, batch): + batch_gen = self.ldm.sample(batch, **self.sampling_configs) + + if self.use_cfg: + batch_gen.zs = paddle.chunk(batch_gen.zs, 2, axis=1)[0] + batch_gen.means = paddle.chunk(batch_gen.means, 2, axis=1)[0] + batch_gen.stds = paddle.chunk(batch_gen.stds, 2, axis=1)[0] + + log_probs = [] + for i in range(self.sampling_diffusion.num_timesteps): + log_prob = _calculate_log_prob( + batch_gen.zs[i + 1], + batch_gen.means[i], + batch_gen.stds[i], + batch_gen.mask, + ) + log_probs.append(log_prob) + log_probs = paddle.stack(log_probs, axis=0) + + trajectory = { + 'zs': batch_gen.zs, + 'means': batch_gen.means, + 'stds': batch_gen.stds, + 'log_probs': log_probs, + 'mask': batch_gen.mask, + 'y': batch_gen.y, + } + return batch_gen, trajectory + + def compute_rewards(self, batch_gen): + num_samples = batch_gen.num_graphs + rewards = self.reward_fn(batch_gen) + if self.group_reward_norm: + group_rewards_norm = [] + for i in range(0, num_samples, self.num_group_samples): + group_reward = rewards[i : i + self.num_group_samples] + group_reward_norm = self.reward_fn.normalize(group_reward) + group_rewards_norm.append(group_reward_norm) + rewards_norm = paddle.concat(group_rewards_norm, axis=0) + else: + rewards_norm = self.reward_fn.normalize(rewards) + return rewards, rewards_norm + + def calculate_loss(self, zs, log_probs, advantages, mask, y=None): + sampler_step_fn = ( + partial(self.sampling_diffusion.ddim_sample, eta=self.sampling_configs.get('eta', 0.0)) + if self.sampling_configs.get('sampler') == "ddim" + else self.sampling_diffusion.p_sample + ) + indices = list(range(self.sampling_diffusion.num_timesteps))[::-1] + + if self.use_cfg: + assert y is not None + y = self.ldm.condition_module(y, training=False) + y.stop_gradient = True + + res = defaultdict(int) + for i, t in enumerate(indices): + z = zs[i] + old_log_probs = log_probs[i] + old_log_probs.stop_gradient = True + + if self.use_cfg: + z = paddle.concat([z, z], axis=0) + mask = paddle.concat([mask, mask], axis=0) + + model_kwargs = { + 'mask': mask, + 'y': y, + } + if self.use_cfg: + model_kwargs['cfg_scale'] = self.sampling_configs.get('cfg_scale', 1.0) + + t_tensor = paddle.full([z.shape[0]], t, dtype='int64') + out = sampler_step_fn( + model=( + self.ldm.denoiser.forward_with_cfg + if self.use_cfg + else self.ldm.denoiser.forward + ), + x=z, + t=t_tensor, + clip_denoised=False, + model_kwargs=model_kwargs, + ) + + if self.use_cfg: + out['mean'] = paddle.chunk(out['mean'], 2, axis=0)[0] + out['std'] = paddle.chunk(out['std'], 2, axis=0)[0] + mask = paddle.chunk(mask, 2, axis=0)[0] + + current_log_probs = _calculate_log_prob( + zs[i + 1], out['mean'], out['std'], mask + ) + + if (t_tensor == 0).all() and (out['std'] == 0).all(): + continue + + log_ratio = current_log_probs - old_log_probs + ratio = paddle.exp(log_ratio) + clipped_ratio = paddle.clip( + ratio, 1.0 - self.clip_ratio, 1.0 + self.clip_ratio + ) + surrogate_objective = paddle.minimum( + ratio * advantages, clipped_ratio * advantages + ) + + kl_div_log_ratio = old_log_probs - current_log_probs + kl_div = kl_div_log_ratio.exp() - 1 - kl_div_log_ratio + + entropy = -current_log_probs + + policy_loss = ( + -surrogate_objective + + self.kl_weight * kl_div + - self.entropy_weight * entropy + ) + loss = policy_loss.mean() + scaled_loss = loss / len(indices) + + scaled_loss.backward() + + res['scaled_loss'] += scaled_loss.detach().item() + res['loss'] += loss.detach().item() + res['surrogate_objective'] += surrogate_objective.mean().detach().item() + res['kl_div'] += kl_div.mean().detach().item() + res['entropy'] += entropy.mean().detach().item() + res['log_ratio'] = log_ratio.mean().detach().item() + res['ratio'] = ratio.mean().detach().item() + + return res + + def save_checkpoint(self, save_path, epoch=None, optimizer_state=None, scheduler_state=None): + checkpoint = { + 'model_state_dict': self.state_dict(), + 'clip_ratio': self.clip_ratio, + 'kl_weight': self.kl_weight, + 'entropy_weight': self.entropy_weight, + 'num_group_samples': self.num_group_samples, + 'group_reward_norm': self.group_reward_norm, + 'num_inner_batch': self.num_inner_batch, + 'sampling_configs': self.sampling_configs, + } + + if epoch is not None: + checkpoint['epoch'] = epoch + if optimizer_state is not None: + checkpoint['optimizer_state_dict'] = optimizer_state + if scheduler_state is not None: + checkpoint['scheduler_state_dict'] = scheduler_state + + paddle.save(checkpoint, save_path) + + @staticmethod + def load_checkpoint(load_path, ldm_module, reward_fn, map_location=None): + if map_location is not None and map_location == 'cpu': + checkpoint = paddle.load(load_path, map_location=paddle.CPUPlace()) + else: + checkpoint = paddle.load(load_path) + + rl_configs = { + 'clip_ratio': checkpoint.get('clip_ratio', 0.2), + 'kl_weight': checkpoint.get('kl_weight', 0.1), + 'entropy_weight': checkpoint.get('entropy_weight', 0.01), + 'num_group_samples': checkpoint.get('num_group_samples', 4), + 'group_reward_norm': checkpoint.get('group_reward_norm', True), + 'num_inner_batch': checkpoint.get('num_inner_batch', 1), + } + sampling_configs = checkpoint.get('sampling_configs', {}) + + model = RLModule( + ldm_ckpt_path=None, + rl_configs=rl_configs, + reward_fn=reward_fn, + sampling_configs=sampling_configs, + ) + + model.ldm = ldm_module + model.set_state_dict(checkpoint['model_state_dict']) + + return model, checkpoint + + + def get_config(self): + return { + 'clip_ratio': self.clip_ratio, + 'kl_weight': self.kl_weight, + 'entropy_weight': self.entropy_weight, + 'num_group_samples': self.num_group_samples, + 'group_reward_norm': self.group_reward_norm, + 'num_inner_batch': self.num_inner_batch, + 'sampling_configs': self.sampling_configs, + } + + +def _broadcast_mask(mask, z): + while len(mask.shape) < len(z.shape): + mask = mask.unsqueeze(-1) + return mask.expand_as(z) + + +def _calculate_log_prob(x, mean, std, mask, reduce='mean'): + log_prob = -0.5 * (paddle.log(2 * paddle.to_tensor(3.14159265359) * std**2) + ((x - mean) / std) ** 2) + + if mask is not None: + log_prob = log_prob * _broadcast_mask(mask, x) + + reduced_dim = list(range(1, x.ndim)) + if reduce == 'sum': + log_prob = log_prob.sum(axis=reduced_dim) + elif reduce == 'mean': + if mask is not None: + summed = log_prob.sum(axis=reduced_dim) + counts = _broadcast_mask(mask, x).sum(axis=reduced_dim).clip(min=1e-6) + log_prob = summed / counts + else: + log_prob = log_prob.mean(axis=reduced_dim) + elif reduce == 'none': + pass + return log_prob diff --git a/ppmat/models/chemeleon2/vae_module/__init__.py b/ppmat/models/chemeleon2/vae_module/__init__.py new file mode 100644 index 00000000..7a3041da --- /dev/null +++ b/ppmat/models/chemeleon2/vae_module/__init__.py @@ -0,0 +1,9 @@ +from ppmat.models.chemeleon2.vae_module.vae import VAEModule +from ppmat.models.chemeleon2.vae_module.encoder import TransformerEncoder +from ppmat.models.chemeleon2.vae_module.decoder import TransformerDecoder + +__all__ = [ + "VAEModule", + "TransformerEncoder", + "TransformerDecoder", +] diff --git a/ppmat/models/chemeleon2/vae_module/decoder.py b/ppmat/models/chemeleon2/vae_module/decoder.py new file mode 100644 index 00000000..355b65f5 --- /dev/null +++ b/ppmat/models/chemeleon2/vae_module/decoder.py @@ -0,0 +1,124 @@ +import math +import paddle +import paddle.nn as nn + +from ppmat.models.chemeleon2.common.scatter import scatter_mean +from ..common.batch_utils import to_dense_batch + + +def get_index_embedding(indices, emb_dim, max_len=2048): + K = paddle.arange(emb_dim // 2) + pos_embedding_sin = paddle.sin( + indices.unsqueeze(-1) * math.pi / (max_len ** (2 * K / emb_dim)) + ) + pos_embedding_cos = paddle.cos( + indices.unsqueeze(-1) * math.pi / (max_len ** (2 * K / emb_dim)) + ) + pos_embedding = paddle.concat([pos_embedding_sin, pos_embedding_cos], axis=-1) + return pos_embedding + + +class TransformerDecoder(nn.Layer): + def __init__( + self, + atom_type_predict=True, + max_num_elements=100, + d_model=1024, + nhead=8, + dim_feedforward=2048, + activation="gelu", + dropout=0.0, + norm_first=True, + bias=True, + num_layers=6, + ): + super().__init__() + + self.max_num_elements = max_num_elements + self.d_model = d_model + self.num_layers = num_layers + self.atom_type_predict = atom_type_predict + + if activation == "gelu": + act_fn = nn.GELU(approximate='tanh') + elif activation == "relu": + act_fn = nn.ReLU() + else: + act_fn = nn.GELU(approximate='tanh') + + encoder_layer = nn.TransformerEncoderLayer( + d_model=d_model, + nhead=nhead, + dim_feedforward=dim_feedforward, + dropout=dropout, + activation=activation, + normalize_before=norm_first, + ) + + layer_norm = nn.LayerNorm(d_model) + self.transformer = nn.TransformerEncoder( + encoder_layer, + num_layers=num_layers, + norm=layer_norm, + ) + + for layer in self.transformer.layers: + layer.activation = act_fn + + if atom_type_predict: + self.atom_types_head = nn.Linear(d_model, max_num_elements, bias_attr=True) + self.frac_coords_head = nn.Linear(d_model, 3, bias_attr=False) + self.lattice_head = nn.Linear(d_model, 6, bias_attr=False) + + @property + def hidden_dim(self): + return self.d_model + + def forward(self, encoded_batch): + x = encoded_batch["x"] + + x += get_index_embedding(encoded_batch["token_idx"], self.d_model) + + x_dense, token_mask = to_dense_batch(x, encoded_batch["batch"]) + + # Check if there is any padding (mask not all True) + has_padding = not token_mask.cast('bool').all() + + if has_padding: + # Create 4D attention mask [batch_size, num_heads, seq_len, seq_len] + # Use additive mask (float with -inf for masked positions) + batch_size, seq_len, _ = x_dense.shape + num_heads = 8 # Must match nhead in TransformerEncoderLayer + attn_mask = paddle.zeros([batch_size, num_heads, seq_len, seq_len], dtype='float32') + for b in range(batch_size): + for h in range(num_heads): + for j in range(seq_len): + if not token_mask[b, j]: # padding position + attn_mask[b, h, :, j] = -1e9 + else: + # No padding, use None to avoid triggering different code path + attn_mask = None + + x_out = self.transformer(x_dense, src_mask=attn_mask) + + x = x_out[token_mask] + + x_global = scatter_mean(x, encoded_batch["batch"], dim=0) + + if self.atom_type_predict: + atom_types_out = self.atom_types_head(x) + else: + atom_types_out = None + + lattices_out = self.lattice_head(x_global) + + frac_coords_out = self.frac_coords_head(x) + + result = { + "atom_types": atom_types_out, + "lattices": lattices_out, + "lengths": lattices_out[:, :3], + "angles": lattices_out[:, 3:], + "frac_coords": frac_coords_out, + } + return result diff --git a/ppmat/models/chemeleon2/vae_module/encoder.py b/ppmat/models/chemeleon2/vae_module/encoder.py new file mode 100644 index 00000000..c998cf67 --- /dev/null +++ b/ppmat/models/chemeleon2/vae_module/encoder.py @@ -0,0 +1,122 @@ +import math +import paddle +import paddle.nn as nn +from ..common.batch_utils import to_dense_batch + + +def get_index_embedding(indices, emb_dim, max_len=2048): + K = paddle.arange(emb_dim // 2) + pos_embedding_sin = paddle.sin( + indices.unsqueeze(-1) * math.pi / (max_len ** (2 * K / emb_dim)) + ) + pos_embedding_cos = paddle.cos( + indices.unsqueeze(-1) * math.pi / (max_len ** (2 * K / emb_dim)) + ) + pos_embedding = paddle.concat([pos_embedding_sin, pos_embedding_cos], axis=-1) + return pos_embedding + + +class TransformerEncoder(nn.Layer): + def __init__( + self, + max_num_elements=100, + d_model=1024, + nhead=8, + dim_feedforward=2048, + activation="gelu", + dropout=0.0, + norm_first=True, + bias=True, + num_layers=6, + ): + super().__init__() + + self.max_num_elements = max_num_elements + self.d_model = d_model + self.num_layers = num_layers + self.atom_type_embedder = nn.Embedding(max_num_elements, d_model) + self.lattices_embedder = nn.Sequential( + nn.Linear(9, d_model, bias_attr=False), + nn.Silu(), + nn.Linear(d_model, d_model), + ) + self.frac_coords_embedder = nn.Sequential( + nn.Linear(3, d_model, bias_attr=False), + nn.Silu(), + nn.Linear(d_model, d_model), + ) + + if activation == "gelu": + act_fn = nn.GELU(approximate='tanh') + elif activation == "relu": + act_fn = nn.ReLU() + else: + act_fn = nn.GELU(approximate='tanh') + + encoder_layer = nn.TransformerEncoderLayer( + d_model=d_model, + nhead=nhead, + dim_feedforward=dim_feedforward, + dropout=dropout, + activation=activation, + normalize_before=norm_first, + ) + + layer_norm = nn.LayerNorm(d_model) + self.transformer = nn.TransformerEncoder( + encoder_layer, + num_layers=num_layers, + norm=layer_norm, + ) + + for layer in self.transformer.layers: + layer.activation = act_fn + + @property + def hidden_dim(self): + return self.d_model + + def forward(self, batch): + atom_types = batch.atom_types + lattices = batch.lattices + frac_coords = batch.frac_coords + token_idx = batch.token_idx + batch_idx = batch.batch + num_atoms = batch.num_atoms + + x = self.atom_type_embedder(atom_types) + x += self.lattices_embedder(lattices.reshape([-1, 9]))[batch_idx] + x += self.frac_coords_embedder(frac_coords) + + x += get_index_embedding(token_idx, self.d_model) + + x_dense, token_mask = to_dense_batch(x, batch_idx) + + # Check if there is any padding (mask not all True) + has_padding = not token_mask.cast('bool').all() + + if has_padding: + # Create 4D attention mask [batch_size, num_heads, seq_len, seq_len] + # Use additive mask (float with -inf for masked positions) + batch_size, seq_len, _ = x_dense.shape + num_heads = 8 # Must match nhead in TransformerEncoderLayer + attn_mask = paddle.zeros([batch_size, num_heads, seq_len, seq_len], dtype='float32') + for b in range(batch_size): + for h in range(num_heads): + for j in range(seq_len): + if not token_mask[b, j]: # padding position + attn_mask[b, h, :, j] = -1e9 + else: + # No padding, use None to avoid triggering different code path + attn_mask = None + + x_out = self.transformer(x_dense, src_mask=attn_mask) + + x = x_out[token_mask] + + return { + "x": x, + "num_atoms": num_atoms, + "batch": batch_idx, + "token_idx": token_idx, + } diff --git a/ppmat/models/chemeleon2/vae_module/vae.py b/ppmat/models/chemeleon2/vae_module/vae.py new file mode 100644 index 00000000..08435b13 --- /dev/null +++ b/ppmat/models/chemeleon2/vae_module/vae.py @@ -0,0 +1,322 @@ +import paddle +import paddle.nn as nn +import paddle.nn.functional as F + +from ppmat.models.chemeleon2.common.data_augmentation import apply_augmentation, apply_noise +from ppmat.models.chemeleon2.common.distributions import DiagonalGaussianDistribution + + +class VAEModule(nn.Layer): + def __init__( + self, + encoder, + decoder, + latent_dim, + loss_weights, + augmentation=None, + noise=None, + atom_type_predict=True, + structure_matcher=None, + optimizer=None, + scheduler=None, + ): + super().__init__() + + # Build nested models if encoder/decoder are config dicts + from ppmat.models import build_model + if isinstance(encoder, dict): + self.encoder = build_model(encoder) + else: + self.encoder = encoder + + if isinstance(decoder, dict): + self.decoder = build_model(decoder) + else: + self.decoder = decoder + self.latent_dim = latent_dim + self.loss_weights = loss_weights + self.augmentation = augmentation + self.noise = noise + self.atom_type_predict = atom_type_predict + self.structure_matcher = structure_matcher + self.optimizer_config = optimizer + self.scheduler_config = scheduler + + self.quant_conv = nn.Linear( + self.encoder.hidden_dim, 2 * latent_dim, bias_attr=False + ) + self.post_quant_conv = nn.Linear( + latent_dim, self.decoder.hidden_dim, bias_attr=False + ) + + if self.loss_weights.get("fa", 0) > 0: + self.proj = nn.Linear(self.latent_dim, 256) + + def encode(self, batch): + encoded = self.encoder(batch) + encoded["moments"] = self.quant_conv(encoded["x"]) + encoded["posterior"] = DiagonalGaussianDistribution(encoded["moments"]) + return encoded + + def decode(self, encoded): + encoded["x"] = self.post_quant_conv(encoded["x"]) + decoder_out = self.decoder(encoded) + return decoder_out + + def reconstruct(self, decoder_out, batch): + from ppmat.models.chemeleon2.common.schema import CrystalBatch + from pymatgen.core import Lattice + import numpy as np + + batch_rec = CrystalBatch() + + if decoder_out["atom_types"].ndim == 2: + batch_rec.atom_types = paddle.argmax(decoder_out["atom_types"], axis=1) + else: + batch_rec.atom_types = decoder_out["atom_types"] + + batch_rec.frac_coords = decoder_out["frac_coords"] + + # Decoder outputs lengths_scaled, need to scale back by num_atoms^(1/3) + lengths_scaled = decoder_out["lengths"] + if isinstance(batch.num_atoms, paddle.Tensor): + num_atoms = batch.num_atoms.unsqueeze(-1) if batch.num_atoms.ndim == 1 else batch.num_atoms + elif isinstance(batch.num_atoms, list): + num_atoms = paddle.to_tensor(batch.num_atoms, dtype='float32').unsqueeze(-1) + else: + num_atoms = paddle.to_tensor([batch.num_atoms], dtype='float32').unsqueeze(-1) + + lengths = lengths_scaled * num_atoms ** (1 / 3) + batch_rec.lengths = lengths + batch_rec.lengths_scaled = lengths_scaled + + angles_radians = decoder_out["angles"] + angles_degrees = paddle.rad2deg(angles_radians) + batch_rec.angles = angles_degrees + batch_rec.angles_radians = angles_radians + + lengths_np = lengths.cpu().numpy() + angles_np = angles_degrees.cpu().numpy() + + lattices_list = [] + for i in range(lengths_np.shape[0]): + lattice = Lattice.from_parameters( + lengths_np[i, 0], lengths_np[i, 1], lengths_np[i, 2], + angles_np[i, 0], angles_np[i, 1], angles_np[i, 2] + ) + lattices_list.append(lattice.matrix) + + lattices_array = np.stack(lattices_list, axis=0) + batch_rec.lattices = paddle.to_tensor(lattices_array, dtype='float32') + + batch_rec.num_atoms = batch.num_atoms + batch_rec.batch = batch.batch + batch_rec.token_idx = batch.token_idx + batch_rec.num_nodes = batch.num_nodes + batch_rec.num_graphs = batch.num_graphs + return batch_rec + + def _dict_to_crystal_batch(self, batch): + """Convert dictionary format to CrystalBatch format. + + Args: + batch: Dict with 'structure_array' key containing structure data + + Returns: + CrystalBatch: Converted batch object + """ + from ppmat.models.chemeleon2.common.schema import CrystalBatch + from ppmat.utils.crystal import lattice_params_to_matrix_paddle + + structure_array = batch["structure_array"] + num_atoms = structure_array["num_atoms"] + batch_size = num_atoms.shape[0] + total_atoms = num_atoms.sum().item() + + # Create CrystalBatch from structure_array + crystal_batch = CrystalBatch() + crystal_batch.atom_types = structure_array["atom_types"] + crystal_batch.frac_coords = structure_array["frac_coords"] + crystal_batch.num_atoms = num_atoms + crystal_batch.batch = paddle.repeat_interleave( + paddle.arange(batch_size), repeats=num_atoms + ) + + # Handle lattice - convert lengths + angles to lattice matrix + if "lattice" in structure_array: + crystal_batch.lattices = structure_array["lattice"] + else: + crystal_batch.lattices = lattice_params_to_matrix_paddle( + structure_array["lengths"], structure_array["angles"] + ) + + # Store original lengths and angles (scaled/radians as needed) + crystal_batch.lengths = structure_array["lengths"] + crystal_batch.angles = structure_array["angles"] + crystal_batch.angles_radians = structure_array["angles"] * 3.141592653589793 / 180.0 + + # Calculate lengths_scaled + num_atoms_tensor = structure_array["num_atoms"] + if num_atoms_tensor.ndim == 1: + num_atoms_tensor = num_atoms_tensor.unsqueeze(-1) + crystal_batch.lengths_scaled = structure_array["lengths"] / ( + num_atoms_tensor ** (1 / 3) + ) + + # Calculate cart_coords: cart_coords = frac_coords @ lattice.T + # lattice shape: [batch_size, 3, 3], frac_coords shape: [total_atoms, 3] + cart_coords = paddle.matmul(crystal_batch.frac_coords, crystal_batch.lattices[crystal_batch.batch]) + crystal_batch.cart_coords = cart_coords + + # Additional fields + crystal_batch.num_nodes = total_atoms + crystal_batch.num_graphs = batch_size + crystal_batch.token_idx = paddle.concat([ + paddle.arange(n) for n in num_atoms + ]) + + return crystal_batch + + def forward(self, batch): + """Forward pass for training compatibility. + + This method converts the standard dictionary format to CrystalBatch format + and then calls calculate_loss. This provides compatibility with the + standard training framework. + + Args: + batch: Input batch data (dict with 'structure_array' key) + + Returns: + dict: Contains 'loss_dict' with training losses (tensors for backward pass) + """ + # Convert dict format to CrystalBatch format + crystal_batch = self._dict_to_crystal_batch(batch) + loss_dict = self.calculate_loss(crystal_batch, training=True) + + # The framework needs loss_dict with tensor values for backward pass + # Add 'loss' key for framework compatibility + loss_dict["loss"] = loss_dict.get("total_loss", paddle.to_tensor([0.0])) + + return {"loss_dict": loss_dict} + + def calculate_loss(self, batch, training=True): + if training and self.augmentation is not None: + translate = self.augmentation.get('translate', False) + rotate = self.augmentation.get('rotate', False) + batch = apply_augmentation(batch, translate=translate, rotate=rotate) + + if training and self.noise is not None: + ratio = self.noise.get('ratio', 0.0) + corruption_scale = self.noise.get('corruption_scale', 0.1) + if ratio > 0: + batch = apply_noise(batch, ratio=ratio, corruption_scale=corruption_scale) + + # Directly call encode and decode to avoid recursion with new forward method + encoded = self.encode(batch) + z = encoded["posterior"].sample() + encoded["x"] = z + encoded["z"] = z + decoder_out = self.decode(encoded) + + loss_atom_types = 0 + if self.atom_type_predict: + loss_atom_types = F.cross_entropy( + decoder_out["atom_types"], batch.atom_types + ) + loss_lengths = F.mse_loss(decoder_out["lengths"], batch.lengths_scaled) + loss_angles = F.mse_loss(decoder_out["angles"], batch.angles_radians) + loss_frac_coords = F.mse_loss( + decoder_out["frac_coords"], batch.frac_coords + ) + + loss_kl = encoded["posterior"].kl().mean() + + fa_loss = 0 + if self.loss_weights.get("fa", 0) > 0: + z = self.proj(encoded["z"]) + mace_features = batch.mace_features + z_norm = F.normalize(z, axis=-1) + mace_features_norm = F.normalize(mace_features, axis=-1) + z_cos_sim = paddle.einsum("ij,kj->ik", z_norm, z_norm) + mace_cos_sim = paddle.einsum( + "ij,kj->ik", mace_features_norm, mace_features_norm + ) + diff = paddle.abs(z_cos_sim - mace_cos_sim) + fa_loss_1 = F.relu(diff - 0.25).mean() + fa_loss_2 = F.relu(1 - 0.5 - F.cosine_similarity(mace_features, z)).mean() + fa_loss = fa_loss_1 + fa_loss_2 + + loss = ( + self.loss_weights.get("atom_types", 1.0) * loss_atom_types + + self.loss_weights.get("lengths", 1.0) * loss_lengths + + self.loss_weights.get("angles", 1.0) * loss_angles + + self.loss_weights.get("frac_coords", 1.0) * loss_frac_coords + + self.loss_weights.get("kl", 1.0) * loss_kl + + self.loss_weights.get("fa", 0.0) * fa_loss + ) + + return { + "total_loss": loss.mean(), + "loss_atom_types": loss_atom_types, + "loss_lengths": loss_lengths, + "loss_angles": loss_angles, + "loss_frac_coords": loss_frac_coords, + "loss_kl": loss_kl, + "fa_loss": fa_loss, + } + + def save_checkpoint(self, save_path, epoch=None, optimizer_state=None, scheduler_state=None): + checkpoint = { + 'model_state_dict': self.state_dict(), + 'latent_dim': self.latent_dim, + 'loss_weights': self.loss_weights, + 'augmentation': self.augmentation, + 'noise': self.noise, + 'atom_type_predict': self.atom_type_predict, + } + + if epoch is not None: + checkpoint['epoch'] = epoch + if optimizer_state is not None: + checkpoint['optimizer_state_dict'] = optimizer_state + if scheduler_state is not None: + checkpoint['scheduler_state_dict'] = scheduler_state + + paddle.save(checkpoint, save_path) + + @staticmethod + def load_checkpoint(load_path, encoder, decoder, map_location=None): + if map_location is not None and map_location == 'cpu': + checkpoint = paddle.load(load_path, map_location=paddle.CPUPlace()) + else: + checkpoint = paddle.load(load_path) + + latent_dim = checkpoint['latent_dim'] + loss_weights = checkpoint['loss_weights'] + augmentation = checkpoint.get('augmentation', None) + noise = checkpoint.get('noise', None) + atom_type_predict = checkpoint.get('atom_type_predict', True) + + model = VAEModule( + encoder=encoder, + decoder=decoder, + latent_dim=latent_dim, + loss_weights=loss_weights, + augmentation=augmentation, + noise=noise, + atom_type_predict=atom_type_predict, + ) + + model.set_state_dict(checkpoint['model_state_dict']) + + return model, checkpoint + + def get_config(self): + return { + 'latent_dim': self.latent_dim, + 'loss_weights': self.loss_weights, + 'augmentation': self.augmentation, + 'noise': self.noise, + 'atom_type_predict': self.atom_type_predict, + } diff --git a/structure_generation/configs/chemeleon2/README.md b/structure_generation/configs/chemeleon2/README.md new file mode 100644 index 00000000..3624f331 --- /dev/null +++ b/structure_generation/configs/chemeleon2/README.md @@ -0,0 +1,84 @@ +# Chemeleon2 配置文件说明 + +本目录包含 Chemeleon2 模型训练和采样的配置文件。 + +## 配置文件列表 + +### 训练配置 + +1. train_vae.yaml - VAE模块的训练 +2. train_ldm.yaml - LDM模块的训练 +3. sample.yaml - 采样生成配置文件 + + + +## 使用方法 + +### 训练 VAE + +```bash +python structure_generation/train.py \ + --config structure_generation/configs/chemeleon2/train_vae.yaml +``` + + +### 训练 LDM + +```bash +python structure_generation/train.py \ + --config structure_generation/configs/chemeleon2/train_ldm.yaml +``` + +## 验证VAE权重 +```bash +python structure_generation/train.py \ + -c structure_generation/configs/chemeleon2/train_vae.yaml \ + Global.do_eval=False \ + Global.do_train=False \ + Global.do_test=True \ + Trainer.pretrained_model_path=output/chemeleon2_vae/checkpoints/latest.pdparams +``` + +## 验证LDM权重 +```bash +python structure_generation/train.py \ + -c structure_generation/configs/chemeleon2/train_ldm.yaml \ + Global.do_eval=False \ + Global.do_train=False \ + Global.do_test=True \ + Trainer.pretrained_model_path=output/chemeleon2_ldm/checkpoints/latest.pdparams +``` + + +### 验证VAE权重 + +```bash +python structure_generation/train.py \ + -c structure_generation/configs/chemeleon2/train_vae.yaml \ + Global.do_eval=True \ + Global.do_train=False \ + Global.do_test=False \ + Trainer.pretrained_model_path=output/chemeleon2_vae/checkpoints/latest.pdparams +``` + +### 验证LDM权重 +```bash +python structure_generation/train.py \ + -c structure_generation/configs/chemeleon2/train_ldm.yaml \ + Global.do_eval=True \ + Global.do_train=False \ + Global.do_test=False \ + Trainer.pretrained_model_path=output/chemeleon2_ldm/checkpoints/latest.pdparams +``` + + + +### 晶体采样 + +```bash +# vae的模型地址在 yaml里约定 +python structure_generation/sample.py \ + --config structure_generation/configs/chemeleon2/sample.yaml \ + --checkpoint_path test-for-weight/converted_weights/ldm_paddle.pdparams +``` + diff --git a/structure_generation/configs/chemeleon2/sample.yaml b/structure_generation/configs/chemeleon2/sample.yaml new file mode 100644 index 00000000..f5eb13dd --- /dev/null +++ b/structure_generation/configs/chemeleon2/sample.yaml @@ -0,0 +1,69 @@ +Model: + __class_name__: LDMModule + __init_params__: + normalize_latent: true + vae_ckpt_path: test-for-weight/converted_weights/vae_paddle.pdparams + denoiser: + __class_name__: DiT + __init_params__: + hidden_dim: 768 + num_layers: 12 + num_heads: 12 + mlp_ratio: 4.0 + input_dim: 8 + learn_sigma: true + diffusion_configs: + timestep_respacing: "" + noise_schedule: linear + use_kl: false + sigma_small: false + predict_xstart: false + learn_sigma: true + rescale_learned_sigmas: false + diffusion_steps: 1000 + vae: + __class_name__: VAEModule + __init_params__: + encoder: + __class_name__: TransformerEncoder + __init_params__: + max_num_elements: 100 + d_model: 512 + nhead: 8 + dim_feedforward: 2048 + activation: gelu + dropout: 0.0 + num_layers: 8 + decoder: + __class_name__: TransformerDecoder + __init_params__: + max_num_elements: 100 + d_model: 512 + nhead: 8 + dim_feedforward: 2048 + activation: gelu + dropout: 0.0 + num_layers: 8 + latent_dim: 8 + loss_weights: {} + +Sample: + data: + dataset: + __class_name__: MP20Dataset + __init_params__: + path: "./data/mp_20/test.csv" + build_structure_cfg: + format: cif_str + num_cpus: 10 + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: false + drop_last: false + batch_size: 10 + build_structure_cfg: + format: array + niggli: false + post_transforms: null + metrics: null diff --git a/structure_generation/configs/chemeleon2/train_ldm.yaml b/structure_generation/configs/chemeleon2/train_ldm.yaml new file mode 100644 index 00000000..21e52335 --- /dev/null +++ b/structure_generation/configs/chemeleon2/train_ldm.yaml @@ -0,0 +1,166 @@ +# Chemeleon2 LDM Training Configuration (Standard Format) + +Global: + do_train: True + do_eval: True + do_test: False + +Model: + __class_name__: LDMModule + __init_params__: + normalize_latent: true + denoiser: + __class_name__: DiT + __init_params__: + hidden_dim: 768 + num_layers: 12 + num_heads: 12 + mlp_ratio: 4.0 + input_dim: 8 + learn_sigma: true + diffusion_configs: + timestep_respacing: "" + noise_schedule: linear + use_kl: false + sigma_small: false + predict_xstart: false + learn_sigma: true + rescale_learned_sigmas: false + diffusion_steps: 1000 + vae: + __class_name__: VAEModule + __init_params__: + encoder: + __class_name__: TransformerEncoder + __init_params__: + max_num_elements: 100 + d_model: 512 + nhead: 8 + dim_feedforward: 2048 + activation: gelu + dropout: 0.0 + num_layers: 8 + decoder: + __class_name__: TransformerDecoder + __init_params__: + max_num_elements: 100 + d_model: 512 + nhead: 8 + dim_feedforward: 2048 + activation: gelu + dropout: 0.0 + num_layers: 8 + latent_dim: 8 + loss_weights: {} + vae_ckpt_path: null + +Dataset: + train: + dataset: + __class_name__: MP20Dataset + __init_params__: + path: "./data/mp_20/train.csv" + property_names: [] + build_structure_cfg: + format: cif_str + num_cpus: 10 + build_graph_cfg: null + cache_path: "./data/mp_20_cache/train" + transforms: [] + num_workers: 4 + use_shared_memory: False + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: True + drop_last: True + batch_size: 32 + val: + dataset: + __class_name__: MP20Dataset + __init_params__: + path: "./data/mp_20/val.csv" + property_names: [] + build_structure_cfg: + format: cif_str + num_cpus: 10 + build_graph_cfg: null + cache_path: "./data/mp_20_cache/val" + transforms: [] + num_workers: 4 + use_shared_memory: False + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: False + drop_last: False + batch_size: 32 + test: + dataset: + __class_name__: MP20Dataset + __init_params__: + path: "./data/mp_20/test.csv" + property_names: [] + build_structure_cfg: + format: cif_str + num_cpus: 10 + build_graph_cfg: null + cache_path: "./data/mp_20_cache/test" + transforms: [] + num_workers: 4 + use_shared_memory: False + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: False + drop_last: False + batch_size: 32 + +Trainer: + max_epochs: 100 + seed: 42 + output_dir: ./output/chemeleon2_ldm + save_freq: 500 + log_freq: 20 + start_eval_epoch: 1 + eval_freq: 100 + pretrained_model_path: null + resume_from_checkpoint: null + use_amp: False + amp_level: 'O1' + eval_with_no_grad: True + gradient_accumulation_steps: 1 + best_metric_indicator: 'eval_loss' + name_for_best_metric: 'loss' + greater_is_better: False + compute_metric_during_train: False + metric_strategy_during_eval: 'epoch' + use_visualdl: False + use_wandb: False + use_tensorboard: False + +Optimizer: + __class_name__: AdamW + __init_params__: + beta1: 0.9 + beta2: 0.999 + lr: 0.0001 + weight_decay: 0.0 + +Sample: + data: + dataset: + __class_name__: NumAtomsCrystalDataset + __init_params__: + total_num: 10 + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: False + drop_last: False + batch_size: 10 + build_structure_cfg: + format: array + niggli: False + model_sample_params: + num_inference_steps: 50 diff --git a/structure_generation/configs/chemeleon2/train_rl.yaml b/structure_generation/configs/chemeleon2/train_rl.yaml new file mode 100644 index 00000000..8f8e3b8b --- /dev/null +++ b/structure_generation/configs/chemeleon2/train_rl.yaml @@ -0,0 +1,98 @@ +task_name: train_rl_chemeleon2 + +seed: 0 + +data: + dataset_name: mp_20 + data_path: data/mp_20 + batch_size: 8 + num_workers: 4 + shuffle: True + max_num_atoms: 100 + +model: + name: chemeleon2_rl + + ldm_ckpt_path: null + + rl_configs: + clip_ratio: 0.2 + kl_weight: 0.1 + entropy_weight: 0.01 + num_group_samples: 4 + group_reward_norm: True + num_inner_batch: 1 + + sampling_configs: + sampler: ddim + sampling_steps: 50 + eta: 1.0 + cfg_scale: 2.0 + return_trajectory: False + + reward_components: + - name: creativity + weight: 1.0 + config: + distance_metric: amd + reference_set: training_set + + - name: energy + weight: 0.5 + config: + calculator: mace + model_path: null + + - name: structure_diversity + weight: 0.3 + config: + distance_metric: mmd + kernel: rbf + + - name: composition_diversity + weight: 0.2 + config: + featurizer: composition + + - name: predictor + weight: 0.5 + config: + predictor_path: null + target_property: formation_energy + +optimizer: + name: AdamW + learning_rate: 1e-5 + weight_decay: 0.0 + beta1: 0.9 + beta2: 0.999 + epsilon: 1e-8 + +scheduler: + name: constant + warmup_steps: 0 + +training: + max_steps: 5000 + log_interval: 1 + eval_interval: 5 + save_interval: 100 + grad_clip: 1.0 + + checkpoint: + monitor: val_reward + mode: max + save_top_k: 3 + + early_stopping: + monitor: val_reward + patience: 200 + mode: max + +output: + save_dir: output/chemeleon2/rl + checkpoint_dir: checkpoints + log_dir: logs + +ckpt_path: null +resume_from: null diff --git a/structure_generation/configs/chemeleon2/train_vae.yaml b/structure_generation/configs/chemeleon2/train_vae.yaml new file mode 100644 index 00000000..4e039de1 --- /dev/null +++ b/structure_generation/configs/chemeleon2/train_vae.yaml @@ -0,0 +1,134 @@ +# Chemeleon2 VAE Training Configuration (Standard Format) + +Global: + do_train: True + do_eval: True + do_test: False + +Model: + __class_name__: VAEModule + __init_params__: + encoder: + __class_name__: TransformerEncoder + __init_params__: + max_num_elements: 100 + d_model: 512 + nhead: 8 + dim_feedforward: 2048 + activation: gelu + dropout: 0.0 + num_layers: 8 + decoder: + __class_name__: TransformerDecoder + __init_params__: + max_num_elements: 100 + d_model: 512 + nhead: 8 + dim_feedforward: 2048 + activation: gelu + dropout: 0.0 + num_layers: 8 + latent_dim: 8 + loss_weights: + atom_types: 1.0 + lengths: 1.0 + angles: 10.0 + frac_coords: 10.0 + kl: 1e-5 + fa: 0.0 + augmentation: + translate: False + rotate: False + +Dataset: + train: + dataset: + __class_name__: MP20Dataset + __init_params__: + path: "./data/mp_20/train.csv" + property_names: [] + build_structure_cfg: + format: cif_str + num_cpus: 10 + build_graph_cfg: null + cache_path: "./data/mp_20_cache/train" + transforms: [] + num_workers: 4 + use_shared_memory: False + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: True + drop_last: True + batch_size: 32 + val: + dataset: + __class_name__: MP20Dataset + __init_params__: + path: "./data/mp_20/val.csv" + property_names: [] + build_structure_cfg: + format: cif_str + num_cpus: 10 + build_graph_cfg: null + cache_path: "./data/mp_20_cache/val" + transforms: [] + num_workers: 4 + use_shared_memory: False + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: False + drop_last: False + batch_size: 32 + test: + dataset: + __class_name__: MP20Dataset + __init_params__: + path: "./data/mp_20/test.csv" + property_names: [] + build_structure_cfg: + format: cif_str + num_cpus: 10 + build_graph_cfg: null + cache_path: "./data/mp_20_cache/test" + transforms: [] + num_workers: 4 + use_shared_memory: False + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: False + drop_last: False + batch_size: 32 + +Trainer: + max_epochs: 100 + seed: 42 + output_dir: ./output/chemeleon2_vae + save_freq: 500 + log_freq: 20 + start_eval_epoch: 1 + eval_freq: 100 + pretrained_model_path: null + resume_from_checkpoint: null + use_amp: False + amp_level: 'O1' + eval_with_no_grad: True + gradient_accumulation_steps: 1 + best_metric_indicator: 'eval_loss' + name_for_best_metric: 'loss' + greater_is_better: False + compute_metric_during_train: False + metric_strategy_during_eval: 'epoch' + use_visualdl: False + use_wandb: False + use_tensorboard: False + +Optimizer: + __class_name__: AdamW + __init_params__: + beta1: 0.9 + beta2: 0.999 + lr: 0.0001 + weight_decay: 0.0 diff --git a/test/chemeleon2/README.md b/test/chemeleon2/README.md new file mode 100644 index 00000000..354f1c6c --- /dev/null +++ b/test/chemeleon2/README.md @@ -0,0 +1,105 @@ +# 本目录是针对 chemeleon2 的 单元测试说明 + +## 目录说明 +- test_loss 是测试paddle版本 和 原始版本的精度差异的目录【可以在当前项目运行】 +- raw_infer_data 是原始的 hspark1212/chemeleon2 commit-id:63769eb4b962278cb0576da1ac7e0cc10d36a3e8 版本项目中，增加的测试脚本， 可以产出工程里的输入和输出数据，方便paddle版本对比【只能在chemeleon2原始项目中运行】。 + + + +## test_loss 目录介绍 + +### test_model_loss_with_raw.py 脚本 +可以通过pytest进行测试，需要提起准备好权重文件，权重文件在：https://aistudio.baidu.com/modelsdetail/43734?modelId=43734 + +## raw_infer_data 目录介绍 + +### create_input_output_npz.py 脚本 + +create_input_output_npz.py 文件需要在 hspark1212/chemeleon2 commit-id:63769eb4b962278cb0576da1ac7e0cc10d36a3e8 项目中运行； + +运行后，会产生一系列的npz文件，其中 input 字样的是输入数据；output 字样是输出数据； + +我们可以拿同样的 input的输入、同样的随机数 让paddle 框架来生成晶体结构获取结果，和原版的output输出的值之间进行对比，观测精度差异。 + +npz的格式化输出脚本如下： + +```python + +import numpy as np + + +def to_python_list(arr, max_dim=4): + """将 numpy 数组转换为 Python 列表格式""" + if arr.ndim == 0: + return arr.item() + if arr.ndim == 1: + return arr.tolist() + # 递归处理多维数组 + return to_python_list(arr, max_dim - 1) if max_dim > 1 else arr.tolist() + + +def print_array(arr, max_dim=4, indent=0): + """以 Python 数组风格打印数组""" + if isinstance(arr, np.ndarray): + if arr.ndim == 0: + print(str(arr.item()), end='') + elif arr.ndim == 1: + print('[', end='') + for i, val in enumerate(arr): + if i < min(4, len(arr)): + print_array(val, max_dim - 1, 0) + if i < min(4, len(arr)) - 1: + print(', ', end='') + if len(arr) > 4: + print(', ...]', end='') + else: + print(']', end='') + else: + print('[', end='') + for i in range(min(4, len(arr))): + print_array(arr[i], max_dim - 1, 0) + if i < min(4, len(arr)) - 1: + print(', ', end='') + if len(arr) > 4: + print(', ...]', end='') + else: + print(']', end='') + else: + # 标量 - 保留小数点后8位 + if isinstance(arr, (np.integer, np.floating)): + val = float(arr) + # 整数不显示小数点 + if val == int(val): + print(int(val), end='') + else: + # 保留小数点后8位 + print(f"{val:.8f}", end='') + else: + print(arr, end='') + + +def print_npz(file_path, max_dim=4): + data = np.load(file_path) + + print(f"文件: {file_path}") + print("=" * 50) + + for key in data.keys(): + arr = data[key] + print(f"\n[{key}]") + print(f" shape: {arr.shape}, dtype: {arr.dtype}") + print(f" 值: ", end='') + print_array(arr, max_dim) + print() + + +if __name__ == "__main__": + # 手动指定 npz 文件路径 + npz_file = "outputs/precision_test/vae_decoder_input_z_full.npz" + + # 可选: 设置最大输出维度，默认 4 + max_output_dim = 4 + + print_npz(npz_file, max_output_dim) + +``` \ No newline at end of file diff --git a/test/chemeleon2/__init__.py b/test/chemeleon2/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/test/chemeleon2/conftest.py b/test/chemeleon2/conftest.py new file mode 100644 index 00000000..e69de29b diff --git a/test/chemeleon2/raw_infer_data/create_input_output_npz.py b/test/chemeleon2/raw_infer_data/create_input_output_npz.py new file mode 100644 index 00000000..41bf67aa --- /dev/null +++ b/test/chemeleon2/raw_infer_data/create_input_output_npz.py @@ -0,0 +1,835 @@ +""" +模型精度测试 - 用于框架迁移对比 + +功能： +1. 测试 VAE Encoder/Decoder 的输出 +2. 测试 LDM Denoiser 的输出 +3. 保存中间结果用于与其他框架对比 +4. 固定随机种子保证可复现 +""" + +import sys +import os +import json +import random +from pathlib import Path +from datetime import datetime + +import torch +import numpy as np +from monty.serialization import loadfn + + +def save_full_tensor_data(tensor, name, output_dir): + """ + 保存完整的tensor数据到NPZ文件 + + Args: + tensor: PyTorch张量 + name: 张量名称 + output_dir: 输出目录 + """ + npz_file = output_dir / f"{name}_full.npz" + np.savez_compressed(npz_file, data=tensor.cpu().numpy()) + return str(npz_file) + + +# 添加项目根目录到路径 +project_root = Path(__file__).parent.parent +sys.path.insert(0, str(project_root)) + +# 切换到项目根目录，确保相对路径正确 +os.chdir(project_root) + +from src.vae_module.vae_module import VAEModule +from src.ldm_module.ldm_module import LDMModule +from src.data.schema import CrystalBatch +from src.data.dataset_util import pmg_structure_to_pyg_data + + +class Logger: + """同时输出到控制台和文件的日志器""" + + def __init__(self, filepath): + self.terminal = sys.stdout + self.log = open(filepath, 'w', encoding='utf-8') + + def write(self, message): + self.terminal.write(message) + self.log.write(message) + + def flush(self): + self.terminal.flush() + self.log.flush() + + def close(self): + self.log.close() + + +# 全局日志器 +_logger = None + + +def set_logger(filepath): + """设置全局日志器""" + global _logger + _logger = Logger(filepath) + sys.stdout = _logger + + +def close_logger(): + """关闭日志器""" + global _logger + if _logger: + sys.stdout = _logger.terminal + _logger.close() + _logger = None + + +def load_fixed_inputs(npz_path): + """从保存的npz文件加载固定输入 + + Args: + npz_path: npz文件路径 + + Returns: + dict: 包含各个模块输入数据的字典 + """ + data = np.load(npz_path, allow_pickle=True) + + fixed_inputs = {} + for key in data.keys(): + module_data = data[key].item() + if 'input' in module_data: + fixed_inputs[key] = module_data['input'] + + return fixed_inputs + + +def test_single_module(module_name, vae_path, ldm_path, structure_path, device="cuda", use_fixed_input=False, + fixed_inputs=None): + """单独测试某个模块 + + Args: + module_name: 'vae_encoder', 'vae_decoder', 'vae_full', 'ldm_denoiser', 'ldm_sampling' + use_fixed_input: 是否使用固定输入 + fixed_inputs: 固定输入数据字典 + + Returns: + 测试结果 + """ + print(f"\n{'=' * 80}") + print(f"Testing: {module_name}") + print(f"{'=' * 80}") + + if use_fixed_input and fixed_inputs and module_name in fixed_inputs: + print(f"Using fixed input data") + fixed_input = fixed_inputs[module_name] + else: + fixed_input = None + if use_fixed_input: + print(f"Warning: Fixed input not found, using random/default input") + + if module_name == 'vae_encoder': + return test_vae_encoder(vae_path, structure_path, device, fixed_input) + elif module_name == 'vae_decoder': + return test_vae_decoder(vae_path, structure_path, device) + elif module_name == 'vae_full': + return test_vae_full(vae_path, structure_path, device) + elif module_name == 'ldm_denoiser': + return test_ldm_denoiser(ldm_path, vae_path, device, fixed_input) + elif module_name == 'ldm_sampling': + return test_ldm_sampling(ldm_path, vae_path, structure_path, device) + else: + raise ValueError(f"未知模块: {module_name}") + + +def set_seed(seed=42): + """固定随机种子，确保完全可复现""" + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + # 设置CUDA确定性模式 + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + + +def import_datetime(): + """返回当前时间字符串""" + return datetime.now().strftime("%Y-%m-%d %H:%M:%S") + + +def print_tensor_info(tensor, name, max_dims=4): + """打印张量信息，只显示前max_dims个值""" + print(f"\n{name}:") + print(f" shape: {tuple(tensor.shape)}") + print(f" dtype: {tensor.dtype}") + print(f" device: {tensor.device}") + print(f" mean: {tensor.mean().item():.8f}") + print(f" std: {tensor.std().item():.8f}") + print(f" min: {tensor.min().item():.8f}") + print(f" max: {tensor.max().item():.8f}") + + # 只显示前4个值 + flat = tensor.flatten() + show_n = min(max_dims, len(flat)) + values_str = " ".join(f"{flat[i].item():.8f}" for i in range(show_n)) + print(f" first {show_n}: {values_str}") + + +def save_tensor_to_dict(tensor, prefix="", output_dir=None, tensor_name=None, save_full=True): + """将张量信息保存为字典（JSON兼容格式） + + Args: + tensor: PyTorch张量 + prefix: 键名前缀 + output_dir: 输出目录，如果提供则保存完整数据 + tensor_name: 张量名称，用于保存完整数据文件 + save_full: 是否保存完整数据到NPZ文件 + + 对于多维张量，JSON中只保存前4个值，但完整数据保存到NPZ文件 + """ + flat = tensor.flatten() + show_n = min(4, len(flat)) + + # 获取前4个值的详细列表 + first_values = [float(flat[i].item()) for i in range(show_n)] + + result = { + f"{prefix}shape": list(tensor.shape), + f"{prefix}dtype": str(tensor.dtype), + f"{prefix}device": str(tensor.device), + f"{prefix}first_4_values": first_values, + f"{prefix}total_elements": int(tensor.numel()), + } + + # 只对浮点数类型计算统计值 + if tensor.dtype in [torch.float32, torch.float64, torch.float16]: + result[f"{prefix}mean"] = float(f"{tensor.mean().item():.8f}") + result[f"{prefix}std"] = float(f"{tensor.std().item():.8f}") + result[f"{prefix}min"] = float(f"{tensor.min().item():.8f}") + result[f"{prefix}max"] = float(f"{tensor.max().item():.8f}") + else: + # 对于整数类型，只记录min和max + result[f"{prefix}min"] = int(tensor.min().item()) + result[f"{prefix}max"] = int(tensor.max().item()) + + # 保存完整数据到NPZ文件 + if save_full and output_dir is not None and tensor_name is not None: + full_data_path = save_full_tensor_data(tensor, tensor_name, output_dir) + result[f"{prefix}full_data_file"] = full_data_path + result[f"{prefix}full_data_available"] = True + else: + result[f"{prefix}full_data_available"] = False + + return result + + +def test_vae_encoder(vae_path, structure_path, device="cuda", fixed_input=None, output_dir=None): + """测试 VAE Encoder + + Args: + fixed_input: 如果提供，使用固定的输入数据而不是从结构加载 + output_dir: 输出目录，用于保存完整tensor数据 + """ + # 设置随机种子确保可复现 + set_seed(42) + + print("\n" + "=" * 80) + print("VAE Encoder Test") + print("=" * 80) + + # 加载模型 + vae = VAEModule.load_from_checkpoint(vae_path, weights_only=False) + vae.to(device) + vae.eval() + + if fixed_input is not None: + # 使用固定输入 + print("\n使用固定输入数据") + # 重建 CrystalBatch + from torch_geometric.data import Data + + data = Data( + atom_types=torch.from_numpy(fixed_input['atom_types']).long(), + frac_coords=torch.from_numpy(fixed_input['frac_coords']).float(), + cart_coords=torch.from_numpy(fixed_input.get('cart_coords', fixed_input['frac_coords'])).float(), + lattices=torch.from_numpy(fixed_input['lattices']).float(), + num_atoms=torch.from_numpy(fixed_input['num_atoms']).long(), + lengths=torch.from_numpy(fixed_input.get('lengths', np.zeros((1, 3)))).float(), + lengths_scaled=torch.from_numpy(fixed_input.get('lengths_scaled', np.zeros((1, 3)))).float(), + angles=torch.from_numpy(fixed_input.get('angles', np.zeros((1, 3)))).float(), + angles_radians=torch.from_numpy(fixed_input.get('angles_radians', np.zeros((1, 3)))).float(), + token_idx=torch.arange(len(fixed_input['atom_types']), dtype=torch.long), + pos=torch.from_numpy(fixed_input.get('cart_coords', fixed_input['frac_coords'])).float(), + ) + batch = CrystalBatch.from_data_list([data]) + batch = batch.to(device) + test_structure_info = "固定输入（从保存的数据加载）" + else: + # 加载测试结构 + structures = loadfn(structure_path) + test_structure = structures[0] # 使用第一个结构 + + print(f"\nTest structure: {test_structure.composition}") + print(f" num_atoms: {len(test_structure)}") + print(f" space_group: {test_structure.get_space_group_info()}") + + # 准备输入 + data = pmg_structure_to_pyg_data(test_structure) + batch = CrystalBatch.from_data_list([data]) + batch = batch.to(device) + test_structure_info = str(test_structure.composition) + + # 保存输入数据的详细描述（JSON格式） + input_description = { + "input_source": "fixed_structure" if fixed_input is None else "fixed_input_data", + "structure_info": test_structure_info, + "batch_info": { + "num_graphs": int(batch.num_graphs), + "total_atoms": int(batch.num_nodes), + }, + "input_tensors": {}, + } + + # 记录每个输入张量的信息并保存完整数据 + for attr_name in ['atom_types', 'frac_coords', 'lattices', 'num_atoms', 'lengths', 'lengths_scaled', 'angles']: + if hasattr(batch, attr_name): + attr_tensor = getattr(batch, attr_name) + # 保存完整数据到NPZ + tensor_name = f"vae_encoder_input_{attr_name}" + input_description["input_tensors"][attr_name] = save_tensor_to_dict( + attr_tensor, "", + output_dir=output_dir, + tensor_name=tensor_name, + save_full=True + ) + # 对于小张量，也在JSON中保存完整值 + if attr_tensor.numel() <= 20: + input_description["input_tensors"][attr_name]["all_values"] = attr_tensor.flatten().cpu().tolist() + + # Encoder 前向传播 + with torch.no_grad(): + # 获取 encoder 输出（返回字典） + encoder_dict = vae.encoder(batch) # Dict with keys: x, num_atoms, batch, token_idx + encoder_output = encoder_dict["x"] # (N, H) + print_tensor_info(encoder_output, "Encoder Output (x)") + + # 量化卷积 + h = vae.quant_conv(encoder_output) # (N, 2*L) + print_tensor_info(h, "After Quant Conv") + + # 分离 mean 和 logvar + mean, logvar = torch.chunk(h, 2, dim=-1) + print_tensor_info(mean, "Latent Mean") + print_tensor_info(logvar, "Latent Logvar") + + # 采样 - 使用numpy生成固定的eps + std = torch.exp(0.5 * logvar) + np.random.seed(42) + eps_np = np.random.randn(*std.shape).astype(np.float32) + eps = torch.from_numpy(eps_np).to(device) + z = mean + eps * std + print_tensor_info(z, "Latent Z (Sampled)") + + # 保存为字典，用于JSON输出，同时保存完整数据 + result = { + "module_name": "vae_encoder", + "description": "VAE编码器：输入晶体结构，输出潜在向量的均值和方差", + "input": input_description, + "output": { + "encoder_output": save_tensor_to_dict( + encoder_output, "", output_dir, "vae_encoder_output", True + ), + "quant_conv_output": save_tensor_to_dict( + h, "", output_dir, "vae_encoder_quant_conv", True + ), + "latent_mean": save_tensor_to_dict( + mean, "", output_dir, "vae_encoder_latent_mean", True + ), + "latent_logvar": save_tensor_to_dict( + logvar, "", output_dir, "vae_encoder_latent_logvar", True + ), + "latent_z_sampled": save_tensor_to_dict( + z, "", output_dir, "vae_encoder_latent_z", True + ), + }, + } + + return result + + +def test_vae_decoder(vae_path, structure_path, device="cuda", output_dir=None): + """测试 VAE Decoder + + Args: + output_dir: 输出目录，用于保存完整tensor数据 + """ + # 设置随机种子确保可复现 + set_seed(42) + + print("\n" + "=" * 80) + print("VAE Decoder Test") + print("=" * 80) + + # 加载模型 + vae = VAEModule.load_from_checkpoint(vae_path, weights_only=False) + vae.to(device) + vae.eval() + + # 加载测试结构 + structures = loadfn(structure_path) + test_structure = structures[0] + data = pmg_structure_to_pyg_data(test_structure) + batch = CrystalBatch.from_data_list([data]) + batch = batch.to(device) + + # 准备输入描述 + input_description = { + "input_source": "encoded_from_structure", + "structure_info": str(test_structure.composition), + "description": "通过VAE编码器编码后采样得到的潜在向量z", + } + + # 完整的编码-解码流程 + with torch.no_grad(): + # Encode + encoded = vae.encoder(batch) # Dict with keys: x, num_atoms, batch, token_idx + encoded["moments"] = vae.quant_conv(encoded["x"]) + + # Use DiagonalGaussianDistribution to sample + from src.vae_module.vae_module import DiagonalGaussianDistribution + encoded["posterior"] = DiagonalGaussianDistribution(encoded["moments"]) + z = encoded["posterior"].sample() + encoded["x"] = z + encoded["z"] = z + + # 添加输入张量信息 + input_description["input_tensor"] = save_tensor_to_dict(z, "") + + print(f"\nInput latent z:") + print_tensor_info(z, "Decoder input") + + # Decode + z_decoder = vae.post_quant_conv(z) + encoded["x"] = z_decoder + decoder_out = vae.decoder(encoded) + + # 输出重建结果 + print(f"\nDecoder output:") + print(f" keys: {decoder_out.keys()}") + + for key in ["atom_types", "frac_coords", "lengths", "angles"]: + if key in decoder_out: + print_tensor_info(decoder_out[key], f"Reconstructed {key}") + + # 保存为字典（JSON格式），同时保存完整数据 + output_dict = { + "decoder_input_z": save_tensor_to_dict(z, "", output_dir, "vae_decoder_input_z", True), + "decoder_post_quant": save_tensor_to_dict(z_decoder, "", output_dir, "vae_decoder_post_quant", True), + } + + for key in ["atom_types", "frac_coords", "lengths", "angles"]: + if key in decoder_out: + output_dict[f"reconstructed_{key}"] = save_tensor_to_dict( + decoder_out[key], "", output_dir, f"vae_decoder_{key}", True + ) + + result = { + "module_name": "vae_decoder", + "description": "VAE解码器：输入潜在向量z，输出重建的晶体结构", + "input": input_description, + "output": output_dict, + } + + return result + + +def test_vae_full(vae_path, structure_path, device="cuda", output_dir=None): + """测试完整的 VAE 编码-解码""" + # 设置随机种子确保可复现 + set_seed(42) + + print("\n" + "=" * 80) + print("VAE Full Test (Encode + Decode)") + print("=" * 80) + + # 加载模型 + vae = VAEModule.load_from_checkpoint(vae_path, weights_only=False) + vae.to(device) + vae.eval() + + # 加载测试结构 + structures = loadfn(structure_path) + test_structure = structures[0] + data = pmg_structure_to_pyg_data(test_structure) + batch = CrystalBatch.from_data_list([data]) + batch = batch.to(device) + + print(f"\nOriginal structure:") + print(f" composition: {test_structure.composition}") + print(f" num_atoms: {len(test_structure)}") + print(f" lattice: {test_structure.lattice.abc}") + + with torch.no_grad(): + # Forward pass + decoder_out, encoded = vae(batch) + + # 获取均值和方差 + mean = encoded["posterior"].mean + logvar = encoded["posterior"].logvar + + print(f"\nReconstruction error:") + # 计算重建误差 + if hasattr(batch, 'frac_coords') and "frac_coords" in decoder_out: + coord_error = torch.abs(batch.frac_coords - decoder_out["frac_coords"]).mean() + print(f" coord MAE: {coord_error.item():.6f}") + + if hasattr(batch, 'lengths') and "lengths" in decoder_out: + lattice_error = torch.abs(batch.lengths_scaled - decoder_out["lengths"]).mean() + print(f" lattice MAE: {lattice_error.item():.6f}") + + # 保存为字典（JSON格式） + input_description = { + "structure_composition": str(test_structure.composition), + "num_atoms": len(test_structure), + "lattice_abc": list(test_structure.lattice.abc), + "lattice_angles": list(test_structure.lattice.angles), + "space_group": test_structure.get_space_group_info()[0], + } + + metrics = {} + if hasattr(batch, 'frac_coords') and "frac_coords" in decoder_out: + metrics["coord_mae"] = float(f"{coord_error.item():.8f}") + if hasattr(batch, 'lengths') and "lengths" in decoder_out: + metrics["lattice_mae"] = float(f"{lattice_error.item():.8f}") + + result = { + "module_name": "vae_full", + "description": "完整VAE：编码-采样-解码，计算重建误差", + "input": input_description, + "output": { + "latent_mean": save_tensor_to_dict(mean, "", output_dir, "vae_full_latent_mean", True), + "latent_logvar": save_tensor_to_dict(logvar, "", output_dir, "vae_full_latent_logvar", True), + "reconstruction_metrics": metrics, + }, + } + + return result + + +def test_ldm_denoiser(ldm_path, vae_path, device="cuda", fixed_input=None, output_dir=None): + """测试 LDM Denoiser + + Args: + fixed_input: dict with keys 'z' and 't' for fixed input + output_dir: 输出目录，用于保存完整tensor数据 + """ + # 设置随机种子确保可复现 + set_seed(42) + + print("\n" + "=" * 80) + print("LDM Denoiser Test") + print("=" * 80) + + # 加载模型 + ldm = LDMModule.load_from_checkpoint(ldm_path, vae_ckpt_path=vae_path, weights_only=False) + ldm.to(device) + ldm.eval() + + # 创建固定的测试输入（不使用随机） + batch_size = 2 + num_atoms = 10 + latent_dim = 8 + + # 固定的输入数据（基于seed=42生成的固定值） + z = torch.randn(batch_size, num_atoms, latent_dim, device=device) + # 固定时间步 + t = torch.tensor([13, 25], device=device, dtype=torch.long) + + print_tensor_info(z, "Input Latent Z") + print(f"\nTimesteps: {t.cpu().numpy()}") + + # 准备详细的输入描述，保存完整数据 + input_description = { + "description": "固定的噪声潜在向量z和扩散时间步t", + "generation_method": "torch.randn with seed=42", + "input_z": save_tensor_to_dict(z, "", output_dir, "ldm_denoiser_input_z", True), + "input_t": { + "values": t.cpu().tolist(), + "shape": list(t.shape), + "dtype": str(t.dtype), + "description": "扩散过程的时间步，范围[0, 1000)", + }, + "batch_size": batch_size, + "num_atoms_per_structure": num_atoms, + "latent_dim": latent_dim, + } + + # Denoiser 前向传播 + with torch.no_grad(): + # 准备 mask (布尔类型) + mask = torch.ones(batch_size, num_atoms, dtype=torch.bool, device=device) + + # 调用 denoiser + noise_pred = ldm.denoiser(z, t, mask=mask) + print_tensor_info(noise_pred, "Denoiser Output (Predicted Noise)") + + # 保存为字典（JSON格式），保存完整数据 + result = { + "module_name": "ldm_denoiser", + "description": "LDM去噪器：输入噪声潜在向量z和时间步t，预测噪声", + "input": input_description, + "output": { + "predicted_noise": save_tensor_to_dict( + noise_pred, "", output_dir, "ldm_denoiser_output_noise", True + ), + }, + } + + return result + + +def test_ldm_sampling(ldm_path, vae_path, structure_path, device="cuda", output_dir=None): + """测试 LDM 完整采样流程 + + Args: + output_dir: 输出目录，用于保存完整tensor数据 + """ + # 设置随机种子确保可复现 + set_seed(42) + + print("\n" + "=" * 80) + print("LDM Sampling Test") + print("=" * 80) + + # 加载模型 + ldm = LDMModule.load_from_checkpoint(ldm_path, vae_ckpt_path=vae_path, weights_only=False) + ldm.to(device) + ldm.eval() + + # 加载参考结构（用于确定原子数分布） + structures = loadfn(structure_path) + ref_structure = structures[0] + + print(f"\nReference structure: {ref_structure.composition}") + print(f" num_atoms: {len(ref_structure)}") + + # 创建 batch（用于采样） + data = pmg_structure_to_pyg_data(ref_structure) + batch = CrystalBatch.from_data_list([data]) + batch = batch.to(device) + + # 准备输入描述 + input_description = { + "reference_structure": str(ref_structure.composition), + "num_atoms": len(ref_structure), + "sampling_config": { + "sampler": "ddim", + "sampling_steps": 10, + "cfg_scale": 2.0, + }, + "description": "从纯噪声开始，通过迭代去噪生成新的晶体结构", + } + + # 采样 + with torch.no_grad(): + # 使用较少的采样步数加快测试 + batch_gen = ldm.sample( + batch, + sampler="ddim", + sampling_steps=10, + cfg_scale=2.0, + ) + + print(f"\nGenerated structures:") + print(f" batch_size: {batch_gen.num_graphs}") + if hasattr(batch_gen, 'frac_coords'): + print_tensor_info(batch_gen.frac_coords, "Generated Coords") + if hasattr(batch_gen, 'lattices'): + print_tensor_info(batch_gen.lattices, "Generated Lattices") + + # 解码为实际结构 + try: + gen_structures = batch_gen.to_structure() # 使用 to_structure() 而不是 to_structures() + print(f"\nDecoded {len(gen_structures)} structures") + for i, struct in enumerate(gen_structures[:3]): # 显示前3个 + print(f" [{i}] {struct.composition}") + except Exception as e: + print(f"\nDecode failed: {e}") + + # 保存为字典（JSON格式），保存完整数据 + output_dict = {} + if hasattr(batch_gen, 'frac_coords'): + output_dict["generated_frac_coords"] = save_tensor_to_dict( + batch_gen.frac_coords, "", output_dir, "ldm_sampling_coords", True + ) + if hasattr(batch_gen, 'lattices'): + output_dict["generated_lattices"] = save_tensor_to_dict( + batch_gen.lattices, "", output_dir, "ldm_sampling_lattices", True + ) + + # 添加生成结构的组成信息 + if 'gen_structures' in locals(): + output_dict["generated_structures"] = [ + {"composition": str(s.composition)} for s in gen_structures[:5] + ] + + result = { + "module_name": "ldm_sampling", + "description": "LDM完整采样：从噪声生成完整的晶体结构", + "input": input_description, + "output": output_dict, + } + + return result + + +def main(): + """主函数""" + # 先创建输出目录 + output_dir = Path("outputs/precision_test") + output_dir.mkdir(parents=True, exist_ok=True) + + # 设置日志输出 + log_file = output_dir / "precision_test_log.txt" + set_logger(log_file) + + print("\n" + "* " * 40) + print(" " * 25 + "Chemeleon2 Precision Test") + print("* " * 40) + + # 设置随机种子 + set_seed(42) + print(f"\nRandom seed: 42") + print(f"Log file: {log_file}") + + # 模型路径 + vae_path = "checkpoints/v0.0.1/alex_mp_20/vae/dng_j1jgz9t0_v1.ckpt" + ldm_path = "checkpoints/v0.0.1/alex_mp_20/ldm/ldm_rl_dng_tuor5vgd.ckpt" + structure_path = "outputs/alex-mp-20/generated_structures.json.gz" + + # 如果测试结构不存在，使用备用路径 + if not Path(structure_path).exists(): + structure_path = "outputs/test_three/generated_structures.json.gz" + + if not Path(structure_path).exists(): + structure_path = "benchmarks/dng/chemeleon2_rl_dng_mp_20.json.gz" + + if not Path(structure_path).exists(): + print(f"\nError: Structure file not found: {structure_path}") + print(" Please provide a valid structure file") + close_logger() + return False + + # 设备 + device = "cuda" if torch.cuda.is_available() else "cpu" + print(f"Device: {device}") + + # 保存所有结果 + all_results = {} + + try: + # 1. 测试 VAE Encoder + results = test_vae_encoder(vae_path, structure_path, device, output_dir=output_dir) + all_results["vae_encoder"] = results + + # 2. 测试 VAE Decoder + results = test_vae_decoder(vae_path, structure_path, device, output_dir=output_dir) + all_results["vae_decoder"] = results + + # 3. 测试完整 VAE + results = test_vae_full(vae_path, structure_path, device, output_dir=output_dir) + all_results["vae_full"] = results + + # 4. 测试 LDM Denoiser + results = test_ldm_denoiser(ldm_path, vae_path, device, output_dir=output_dir) + all_results["ldm_denoiser"] = results + + # 5. 测试 LDM 采样 + results = test_ldm_sampling(ldm_path, vae_path, structure_path, device, output_dir=output_dir) + all_results["ldm_sampling"] = results + + # 保存结果为JSON文件 + output_json = output_dir / "pytorch_results.json" + output_txt = output_dir / "pytorch_results_readable.txt" # 也保存易读的文本版本 + + # 构建完整的JSON结构 + full_results = { + "test_info": { + "framework": "PyTorch", + "test_name": "Chemeleon2 精度测试", + "test_time": datetime.now().strftime('%Y-%m-%d %H:%M:%S'), + "random_seed": 42, + "device": device, + "description": "用于框架迁移的精度对比测试，包含VAE和LDM的各个模块", + }, + "model_info": { + "vae_checkpoint": vae_path, + "ldm_checkpoint": ldm_path, + "structure_file": structure_path, + }, + "modules": all_results, + } + + # 保存JSON文件 + with open(output_json, 'w', encoding='utf-8') as f: + json.dump(full_results, f, indent=2, ensure_ascii=False) + + # 保存易读的文本版本 + with open(output_txt, 'w', encoding='utf-8') as f: + f.write("=" * 80 + "\n") + f.write("Chemeleon2 精度测试结果 - PyTorch (易读版本)\n") + f.write("=" * 80 + "\n") + f.write(f"测试时间: {full_results['test_info']['test_time']}\n") + f.write(f"随机种子: {full_results['test_info']['random_seed']}\n") + f.write(f"设备: {full_results['test_info']['device']}\n") + f.write("=" * 80 + "\n\n") + + for module_name, module_data in all_results.items(): + f.write(f"\n{'=' * 80}\n") + f.write(f"模块: {module_data.get('module_name', module_name)}\n") + f.write(f"描述: {module_data.get('description', '')}\n") + f.write('=' * 80 + "\n\n") + + # 写入输入信息 + if 'input' in module_data: + f.write("输入信息:\n") + f.write(json.dumps(module_data['input'], indent=2, ensure_ascii=False)) + f.write("\n\n") + + # 写入输出信息 + if 'output' in module_data: + f.write("输出信息:\n") + f.write(json.dumps(module_data['output'], indent=2, ensure_ascii=False)) + f.write("\n") + + print("\n" + "=" * 80) + print("Done!") + print("=" * 80) + print(f"\nJSON saved to: {output_json}") + print(f"Text saved to: {output_txt}") + print(f"Log saved to: {output_dir}/precision_test_log.txt") + print("\nYou can compare these results with other frameworks (e.g. Paddle)") + print(f"\nComparison methods:") + print(f" 1. JSON format (recommended):") + print(f" python compare_precision.py {output_json} paddle_results.json") + print(f" 2. Text format:") + print(f" diff {output_txt} paddle_results_readable.txt") + + close_logger() + return True + + except Exception as e: + print(f"\nError: Test failed - {str(e)}") + import traceback + traceback.print_exc() + close_logger() + return False + + +if __name__ == "__main__": + success = main() + sys.exit(0 if success else 1) diff --git a/test/chemeleon2/test_loss/__init__.py b/test/chemeleon2/test_loss/__init__.py new file mode 100644 index 00000000..51ec7110 --- /dev/null +++ b/test/chemeleon2/test_loss/__init__.py @@ -0,0 +1,5 @@ +""" +Test module for chemeleon2 configuration tests. + +This module contains tests for configuration loading, validation, and parsing. +""" diff --git a/test/chemeleon2/test_loss/test_model_loss_with_raw.py b/test/chemeleon2/test_loss/test_model_loss_with_raw.py new file mode 100644 index 00000000..6c5be9e9 --- /dev/null +++ b/test/chemeleon2/test_loss/test_model_loss_with_raw.py @@ -0,0 +1,658 @@ +#!/usr/bin/env python3 +""" +VAE 和 LDM 模型的集成验证脚本。 + +""" + +import sys +from pathlib import Path + +parent_dir = Path(__file__).parent.parent.parent.parent +sys.path.insert(0, str(parent_dir)) +sys.path.insert(0, str(parent_dir / 'chemeleon2')) + +import paddle +import numpy as np +import json + +from ppmat.models.chemeleon2.vae_module.vae import VAEModule +from ppmat.models.chemeleon2.vae_module.encoder import TransformerEncoder +from ppmat.models.chemeleon2.vae_module.decoder import TransformerDecoder +from ppmat.models.chemeleon2.ldm_module.dit import DiT +from ppmat.models.chemeleon2.common.schema import CrystalBatch + +import os +HOME_DIR = Path(os.path.expanduser("~")) + +DEFAULT_WEIGHT_DIR = str(HOME_DIR / '.paddlemat/models/chemeleon2') +DEFAULT_OUTPUT_DIR = str(parent_dir) + "/outputs/" +DEFAULT_OUTPUT_FILE = DEFAULT_OUTPUT_DIR + 'chemeleon2_diff_lose.json' +DEFAULT_PRECISION_THRESHOLD = 1e-5 + + +# VAE 编码器输入数据 +INPUT_ATOM_TYPES = np.array( + [66, 67, 68, 68, 68, 69, 52, 52, 52, 33, 15, 15] +) + +INPUT_FRAC_COORDS = np.array( + [[0.16941863, 0.25193024, 0.45317790], + [0.49152774, 0.45137846, 0.11796811], + [0.48285794, 0.91339254, 0.14510530], + [0.82677758, 0.59025151, 0.79817736], + [0.81833100, 0.11027566, 0.76181233], + [0.15705948, 0.77731586, 0.42618808], + [0.33342591, 0.85136592, 0.77985525], + [0.65449405, 0.51935136, 0.45445091], + [0.98995161, 0.17083004, 0.11222821], + [0.33672857, 0.33988437, 0.77312100], + [0.66850150, 0.00315657, 0.45171800], + [0.99026763, 0.65868866, 0.12688932]] +) + +INPUT_LATTICES = np.array( + [[[4.07097721, 0.0, -0.91007066], + [-0.46107024, 8.27997398, -2.12557125], + [0.0, 0.0, 9.44053841]]] +) + +INPUT_NUM_ATOMS = np.array([12]) + +INPUT_LENGTHS = np.array( + [[4.17146063, 8.56087685, 9.44053841]] +) + +INPUT_ANGLES = np.array( + [[104.37628174, 102.60133362, 89.90788269]] +) + +# VAE 解码器输入数据（潜在变量 z） +INPUT_Z = np.array( + [[1.96778619, 0.52110022, -0.37526771, -1.06078792, -0.85598081, 1.11346722, 0.22077970, 0.96572930], + [0.64574736, -0.43119499, 0.03418273, -0.00645405, 0.89249128, 0.37275982, 2.41711688, 0.11571212], + [0.42960018, -0.43622142, 1.44103599, 0.11944307, -1.84233403, -0.07848608, 1.15001416, -0.54063123], + [0.91394019, 1.63482583, 0.57408303, 1.07989991, -2.48861003, 1.02283752, 1.02458203, 0.99812806], + [0.66510350, 1.63272715, -0.96877307, 0.91754431, -2.04141307, 2.00096869, 0.64047354, 0.05152281], + [2.28735399, 0.36919677, 1.11381781, -1.00699818, 0.64713156, 0.05649607, -0.62192798, 0.97391701], + [-1.52237606, 1.39929760, 1.43075931, -0.46554267, -0.68277979, -0.56380445, 1.09997213, 0.73207515], + [-1.66961420, 0.57721722, 0.25194433, 0.39816928, -0.65705436, -0.12071487, 0.88596821, 1.20686293], + [-0.70767230, -0.25845712, -1.09743047, 1.77922177, -0.84514201, 2.08224416, 0.54471433, -0.14093801], + [-1.00094223, 1.52224243, -0.10302663, -0.58021611, 1.42938161, -0.58981705, -1.16647685, 0.97821933], + [-0.58995736, 0.75216556, -1.43728018, 0.37418503, -0.30776960, 2.80792618, -0.79698497, -1.00396240], + [-0.84623051, -0.31481862, 0.54803443, 1.99706137, -0.12883027, 1.18385315, -0.81202215, -0.88749123]] +) + +# LDM 去噪器输入数据（扩散潜在变量 z） +INPUT_LDM_Z = np.array( + [[[0.19401880, 2.16137362, -0.17205022, 0.84906012, -1.92439902, 0.65298551, -0.64944082, -0.81752473], + [0.52796447, -1.27534986, -1.66212630, -0.30331373, -0.09256987, 0.19923715, -1.12043285, 1.85765862], + [-0.71451885, 0.68810511, 0.79683083, -0.03340188, 1.49173188, -0.51650929, -0.25409597, 1.47461557], + [-0.32603732, -1.15996265, 2.35513091, -0.69244707, 0.18374282, -1.18350995, -1.80286717, -1.58075690], + [0.83866954, 1.41918027, 0.64693671, 0.42527241, -1.58924079, 0.62234497, 1.68980360, -0.66480386], + [0.94254267, 0.07832550, 0.08465634, -0.14082992, 0.33156055, -0.58897614, -1.07228947, 0.09539576], + [-0.33469191, -0.52579743, -0.87762552, 0.39383137, 0.16395937, -0.19768225, 1.01041365, -1.34824479], + [-0.34977224, -0.64426798, 0.44678342, -0.53710973, 1.24231851, -0.81459534, 0.25015041, -0.42725861], + [1.10436928, -1.10279870, 0.55432665, -1.28465545, -0.38157833, 0.51394576, 0.10019008, 0.25862604], + [0.36168072, 2.27866960, 0.02334509, 1.58275771, -1.15917921, 0.94839233, -0.45734766, 0.76054770]], + [[-0.57868302, -0.70502084, -0.72338772, -0.50706196, -0.43984994, -0.41817018, 0.17413868, 0.44268036], + [0.50689828, -1.21680868, -0.27187300, 0.27654943, -1.43981659, -0.64632124, 0.07486922, 0.19387875], + [0.59601170, 0.23220330, 1.14146543, -0.68170702, -1.65314484, 0.00603564, 1.38148701, 1.27042663], + [0.02323810, -1.30014515, -0.75094134, 0.37562433, -0.54744226, -0.03964127, -0.77786469, -2.50188589], + [0.70001656, -0.09377469, -0.21625695, 0.44839421, -0.31519616, 0.02163736, 0.62534708, 0.24658130], + [0.74856061, -0.11692451, -0.10216469, -0.50108075, -0.50488758, -1.20719242, -0.24375997, -0.67842638], + [0.19728611, 0.97822028, -0.02866771, 1.68258953, 1.09085572, -0.99209458, -0.67126238, 1.71963453], + [2.46055436, -0.61983937, 1.27138603, -0.27986664, 0.43596676, 0.42602390, 1.06455433, -2.02799630], + [-0.63258272, 2.11064816, -0.09474602, 0.23587526, -0.73007232, -1.68571997, 0.91141981, 0.78854555], + [-0.62873089, 2.15955496, 1.16424942, -0.42566288, 0.23932022, -1.27767038, -0.12064690, -0.60658425]]] +) + + +# ============================================================================ +# 硬编码的参考输出数据 +# 这些数据是从原版项目中通过固定随机数种子 (seed=42) 推理计算得出的 +# 具体的复现代码参考当前目录的 README.md +# 输出的数据进行了裁剪，每个维度如果大于4维，则最多只取4维，避免太多数据 + +REFERENCE_OUTPUTS = { + 'vae_encoder_output': np.array([ + [-0.17990004, -0.04645369, -0.29342285, -0.05779608], + [-2.15766120, 0.05053813, 0.05233471, -0.08952494], + [0.32678959, 0.10300802, 0.28084213, -0.10772780], + [0.75947309, 0.17603433, 0.14910318, -0.12478722] + ]), + 'vae_encoder_quant_conv': np.array([ + [1.80965030, 0.52404141, -0.37543607, -1.06138432], + [0.62111264, -0.43615344, 0.03323286, -0.00297468], + [0.57471651, -0.43998951, 1.44004822, 0.11691619], + [0.44248250, 1.63856602, 0.57267267, 1.08317006] + ]), + 'vae_encoder_latent_mean': np.array([ + [1.80965030, 0.52404141, -0.37543607, -1.06138432], + [0.62111264, -0.43615344, 0.03323286, -0.00297468], + [0.57471651, -0.43998951, 1.44004822, 0.11691619], + [0.44248250, 1.63856602, 0.57267267, 1.08317006] + ]), + 'vae_encoder_latent_logvar': np.array([ + [-3.96147633, -11.50236607, -11.20673370, -10.69168186], + [-3.86336756, -10.93966484, -11.33978176, -11.27313232], + [-2.90164447, -10.99563599, -10.80719090, -11.23406792], + [-2.86538887, -10.86329651, -11.36846542, -10.63086319] + ]), + 'vae_decoder_atom_types': np.array([ + [-9.97118473, -12.22323322, -9.96976662, -8.63844585], + [-10.63955688, -17.72699165, -10.35566330, -10.19803047], + [-9.42540455, -1.47579587, -9.59889793, -11.58701992], + [-9.40364838, -2.24563265, -9.89963722, -12.12908268] + ]), + 'vae_decoder_frac_coords': np.array([ + [0.17350805, 0.25384152, 0.45326898], + [0.49485329, 0.45135695, 0.12035841], + [0.48613694, 0.91251493, 0.14744006], + [0.82910115, 0.58979881, 0.79950404] + ]), + 'vae_decoder_lengths': np.array([[1.81244886, 3.73297477, 4.11331844]]), + 'vae_decoder_angles': np.array([[1.82423151, 1.79026175, 1.56918395]]), + 'ldm_denoiser_output_noise': np.array([ + [[-0.38310128, 1.03488433, 0.64289629, 0.23227391], + [-0.52038497, -0.62339407, -0.72368735, 0.30327153], + [-0.47714779, 0.58150470, 0.24278848, 0.27084529], + [0.22972283, 0.21897136, 0.25752124, -0.38142803]], + [[0.02938975, -0.40452722, -0.57233673, 1.40531123], + [-0.13082200, -1.04666209, -0.68292052, 0.29126278], + [0.15657395, -0.48006776, 0.57248652, -0.19115604], + [0.08723904, -1.17863417, 0.76299840, -0.34876239]] + ]), +} + + +def get_first_n_elements(arr, max_dims=(4, 4, 4, 4)): + """根据维度从数组中提取前 N 个元素。""" + slice_shape = tuple(min(d, s) for d, s in zip(max_dims, arr.shape)) + slices = tuple(slice(0, s) for s in slice_shape) + return arr[slices] + + + +def run_test_vae_encoder(reference_data, weight_dir, precision_threshold): + """测试 VAE 编码器并与参考数据进行比较。""" + print("\n" + "="*80) + print("Testing VAE Encoder") + print("="*80) + + # 初始化模型 + encoder = TransformerEncoder( + max_num_elements=100, + d_model=512, + nhead=8, + dim_feedforward=2048, + activation='gelu', + dropout=0.0, + num_layers=8 + ) + + decoder = TransformerDecoder( + max_num_elements=100, + d_model=512, + nhead=8, + dim_feedforward=2048, + activation='gelu', + dropout=0.0, + num_layers=8 + ) + + vae = VAEModule( + encoder=encoder, + decoder=decoder, + latent_dim=8, + loss_weights={} + ) + + # 加载权重 + weight_path = Path(weight_dir) / 'vae_paddle.pdparams' + print(f"Loading weights from: {weight_path}") + full_state_dict = paddle.load(str(weight_path)) + vae.set_state_dict(full_state_dict) + vae.eval() + + # 使用硬编码数据创建 batch + batch = CrystalBatch() + batch.atom_types = paddle.to_tensor(INPUT_ATOM_TYPES, dtype='int64') + batch.frac_coords = paddle.to_tensor(INPUT_FRAC_COORDS, dtype='float32') + batch.lattices = paddle.to_tensor(INPUT_LATTICES, dtype='float32') + batch.num_atoms = paddle.to_tensor(INPUT_NUM_ATOMS, dtype='int64') + batch.batch = paddle.zeros([12], dtype='int64') + batch.token_idx = paddle.arange(12, dtype='int64') + batch.num_graphs = 1 + + # 运行推理 + with paddle.no_grad(): + encoded = vae.encode(batch) + encoder_output = encoded['x'] + quant_conv_output = encoded['moments'] + latent_mean = encoded['posterior'].mean + latent_logvar = encoded['posterior'].logvar + latent_z = encoded['posterior'].sample() + + # 收集输出以进行比较 + outputs = { + 'vae_encoder_output': encoder_output.numpy(), + 'vae_encoder_quant_conv': quant_conv_output.numpy(), + 'vae_encoder_latent_mean': latent_mean.numpy(), + 'vae_encoder_latent_logvar': latent_logvar.numpy(), + } + + # 与参考数据比较 + results = [] + all_passed = True + + for key, output in outputs.items(): + if key not in reference_data: + print(f"WARNING {key}: No reference data found") + continue + + ref_array = reference_data[key] + output_slice = get_first_n_elements(output) + + # 展平以进行比较 + ref_flat = ref_array.flatten() + output_flat = output_slice.flatten() + + min_len = min(len(ref_flat), len(output_flat)) + ref_flat = ref_flat[:min_len] + output_flat = output_flat[:min_len] + + diff = np.abs(ref_flat - output_flat) + max_diff = diff.max() + mean_diff = diff.mean() + + passed = max_diff < precision_threshold + + if not passed: + all_passed = False + + status = "PASS" if passed else "FAIL" + print(f"{status} {key}:") + print(f" Max diff: {max_diff:.6e}, Mean diff: {mean_diff:.6e}") + print(f" Shape: {output.shape}") + + if not passed: + max_idx = diff.argmax() + print(f" Ref value: {ref_flat[max_idx]:.8f}, Output value: {output_flat[max_idx]:.8f}") + + results.append({ + 'name': key, + 'max_diff': float(max_diff), + 'mean_diff': float(mean_diff), + 'passed': bool(passed) + }) + + return all_passed, results + + +def run_test_vae_decoder(reference_data, weight_dir, precision_threshold): + """测试 VAE 解码器并与参考数据进行比较。""" + print("\n" + "="*80) + print("Testing VAE Decoder") + print("="*80) + + # 初始化模型 + encoder = TransformerEncoder( + max_num_elements=100, + d_model=512, + nhead=8, + dim_feedforward=2048, + activation='gelu', + dropout=0.0, + num_layers=8 + ) + + decoder = TransformerDecoder( + max_num_elements=100, + d_model=512, + nhead=8, + dim_feedforward=2048, + activation='gelu', + dropout=0.0, + num_layers=8 + ) + + vae = VAEModule( + encoder=encoder, + decoder=decoder, + latent_dim=8, + loss_weights={} + ) + + # 加载权重 + weight_path = Path(weight_dir) / 'vae_paddle.pdparams' + print(f"Loading weights from: {weight_path}") + full_state_dict = paddle.load(str(weight_path)) + vae.set_state_dict(full_state_dict) + vae.eval() + + z_t = paddle.to_tensor(INPUT_Z, dtype='float32') + batch = CrystalBatch() + batch.atom_types = paddle.zeros([12], dtype='int64') + batch.frac_coords = paddle.zeros([12, 3], dtype='float32') + batch.lattices = paddle.zeros([1, 3, 3], dtype='float32') + batch.num_atoms = paddle.to_tensor(INPUT_NUM_ATOMS, dtype='int64') + batch.batch = paddle.zeros([12], dtype='int64') + batch.token_idx = paddle.arange(12, dtype='int64') + batch.num_graphs = 1 + + # 运行推理 + with paddle.no_grad(): + encoded = { + 'x': z_t, + 'z': z_t, + 'batch': batch.batch, + 'token_idx': batch.token_idx, + 'num_atoms': batch.num_atoms + } + decoded = vae.decode(encoded) + atom_types = decoded['atom_types'] + frac_coords = decoded['frac_coords'] + lengths = decoded['lengths'] + angles = decoded['angles'] + + # 收集输出以进行比较 + outputs = { + 'vae_decoder_atom_types': atom_types.numpy(), + 'vae_decoder_frac_coords': frac_coords.numpy(), + 'vae_decoder_lengths': lengths.numpy(), + 'vae_decoder_angles': angles.numpy(), + } + + # 与参考数据比较 + results = [] + all_passed = True + + for key, output in outputs.items(): + if key not in reference_data: + print(f"WARNING {key}: No reference data found") + continue + + ref_array = reference_data[key] + output_slice = get_first_n_elements(output) + + # 展平以进行比较 + ref_flat = ref_array.flatten() + output_flat = output_slice.flatten() + + min_len = min(len(ref_flat), len(output_flat)) + ref_flat = ref_flat[:min_len] + output_flat = output_flat[:min_len] + + diff = np.abs(ref_flat - output_flat) + max_diff = diff.max() + mean_diff = diff.mean() + + passed = max_diff < precision_threshold + + if not passed: + all_passed = False + + status = "PASS" if passed else "FAIL" + print(f"{status} {key}:") + print(f" Max diff: {max_diff:.6e}, Mean diff: {mean_diff:.6e}") + print(f" Shape: {output.shape}") + + if not passed: + max_idx = diff.argmax() + print(f" Ref value: {ref_flat[max_idx]:.8f}, Output value: {output_flat[max_idx]:.8f}") + + results.append({ + 'name': key, + 'max_diff': float(max_diff), + 'mean_diff': float(mean_diff), + 'passed': bool(passed) + }) + + return all_passed, results + + +def run_test_ldm_denoiser(reference_data, weight_dir, precision_threshold): + """测试 LDM 去噪器并与参考数据进行比较。""" + print("\n" + "="*80) + print("Testing LDM Denoiser") + print("="*80) + + # 初始化模型 + denoiser = DiT( + hidden_dim=768, + num_layers=12, + num_heads=12, + mlp_ratio=4.0, + input_dim=8, + learn_sigma=True + ) + + # 加载权重 + weight_path = Path(weight_dir) / 'ldm_paddle.pdparams' + print(f"Loading weights from: {weight_path}") + full_state_dict = paddle.load(str(weight_path)) + + denoiser_state_dict = {} + for key, value in full_state_dict.items(): + if key.startswith('denoiser.'): + new_key = key.replace('denoiser.', '') + denoiser_state_dict[new_key] = value + + denoiser.set_state_dict(denoiser_state_dict) + denoiser.eval() + + z_t = paddle.to_tensor(INPUT_LDM_Z, dtype='float32') + t = paddle.to_tensor([13, 25], dtype='int64') + batch_size, num_atoms, _ = INPUT_LDM_Z.shape + mask = paddle.ones([batch_size, num_atoms], dtype='bool') + + # 运行推理 + with paddle.no_grad(): + noise_pred = denoiser(z_t, t, mask=mask) + + # 收集输出以进行比较 + outputs = { + 'ldm_denoiser_output_noise': noise_pred.numpy(), + } + + # 与参考数据比较 + results = [] + all_passed = True + + for key, output in outputs.items(): + if key not in reference_data: + print(f"WARNING {key}: No reference data found") + continue + + ref_array = reference_data[key] + output_slice = get_first_n_elements(output) + + # 展平以进行比较 + ref_flat = ref_array.flatten() + output_flat = output_slice.flatten() + + min_len = min(len(ref_flat), len(output_flat)) + ref_flat = ref_flat[:min_len] + output_flat = output_flat[:min_len] + + diff = np.abs(ref_flat - output_flat) + max_diff = diff.max() + mean_diff = diff.mean() + + passed = max_diff < precision_threshold + + if not passed: + all_passed = False + + status = "PASS" if passed else "FAIL" + print(f"{status} {key}:") + print(f" Max diff: {max_diff:.6e}, Mean diff: {mean_diff:.6e}") + print(f" Shape: {output.shape}") + + if not passed: + max_idx = diff.argmax() + print(f" Ref value: {ref_flat[max_idx]:.8f}, Output value: {output_flat[max_idx]:.8f}") + + results.append({ + 'name': key, + 'max_diff': float(max_diff), + 'mean_diff': float(mean_diff), + 'passed': bool(passed) + }) + + return all_passed, results + + +def main(weight_dir=None, output_file=None, precision_threshold=None, save_result=True): + """ + 运行所有验证测试的主函数。 + + 参数: + weight_dir: 包含转换权重的目录（默认: DEFAULT_WEIGHT_DIR） + output_file: 结果的输出 JSON 文件路径（默认: DEFAULT_OUTPUT_FILE） + precision_threshold: 比较的精度阈值（默认: DEFAULT_PRECISION_THRESHOLD） + save_result: 是否保存结果到文件（默认: True） + + 返回: + int: 如果所有测试通过则返回 0，否则返回 1 + """ + # 如果未提供则使用默认值 + if weight_dir is None: + weight_dir = DEFAULT_WEIGHT_DIR + if output_file is None: + output_file = DEFAULT_OUTPUT_FILE + if precision_threshold is None: + precision_threshold = DEFAULT_PRECISION_THRESHOLD + + # 转换为 Path 对象 + weight_dir = Path(weight_dir) + output_file = Path(output_file) + + print("="*80) + print("Integrated VAE & LDM Validation (Hardcoded Input & Reference Data)") + print(f"Precision Threshold: {precision_threshold:.1e}") + print(f"Weight Directory: {weight_dir}") + print(f"Output File: {output_file}") + print("="*80) + + # 使用硬编码的参考数据 + print("\n使用硬编码的参考输出数据（来自原版项目，seed=42）") + reference_data = REFERENCE_OUTPUTS + print(f"Loaded {len(reference_data)} reference outputs") + + # 运行所有测试 + all_results = [] + global_passed = True + + # 测试 VAE 编码器 + passed, results = run_test_vae_encoder(reference_data, weight_dir, precision_threshold) + all_results.extend(results) + global_passed = global_passed and passed + + # 测试 VAE 解码器 + passed, results = run_test_vae_decoder(reference_data, weight_dir, precision_threshold) + all_results.extend(results) + global_passed = global_passed and passed + + # 测试 LDM 去噪器 + passed, results = run_test_ldm_denoiser(reference_data, weight_dir, precision_threshold) + all_results.extend(results) + global_passed = global_passed and passed + + # 打印摘要 + print("\n" + "="*80) + print("SUMMARY") + print("="*80) + + passed_count = sum(1 for r in all_results if r['passed']) + total_count = len(all_results) + + print(f"Passed: {passed_count}/{total_count}") + print(f"Threshold: {precision_threshold:.1e}") + + if global_passed: + print("ALL TESTS PASSED") + else: + print("FAILED TESTS:") + for r in all_results: + if not r['passed']: + print(f" - {r['name']}: Max diff = {r['max_diff']:.6e}") + + if save_result: + output_file.parent.mkdir(parents=True, exist_ok=True) + with open(output_file, 'w') as f: + json.dump({ + 'threshold': precision_threshold, + 'global_passed': global_passed, + 'passed_count': passed_count, + 'total_count': total_count, + 'results': all_results + }, f, indent=2) + + print(f"\nResults saved to: {output_file}") + print("="*80) + + return 0 if global_passed else 1 + +# 这里硬编码的矩阵，实际是从 原版的项目中，通过固定42随机数，推理计算出来，具体的复现代码参考 当前目录的README.md +def test_with_fix_random(weight_dir=None, output_file=None, precision_threshold=None): + """ + 固定随机种子的主测试入口点。 + + 参数: + weight_dir: 包含转换权重的目录（默认: DEFAULT_WEIGHT_DIR） + output_file: 结果的输出 JSON 文件路径（默认: DEFAULT_OUTPUT_FILE） + precision_threshold: 比较的精度阈值（默认: DEFAULT_PRECISION_THRESHOLD） + + 返回: + int: 如果所有测试通过则返回 0，否则返回 1 + """ + paddle.seed(42) + np.random.seed(42) + + exit_code = main( + weight_dir=weight_dir, + output_file=output_file, + precision_threshold=precision_threshold + ) + return exit_code + + +# 一般没有必要，直接从main开始执行，执行单元测试 test_with_fix_random 即可 +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser( + description='Integrated validation script for VAE and LDM models' + ) + + parser.add_argument( + '--weight-dir', + type=str, + default=None, + help='Directory containing converted weights' + ) + + parser.add_argument( + '--output-file', + type=str, + default=None, + help='Output JSON file path for results' + ) + + parser.add_argument( + '--threshold', + type=float, + default=None, + help='Precision threshold' + ) + + args = parser.parse_args() + + sys.exit( + test_with_fix_random( + weight_dir=args.weight_dir, + output_file=args.output_file, + precision_threshold=args.threshold + ) + )