Initial commit

compare_models.py

@@ -0,0 +1,130 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+from torch_geometric.nn.models import SchNet, DimeNetPlusPlus
+
+import ase
+import ase.io
+import re
+from sklearn.metrics import r2_score
+import numpy as np
+
+import sys
+
+from tqdm import tqdm
+
+import argparse
+
+from utils_model import ModellDimeNet
+
+def get_model_and_optimizer(model_type):
+    if model_type == 'SchNet':
+        model = SchNet()
+    elif model_type == 'DimeNet':
+        model =  ModelDimeNet()
+
+    optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
+
+    return model, optimizer
+
+
+def train_epoch(model, optimizer, geoms, energies_n, mean_grad=32):
+    j = 0
+    model.train();
+
+    for geom, energy in zip(geoms, energies_n.clone().detach()):
+        if j == 0:
+            optimizer.zero_grad()
+        coords = torch.tensor(geom.get_positions(), dtype=torch.float32)
+        atoms = torch.tensor(geom.get_atomic_numbers())
+        batch = torch.zeros_like(atoms)
+        
+        en = energy.clone().detach()
+        pred = model(atoms, coords, batch)
+        loss = F.huber_loss(pred.squeeze(), en)
+        (loss / mean_grad).backward(); j += 1
+        if j == mean_grad - 1:
+            optimizer.step()
+            j = 0
+
+def test_epoch(model, optimizer, geoms, energies_n):
+    all_loss = 0
+    all_mols = 0
+    all_preds = []
+    all_trues = []
+    model.eval();
+    
+    for geom, energy in zip(geoms, energies_n.clone().detach()):
+        coords = torch.tensor(geom.get_positions(), dtype=torch.float32)
+        atoms = torch.tensor(geom.get_atomic_numbers())
+        batch = torch.zeros_like(atoms)
+        
+        en = energy.clone().detach()
+        with torch.no_grad():
+            pred = model(atoms, coords, batch)
+            all_preds.append(pred.item())
+            all_trues.append(en.item())
+        all_loss += F.l1_loss(pred.squeeze(), en).item()
+        all_mols += 1
+    return {
+        'r2_score': r2_score(np.array(all_trues), np.array(all_preds)),
+        'mae': all_loss / all_mols,
+    }
+
+def train(model, optimizer, geoms, energies_n, n_epochs=100):
+    best_r2score = -1e100
+    best_mae = 1e100
+    
+    for i in tqdm(range(n_epochs)):
+        train_epoch(model, optimizer, geoms, energies_n)
+
+        metrics = test_epoch(model, optimizer, geoms, energies_n)
+
+        if best_r2score < metrics['r2_score']:
+            best_r2score = metrics['r2_score']
+        if best_mae > metrics['mae']:
+            best_mae = metrics['mae']
+
+    return best_r2score, best_mae
+
+def main(trajectory_file, model_type):
+    geoms = ase.io.read(trajectory_file, format='xyz', index=':')
+    
+    with open(trajectory_file) as f:
+        cont = f.read()
+    
+    energies = []
+    lines = cont.split('\n'); i = 0
+    while i < len(lines):
+        try:
+            n = int(lines[i].strip())
+        except ValueError:
+            break
+        comment = lines[i+1]
+        energy = float(re.findall('energy\\:\\s+(-?\\d*\.\\d*)', comment)[0])
+        energies.append(energy)
+        i += n + 2
+    
+    energies = torch.tensor(energies) 
+    energies_n = (energies  - energies.min()) * 627.5     
+
+    model, optimizer = get_model_and_optimizer(model_type)
+
+    best_r2score, best_mae = train(model, optimizer, geoms, energies_n)
+
+    print(f'R2_score: {best_r2score:.4f}')
+    print(f'MAE: {best_mae:.3f} kcal/mol')
+
+avaliable_models = ['SchNet', 'DimeNet']
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Обработчик файлов с различными моделями")
+    parser.add_argument("filename", help="Путь к обрабатываемому файлу")
+    parser.add_argument("model", 
+                        choices=avaliable_models,
+                        help=f"Выбор модели из доступных: {', '.join(avaliable_models)}")
+    
+    args = parser.parse_args()
+    
+    main(args.filename, args.model)
+

