CrysBFN

https://github.com/wu-han-lin/crysbfn
Admin
CrysBFN is a periodic Bayesian flow model for material structure generation and prediction using...
Tokens:11,027
Snippets:95
Trust Score:5.4
Update:5 months ago
Show doc for...
Tokens:
Raw
### Set up Environment Variables

Source: https://context7.com/wu-han-lin/crysbfn/llms.txt

Create a .env file with required project paths before installing dependencies. This sets up the necessary environment variables for the project.

```bash
cat > .env << 'EOF' 
PROJECT_ROOT=/path/to/CrysBFN
HYDRA_JOBS=/path/to/CrysBFN/hydra
WABDB_DIR=/path/to/CrysBFN/wandb
EOF

conda env create -f environment.yml
conda activate crysbfn
```

--------------------------------

### Load and Setup CrysBFN DataModule

Source: https://context7.com/wu-han-lin/crysbfn/llms.txt

Loads and preprocesses crystal structure datasets using `CrystDataModule` and Hydra for configuration. This involves instantiating the datamodule from a configuration file and setting up the data loaders for training, validation, and testing. The example shows how to iterate through a batch to inspect its contents.

```python
from crysbfn.pl_data.dataset import CrystDataset
from crysbfn.pl_data.datamodule import CrystDataModule
import hydra

# Using Hydra config
datamodule = hydra.utils.instantiate(
    cfg.data.datamodule,
    _recursive_=False
)
datamodule.setup('fit')

# Access data loaders
train_loader = datamodule.train_dataloader()
val_loader = datamodule.val_dataloader()
test_loader = datamodule.test_dataloader()

# Iterate over batches
for batch in train_loader:
    print(f"Batch size: {batch.num_atoms.shape[0]}")
    print(f"Total atoms: {batch.frac_coords.shape[0]}")
    print(f"Atom types: {batch.atom_types.unique()}")
    print(f"Lattice lengths: {batch.lengths}")
    print(f"Lattice angles: {batch.angles}")
    break
```

--------------------------------

### Launch De Novo Generation Training

Source: https://github.com/wu-han-lin/crysbfn/blob/main/README.md

Execute this script to start a de novo generation task training experiment. The first run on a dataset may take longer due to data caching.

```bash
bash ./scripts/gen_scripts/mp20_exps.sh
```

--------------------------------

### Run Evaluation with Checkpoint

Source: https://github.com/wu-han-lin/crysbfn/blob/main/README.md

Navigate to the project root and run the evaluation script after configuring the MODEL_PATH. This verifies your installation using a provided checkpoint.

```bash
cd ..
bash scripts/csp_scripts/mp20_eval.sh
```

--------------------------------

### Convert Generated Data to Pymatgen Structures

Source: https://context7.com/wu-han-lin/crysbfn/llms.txt

Converts the raw output from the CrysBFN model into pymatgen Structure objects. This involves iterating through the generated data and creating `pymatgen.core.structure.Structure` and `pymatgen.core.lattice.Lattice` objects. Ensure pymatgen is installed.

```python
from pymatgen.core.structure import Structure
from pymatgen.core.lattice import Lattice

structures = []
atom_offset = 0
for i in range(len(num_atoms)):
    n = num_atoms[i].item()
    lattice = Lattice.from_parameters(
        *lengths[i].tolist(),
        *angles[i].tolist()
    )
    structure = Structure(
        lattice=lattice,
        species=atom_types[atom_offset:atom_offset+n].tolist(),
        coords=frac_coords[atom_offset:atom_offset+n].tolist(),
        coords_are_cartesian=False
    )
    structures.append(structure)
    atom_offset += n
```

--------------------------------

### Compute Crystal Validity, Diversity, and Coverage Metrics

Source: https://context7.com/wu-han-lin/crysbfn/llms.txt

Utilizes CrysBFN evaluation classes to compute validity, diversity, and coverage metrics for generated crystals. This involves creating `Crystal` objects from generated data and then using `GenEval` or `RecEval` for specific task evaluations. Ensure `p_tqdm` is installed for parallel processing.

