diffcrysgen 0.1.0__py3-none-any.whl

diffcrysgen/analyze_generated_structures.py

@@ -0,0 +1,88 @@
+# diffcrysgen/analyze_generated_structures.py
+
+import numpy as np
+import pandas as pd
+import spglib
+from diffcrysgen.invert_pcr import InvertPCR
+from ase import Atoms
+import ase.io
+import os
+
+def analyze_ircr_data(generated_ircr_path="generated_ircr.npy", output_csv="generated_material_data.csv", min_valid_distance=0.5, save_cifs=False, cif_dir="cif-files"):
+    """
+    Analyzes generated IRCR data and extracts crystallographic info.
+
+    Args:
+        generated_ircr_path (str): Path to the generated IRCR .npy file
+        output_csv (str): Path to save the resulting CSV
+        min_valid_distance (float): Distance threshold for validity check in Angstrom
+        save_cifs (bool): If True, saves CIF files into subfolders
+        cif_dir (str): Base directory to store valid/invalid CIFs
+    """
+    data = np.load(generated_ircr_path)
+    Nmat = data.shape[0]
+    print(f"Loaded {Nmat} materials from {generated_ircr_path}")
+
+    results = {
+        "ID": [], "Formula": [], "type": [], "Natoms": [], "min-pairwise-distance": [],
+        "Zmax": [], "SPG-symbol": [], "SPG-number": [], "validity": [],
+        "a": [], "b": [], "c": [], "alpha": [], "beta": [], "gamma": []
+    }
+
+    if save_cifs:
+        os.makedirs(os.path.join(cif_dir, "valid"), exist_ok=True)
+        os.makedirs(os.path.join(cif_dir, "invalid"), exist_ok=True)
+
+    for i in range(Nmat):
+        mat = data[i, :]
+        ipcr = InvertPCR(mat, 94, 20)
+
+        try:
+            atoms = ipcr.get_atoms_object()
+            formula = atoms.get_chemical_formula()
+            atomic_numbers = atoms.get_atomic_numbers()
+            unique_elements = len(set(atomic_numbers))
+
+            mat_type = {1: "elemental", 2: "binary", 3: "ternary"}.get(unique_elements, "complex")
+            a, b, c = atoms.cell.lengths()
+            alpha, beta, gamma = atoms.cell.angles()
+
+            structure = (atoms.get_cell(), atoms.get_scaled_positions(), atomic_numbers)
+            spg_symbol, spg_number = spglib.get_spacegroup(structure, symprec=0.1).split(" ")
+            spg_number = int(spg_number.strip("()"))
+
+            distances = ipcr.get_distances()
+            min_dist = min(distances) if distances else 3.0
+            is_valid = min_dist >= min_valid_distance
+            mat_id = f"smps-{i+1}"
+
+            if save_cifs:
+                cif_path = os.path.join(cif_dir, "valid" if is_valid else "invalid", f"{mat_id}.cif")
+                ase.io.write(cif_path, atoms)
+
+            # Append data
+            results["ID"].append(mat_id)
+            results["Formula"].append(formula)
+            results["type"].append(mat_type)
+            results["Natoms"].append(len(atomic_numbers))
+            results["min-pairwise-distance"].append(min_dist)
+            results["Zmax"].append(max(atomic_numbers))
+            results["SPG-symbol"].append(spg_symbol)
+            results["SPG-number"].append(spg_number)
+            results["validity"].append("valid" if is_valid else "invalid")
+            results["a"].append(a)
+            results["b"].append(b)
+            results["c"].append(c)
+            results["alpha"].append(alpha)
+            results["beta"].append(beta)
+            results["gamma"].append(gamma)
+
+        except Exception as e:
+            #print(f"[Material {i}] Error: {e}")
+            print("spglib failed to assign spg")
+
+    df = pd.DataFrame(results)
+    df.to_csv(output_csv, index=False)
+    print(f"Saved extracted data to: {output_csv}")
+    return df
+

diffcrysgen/invert_pcr.py