example_eval.py

@@ -0,0 +1,102 @@
+import torch
+import argparse
+
+import numpy as np
+
+from utils_data import read_data, MolDataset, collate_mol, get_train_test_data
+from utils_model import ModelDimeNet
+
+def main(denormalize, checkpoint_path, data_filename):
+    model = ModelDimeNet()
+    model.load_state_dict(torch.load(checkpoint_path, weights_only=True))
+    
+    all_numbers, all_coords, energies, groups = read_data(data_filename)
+    ds_all = MolDataset(all_numbers, all_coords, energies, normalize=denormalize)
+
+    ds_train, ds_test = get_train_test_data(ds_all, groups, 'finetune')
+
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+    model.to(device);
+    
+    ens = []
+
+    for numbers, coords, energy in ds_test:
+        if len(ens) in ds_train.indices:
+            ens.append(energy)
+        else:
+            coords = torch.tensor(coords, dtype=torch.float32).to(device)
+            atoms = torch.tensor(numbers).to(device)
+            batch = torch.zeros_like(atoms).to(device)
+
+            with torch.no_grad():
+                ens.append(model(atoms, coords, batch).item())
+
+
+    if denormalize:
+        ens = ens.sign() * ens.abs() ** 10
+
+    ensa = np.array(ens[1:])
+
+    n_atoms = 3
+    n_modes = n_atoms * 3 - 6
+    modes_i = []
+
+    for i in range(n_modes):
+        modes_i.append(ensa[0:][16*i:16*i+16])
+    
+    all_a = []
+
+    m = 0
+    for i in range(n_modes):
+        for j in range(i + 1, n_modes):
+            for k in range(16):
+                all_a.append(ensa[16*n_modes+m:16*n_modes+m+16] - modes_i[j] - modes_i[i][k])
+                m += 16
+
+    arr = np.concatenate([modes_i, all_a]).reshape((-1))
+
+    with open('template.rst', 'r') as f:
+        content = f.read()
+
+    fo = open('filled.rst', 'w')
+
+    i = 0
+    for line in content.split('\n'):
+        if i < len(arr):
+            new_line = line.replace('{}', f'{arr[i]:.10f}')
+            if line != new_line:
+                i += 1
+        else:
+            new_line = line
+        fo.write(new_line + '\n')
+
+    del fo
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    
+    # Обязательные аргументы
+
+    parser.add_argument('checkpoint_path', 
+                        type=str,
+                       )
+    parser.add_argument('data_filename', 
+                        type=str,
+                       )
+    
+    # Флаги (булевые параметры)
+    parser.add_argument('--denormalize', 
+                        action='store_true',
+                       )
+    
+    
+    args = parser.parse_args()
+    
+    
+    # Вызов основной функции
+    main(
+        denormalize=args.denormalize,
+        checkpoint_path=args.checkpoint_path,
+        data_filename=args.data_filename
+    )

example_plot.py

@@ -0,0 +1,28 @@
+import pickle
+import matplotlib.pyplot as plt
+import argparse
+
+def main(picklefile):
+    with open(picklefile, 'rb') as f:
+        all_metrics = pickle.load(f)
+    
+    plt.plot([i for i in range(len(all_metrics))], [a[1] for a in all_metrics])
+    plt.grid(True)
+    plt.ylim(0.0, 0.05)
+    plt.xlabel('Эпоха')
+    plt.ylabel('MAE$')
+    plt.show()
+    
+    plt.plot([i for i in range(len(all_metrics))], [a[0] for a in all_metrics])
+    plt.grid(True)
+    plt.ylim(0.0, 1.0)
+    plt.xlabel('Эпоха')
+    plt.ylabel('$R^2$')
+    plt.show()
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='Process a pickle file')
+    parser.add_argument('picklefile', type=str, help='Path to pickle file')
+    args = parser.parse_args()
+    
+    main(args.picklefile)

train.py