```python
from crysbfn.evaluate.compute_metrics import Crystal, GenEval, RecEval
from p_tqdm import p_map

# Create Crystal objects from generated data
crys_array_list = [
    {
        'frac_coords': frac_coords[i],
        'atom_types': atom_types[i],
        'lengths': lengths[i],
        'angles': angles[i]
    }
    for i in range(num_samples)
]

# Convert to Crystal objects (parallelized)
pred_crystals = p_map(lambda x: Crystal(x), crys_array_list)

# Check validity of each crystal
for crys in pred_crystals:
    print(f"Composition valid: {crys.comp_valid}")
    print(f"Structure valid: {crys.struct_valid}")
    print(f"Overall valid: {crys.valid}")
    if crys.constructed:
        print(f"Density: {crys.structure.density}")

# For generation task evaluation
gen_evaluator = GenEval(
    pred_crys=pred_crystals,
    gt_crys=ground_truth_crystals,
    n_samples=1000,
    eval_model_name='mp20'  # or 'carbon', 'perovskite'
)
gen_metrics = gen_evaluator.get_metrics()
print(f"Validity: {gen_metrics['valid']:.4f}")
print(f"Coverage Recall: {gen_metrics['cov_r']:.4f}")
print(f"Coverage Precision: {gen_metrics['cov_p']:.4f}")

# For reconstruction/CSP task evaluation
rec_evaluator = RecEval(
    pred_crys=pred_crystals,
    gt_crys=ground_truth_crystals,
    stol=0.5,        # Site tolerance
    angle_tol=10,    # Angle tolerance in degrees
    ltol=0.3         # Length tolerance
)
recon_metrics = rec_evaluator.get_metrics()
print(f"Match rate: {recon_metrics['match_rate']:.4f}")
print(f"RMS distance: {recon_metrics['rms_dist']:.4f}")
```

--------------------------------

### Prepare Initial Sample Batch

Source: https://github.com/wu-han-lin/crysbfn/blob/main/toy_example/Gaussian_BFN.ipynb

Selects and clones a random subset of samples from `samples_0` to be used as the initial data points for a batch.

```python
x_0 = samples_0.detach().clone()[torch.randperm(len(samples_0))]
```

--------------------------------

### Download and Run CSP Evaluation

Source: https://context7.com/wu-han-lin/crysbfn/llms.txt

Steps to download a pretrained checkpoint, update the model path, and run the CSP evaluation script. Ensure the MODEL_PATH variable points to the unzipped checkpoint directory.

```bash
# Download checkpoint
cd hydra
# Download mp20_csp_s10.zip from:
# https://drive.google.com/drive/folders/1W5kGiZYFRJZiyKyTwCdcPk9lbjTsTCO-
unzip mp20_csp_s10.zip

# Update evaluation script with model path
MODEL_PATH=/path/to/hydra/mp20_csp_s10

# Run CSP evaluation
python scripts/evaluate_vmbfn.py \
    --model_path $MODEL_PATH \
    --tasks csp \
    --num_batches_to_csp 10 \
    --batch_size 1000 \
    --n_step_each 1000 \
    --samp_acc_factor 1 \
    --label end_back
```

--------------------------------

### Create Toy Data Distribution with PyTorch

Source: https://github.com/wu-han-lin/crysbfn/blob/main/toy_example/Gaussian_BFN.ipynb

Sets up random seeds and defines mixture models for generating toy data samples using PyTorch distributions. Includes plotting of generated samples.

```python
import torch
import numpy as np
import torch.nn as nn
from torch.distributions import Normal, Categorical
from torch.distributions.multivariate_normal import MultivariateNormal
from torch.distributions.mixture_same_family import MixtureSameFamily
import matplotlib.pyplot as plt
import torch.nn.functional as F
import random
import os

def setup_seed(seed):
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    random.seed(seed)
    torch.backends.cudnn.deterministic = True

setup_seed(0)

D = 10.
M = D+5
VAR = 0.1
DOT_SIZE = 4
COMP = 3

initial_mix = Categorical(torch.tensor([0.1,0.7]))

initial_comp = MultivariateNormal(torch.tensor([[-1.0, 1.0], [1.0, -1.0]]).float(), VAR * torch.stack([torch.eye(2) for i in range(2)]))

initial_model = MixtureSameFamily(initial_mix, initial_comp)
samples_1 = initial_model.sample([1000])

target_mix = Categorical(torch.tensor([0.7,0.1]))

target_comp = MultivariateNormal(torch.tensor([[0.5, 0.5], [0.5, 0.5]]).float(), VAR * torch.stack([torch.eye(2) for i in range(2)]))
target_model = MixtureSameFamily(target_mix, target_comp)
samples_0 = target_model.sample([1000])

plt.figure(figsize=(4,4))
plt.title(r'Samples from data distribution')
plt.scatter(samples_1[:, 0].cpu().numpy(), samples_1[:, 1].cpu().numpy(), alpha=0.1)
plt.xlim(-5.,5.)
plt.ylim(-5.,5.)
plt.tight_layout()
```

--------------------------------

### CrysBFN Model Initialization and Loss Calculation

Source: https://context7.com/wu-han-lin/crysbfn/llms.txt

Initialize the CrysBFN model with specified hyperparameters for different modalities and compute the loss for a single training step. Ensure batch data is correctly formatted.