@@ -0,0 +1,163 @@
+
+from ase.spacegroup import crystal
+import numpy as np
+import ase.io 
+
+
+# Dictionary containing atomic number with respective element
+
+atomic_name = {"1":"H", "2":"He", "3":"Li", "4":"Be", "5":"B", "6":"C", "7":"N", "8":"O", "9":"F", "10":"Ne",\
+               "11":"Na", "12":"Mg", "13":"Al", "14":"Si", "15":"P", "16":"S", "17":"Cl", "18":"Ar", "19": "K",\
+               "20":"Ca", "21":"Sc", "22":"Ti", "23":"V", "24":"Cr", "25":"Mn", "26":"Fe", "27":"Co", "28":"Ni",\
+               "29":"Cu", "30":"Zn", "31":"Ga", "32":"Ge", "33":"As","34":"Se", "35":"Br","36":"Kr", "37":"Rb",\
+               "38":"Sr", "39":"Y", "40":"Zr", "41":"Nb", "42":"Mo", "43":"Tc", "44":"Ru", "45":"Rh", "46":"Pd",\
+               "47":"Ag", "48":"Cd", "49":"In", "50":"Sn", "51":"Sb", "52":"Te", "53":"I", "54":"Xe", "55":"Cs", "56":"Ba",\
+               "57":"La", "58":"Ce", "59":"Pr", "60":"Nd", "61":"Pm", "62":"Sm", "63":"Eu", "64":"Gd", "65":"Tb", "66":"Dy",\
+               "67":"Ho", "68":"Er", "69":"Tm", "70":"Yb", "71":"Lu", "72":"Hf", "73":"Ta", "74":"W", "75":"Re", "76":"Os",\
+               "77":"Ir", "78":"Pt", "79":"Au", "80":"Hg", "81":"Tl", "82":"Pb", "83":"Bi", "84":"Po", "85":"At", "86":"Rn",\
+               "87":"Fr", "88":"Ra", "89":"Ac", "90":"Th", "91":"Pa", "92":"U", "93":"Np", "94":"Pu"}
+               
+
+class InvertPCR :
+    
+    """ Here we will extract necessary informations from point clound reps (PCR) in order to construct cif file."""
+    
+    def __init__(self,PCR,z_max,n_sites) :
+        self.PCR = PCR
+        self.z_max = z_max
+        self.n_sites = n_sites
+        
+    def get_element_matrix(self):
+        ele_mat = self.PCR[0:self.z_max,:]
+        ele_mat[ele_mat < 0.5] = 0
+        return ele_mat
+    
+    def get_lattice_matrix(self):
+        lat_mat = self.PCR[self.z_max:self.z_max+2,:]
+        return lat_mat
+    
+    def get_lattice_parameters(self):
+        lattice_matrix = self.get_lattice_matrix()
+        lattice_parameters = list(lattice_matrix.flatten())
+        a = lattice_parameters[0]
+        b = lattice_parameters[1]
+        c = lattice_parameters[2]
+        alpha = lattice_parameters[3]
+        beta = lattice_parameters[4]
+        gamma = lattice_parameters[5]
+        scaled_par = [a,b,c,alpha,beta,gamma]
+        return scaled_par
+    
+    def get_basis_matrix(self):
+        basis_mat = self.PCR[self.z_max+2:self.z_max+2+self.n_sites,:]
+        return basis_mat
+    
+    def get_site_matrix(self):
+        site_mat = self.PCR[self.z_max+2+self.n_sites:self.z_max+2+2*self.n_sites,:]
+        site_mat[site_mat < 0.5] = 0
+        return site_mat
+
+    def get_property_matrix(self):
+        prop_mat = self.PCR[self.z_max+2+2*self.n_sites:self.z_max+2+2*self.n_sites+8,:]
+        return prop_mat
+    
+    def get_unique_atomic_numbers(self):
+        element_matrix = self.get_element_matrix()
+        z_unique = []
+        for i in range(3):
+            col = list(element_matrix[:,i])
+            max_value = max(col)
+            if max_value == 0 :
+                z_unique.append(0)
+            if max_value != 0 :
+                z_unique.append(col.index(max_value)+1)
+        return z_unique
+    
+    def get_unique_elements(self):
+        z_unique = self.get_unique_atomic_numbers()
+        unique_element = []
+        for z in z_unique :
+            if z == 0 :
+                unique_element.append("-")
+            else :
+                unique_element.append(atomic_name[str(z)])
+        return unique_element
+
+    
+    def get_atomic_numbers(self):
+        element_matrix = self.get_element_matrix()
+        site_matrix = self.get_site_matrix()
+        z_unique = self.get_unique_atomic_numbers()
+        z_total = []
+        for i in range(self.n_sites):
+            row = list(site_matrix[i,:])
+            val = max(row)
+            if val == 0 :
+                z_total.append(0)
+            if val != 0 :
+                z_total.append(z_unique[row.index(val)])
+        return z_total
+    
+    def get_elements_basis(self):
+        basis = self.get_basis_matrix()
+        z_total = self.get_atomic_numbers()
+        final_atoms = []
+        final_atoms_index = []
+        final_elements = []
+        for i in range(len(z_total)):
+            z = z_total[i]
+            if z != 0 :
+                final_atoms_index.append(i)
+                final_atoms.append(z)
+                final_elements.append(atomic_name[str(z)])
+        final_basis = basis[final_atoms_index,:]
+   
+        return final_elements, final_basis
+
+    def get_formula(self) :
+        a,_ = self.get_elements_basis()
+        elements = [a[0]]
+        numbers = []
+
+        for i in range(len(a)) :
+            ele = a[i]
+            if ele != elements[len(elements)-1] :
+                elements.append(ele)
+                numbers.append(i)
+
+        numbers.append(len(a))
+
+        final_numbers = []
+        final_numbers.append(numbers[0])
+        for i in range(1,len(numbers)):
+            final_numbers.append(numbers[i]-numbers[i-1])
+
+        formula = elements[0] + str(final_numbers[0])
+        for i in range(1,len(elements)) :
+            formula += elements[i] + str(final_numbers[i])
+
+        return formula
+    
+    def get_atoms_object(self):
+        symbol = self.get_formula()
+        _, basis = self.get_elements_basis()
+        lattice = self.get_lattice_parameters()
+        atoms = crystal(symbols=symbol,
+                        basis=basis,
+                        cellpar=lattice)
+        return atoms
+
+    def get_distances(self):
+        atoms = self.get_atoms_object()
+        positions = atoms.get_positions() 
+        nele = positions.shape[0]
+        distances = []
+        for i in range(nele):
+            for j in range(i+1,nele):
+                pos1 = positions[i,:]
+                pos2 = positions[j,:]
+                vec = pos1 - pos2
+                dist = np.linalg.norm(vec)
+                distances.append(np.around(dist,4))
+        return distances    
+