@@ -0,0 +1,198 @@
+from sklearn.metrics import r2_score
+import numpy as np
+
+import torch
+import torch.nn.functional as F
+from tqdm import tqdm
+import argparse
+
+import pickle
+
+from utils_data import read_data, MolDataset, collate_mol, get_train_test_data
+from utils_model import ModelDimeNet
+
+
+def get_model():
+    model =  ModelDimeNet()
+    return model
+
+    
+def get_optimizer(model, e_start):
+    optimizer = torch.optim.RMSprop(model.parameters(), lr=10 ** -e_start)
+
+    return optimizer
+
+def get_loss(mode):
+    if mode[0] == 'mae':
+        return lambda pred, en: (pred - en).abs().mean()
+
+    if mode[0] == 'adaptive':
+        return lambda pred, en: ((pred - en).abs() / (en.abs() + 1e-5) ** mode[1]).mean()
+
+def train_epoch(model, optimizer, dl_train, loss_fn, device):
+    model.train()
+    
+    for atoms, coords, energy, batch in dl_train:
+        optimizer.zero_grad()
+
+        atoms = atoms.to(device)
+        coords = coords.to(device)
+        energy = energy.to(device)
+        batch = batch.to(device)
+        
+        en = energy.squeeze()
+        pred = model(atoms, coords, batch).squeeze() 
+        loss = loss_fn(pred, en)
+        loss.backward()
+        optimizer.step()   
+
+
+def test_epoch(model, optimizer, dl_test, device):
+    all_loss = 0
+    all_mols = 0
+    all_preds = []
+    all_trues = []
+    model.eval()
+    
+    for atoms, coords, energy, batch in dl_test:  
+        atoms = atoms.to(device)
+        coords = coords.to(device)
+        energy = energy.to(device)
+        batch = batch.to(device)
+        
+        en = energy.squeeze()
+        with torch.no_grad():
+            pred = model(atoms, coords, batch).squeeze()
+            all_preds.append(pred.cpu().numpy())
+            all_trues.append(en.cpu().numpy())
+        all_loss += F.l1_loss(pred.squeeze(), en).item() * len(pred)
+        all_mols += len(pred)
+
+    all_trues = np.concatenate(all_trues)
+    all_preds = np.concatenate(all_preds)
+
+    return {
+        'r2_score': r2_score(np.array(all_trues), np.array(all_preds)),
+        'mae': all_loss / all_mols,
+    }
+
+def refresh_lr(optimizer, i, n, e_start, downscale=2.0):
+    for g in optimizer.param_groups:
+        g['lr'] = 10 ** -(e_start + i / n * downscale)
+
+    return 10 ** -(e_start + i / n * downscale)
+
+
+def train(n_epoch, model, optimizer, loss_fn, e_start, dl_train, dl_test, device, checkpoint_prefix):
+    all_metrics = []
+    new_lr = e_start
+
+    for i in tqdm(range(n_epoch)):
+        train_epoch(model, optimizer, dl_train, loss_fn, device)
+
+        metrics = test_epoch(model, optimizer, dl_test, device)
+
+        cur_lr = new_lr
+        new_lr = refresh_lr(optimizer, i, n_epoch, e_start)
+        
+
+        all_metrics.append((
+            metrics['r2_score'],
+            metrics['mae'], 
+            cur_lr, 
+        ))
+
+        torch.save(model.state_dict(), checkpoint_prefix + '.ckpt')
+
+    return all_metrics
+
+def main(loss_mode, normalize, pretrain, checkpoint_prefix, data_filename):
+    all_numbers, all_coords, energies, groups = read_data(data_filename)
+    ds_all = MolDataset(all_numbers, all_coords, energies, normalize=normalize)
+
+    loss_fn = get_loss(loss_mode)
+    
+    model = get_model()
+
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    model = model.to(device)
+
+    # Pretraining
+    if pretrain:
+        e_start = 4
+    
+        ds_train, ds_test = get_train_test_data(ds_all, groups, 'pretrain')
+        dl_train = torch.utils.data.DataLoader(ds_train, batch_size=32, shuffle=True, collate_fn=collate_mol)
+        dl_test = torch.utils.data.DataLoader(ds_test, batch_size=32, shuffle=False, collate_fn=collate_mol)
+    
+        optimizer = get_optimizer(model, e_start=e_start)
+
+        all_metrics = train(100, model, optimizer, loss_fn, e_start, dl_train, dl_test, device, checkpoint_prefix + '_pretrain_model')
+        with open(checkpoint_prefix + '_pretrain_metrics.pkl', 'wb') as f:
+            pickle.dump(all_metrics, f)
+
+    # Fine-tuting
+    e_start = 5
+
+    ds_train, ds_test = get_train_test_data(ds_all, groups, 'finetune')
+    dl_train = torch.utils.data.DataLoader(ds_train, batch_size=32, shuffle=True, collate_fn=collate_mol)
+    dl_test = torch.utils.data.DataLoader(ds_test, batch_size=32, shuffle=False, collate_fn=collate_mol)
+
+    optimizer = get_optimizer(model, e_start=e_start)
+
+    all_metrics = train(100, model, optimizer, loss_fn, e_start, dl_train, dl_test, device, checkpoint_prefix + '_finetune_model')
+    with open(checkpoint_prefix + '_finetune_metrics.pkl', 'wb') as f:
+        pickle.dump(all_metrics, f)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='Параметры для обучения модели')
+    
+    # Обязательные аргументы
+    parser.add_argument('loss_mode', 
+                        choices=['mae', 'adaptive'],
+                        help="Режим потерь: 'mae' или 'adaptive'")
+    parser.add_argument('checkpoint_prefix', 
+                        type=str,
+                        help="Префикс для чекпоинтов")
+    parser.add_argument('data_filename', 
+                        type=str,
+                        help="Путь к файлу с датасетом")
+    
+    # Флаги (булевые параметры)
+    parser.add_argument('--normalize', 
+                        action='store_true',
+                        help="Применить нормализацию (только для loss_mode='mae')")
+    parser.add_argument('--pretrain', 
+                        action='store_true',
+                        help="Использовать предобучение")
+    
+    # Параметр только для adaptive режима
+    parser.add_argument('--loss_k', 
+                        type=float,
+                        default=None,
+                        help="Коэффициент k для adaptive loss (требуется при loss_mode='adaptive')")
+    
+    args = parser.parse_args()
+    
+    # Проверка совместимости параметров
+    if args.loss_mode == 'adaptive':
+        if args.normalize:
+            raise ValueError("Параметр --normalize несовместим с loss_mode='adaptive'")
+        if args.loss_k is None:
+            raise ValueError("Для adaptive loss требуется параметр --loss_k")
+        # Формируем кортеж для adaptive режима
+        loss_mode_arg = ('adaptive', args.loss_k)
+    else:  # loss_mode == 'mae'
+        if args.loss_k is not None:
+            raise ValueError("Параметр --loss_k можно использовать только с loss_mode='adaptive'")
+        loss_mode_arg = ('mae', )
+    
+    # Вызов основной функции
+    main(
+        loss_mode=loss_mode_arg,
+        normalize=args.normalize,
+        pretrain=args.pretrain,
+        checkpoint_prefix=args.checkpoint_prefix,
+        data_filename=args.data_filename,
+    )