```python
from crysbfn.pl_modules.crysbfn import CrysBFN
import torch

# Initialize the model
model = CrysBFN(
    hparams=cfg,                    # Hydra config object
    device="cuda",
    beta1_type=0.5,                 # Accuracy parameter for discrete (atom type) flow
    beta1_coord=1e3,                # Accuracy parameter for circular (coordinate) flow
    sigma1_lattice=0.0316,          # Noise parameter for continuous (lattice) flow
    K=88,                           # Number of atom types
    t_min=0.0001,                   # Minimum time step
    dtime_loss=True,                # Use discrete-time loss
    dtime_loss_steps=1000,          # Number of diffusion steps
    pred_mean=True,                 # Predict mean directly
    end_back=True,                  # Use end-back sampling strategy
    sim_cir_flow=True,              # Simulate circular flow
    sch_type='linear'               # Schedule type for accuracy
)

# Compute loss for one training step
type_loss, lattice_loss, coord_loss = model.loss_one_step(
    t=None,                         # Time will be sampled uniformly
    atom_type=batch.atom_types,     # One-hot encoded atom types [N, K]
    frac_coords=batch.frac_coords,  # Fractional coordinates [N, 3]
    lengths=batch.lengths,          # Lattice lengths [B, 3]
    angles=batch.angles,            # Lattice angles [B, 3]
    num_atoms=batch.num_atoms,      # Number of atoms per crystal [B]
    segment_ids=batch.batch,        # Batch indices for atoms
    edge_index=batch.edge_index     # Graph connectivity
)

total_loss = type_loss + lattice_loss + coord_loss
```

--------------------------------

### Generate and Evaluate Samples

Source: https://github.com/wu-han-lin/crysbfn/blob/main/README.md

After training, modify the MODEL_PATH variable in the script to point to your training experiment's hydra directory. Then, run this script to generate and evaluate samples.

```bash
bash scripts/csp_scripts/eval_mp20.sh
```

--------------------------------

### Launch Crystal Structure Prediction Training

Source: https://github.com/wu-han-lin/crysbfn/blob/main/README.md

Use this script to initiate a crystal structure prediction task training experiment.

```bash
bash ./scripts/csp_scripts/mp20_exps.sh
```

--------------------------------

### Access Individual Samples from Dataset

Source: https://context7.com/wu-han-lin/crysbfn/llms.txt

Demonstrates how to access and print properties of an individual sample from a dataset.

```python
sample = dataset[0]
print(f"Fractional coords shape: {sample.frac_coords.shape}")
print(f"Atom types: {sample.atom_types}")
print(f"Number of atoms: {sample.num_atoms}")
```

--------------------------------

### Load Trained CrysBFN Model and Generate Samples

Source: https://context7.com/wu-han-lin/crysbfn/llms.txt

Loads a trained CrysBFN model from a checkpoint and generates new crystal structures. Ensure the model checkpoint and configuration are correctly specified. The generation process requires PyTorch and can be configured with parameters like `num_samples` and `sample_steps`.

```python
model = CrysBFN_PL_Model.load_from_checkpoint(
    "path/to/checkpoint.ckpt",
    hparams=cfg
).to("cuda")
model.eval()

# Generate samples
with torch.no_grad():
    output_dict = model.sample(
        num_samples=256,           # Number of crystals to generate
        sample_steps=1000,         # Number of denoising steps
        show_bar=True,             # Show progress bar
        samp_acc_factor=1.0        # Sampling accuracy factor
    )

# Extract generated structures
frac_coords = output_dict['frac_coords']    # [total_atoms, 3]
atom_types = output_dict['atom_types']       # [total_atoms]
lengths = output_dict['lengths']             # [num_samples, 3]
angles = output_dict['angles']               # [num_samples, 3]
num_atoms = output_dict['num_atoms']         # [num_samples]
```

--------------------------------

### Evaluate New Material Generation Models

Source: https://github.com/wu-han-lin/crysbfn/blob/main/data/README.md

Use this script to evaluate new material generation models. It computes metrics for reconstruction, generation, and property optimization tasks. Ensure that the specified root path contains the necessary PyTorch pickle files (`eval_recon.pt`, `eval_gen.pt`, `eval_opt.pt`).

```python
python scripts/compute_metrics.py --root_path ROOT_PATH --tasks recon gen opt
```

--------------------------------

### Initialize and Use CSPNet Neural Network

Source: https://context7.com/wu-han-lin/crysbfn/llms.txt

Initializes the CSPNet model with specified parameters and performs a forward pass with sample data. Ensure CUDA is available for GPU acceleration.