diffcrysgen/model.py

@@ -0,0 +1,316 @@
+
+
+# Importing necessary libraries
+
+import torch
+import numpy as np
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+# Sinusoidal Positional Embedding
+
+class PositionalEmbedding(torch.nn.Module):
+    """
+    Given a noise label, it provides the corresponding positional embedding vector.
+    """
+    def __init__(self, emb_dim, max_positions = 10000):
+        """
+        Parameter : emb_dim (Type : Int) : Embedding Dimension
+        """
+        super(PositionalEmbedding, self).__init__()
+        self.emb_dim = emb_dim
+        self.max_positions = max_positions
+
+    def forward(self, x):
+        """
+        Input : x [Tensor object], representing noise levels for a batch of input images
+        """
+
+        # Precompute frequencies for efficiency
+        freqs = torch.arange(0, self.emb_dim // 2, dtype=torch.float32, device=x.device)
+        freqs = freqs / (self.emb_dim // 2)
+        freqs = (1 / self.max_positions) ** freqs
+
+        # Outer product between Input tensor and freqs tensor
+        x = torch.tensordot(x, freqs, dims=0)
+        # Sinusoidal embedding
+        # sin block
+        sin_block = torch.sin(x)
+        # cos block
+        cos_block = torch.cos(x)
+        # Concatenate sin and cos block along column dimension
+        embedding = torch.cat([sin_block, cos_block], dim=1)
+        # Finally, concatenate sin and cos block in such a way so that even columns
+        # come from sin block and odd column comes from cos block.
+        sin_unsqueezed = sin_block.unsqueeze(2)
+        cos_unsqueezed = cos_block.unsqueeze(2)
+        final_embedding = torch.reshape(torch.cat([sin_unsqueezed, cos_unsqueezed], dim=2), (embedding.shape[0], -1))
+        return final_embedding
+
+
+# Module for Preconditioning
+# Adapted from "Elucidating the Design Space of Diffusion-Based
+# Generative Models" by Karras et al.
+
+class Precond(nn.Module):
+  def __init__(self,
+               denoise_fn,
+               sigma_min=0,               # minimum supported noise level
+               sigma_max=float("inf"),    # maximum supported noise level
+               sigma_data=0.5,            # expected standard deviation of training data
+               ):
+    super().__init__()
+
+    self.denoise_fn_F = denoise_fn
+    self.sigma_min = sigma_min
+    self.sigma_max = sigma_max
+    self.sigma_data = sigma_data
+
+  def forward(self, x, sigma):
+    x = x.to(torch.float32)
+    sigma = sigma.to(torch.float32).reshape(-1,1,1)
+    dtype = torch.float32
+
+    c_skip = self.sigma_data ** 2 / (sigma ** 2 + self.sigma_data ** 2)
+    c_out = sigma * self.sigma_data / (sigma ** 2 + self.sigma_data ** 2).sqrt()
+    c_in = 1 / (sigma ** 2 + self.sigma_data ** 2).sqrt()
+    c_noise = sigma.log() / 4
+
+    x_in = c_in * x
+    F_x = self.denoise_fn_F((x_in).to(dtype), c_noise.flatten())
+
+    assert F_x.dtype == dtype
+    D_x = c_skip * x + c_out * F_x.to(torch.float32)
+    return D_x
+
+  def round_sigma(self, sigma):
+    return torch.as_tensor(sigma)
+
+
+# EDM Loss
+
+def EDMLoss(denoise_fn, data):
+  P_mean = -1.2
+  P_std = 1.2
+  sigma_data = 0.5
+  rnd_normal = torch.randn(data.shape[0], device=data.device)
+  sigma = (rnd_normal * P_std + P_mean).exp()
+  weight = (sigma ** 2 + sigma_data ** 2) / (sigma * sigma_data) ** 2
+  y = data
+  n = torch.randn_like(y) * sigma.unsqueeze(1).unsqueeze(1)
+  D_yn = denoise_fn(y + n, sigma)
+  target = y
+  loss = weight.unsqueeze(1).unsqueeze(1) * ((D_yn - target) ** 2)
+  return loss
+
+
+# Final Model
+
+class Model(nn.Module):
+  def __init__(self, denoise_fn, P_mean=-1.2, P_std=1.2, sigma_data=0.5):
+    super().__init__()
+    self.P_mean = P_mean
+    self.P_std = P_std
+    self.sigma_data = sigma_data
+    self.denoise_fn_D = Precond(denoise_fn)
+
+  def forward(self, x):
+    loss = EDMLoss(self.denoise_fn_D, x)
+    return loss.mean(-1).mean()
+
+
+# time-encoding layer
+
+def make_te(dim_in, dim_out):
+    return nn.Sequential(nn.Linear(dim_in, dim_out),nn.SiLU(),nn.Linear(dim_out, dim_out))
+
+
+
+#=================================== UNet================================================= 
+
+def weight_standardization(weight: torch.Tensor, eps: float):
+    c_out, c_in, kernel_size = weight.shape
+    weight = weight.view(c_out, -1)
+    var, mean = torch.var_mean(weight, dim=1, keepdim=True)
+    # Standardize weights
+    weight = (weight - mean) / (torch.sqrt(var + eps))
+    return weight.view(c_out, c_in, kernel_size)
+
+
+class WSConv1d(torch.nn.Conv1d):
+    def __init__(self, in_channels, out_channels, kernel_size,
+                 stride=1, padding=0, dilation=1, groups: int = 1, bias: bool = True,
+                 padding_mode: str = 'zeros', eps: float = 1e-5):
+        super(WSConv1d, self).__init__(in_channels, out_channels, kernel_size,
+                                     stride=stride, padding=padding, dilation=dilation,
+                                     groups=groups, bias=bias, padding_mode=padding_mode)
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor):
+        # Apply weight standardization before convolution
+        standardized_weight = weight_standardization(self.weight, self.eps)
+        return F.conv1d(x, standardized_weight, self.bias, self.stride,
+                        self.padding, self.dilation, self.groups)
+
+
+class ResNetBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, groups=1, eps=1e-5):
+        super(ResNetBlock, self).__init__()
+
+        # First WSConv1d layer with GroupNorm
+        self.conv1 = WSConv1d(in_channels, out_channels, kernel_size, stride=stride, padding=kernel_size//2, groups=groups, eps=eps)
+        self.norm1 = nn.GroupNorm(num_groups=1, num_channels=out_channels)
+
+        # Second WSConv1d layer with GroupNorm
+        self.conv2 = WSConv1d(out_channels, out_channels, kernel_size, stride=stride, padding=kernel_size//2, groups=groups, eps=eps)
+        self.norm2 = nn.GroupNorm(num_groups=1, num_channels=out_channels)
+
+        # Shortcut connection to match the shape when in_channels != out_channels
+        self.shortcut = nn.Conv1d(in_channels, out_channels, kernel_size=1, stride=stride, bias=False) if in_channels != out_channels else nn.Identity()
+
+        # activation
+        self.activation = nn.SiLU()
+
+    def forward(self,x):
+        # Save the input for the residual connection
+        identity = x
+
+        # First conv-norm-activation
+        out = self.conv1(x)
+        out = self.norm1(out)
+        out = self.activation(out)
+
+        # Second conv-norm-activation
+        out = self.conv2(out)
+        out = self.norm2(out)
+
+        # Add the residual (shortcut connection)
+        identity = self.shortcut(identity)
+        out += identity
+
+        # Apply ReLU after adding the residual
+        out = self.activation(out)
+        
+        return out
+
+
+# attention block
+# we apply attention block at a resolution level of 17
+
+class AttentionBlock(nn.Module):
+    """
+    MHA expects input of shape (batch_size, seq_length, embedding_dim).
+    Our data shape (batch_size, num_channels, resolution), e.g. (64,128,17).
+    embedding_dim == resolution
+    seq_length = num_channels
+    """
+    def __init__(self, emb_dim, num_heads=1, dropout=0.1):
+        super(AttentionBlock, self).__init__()
+        
+        self.attention = nn.MultiheadAttention(embed_dim=emb_dim, num_heads=num_heads, dropout=dropout)
+        
+        # Layer normalization
+        self.layernorm1 = nn.LayerNorm(emb_dim)
+        self.layernorm2 = nn.LayerNorm(emb_dim)
+        
+        # Feed-forward network (after attention)
+        self.ffn = nn.Sequential(
+            nn.Linear(emb_dim, emb_dim * 4),
+            nn.ReLU(),
+            nn.Linear(emb_dim * 4, emb_dim)
+        )
+       
+    def forward(self, x):
+        
+        # Apply self-attention
+        attn_output, attn_weights = self.attention(x, x, x)  # Query, Key, Value all are x (self-attention)
+        
+        # Residual connection followed by layer normalization : add & norm
+        x = self.layernorm1(x + attn_output)
+        
+        # Feed-forward network
+        ffn_output = self.ffn(x)
+        
+        # Residual connection followed by layer normalization : add & norm
+        x = self.layernorm2(x + ffn_output)
+        
+        return x
+
+
+class UNet(nn.Module):
+    def __init__(self, in_c, out_c, time_emb_dim):
+        super(UNet, self).__init__()
+
+        # map noise labels
+        self.map_noise = PositionalEmbedding(emb_dim=time_emb_dim)
+
+        # encoder path
+        self.te1 = make_te(time_emb_dim, 3)
+        self.b1 = nn.Sequential(ResNetBlock(3, 32),ResNetBlock(32, 32),ResNetBlock(32,32)) # shape = [batch,32,136]
+        self.down1 = nn.Conv1d(32, 32, 4, 2, 1) # shape = [batch,32,68]
+
+        self.te2 = make_te(time_emb_dim, 32)
+        self.b2 = nn.Sequential(ResNetBlock(32, 64),ResNetBlock(64, 64),ResNetBlock(64,64)) # shape = [batch,64,68]
+        self.down2 = nn.Conv1d(64, 64, 4, 2, 1) # shape = [batch,64,34]
+
+        self.te3 = make_te(time_emb_dim, 64)
+        self.b3 = nn.Sequential(ResNetBlock(64, 128),ResNetBlock(128, 128),ResNetBlock(128,128)) # shape = [batch,128,34]
+        self.down3 = nn.Conv1d(128, 128, 4, 2, 1) # shape = [batch,128,17]
+
+        self.attn1 = AttentionBlock(emb_dim=17)
+
+        # Bottleneck
+        self.te_mid = make_te(time_emb_dim, 128)
+        self.b_mid = nn.Sequential(ResNetBlock(128, 64),ResNetBlock(64, 64),ResNetBlock(64,128)) # shape = [batch,128,17]
+
+        self.attn2 = AttentionBlock(emb_dim=17)
+
+
+        # decoder path
+        self.up1 = nn.ConvTranspose1d(128, 128, 4, 2, 1) # shape = [batch,128,34]
+        self.te4 = make_te(time_emb_dim, 256)
+        self.b4 = nn.Sequential(ResNetBlock(256,128),ResNetBlock(128,64),ResNetBlock(64,64)) # shape = [batch,64,34]
+
+        self.up2 = nn.ConvTranspose1d(64, 64, 4, 2, 1) # shape = [batch,64,68]
+        self.te5 = make_te(time_emb_dim, 128)
+        self.b5 = nn.Sequential(ResNetBlock(128, 64),ResNetBlock(64, 32),ResNetBlock(32, 32)) # shape = [batch,32,68]
+
+        self.up3 = nn.ConvTranspose1d(32, 32, 4, 2, 1) # shape = [batch,32,136]
+        self.te6 = make_te(time_emb_dim, 64)
+        self.b6 = nn.Sequential(ResNetBlock(64, 32),ResNetBlock(32, 32),ResNetBlock(32, 32)) # shape = [batch,32,136]
+
+        # output
+        self.conv_out = nn.Conv1d(32, 3, 3, 1, 1) # shape = [batch,3,136]
+
+    def forward(self, x, noise_labels):
+        t = self.map_noise(noise_labels)
+        n = len(x)
+        out1 = self.b1(x + self.te1(t).reshape(n, -1, 1))  
+        out2 = self.b2(self.down1(out1) + self.te2(t).reshape(n, -1, 1))  
+        out3 = self.b3(self.down2(out2) + self.te3(t).reshape(n, -1, 1))  
+
+        out_mid = self.b_mid(self.attn1(self.down3(out3)) + self.te_mid(t).reshape(n, -1, 1)) 
+
+        # self-attention
+        out_mid = self.attn2(out_mid)
+
+        out4 = torch.cat((out3,self.up1(out_mid)), dim=1) 
+        out4 = self.b4(out4 + self.te4(t).reshape(n, -1, 1))  
+        out5 = torch.cat((out2, self.up2(out4)), dim=1) 
+        out5 = self.b5(out5 + self.te5(t).reshape(n, -1, 1)) 
+
+        out6 = torch.cat((out1, self.up3(out5)), dim=1)  
+        out6 = self.b6(out6 + self.te6(t).reshape(n, -1, 1))  
+
+        out = self.conv_out(out6) 
+        return out
+
+
+
+
+
+
+
+
+

diffcrysgen/sampler.py

@@ -0,0 +1,116 @@
+
+import torch
+import numpy as np
+import tqdm
+import time
+from diffcrysgen.model import Model, UNet
+from diffcrysgen.utils import *
+
+
+# Adapted from "Elucidating the Design Space of Diffusion-Based
+# Generative Models" by Karras et al.
+
+
+def round_sigma(sigma):
+    return torch.as_tensor(sigma)
+
+
+def generate_samples(num_samples: int = 100, batch_size: int = 1000, model_path: str = "assets/saved-model/sdm.pt"):
+    """
+    Generate IRCR samples using our pre-trained version.
+    
+    Args:
+        num_samples (int): Total number of samples to generate.
+        batch_size (int): Batch size for sampling.
+        model_path (str): Path to the saved model (.pt file).
+
+    Returns:
+        np.ndarray: Generated IRCR array of shape (num_samples, F, C)
+    """
+    # Load scaler
+    scaler = load_saved_diffusion_scaler()
+    print("Scaler loaded.")
+
+    # Load device
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print("Using device:", device)
+
+    # Load the model
+    denoise_fn = UNet(in_c=3, out_c=3, time_emb_dim=256).to(device)
+    model = Model(denoise_fn=denoise_fn).to(device)
+    model.load_state_dict(torch.load(model_path, map_location=device, weights_only=True))
+    model.eval()
+    print("Model loaded from:", model_path)
+
+    # Sampling hyperparameters
+    num_steps = 100
+    sigma_min = 0.002
+    sigma_max = 80
+    rho = 7
+    S_churn = 1
+    S_min = 0
+    S_max = float("inf")
+    S_noise = 1
+
+    net = model.denoise_fn_D
+    sigma_min = max(sigma_min, net.sigma_min)
+    sigma_max = min(sigma_max, net.sigma_max)
+
+    # Time steps
+    step_indices = torch.arange(num_steps, dtype=torch.float64, device=device)
+    t_steps = (sigma_max ** (1 / rho) + step_indices / (num_steps - 1) * (sigma_min ** (1 / rho) - sigma_max ** (1 / rho))) ** rho
+    t_steps = torch.cat([round_sigma(t_steps), torch.zeros_like(t_steps[:1])])  # append t_N = 0
+
+    img_channels = 3
+    img_resolution = 136
+
+    total_loops = int(np.ceil(num_samples / batch_size))
+    final_data = []
+
+    start_time = time.time()
+
+    with torch.no_grad():
+        for loop in range(total_loops):
+            curr_batch_size = min(batch_size, num_samples - loop * batch_size)
+            latents = torch.randn([curr_batch_size, img_channels, img_resolution], device=device)
+            x_next = latents * t_steps[0]
+
+            # Main Sampling Loop
+            for i, (t_cur, t_next) in tqdm.tqdm(list(enumerate(zip(t_steps[:-1], t_steps[1:]))), desc=f"Sampling {loop+1}/{total_loops}", leave=False):
+                x_cur = x_next
+
+                # Increase Noise temporarily
+                gamma = min(S_churn / num_steps, np.sqrt(2) - 1) if S_min <= t_cur <= S_max else 0
+                t_hat = round_sigma(t_cur + gamma * t_cur)
+
+                t_cur = t_cur.repeat(curr_batch_size)
+                t_next = t_next.repeat(curr_batch_size)
+                t_hat = t_hat.repeat(curr_batch_size)
+
+                x_hat = x_cur + (t_hat[:, None, None] ** 2 - t_cur[:, None, None] ** 2).sqrt() * S_noise * torch.randn_like(x_cur)
+
+                denoised = net(x_hat, t_hat).float()
+                d_cur = (x_hat - denoised) / t_hat[:, None, None]
+                x_next = x_hat + (t_next[:, None, None] - t_hat[:, None, None]) * d_cur
+
+                if i < num_steps - 1:
+                    denoised = net(x_next, t_next).float()
+                    d_prime = (x_next - denoised) / t_next[:, None, None]
+                    x_next = x_hat + (t_next[:, None, None] - t_hat[:, None, None]) * (0.5 * d_cur + 0.5 * d_prime)
+
+            x_next = x_next.permute(0, 2, 1).cpu().numpy()  # shape: (batch, F, C)
+            x_next = (x_next + 1) / 2
+            x_next = inv_minmax(x_next, scaler)
+            final_data.append(x_next)
+
+            print(f"Sampling batch {loop + 1}/{total_loops} done.")
+
+    elapsed = (time.time() - start_time) / 60
+    print(f"Total sampling time: {elapsed:.2f} minutes")
+
+    generated_array = np.concatenate(final_data, axis=0)[:num_samples]
+    print("Generated IRCR shape:", generated_array.shape)
+    return generated_array
+
+
+

diffcrysgen/trainer.py

@@ -0,0 +1,128 @@
+
+
+import time
+import torch
+import numpy as np
+import pandas as pd
+from torch.utils.data import DataLoader
+from torch.optim.lr_scheduler import ReduceLROnPlateau as RLROP
+from tqdm.auto import tqdm
+from model import UNet, Model
+from utils import minmax
+
+def load_data(train_path, test_path, batch_size=64):
+    # Original IRCR has shape [batch,142,3] : IRCR={E,L,C,O,P}
+    # We are not considering property matrix (P) for unconditional generation.
+    train_data = np.load(train_path)[:, :136, :]
+    test_data = np.load(test_path)[:, :136, :]
+
+    print(f"Training data shape: {train_data.shape}")
+    print(f"Test data shape: {test_data.shape}")
+
+    # Normalize [0, 1]
+    train_scaled, train_scaler = minmax(train_data)
+    test_scaled, test_scaler = minmax(test_data)
+
+    # Normalize to [-1, 1]
+    train_scaled = 2 * train_scaled - 1
+    test_scaled = 2 * test_scaled - 1
+
+    # Convert to torch tensors (N, 3, F)
+    train_tensor = torch.from_numpy(train_scaled).float().permute(0, 2, 1)
+    test_tensor = torch.from_numpy(test_scaled).float().permute(0, 2, 1)
+
+    print(f"Transformed training shape: {train_tensor.shape}")
+    print(f"Transformed test shape: {test_tensor.shape}")
+    print(f"Train value range: {train_tensor.min().item()} to {train_tensor.max().item()}")
+
+    train_loader = DataLoader(train_tensor, batch_size=batch_size)
+    test_loader = DataLoader(test_tensor, batch_size=batch_size)
+
+    return train_loader, test_loader, train_scaler, test_scaler
+
+
+def train_diffcrysgen(
+    train_loader,
+    test_loader,
+    num_epochs=1000,
+    save_path="sdm.pt",
+    initial_lr=5e-4,
+    device=None,
+):
+    if device is None:
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Using device: {device}")
+
+    denoise_fn = UNet(in_c=3, out_c=3, time_emb_dim=256).to(device)
+    model = Model(denoise_fn=denoise_fn).to(device)
+
+    optimizer = torch.optim.Adam(model.parameters(), lr=initial_lr, weight_decay=0)
+    scheduler = RLROP(optimizer, mode="min", factor=0.3, patience=50, verbose=True)
+
+    best_loss = float("inf")
+    train_losses = []
+    test_losses = []
+    lrs = []
+
+    start_time = time.time()
+
+    for epoch in tqdm(range(num_epochs), desc="Training Progress", colour="green"):
+        # Manual LR scheduling based on epochs
+        if epoch == 500:
+            for group in optimizer.param_groups:
+                group["lr"] = 1e-4
+        elif epoch == 800:
+            for group in optimizer.param_groups:
+                group["lr"] = 5e-5
+
+        model.train()
+        train_loss = 0
+        for batch in tqdm(train_loader, leave=False, desc=f"Epoch {epoch+1}/{num_epochs}", colour="blue"):
+            batch = batch.to(device)
+            loss = model(batch)
+            train_loss += loss.item()
+
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+
+        train_loss /= len(train_loader)
+        train_losses.append(train_loss)
+
+        model.eval()
+        test_loss = 0
+        with torch.no_grad():
+            for batch in tqdm(test_loader, leave=False, desc=f"Epoch {epoch+1}/{num_epochs}", colour="blue"):
+                batch = batch.to(device)
+                loss = model(batch)
+                test_loss += loss.item()
+        test_loss /= len(test_loader)
+        test_losses.append(test_loss)
+ 
+        # step the scheduler
+        # scheduler.step(test_loss)
+
+        current_lr = optimizer.param_groups[0]["lr"]
+        lrs.append(current_lr)
+
+        log_msg = f"Epoch {epoch+1}: train_loss: {train_loss:.4f}, test_loss: {test_loss:.4f}, lr: {current_lr:.6f}"
+        if test_loss < best_loss:
+            best_loss = test_loss
+            torch.save(model.state_dict(), save_path)
+            log_msg += " --> Best model ever (stored)"
+        print(log_msg)
+
+        # scheduler.step(val_loss)
+
+    print(f"Training completed in {(time.time() - start_time):.2f}s")
+
+    # Save epoch-wise loss
+    results_df = pd.DataFrame({
+        "epoch": np.arange(num_epochs),
+        "train_loss": train_losses,
+        "test_loss": test_losses,
+        "lr": lrs
+    })
+    results_df.to_csv("epoch-loss.csv", index=False)
+    print(results_df.tail())
+

diffcrysgen/utils.py

@@ -0,0 +1,38 @@
+
+import os
+import joblib
+import numpy as np
+from sklearn.preprocessing import MinMaxScaler
+
+# This is the normalization script adapted from FTCP work.
+# It works for FTCP, IRCR.
+
+def minmax(pcr): 
+    dim0, dim1, dim2 = pcr.shape
+    scaler = MinMaxScaler()
+    pcr_ = np.transpose(pcr, (1, 0, 2))
+    pcr_ = pcr_.reshape(dim1, dim0*dim2)
+    pcr_ = scaler.fit_transform(pcr_.T)
+    pcr_ = pcr_.T
+    pcr_ = pcr_.reshape(dim1, dim0, dim2)
+    pcr_normed = np.transpose(pcr_, (1, 0, 2))
+    return pcr_normed, scaler
+
+def inv_minmax(pcr_normed, scaler):
+    dim0, dim1, dim2 = pcr_normed.shape
+    pcr_ = np.transpose(pcr_normed, (1, 0, 2))
+    pcr_ = pcr_.reshape(dim1, dim0*dim2)
+    pcr_ = scaler.inverse_transform(pcr_.T)
+    pcr_ = pcr_.T
+    pcr_ = pcr_.reshape(dim1, dim0, dim2)
+    pcr = np.transpose(pcr_, (1, 0, 2))
+    return pcr
+
+def load_saved_diffusion_scaler(path=None):
+    if path is None:
+        here = os.path.dirname(__file__)
+        path = os.path.join(here, "..", "assets", "ircr_diffusion_scaler.pkl")
+        path = os.path.abspath(path)
+    return joblib.load(path)
+
+

diffcrysgen-0.1.0.dist-info/METADATA

@@ -0,0 +1,103 @@
+Metadata-Version: 2.4
+Name: diffcrysgen
+Version: 0.1.0
+Summary: "DiffCrysGen is a score-based diffusion model for accelerated design of diverse inorganic crystalline materials."
+Home-page: https://github.com/SouravMal/DiffCrysGen.git
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Development Status :: 3 - Alpha
+Classifier: Topic :: Scientific/Engineering :: Chemistry
+Classifier: Topic :: Scientific/Engineering :: Physics
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch~=2.5.1
+Requires-Dist: numpy~=2.0.1
+Requires-Dist: pandas~=2.2.3
+Requires-Dist: scikit-learn~=1.6.1
+Requires-Dist: tqdm
+Requires-Dist: ase~=3.24.0
+Requires-Dist: spglib~=2.5.0
+Requires-Dist: joblib
+Dynamic: license-file
+
+# DiffCrysGen [![Project Version](https://img.shields.io/badge/version-v0.1.0-blue)](https://github.com/SouravMal/DiffCrysGen)
+
+DiffCrysGen is a score-based diffusion model. It treats the entire materials representation with a single, unified diffusion process, allowing a single denosing neural network to predict a holistic score for the entire noisy crystal data. This unified treatment significantly simplifies the architecture and improves the computational efficiency.
+
+
+<p align="center">
+  <img src="images/logo-DiffCrysGen.png" alt="DiffCrysGen Logo" width="350"/>
+</p>
+
+
+## Generative diffusion framework in DiffCrysGen
+<img src="images/diffusion-schematic.png" alt="DiffCrysGen Schematic" width="550">
+
+---
+
+## Architecture of the denoising neural network
+<img src="images/architecture.png" alt="DiffCrysGen Architecture" width="750">
+
+---
+
+## Installation
+
+### Prerequisites
+
+The package requires specific environments and dependencies.
+Using a virtual environment is highly recommended.
+**Environment Setup using Conda**
+
+```
+conda create -n diffcrysgen python=3.11
+conda activate diffcrysgen
+```
+
+### Install from PyPI
+```
+pip install diffcrysgen
+```
+
+### Install from Source Code
+```
+git clone https://github.com/SouravMal/DiffCrysGen.git
+cd DiffCrysGen
+pip install -e .
+```
+
+## Quick Start
+For a simple walkthrough of generating materials and analyzing them, see the [DiffCrysGen Demo Notebook](./notebooks/DiffCrysGen-demo.ipynb).
+
+## License
+
+This project is licensed under the **MIT License**.
+
+See the [LICENSE](LICENSE) file for details.
+
+Developed by: [Sourav Mal](https://github.com/SouravMal) at Harish-Chandra Research Institute (HRI) (https://www.hri.res.in/), Prayagraj, India.
+
+
+## Citation
+
+Please consider citing our work if you find it helpful:
+
+```bibtex
+@misc{mal2025generativediffusionmodeldiffcrysgen,
+      title={Generative Diffusion Model DiffCrysGen Discovers Rare Earth-Free Magnetic Materials}, 
+      author={Sourav Mal and Nehad Ahmed and Subhankar Mishra and Prasenjit Sen},
+      year={2025},
+      eprint={2510.12329},
+      archivePrefix={arXiv},
+      primaryClass={cond-mat.mtrl-sci},
+      url={https://arxiv.org/abs/2510.12329}, 
+}
+```
+
+
+## Contact
+
+If you have any questions, feel free to reach us at:
+**Sourav Mal** <souravmal492@gmail.com> 

diffcrysgen-0.1.0.dist-info/RECORD

@@ -0,0 +1,12 @@
+diffcrysgen/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+diffcrysgen/analyze_generated_structures.py,sha256=fbybCbvH0o_xv_1v_tSV1OfbDEau7FHPs8oHyCAlf0Q,3475
+diffcrysgen/invert_pcr.py,sha256=Q6j7AVb6iZWR7t3VjB1yFLnuPB9hCc4_0Vp1K72gcak,5944
+diffcrysgen/model.py,sha256=sb42SQJPm9nQNSGMNN88RXwWE0tIsompypFZxmJob3k,11166
+diffcrysgen/sampler.py,sha256=k5r2D0MmWcQRebcg0KBugl1AEg6FPrnvfobZ_uyznP0,4092
+diffcrysgen/trainer.py,sha256=SI988DGpG0rgfixV9Sz9Ue1Lk5qG6uZvC1ka-BP9KOA,4189
+diffcrysgen/utils.py,sha256=FI6bkDGMt7E5_WbL55la--MGkOryz11KBgmVfivBvao,1094
+diffcrysgen-0.1.0.dist-info/licenses/LICENSE,sha256=7WpR2tbfkkPUcuvEnP_B3ijJczHpi-_U_i2ngeJGqSY,1067
+diffcrysgen-0.1.0.dist-info/METADATA,sha256=KhTsxU3Q6uqtvBWe9UrwHrrkUcFyMDvOJrcCt3x-_t4,3209
+diffcrysgen-0.1.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+diffcrysgen-0.1.0.dist-info/top_level.txt,sha256=aN4XuKXo6kncvD9N8mx9SOVKSFfNqBglK7G9jm5u3Fs,12
+diffcrysgen-0.1.0.dist-info/RECORD,,

diffcrysgen-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.9.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

diffcrysgen-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Sourav Mal
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

diffcrysgen-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+diffcrysgen