utils_data.py

@@ -0,0 +1,131 @@
+import re
+import torch
+
+
+def read_data(filename):
+    all_coords = []
+    all_numbers = []
+    
+    with open(filename) as f:
+        cont = f.read()
+    
+    energies = []
+    groups = []
+    lines = cont.split('\n'); i = 0
+    mol_en = None
+    while i < len(lines):
+        try:
+            n = int(lines[i].strip())
+        except ValueError:
+            break
+        comment = lines[i+1]
+        energy = float(re.findall('Energy\\:\\s+(-?\\d*\.\\d*)', comment)[0])
+        g0 = re.findall('Grid: 0', comment)
+        if g0:
+            mol_en = energy
+            grp = ()
+        g1 = re.findall('Grid: (\\d+): (\\d+)', comment)
+        if g1:
+            grp = (g1[0][1], )
+        g2 = re.findall('Grid: (\\d+): (\\d+), (\\d+): (\\d+)', comment)
+        if g2:
+            grp = (g2[0][1], g2[0][3])
+        energies.append(energy - mol_en)
+        groups.append(grp)
+        j = 0
+        all_coords.append([])
+        all_numbers.append([])
+        while j < n:
+            at, x, y, z = list(filter(lambda x: x != '', lines[i+j+2].strip().split(' ')))
+            all_coords[-1].append((float(x), float(y), float(z)))
+            all_numbers[-1].append(int(at))
+            j += 1
+        i += n + 2
+    
+    energies = torch.tensor(energies)
+    
+    return all_numbers, all_coords, energies, groups
+
+class MolDataset(torch.utils.data.Dataset):
+    def __init__(self, all_numbers, all_coords, energies, normalize=False):
+        self.numbers = all_numbers
+        self.coords = all_coords
+        self.energies = energies
+        self.normalize = normalize
+
+    def __len__(self):
+        return len(self.energies)
+
+    def __getitem__(self, ind):
+        energy = self.energies[ind]
+
+        atoms = torch.tensor(self.numbers[ind])
+        coords = torch.tensor(self.coords[ind], dtype=torch.float32)
+
+        if self.normalize:
+            energy = energy.sign() * energy.abs() ** 0.1
+
+        return atoms, coords, energy
+
+def collate_mol(batch):
+    """
+    Collate function for molecular dataset.
+    
+    Args:
+        batch: List of tuples (atoms, coords, energy) from MolDataset
+    
+    Returns:
+        atoms_cat: Concatenated atomic numbers tensor of shape [total_atoms]
+        coords_cat: Concatenated coordinates tensor of shape [total_atoms, 3]
+        energies: Energy tensor of shape [batch_size]
+        batch_tensor: Batch indices tensor of shape [total_atoms]
+    """
+    atoms_list = []
+    coords_list = []
+    energies_list = []
+    batch_indices = []
+    
+    # Process each molecule in the batch
+    for i, (atoms, coords, energy) in enumerate(batch):
+        n_atoms = atoms.size(0)
+        
+        # Store components
+        atoms_list.append(atoms)
+        coords_list.append(coords)
+        energies_list.append(energy)
+        
+        # Create batch indices: [i, i, ..., i] for n_atoms times
+        batch_indices.append(torch.full((n_atoms,), i, dtype=torch.long))
+    
+    # Concatenate all components
+    atoms_cat = torch.cat(atoms_list, dim=0)      # shape: [total_atoms]
+    coords_cat = torch.cat(coords_list, dim=0)    # shape: [total_atoms, 3]
+    energies = torch.stack(energies_list)          # shape: [batch_size]
+    batch_tensor = torch.cat(batch_indices, dim=0) # shape: [total_atoms]
+    
+    return atoms_cat, coords_cat, energies, batch_tensor
+
+def get_train_test_data(ds_all, groups, mode, test_idcs=range(28986, 29803)):
+    grid1_selection = ['1', '3', '5', '7', '8', '9', '10', '12', '14', '16']
+    grid2_selection = ['1', '5', '8', '9', '12', '16']
+
+    assert(mode in ['pretrain', 'finetune'])
+
+    pretrain = mode == 'pretrain'
+
+    train_idces = []
+    
+    for i in range(len(groups)):
+        if len(groups[i]) == 0:
+            if (i in test_idcs) != pretrain: train_idces.append(i)
+        elif len(groups[i]) == 1:
+            if pretrain or groups[i][0] in grid1_selection:
+                if (i in test_idcs) != pretrain: train_idces.append(i)
+        elif len(groups[i]) == 2:
+            if pretrain or groups[i][0] in grid2_selection and groups[i][1] in grid2_selection:
+                if (i in test_idcs) != pretrain: train_idces.append(i)
+
+    ds_train = torch.utils.data.Subset(ds_all, train_idces)
+    ds_test = torch.utils.data.Subset(ds_all, test_idcs)
+
+    return ds_train, ds_test

utils_model.py

@@ -0,0 +1,15 @@
+import torch
+from torch import nn
+import torch_geometric
+from torch_geometric.nn.models import DimeNetPlusPlus
+
+class ModelDimeNet(nn.Module):
+    def __init__(self):
+        super().__init__()
+        
+        self.net = DimeNetPlusPlus(hidden_channels=256, out_channels=256, num_blocks=4, num_spherical=8, num_radial=8, int_emb_size=64, basis_emb_size=64, out_emb_channels=64)
+        self.head = nn.Linear(256, 1)
+
+    def forward(self, atoms, coords, batch):
+        emb = self.net(atoms, coords, batch)
+        return self.head(emb)