```python
from crysbfn.pl_modules.egnn.cspnet import CSPNet
import torch

# Initialize CSPNet
cspnet = CSPNet(
    hidden_dim=256,          # Hidden dimension
    time_dim=128,            # Time embedding dimension
    num_layers=4,            # Number of message passing layers
    max_atoms=100,           # Maximum atoms (vocabulary size)
    act_fn='silu',           # Activation function
    dis_emb='sin',           # Distance embedding type
    num_freqs=10,            # Sinusoidal frequencies
    edge_style='fc',         # Fully connected edges
    cutoff=6.0,              # Distance cutoff for knn
    max_neighbors=20,        # Max neighbors for knn
    period='2*np.pi',        # Periodic boundary
    ln=False,                # Layer normalization
    ip=True,                 # Inner product for lattice
    pred_type=True,          # Predict atom types
    cond_acc=True            # Condition on accuracy
).to("cuda")

# Forward pass
batch_size = 4
num_atoms = torch.tensor([10, 12, 8, 15]).cuda()
total_atoms = num_atoms.sum()

t = torch.randn(batch_size, 128).cuda()                    # Time embedding
atom_types = torch.randn(total_atoms, 256).cuda()          # Atom type embeddings
frac_coords = torch.rand(total_atoms, 3).cuda() * 2 * 3.14159 - 3.14159  # [-pi, pi)
lattices = torch.randn(batch_size, 3, 3).cuda()            # Lattice matrices
node2graph = torch.repeat_interleave(
    torch.arange(batch_size).cuda(),
    num_atoms
)
log_acc = torch.randn(total_atoms, 3).cuda()               # Log accuracy

lattice_out, coord_out, type_out = cspnet(
    t=t,
    atom_types=atom_types,
    frac_coords=frac_coords,
    lattices=lattices,
    num_atoms=num_atoms,
    node2graph=node2graph,
    log_acc=log_acc
)

print(f"Lattice output: {lattice_out.shape}")   # [batch_size, 3, 3]
print(f"Coord output: {coord_out.shape}")       # [total_atoms, 3]
print(f"Type output: {type_out.shape}")         # [total_atoms, max_atoms]
```

--------------------------------

### Download and Unzip Checkpoint

Source: https://github.com/wu-han-lin/crysbfn/blob/main/README.md

Download the provided checkpoint zip file into the hydra directory and unzip it. This is part of the process to use pre-trained checkpoints for verification.

```bash
cd hydra
unzip ./mp20_csp_s10.zip
```

--------------------------------

### Initialize and Train Bayesian Flows Network

Source: https://github.com/wu-han-lin/crysbfn/blob/main/toy_example/Gaussian_BFN.ipynb

Initializes a Bayesian Flows Network with a specified model and sigma, then trains it using a given optimizer and data. Requires PyTorch and CUDA.

```python
indices = torch.randint(low=0, high=samples_1.shape[0], size=(1000,))
x_1 = samples_1.detach().clone()[indices]
x_pairs = torch.stack([x_0, x_1], dim=1)

iterations = 10000
batchsize = 512
input_dim = 2

x_pairs = x_pairs.cuda()

BFN = Baysian_Flows(model=MLP(input_dim, hidden_num=200).cuda(),sigma1=0.02)
optimizer = torch.optim.Adam(BFN.model.parameters(), lr=1e-2)

BFN, loss_curve = train_BFN_continous(BFN, optimizer, x_pairs, batchsize, iterations)
```

--------------------------------

### Create Conda Environment

Source: https://github.com/wu-han-lin/crysbfn/blob/main/README.md

Use Mamba or conda to create the Python environment from the provided YAML file. This may take several minutes to resolve dependencies.

```bash
conda env create -f environment.yml
conda activate crysbfn
```

--------------------------------

### Sampling Crystal Structures

Source: https://context7.com/wu-han-lin/crysbfn/llms.txt

Import the CrysBFN_PL_Model for generating new crystal structures using a trained model. This is the entry point for inference.

```python
from crysbfn.pl_modules.crysbfn_plmodel import CrysBFN_PL_Model
import torch
```

--------------------------------

### Direct Instantiation of CrystDataset

Source: https://context7.com/wu-han-lin/crysbfn/llms.txt

Directly instantiates a `CrystDataset` object for loading crystal structure data from a specified CSV file. This method allows for fine-grained control over dataset parameters such as name, path, property to predict, and preprocessing options like Niggli reduction and graph method. Ensure the data path is correct.

```python
from crysbfn.pl_data.dataset import CrystDataset

# Direct dataset instantiation
dataset = CrystDataset(
    name='mp_20',
    path='data/mp_20/train.csv',
    prop='formation_energy_per_atom',
    niggli=True,
    primitive=False,
    graph_method='crystalnn',
    preprocess_workers=4,
    lattice_scale_method='scale_length'
)
```

--------------------------------

### Run CrysBFN Evaluation Script

Source: https://context7.com/wu-han-lin/crysbfn/llms.txt

Executes the evaluation pipeline for generated crystals using the `evaluate_vmbfn.py` script. This script requires specifying the model path, tasks, and various sampling parameters. Ensure the `scripts` directory is accessible and model paths are correct.

```bash
# Run evaluation script
MODEL_PATH=/path/to/trained/model
LABEL=end_back

# Generate samples and evaluate
python scripts/evaluate_vmbfn.py \
    --model_path $MODEL_PATH \
    --tasks gen \
    --num_batches_to_samples 20 \
    --batch_size 500 \
    --n_step_each 1000 \
    --label $LABEL

# Compute metrics
python scripts/compute_metrics.py \
    --root_path $MODEL_PATH \
    --tasks gen \
    --label $LABEL
```

--------------------------------

### Train Crystal Structure Prediction (CSP) Model

Source: https://context7.com/wu-han-lin/crysbfn/llms.txt

Launch training for the crystal structure prediction task, where the model predicts atomic positions given a composition. This can be done via a script or directly with custom parameters.

```bash
# Launch CSP training
bash ./scripts/csp_scripts/mp20_exps.sh

# Or run directly with custom parameters
python crysbfn/run.py \
    data=mp_20 \
    model=bfn_csp \
    expname=mp20_csp \
    model.BFN.dtime_loss_steps=1000 \
    model.BFN.end_back=true \
    train.pl_trainer.max_epochs=1000
```

--------------------------------

### Train De Novo Crystal Structure Generation Model

Source: https://context7.com/wu-han-lin/crysbfn/llms.txt

Configure and launch training for unconditional crystal structure generation using the MP-20 dataset. Adjust diffusion steps, beta parameters, and EMA settings as needed.

```bash
# Configure and launch training for de novo generation
STEPS=1000
BETA1_COORD=1e3
BETA1_TYPE=0.5
SIGMA1_LATTICE=0.03162277660168379
USE_EMA=true
EMA_DECAY=0.995
DATASET=mp_20

# Launch training
python crysbfn/run.py \
    data=${DATASET} \
    expname=mp20_generation \
    model.BFN.dtime_loss_steps=$STEPS \
    model.BFN.beta1_coord=$BETA1_COORD \
    model.BFN.beta1_type=$BETA1_TYPE \
    model.BFN.sigma1_lattice=$SIGMA1_LATTICE \
    train.ema.enable=$USE_EMA \
    train.ema.decay=$EMA_DECAY \
    logging.gen_check.num_samples=128
```

--------------------------------

### Simple MLP Network Implementation

Source: https://github.com/wu-han-lin/crysbfn/blob/main/toy_example/Gaussian_BFN.ipynb

Defines a basic Multi-Layer Perceptron (MLP) using PyTorch's nn.Module. This network takes input features and a time step, concatenates them, and passes them through several linear layers with tanh activation.

```python
class MLP(nn.Module):
    def __init__(self, input_dim=2, hidden_num=100):
        super().__init__()
        self.fc1 = nn.Linear(input_dim+1, hidden_num, bias=True)
        self.fc2 = nn.Linear(hidden_num, hidden_num, bias=True)
        self.fc3 = nn.Linear(hidden_num, input_dim, bias=True)
        self.act = lambda x: torch.tanh(x)
    
    def forward(self, x_input, t):
        inputs = torch.cat([x_input, t], dim=-1)
        x = self.fc1(inputs)

        x = self.act(x)
        x = self.fc2(x)
        x = self.act(x)
        x = self.fc3(x)

        return x
```

--------------------------------

### CrysBFN Model Configuration

Source: https://context7.com/wu-han-lin/crysbfn/llms.txt

Configuration for the CrysBFN model, specifying parameters like time and hidden dimensions, cost functions, and BFN-specific settings such as beta values, diffusion steps, and sampling strategies.

```yaml
_target_: crysbfn.pl_modules.crysbfn_plmodel.CrysBFN_PL_Model

time_dim: 128
hidden_dim: 256
cost_coord: 1.0
cost_type: 1.0
cost_lattice: 1.0
T_min: -np.pi
T_max: np.pi

BFN:
  _target_: crysbfn.pl_modules.crysbfn.CrysBFN
  device: cuda
  beta1_type: 0.5           # Controls discrete flow accuracy
  beta1_coord: 1e3          # Controls circular flow accuracy
  K: ${data.num_atom_types}
  dtime_loss: true
  dtime_loss_steps: 1000    # Number of diffusion steps
  sigma1_lattice: 0.03162   # sqrt(1e-3)
  pred_mean: true
  end_back: true            # End-back sampling strategy
  sim_cir_flow: true        # Simulated circular flow
  sch_type: linear          # Accuracy schedule type

defaults:
  - decoder: cspnet
```

--------------------------------

### Draw BFN Fit Visualization

Source: https://github.com/wu-han-lin/crysbfn/blob/main/toy_example/Gaussian_BFN.ipynb

Visualizes the fitted data distribution using the trained Bayesian Flows Network. Requires the BFN model and sample data.

```python
draw_plot(
    BFN,
    z1=samples_1.detach().clone(),
    N=100,
)
```

--------------------------------

### Compute Metrics for CSP Tasks

Source: https://context7.com/wu-han-lin/crysbfn/llms.txt

Command to compute standard metrics after running the evaluation. This script requires the root path to the model checkpoint and specifies the tasks and label for metric calculation.

```bash
# Compute metrics
python scripts/compute_metrics.py \
    --root_path $MODEL_PATH \
    --tasks csp \
    --label end_back
```

--------------------------------

### CrysBFN Data Module Configuration

Source: https://context7.com/wu-han-lin/crysbfn/llms.txt

Configuration for the CrysBFN data module, defining batch sizes for training, validation, and testing, number of workers, maximum atoms, and atom types, along with lattice mean and standard deviation.

```yaml
datamodule:
  _target_: crysbfn.pl_data.datamodule.CrystDataModule
  batch_size:
    train: 256
    val: 256
    test: 256
  num_workers: 4

max_atoms: 20
num_atom_types: 88
lattice_mean: [0., 0., 0., 0., 0., 0., 0., 0., 0.]
lattice_std: [5., 5., 5., 5., 5., 5., 5., 5., 5.]
```

--------------------------------

### Bayesian Flow for Continuous Variables

Source: https://context7.com/wu-han-lin/crysbfn/llms.txt

Applies Bayesian flow to continuous variables, such as lattice parameters. Requires a `bfnBase` model instance and specifies the final noise level.

```python
from crysbfn.pl_modules.bfn_base import bfnBase
import torch

class MyBFN(bfnBase):
    def __init__(self):
        super().__init__()
        self.device = "cuda"

model = MyBFN()

# Continuous variable Bayesian flow (for lattice parameters)
x_continuous = torch.randn(32, 9).cuda()      # Lattice parameters
t = torch.rand(32, 1).cuda()                   # Time in [0, 1]
sigma1 = torch.tensor(0.0316).cuda()           # Final noise level

mu, gamma = model.continuous_var_bayesian_flow(
    x=x_continuous,
    t=t,
    sigma1=sigma1
)
# mu: noisy observation, gamma: signal-to-noise ratio
```

--------------------------------

### Configure Model Path for Evaluation

Source: https://github.com/wu-han-lin/crysbfn/blob/main/README.md

Modify the MODEL_PATH variable in the evaluation script to the correct path of your downloaded and unzipped checkpoint. This ensures the evaluation script uses the correct model.

```bash
MODEL_PATH=/data/wuhl/CrysBFN/hydra/mp20_csp_s10 # modify according to your path
```

--------------------------------

### Discrete-Time Loss for Continuous Variables

Source: https://context7.com/wu-han-lin/crysbfn/llms.txt

Calculates the discrete-time loss for continuous variables. Requires the current time index (i), total steps (N), final noise level (sigma1), predicted values, and actual values.

```python
from crysbfn.pl_modules.bfn_base import bfnBase
import torch

class MyBFN(bfnBase):
    def __init__(self):
        super().__init__()
        self.device = "cuda"

model = MyBFN()

# Discrete-time loss for continuous variables
loss_continuous = model.dtime4continuous_loss(
    i=torch.randint(1, 1001, (32, 1)).cuda(),  # Time index
    N=1000,                                      # Total steps
    sigma1=sigma1,
    x_pred=torch.randn(32, 9).cuda(),
    x=x_continuous
)
```

--------------------------------

### Discrete-Time Loss for Circular Variables

Source: https://context7.com/wu-han-lin/crysbfn/llms.txt

Calculates the discrete-time loss for circular variables, such as coordinates. Requires time index, total steps, concentration parameter (alpha_i), predicted values, and actual values. Ensure `torch.special.i0e` and `torch.special.i1e` are imported.

```python
from crysbfn.pl_modules.bfn_base import bfnBase
import torch
from torch.special import i0e, i1e

class MyBFN(bfnBase):
    def __init__(self):
        super().__init__()
        self.device = "cuda"

model = MyBFN()

# Discrete-time loss for circular variables (coordinates)
from torch.special import i0e, i1e
alpha_i = torch.rand(100, 1).cuda() * 1000     # Concentration parameter
loss_circular = model.dtime4circular_loss(
    i=torch.randint(1, 1001, (100, 1)).cuda(),
    N=1000,
    alpha_i=alpha_i,
    x_pred=torch.rand(100, 3).cuda() * 2 * 3.14159 - 3.14159,
    x=torch.rand(100, 3).cuda() * 2 * 3.14159 - 3.14159
)
```

--------------------------------

### Train BFN Continuously

Source: https://github.com/wu-han-lin/crysbfn/blob/main/toy_example/Gaussian_BFN.ipynb

Trains the BFN model using a discrete loss function over a specified number of iterations. It includes optimizer zeroing, batch selection, loss calculation, backpropagation, and optimizer stepping. Prints loss every 1000 iterations and stores the loss curve.

```python
def train_BFN_continous(BFN, optimizer, pairs, batchsize, inner_iters):
    loss_curve = []
    for i in range(inner_iters + 1):
        optimizer.zero_grad()
        indices = torch.randperm(len(pairs))[:batchsize]
        batch = pairs[indices]
        z0 = batch[:, 0].detach().clone()
        z1 = batch[:, 1].detach().clone()
        # z_t, t, target = rectified_flow.get_train_tuple(z0=z0, z1=z1)
        loss = BFN.BFN_loss_discrete(z1)
        loss = loss.mean()
        loss.backward()

        if i % 1000 == 0:
            print(i, "loss:", loss.item())

        optimizer.step()
        loss_curve.append(np.log(loss.item()))  ## to store the loss curve

    return BFN, loss_curve
```

--------------------------------

### Plot Training Loss Curve

Source: https://github.com/wu-han-lin/crysbfn/blob/main/toy_example/Gaussian_BFN.ipynb

Plots the training loss curve over the specified number of iterations. Requires matplotlib and numpy.

```python
plt.plot(np.linspace(0, iterations, iterations+1), loss_curve[:(iterations+1)])
plt.title('Training Loss Curve')
```

--------------------------------

### Bayesian Flow for Discrete Variables

Source: https://context7.com/wu-han-lin/crysbfn/llms.txt

Applies Bayesian flow to discrete variables, such as atom types. Requires specifying the number of types (K) and the initial time step parameter (beta1).

```python
from crysbfn.pl_modules.bfn_base import bfnBase
import torch

class MyBFN(bfnBase):
    def __init__(self):
        super().__init__()
        self.device = "cuda"

model = MyBFN()

# Discrete variable Bayesian flow (for atom types)
K = 88  # Number of atom types
one_hot_x = torch.zeros(100, K).cuda()
one_hot_x[torch.arange(100), torch.randint(0, K, (100,))] = 1
beta1 = torch.tensor(0.5).cuda()

theta = model.discrete_var_bayesian_flow(
    t=t.repeat_interleave(torch.tensor([100//32]*32).cuda(), dim=0)[:100],
    beta1=beta1,
    x=one_hot_x,
    K=K
)
# theta: posterior probability distribution over atom types
```

--------------------------------

### Draw Distribution and Trajectory Plots

Source: https://github.com/wu-han-lin/crysbfn/blob/main/toy_example/Gaussian_BFN.ipynb

Generates plots for the original distribution and generated samples, as well as transport trajectories. It calculates and prints MMD and straightness metrics. Requires `matplotlib.pyplot` and `mmd_rbf` function.

```python
@torch.no_grad()
def draw_plot(BFN, z1, N=None, traj=None):
  z1 = z1.cuda()
  if traj is None:
    samples_nums = z1.shape[0]
    traj = BFN.Sampling_BFN(N=N,samples_nums=samples_nums)
  
  plt.figure(figsize=(4,4))
  plt.xlim(-5.,5.)
  plt.ylim(-5.,5.)
    
  plt.scatter(z1[:, 0].cpu().numpy(), z1[:, 1].cpu().numpy(), label=r'$\\pi_1$', alpha=0.15)
  plt.scatter(traj[-1][:, 0].cpu().numpy(), traj[-1][:, 1].cpu().numpy(), label='Generated', alpha=0.15)
  plt.legend()
  plt.title('Distribution')
  plt.tight_layout()
  traj_particles = torch.stack(traj)
  print(f'traj:{traj_particles.shape}')
  plt.figure(figsize=(4,4))
  plt.xlim(-1.,1.)
  plt.ylim(-3.,3.)
  plt.axis('equal')
  for i in range(100):
    plt.plot(traj_particles[:, i, 0].cpu().numpy(), traj_particles[:, i, 1].cpu().numpy())
  plt.title('Transport Trajectory')
  plt.tight_layout()
  
  mmd = mmd_rbf(z1.cpu().numpy(), traj[-1].cpu().numpy(), gamma=0.1)
  print('MMD:', mmd)

  straightness = 0.0
  straight_v = traj_particles[-1] - traj_particles[0]
  for i in range(N):
    v = (traj_particles[i+1] - traj_particles[i]) * N
    straightness += (v - straight_v).norm(dim=1).mean() 

  print('Straightness:', straightness / N)
```

--------------------------------

### BibTeX Citation for CrysBFN

Source: https://github.com/wu-han-lin/crysbfn/blob/main/README.md

This is the BibTeX entry for citing the CrysBFN paper. Use this in your academic work if you find the repository or paper useful.

```bibtex
@misc{wu2025periodicbayesianflowmaterial,
        title={A Periodic Bayesian Flow for Material Generation},
        author={Hanlin Wu and Yuxuan Song and Jingjing Gong and Ziyao Cao and Yawen Ouyang and Jianbing Zhang and Hao Zhou and Wei-Ying Ma and Jingjing Liu},
        year={2025},
        eprint={2502.02016},
        archivePrefix={arXiv},
        primaryClass={cs.LG},
        url={https://arxiv.org/abs/2502.02016},
  }
```

--------------------------------

### Define Bayesian Flow Networks Class

Source: https://github.com/wu-han-lin/crysbfn/blob/main/toy_example/Gaussian_BFN.ipynb

Initializes the BFN model with a base model and noise parameters. The `corrupt_t_pred` method predicts corrupted data, `get_train_tuple` prepares training data, and `forward` handles model predictions. `Sampling_BFN` generates samples, `BFN_loss` calculates continuous loss, and `BFN_loss_discrete` calculates discrete loss.

```python
class Baysian_Flows(nn.Module):
    def __init__(self, model=None, t_min=0.0001, sigma1=0.02):
        super().__init__()
        self.model = model
        self.t_min = t_min
        self.sigma1 = torch.tensor(sigma1).cuda()
        self.sigma1.requires_grad = False

    def corrupt_t_pred(self, mu, t, gamma):
        eps_pred = self.model(mu, t)
        x_pred = mu / gamma - torch.sqrt((1 - gamma) / gamma) * eps_pred
        return x_pred

    def get_train_tuple(self, z0=None, z1=None):
        t = torch.rand((z1.shape[0], 1)).to(z0.device)
        z_t = t * z1 + (1.0 - t) * z0
        target = z1 - z0
        return z_t, t, target

    def forward(self, t, x):
        t = torch.ones((x.size(0), 1)).cuda() * t
        return self.model(x, t)

    @torch.no_grad()
    def Sampling_BFN(self, N=10, samples_nums=2000):
        mu = torch.zeros((samples_nums, 2)).cuda()
        ro = 1
        if N is None:
            N = self.N
        traj = []  # to store the trajectory
        for i in range(1, N + 1):
            t = torch.ones((samples_nums, 1)).cuda() * (i - 1) / N
            t = torch.clamp(t, min=self.t_min)
            gamma = 1 - torch.pow(self.sigma1, 2 * t)
            x_pred = self.corrupt_t_pred(mu, t, gamma)
            alpha = torch.pow(self.sigma1, -2 * i / N) * (
                1 - torch.pow(self.sigma1, 2 / N)
            )

            y = x_pred + torch.randn_like(x_pred) * torch.sqrt(1 / alpha)
            mu = (ro * mu + alpha * y) / (ro + alpha)
            ro = ro + alpha
            traj.append(x_pred.detach().clone())

        x_final_pred = self.corrupt_t_pred(
            mu, torch.ones((samples_nums, 1)).cuda(), 1 - self.sigma1**2
        )
        traj.append(x_final_pred.detach().clone())
        return traj

    def BFN_loss(self, z1):
        t = torch.rand((z1.shape[0], 1)).to(z1.device)
        sigma = self.sigma1 * torch.ones_like(t)
        gamma = 1 - torch.pow(sigma, 2 * t)
        mu = gamma * z1 + torch.randn_like(z1) * torch.sqrt(gamma * (1 - gamma))
        x_pred = self.corrupt_t_pred(mu, t, gamma)
        loss = (x_pred - z1).view(z1.shape[0], -1).abs().pow(2).sum(dim=1)

        return -torch.log(self.sigma1) * loss * torch.pow(self.sigma1, -2 * t)

    def BFN_loss_discrete(self, z1, N=10):
        i = torch.randint(1, N + 1, (z1.shape[0], 1)).to(z1.device).float()
        t = (i - 1) / N * torch.ones((z1.shape[0], 1)).cuda()
        t = torch.clamp(t, min=self.t_min)
        gamma = 1 - torch.pow(self.sigma1, 2 * t)
        mu = gamma * z1 + torch.randn_like(z1) * torch.sqrt(gamma * (1 - gamma))
        x_pred = self.corrupt_t_pred(mu, t, gamma)
        weight = (
            N
            * (1 - torch.pow(self.sigma1, 2 / N))
            / (2 * torch.pow(self.sigma1, 2 * i / N))
        )

        return weight * (x_pred - z1).view(z1.shape[0], -1).abs().pow(2).sum(dim=1)
```

--------------------------------

### Calculate MMD using RBF Kernel

Source: https://github.com/wu-han-lin/crysbfn/blob/main/toy_example/Gaussian_BFN.ipynb

Implements the Maximum Mean Discrepancy (MMD) calculation using a Radial Basis Function (RBF) kernel. This function is useful for quantitatively evaluating the similarity between two data distributions.

```python
from sklearn import metrics

def mmd_rbf(X, Y, gamma=1.0):
    """MMD using rbf (gaussian) kernel (i.e., k(x,y) = exp(-gamma * ||x-y||^2 / 2))
    Arguments:
        X {[n_sample1, dim]} -- [X matrix]
        Y {[n_sample2, dim]} -- [Y matrix]
    Keyword Arguments:
        gamma {float} -- [kernel parameter] (default: {1.0})
    Returns:
        [scalar] -- [MMD value]
    """
    XX = metrics.pairwise.rbf_kernel(X, X, gamma)
    YY = metrics.pairwise.rbf_kernel(Y, Y, gamma)
    XY = metrics.pairwise.rbf_kernel(X, Y, gamma)
    return XX.mean() + YY.mean() - 2 * XY.mean()
```
Tokens:
Raw