saxs_refinement.py

import argparse
import torch
import os
import numpy as np
from scipy.linalg import svd, inv
import torch.nn.functional as F
import pickle
from chroma import Protein
from chroma import Chroma
from chroma.layers.structure.rmsd import CrossRMSD

from src.plots import plot_metric, save_trajectory, plot_rmsd_ca_vs_completeness, plot_SAXS_profile
from src.profile import compute_profile, chi_square_loss, FitParameters, Profile
from src.structure import FormFactorType
from src.utils import Particles

def main(args):
    torch.manual_seed(args.seed)
    torch.manual_seed(args.seed)
    np.random.seed(args.seed)
    os.makedirs(args.outdir, exist_ok=True)

    print("Saving arguments", flush=True)
    args_dict = vars(args)
    with open(f"{args.outdir}/config.txt", 'w') as file:
        for key in sorted(args_dict.keys()):
            file.write(f"{key}: {args_dict[key]}\n")

    print("Loading Chroma model", flush=True)
    if args.weights_backbone is not None and args.weights_design is not None:
        chroma = Chroma(weights_backbone=args.weights_backbone, weights_design=args.weights_design)
    else:
        chroma = Chroma()
    backbone_network = chroma.backbone_network
    design_network = chroma.design_network
    def multiply_R(Z, C): return backbone_network.noise_perturb.base_gaussian._multiply_R(Z, C)
    def multiply_R_inverse(X, C): return backbone_network.noise_perturb.base_gaussian._multiply_R_inverse(X, C)
    def multiply_covariance(dU, C): return backbone_network.noise_perturb.base_gaussian.multiply_covariance(dU, C)

    print("Loading experimental SAXS profile", flush=True)
    if not os.path.exists(args.dat):
        raise FileNotFoundError(f"File {args.dat} not found")
    exp_profile = Profile(file_name=args.dat, number_of_q_points=args.nb_points, constructor=1)
    if exp_profile is None:
        print("No experimental profile found", flush=True)
    qmin = exp_profile.min_q_
    qmax = exp_profile.max_q_
    q_values = exp_profile.q_
    ff_type = FormFactorType.HEAVY_ATOMS   

    print("Initializing Protein objects", flush=True)
    protein = Protein.from_PDB(args.pdb, device='cuda')
    X_gt, C_gt, S_gt = protein.to_XCS(all_atom=False)  # we use X_gt to compute RMSD
    particle_model = Particles(S_gt)
    particles = particle_model.particles
    mask_gt = (C_gt > 0)[0]
    n_residues = X_gt.shape[1]
    X_gt_full, C_gt_full, S_gt_full = protein.to_XCS(all_atom=True)
    chi_gt, mask_chi = design_network.X_to_chi(X_gt_full, C_gt_full, S_gt_full)

    print("Precomputing preconditioning matrices based on ground truth", flush=True)
    Y = X_gt[:, mask_gt.cpu()].reshape(1, -1, 3)

    def mR_fun(z):
        z = z.reshape(-1, n_residues, 4, 3)
        return multiply_R(z, C_gt.expand(z.shape[0], -1))[:, mask_gt.cpu()].reshape(z.shape[0], -1).permute(1, 0)

    Z_0 = torch.eye(n_residues * 12).cuda()
    mR_mat = mR_fun(Z_0).cpu().numpy()

    U, S, Vh = svd(mR_mat)
    Um1 = inv(U)
    Um1 = torch.tensor(Um1).float().cuda()[None].expand(args.population_size, -1, -1)
    S = torch.tensor(S).float().cuda()
    Vh = torch.tensor(Vh).float().cuda()[None].expand(args.population_size, -1, -1)

    print("Initializing backbone", flush=True)
    C_gt = torch.abs(C_gt).expand(args.population_size, -1)
    S_gt = S_gt.expand(args.population_size, -1)
    chi_gt = chi_gt.expand(args.population_size, -1, -1)
    Y = Y.expand(args.population_size, -1, -1, -1)
    if args.init_gt:
        X = torch.clone(X_gt).expand(args.population_size, -1, -1, -1) + args.eps_init
        assert not torch.isnan(X).any(), "X_gt clone contains NaNs before gradient computation"
        if args.std_dev_init > 1e-8:
            X += args.std_dev_init * torch.randn_like(X)
        Z = multiply_R_inverse(X, C_gt)
        assert not torch.isnan(Z).any(), "Z R-1 X_gt clone contains NaNs before gradient computation"
    else:
        Z = torch.randn(args.population_size, *X_gt.shape[1:]).float().cuda()
        X = multiply_R(Z, C_gt)

    V_m = torch.zeros_like(Z)
    V_p = torch.zeros_like(Z)
    V_s = torch.zeros_like(Z)
    V_c = torch.zeros_like(Z)
    
    def sample_chi(X, t=0):
        print("Sampling side chain angles", flush=True)
        _X = F.pad(X, [0, 0, 0, 10])
        node_h, edge_h, edge_idx, mask_i, mask_ij = design_network.encode(_X, C_gt, t=t)
        permute_idx = design_network.traversal(_X, C_gt)
        _, chi_sample, _, logp_chi, _ = design_network.decoder_chi.decode(
            _X,
            C_gt,
            S_gt,
            node_h,
            edge_h,
            edge_idx,
            mask_i,
            mask_ij,
            permute_idx
        )
        return chi_sample
    if args.use_gt_chi:
        chi_sample = chi_gt
    else:
        chi_sample = torch.clone(sample_chi(X, t=1)).detach() if args.sample_chi_every > 0 else None

    def t_fn(epoch):
        if args.temporal_schedule == 'linear':
            return (-args.t + 0.001) * epoch / args.epochs + args.t
        elif args.temporal_schedule == 'sqrt':
            return (1.0 - 0.001) * (1. - np.sqrt(epoch / args.epochs)) + 0.001
        elif args.temporal_schedule == 'constant':
            return args.t
        else:
            raise NotImplementedError

    def lr_fn(epoch):
        return args.lr_profile

    def resolution_fn(epoch):
        if epoch < args.activate_resolution_drop:
            return args.resolution_cutoff_start
        else:
            a = (args.resolution_cutoff_end - args.resolution_cutoff_start) / (args.epochs - args.activate_resolution_drop)
            b = args.resolution_cutoff_start - args.activate_resolution_drop * a
            return a * epoch + b

    def sampling_rate_fn(epoch):
        if args.sampling_rate_schedule == 'constant':
            return args.sampling_rate_start
        elif args.sampling_rate_schedule == 'linear':
            return args.sampling_rate_start + epoch * (args.sampling_rate_end - args.sampling_rate_start) / args.epochs
        elif args.sampling_rate_schedule == 'exp':
            return np.exp(np.log(args.sampling_rate_start) + epoch * (np.log(args.sampling_rate_end) - np.log(args.sampling_rate_start)) / args.epochs)
        else:
            raise NotImplementedError
        
    def profile_error(exp_profile, model_profile, c1, c2, use_offset):
        loss_p = chi_square_loss(exp_profile, model_profile, c1, c2, use_offset=use_offset)
        norm = 1.0 # hard-coded parameter -- has roughly the same role as the learning rate
        return loss_p / norm

    def get_gradient_Z_m(Z, t, epoch):
        if args.lr_model > 0. and (args.de_activate_model < 0 or epoch < args.de_activate_model):
            with (torch.enable_grad()):
                Z.requires_grad_(True)
                _Z = Z
                if args.preconditioning_model:
                    Um1Y = torch.bmm(Um1, Y.reshape(Y.shape[0], -1, 1)).reshape(Y.shape[0], -1)
                    Sm1Um1Y = Um1Y / S
                    loss_m = ((torch.bmm(Vh, _Z.reshape(Z.shape[0], -1, 1))[:, :Um1.shape[1], 0] - Sm1Um1Y) ** 2).sum()
                else:
                    loss_m = ((multiply_R(_Z, C_gt)[:, mask_gt.cpu()] - Y) ** 2).sum()
                loss_m.backward()
                grad_Z_m = Z.grad
            Z.requires_grad_(False)
        else:
            grad_Z_m = torch.zeros(*Z.shape).float().to(X.device)
            loss_m = torch.tensor([0.]).float().cuda()
        return grad_Z_m, loss_m

    def get_gradient_Z_p(Z, chi_sample, t, epoch):
        if args.lr_profile > 0.:
            with (torch.enable_grad()):
                Z.requires_grad_(True)
                print(f'Z contains NaNs: {torch.isnan(Z).any()}', flush=True)
                _X = multiply_R(Z, C_gt)
                print(f'_X contains NaNs: {torch.isnan(_X).any()}', flush=True)
                if args.sample_chi_every > 0 and epoch + 1 % args.sample_chi_every == 0:
                    if args.use_gt_chi:
                        chi_sample = chi_gt
                    else:
                        chi_sample = sample_chi(_X, t)
                _X_full, _ = design_network.chi_to_X(_X, C_gt, S_gt, chi_sample)
                print(f'_X_full contains NaNs: {torch.isnan(_X_full).any()}', flush=True)
                coords = particle_model.X_to_coords(_X_full)
                print(f'coords contains NaNs: {torch.isnan(coords).any()}', flush=True)
                model_profile = compute_profile(particles=particles, coordinates=coords, min_q=qmin, max_q=qmax, q_values=q_values, ff_type=ff_type)
                loss_p = profile_error(exp_profile, model_profile, args.c1, args.c2, use_offset=False)
                print(f'loss_p: {loss_p}', flush=True)
                loss_p.backward()
                grad_Z_p = Z.grad
                print(f'grad_Z_p: {grad_Z_p}', flush=True)
            Z.requires_grad_(False)
        else:
            grad_Z_p = torch.zeros(*Z.shape).float().to(X.device)
            loss_p = torch.tensor([0.]).float().cuda()
        return grad_Z_p, loss_p, chi_sample

    def get_gradient_Z_s(Z, t, epoch):
        if args.lr_sequence > 0. and epoch >= args.activate_sequence:
            with (torch.enable_grad()):
                Z.requires_grad_(True)
                _X = multiply_R(Z, C_gt)
                _X_input = F.pad(_X, [0, 0, 0, 10])
                out = design_network(_X_input, C_gt, S_gt, t.cuda())
                logp_S = out["logp_S"]
                loss_s = -logp_S.sum()
                loss_s.backward()
                grad_Z_s = Z.grad
            Z.requires_grad_(False)
        else:
            grad_Z_s = torch.zeros(*Z.shape).float().to(X.device)
            loss_s = torch.tensor([0.]).float().cuda()
        return grad_Z_s, loss_s

    def get_gradient_Z_c(Z):
        if args.lr_inter_ca > 0.:
            with (torch.enable_grad()):
                Z.requires_grad_(True)
                _X = multiply_R(Z, C_gt)
                distances = torch.linalg.norm(_X[:, 1:, 1] - _X[:, :-1, 1])
                loss_c = ((distances - 3.8) ** 2).sum()
                loss_c.backward()
                grad_Z_c = Z.grad
            Z.requires_grad_(False)
        else:
            grad_Z_c = torch.zeros(*Z.shape).float().to(X.device)
            loss_c = torch.tensor([0.]).float().cuda()
        return grad_Z_c, loss_c
    
    trajectory = [torch.clone(X_gt[:, mask_gt]).detach().cpu().numpy(),
                  torch.clone(Y[:1]).detach().cpu().numpy(),
                  (torch.clone(X).detach().cpu().numpy(), 'initial state')]
    
    metrics = {'epoch': [], 'rmsd': [], 't': [], 'loss_m': [], 'loss_d': [], 'rmsd_ca': [],
               'resolution': [], 'loss_s': [], 'lr_density': [], 'loss_d_per_sample': [], 'sampling_rate': [],
               'loss_c': []}

    print("--- Optimization starts now ---", flush=True)
    for epoch in range(args.epochs):
        t = torch.tensor(t_fn(epoch)).float().cuda()
        print(f"Epoch {epoch + 1}/{args.epochs}, t={t.item()}", flush=True)
        if args.use_diffusion:
            with torch.no_grad():
                X0 = backbone_network.denoise(X.detach(), C_gt, t)
                print(f'X0 contains NaNs: {torch.isnan(X0).any()}', flush=True)
        else:
            X0 = X
        Z0 = multiply_R_inverse(X0, C_gt)
        print(f'Z0 contains NaNs: {torch.isnan(Z0).any()}', flush=True)

        grad_Z_m, loss_m = get_gradient_Z_m(Z0, t, epoch)
        V_m = args.rho_model * V_m + args.lr_model * grad_Z_m

        grad_Z_p, loss_p, chi_sample = get_gradient_Z_p(Z0, chi_sample, t, epoch)
        chi_sample = torch.clone(chi_sample).detach() if chi_sample is not None else None
        V_p = args.rho_profile * V_p + lr_fn(epoch) * grad_Z_p

        grad_Z_s, loss_s = get_gradient_Z_s(Z0, t, epoch)
        V_s = args.rho_sequence * V_s + args.lr_sequence * grad_Z_s

        grad_Z_c, loss_c = get_gradient_Z_c(Z0)
        V_c = args.rho_inter_ca * V_c + args.lr_inter_ca * grad_Z_c

        Z0 = Z0 - V_m - V_p - V_s - V_c

        # replicate models with lowest density error
        if args.select_best_every > 0 and epoch >= args.activate_replication and epoch % args.select_best_every == 0:
            assert Z0.shape[0] % args.replication_factor == 0, "The population size must be an integer multiple of the replication factor"
            _, indices = torch.topk(loss_p, Z0.shape[0] // args.replication_factor, largest=False)
            Z0 = torch.clone(Z0[indices][:, None].expand(-1, args.replication_factor, -1, -1, -1)).reshape(-1, *Z0.shape[1:])

        if args.use_diffusion:
            tm1 = torch.tensor(t_fn(epoch + 1)).float().cuda()
            alpha, sigma, _, _, _, _ = backbone_network.noise_perturb._schedule_coefficients(tm1)
            X = multiply_R(alpha * Z0 + sigma * torch.randn_like(Z0), C_gt)
        else:
            X = multiply_R(Z0, C_gt)

        if (epoch + 1) % args.log_every == 0:
            rmsds = []
            rmsds_cas = []
            for i in range(args.population_size):
                rmsd, _ = CrossRMSD().pairedRMSD(
                    torch.clone(X[i, mask_gt]).cpu().reshape(1, -1, 3),
                    torch.clone(X_gt[0, mask_gt]).cpu().reshape(1, -1, 3),
                    compute_alignment=True
                )
                rmsd_ca, _ = CrossRMSD().pairedRMSD(
                    torch.clone(X[i, mask_gt, 1, :]).cpu().reshape(1, -1, 3),
                    torch.clone(X_gt[0, mask_gt, 1, :]).cpu().reshape(1, -1, 3),
                    compute_alignment=True
                )
                rmsds.append(rmsd.item())
                rmsds_cas.append(rmsd_ca.item())
            idx_best = np.argmin(rmsds)
            rmsd_best = rmsds[idx_best]
            rmsd_ca_best = rmsds_cas[idx_best]

            rmsd, X_aligned = CrossRMSD().pairedRMSD(
                torch.clone(X[idx_best, mask_gt]).cpu().reshape(1, -1, 3),
                torch.clone(X_gt[0, mask_gt]).cpu().reshape(1, -1, 3),
                compute_alignment=True
            )
            X_aligned = X_aligned.reshape(1, -1, 4, 3)

            metrics['epoch'].append(epoch)
            metrics['t'].append(t.item())
            metrics['loss_m'].append(loss_m.item())
            metrics['loss_p'].append(loss_p.item())
            metrics['loss_s'].append(loss_s.item())
            metrics['loss_c'].append(loss_c.item())
            metrics['resolution'].append(resolution_fn(epoch))
            metrics['sampling_rate'].append(sampling_rate_fn(epoch))
            metrics['lr_density'].append(lr_fn(epoch))
            metrics['rmsd'].append(rmsds)
            metrics['rmsd_ca'].append(rmsds_cas)
            trajectory.append((torch.clone(X[idx_best][None]).detach().cpu().numpy(), 'x-update'))
            print(f"Epoch {epoch + 1}/{args.epochs}, Loss Model: {loss_m.item():.4e}, Loss Profile: {loss_p.item():.4e}, RMSD: {rmsd_best:.2e}, RMSD CA: {rmsd_ca_best:.2e}", flush=True)

    C_gt = C_gt[0:1]
    S_gt = S_gt[0:1]
    X_full, _ = design_network.chi_to_X(X[idx_best][None], C_gt, S_gt, chi_sample)

    print(f"Saving {args.outdir}/metrics.pkl" , flush=True)
    with open(f"{args.outdir}/metrics.pkl", 'wb') as file:
        pickle.dump(metrics, file)

    print(f"Saving {args.outdir}/{args.outdir.split('/')[-1]}.pdb", flush=True)
    protein_out = Protein.from_XCS(X_full, C_gt, S_gt)
    protein_out.to_PDB(f"{args.outdir}/{args.outdir.split('/')[-1]}.pdb")

    for key in metrics.keys():
        if key != 'epoch' and key != 'loss_d_per_sample':
            print(f"Saving {args.outdir}/{key}.png", flush=True)
            plot_metric(metrics, key, f"{args.outdir}/{key}.png")

    print(f"Saving {args.outdir}/{args.outdir.split('/')[-1]}.mp4", flush=True)
    save_trajectory(trajectory, f"{args.outdir}/{args.outdir.split('/')[-1]}.mp4")

    print(f"Saving {args.outdir}/rmsd_ca_vs_completeness.png", flush=True)
    plot_rmsd_ca_vs_completeness(X_gt, X[idx_best][None], mask_gt.cpu(), f"{args.outdir}/rmsd_ca_vs_completeness.png")

    print(f"Saving {args.outdir}/saxs_profile.png", flush=True)
    model_profile = compute_profile(particles=particles, coordinates=particle_model.X_to_coords(X_full), min_q=qmin, max_q=qmax, q_values=q_values, ff_type=ff_type)
    chi_square, fit_profile = chi_square_loss(exp_profile, model_profile, c1=args.c1, c2=args.c2, use_offset=False)
    fitted_params = FitParameters(chi_square, args.c1, args.c2)
    plot_SAXS_profile(profile=fit_profile, exp_profile=exp_profile, chi_square=chi_square, fitted_params=fitted_params, name=f"{args.outdir}/saxs_profile.png")

if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    # required parameters
    parser.add_argument('--outdir', type=str, required=True, help="Path to output directory.")
    parser.add_argument('--dat', type=str, required=True, help="Path to experimental SAXS profile dat file.")
    parser.add_argument('--pdb', type=str, required=True, help="Path to deposited PDB file.")

    # I/O parameters
    parser.add_argument('--weights-backbone', type=str, default=None, help="Path to Chroma weights (backbone).")
    parser.add_argument('--weights-design', type=str, default=None, help="Path to Chroma weights (design).")

    # optimization parameters
    parser.add_argument('--epochs', type=int, default=4000, help="Number of epochs.")
    parser.add_argument('--population-size', type=int, default=16, help="Number of atomic models to simultaneously optimize.")
    parser.add_argument('--lr-model', type=float, default=1e-2, help="Learning rate for the model loss.")
    parser.add_argument('--rho-model', type=float, default=0.9, help="Momentum for the model loss.")
    parser.add_argument('--lr-sequence', type=float, default=1e-5, help="Learning rate for the sequence loss.")
    parser.add_argument('--rho-sequence', type=float, default=0.9, help="Momentum for the sequence loss.")
    parser.add_argument('--lr-profile', type=float, default=1e-5, help="Learning rate for the profile loss.")
    parser.add_argument('--rho-profile', type=float, default=0.9, help="Momentum for the profile loss.")
    parser.add_argument('--lr-inter-ca', type=float, default=0.0, help="Learning rate for the inter-CA loss.")
    parser.add_argument('--rho-inter-ca', type=float, default=0.9, help="Momentum for the inter-CA loss.")
    parser.add_argument('--preconditioning-model', type=int, default=1, help="Flag to use preconditioning on the model loss.")
    parser.add_argument('--de-activate-model', type=int, default=-1, help="Number of epochs before de-activating the model loss (-1 to always activate).")
    parser.add_argument('--activate-sequence', type=int, default=3000, help="Number of epochs before activating sequence loss.")

    # SAXS parameters 
    parser.add_argument('--nb-points', type=int, default=-1, help="Number of q points in the experimental profile to use (-1 for all points).")
    parser.add_argument('--c1', type=float, default=1.0, help="c1 parameter used for the fitting of the experimental SAXS profile.")
    parser.add_argument('--c2', type=float, default=0.0, help="c2 parameter used for the fitting of the experimental SAXS profile.")

    # diffusion parameters
    parser.add_argument('--use-diffusion', type=int, default=1, help="Flag to use the diffusion model.")
    parser.add_argument('--temporal-schedule', type=str, default='sqrt', choices=['sqrt', 'linear', 'constant'], help="Type of temporal schedule.")
    parser.add_argument('--t', type=float, default=1.0, help="Initial diffusion time (between 0 and 1).")

    # random sampling
    parser.add_argument('--sampling-rate-schedule', type=str, default='constant', choices=['constant', 'linear', 'exp'], help='Type of schedule for the sampling rate.')
    parser.add_argument('--sampling-rate-start', type=float, default=0.1, help='Initial sampling rate.')
    parser.add_argument('--sampling-rate-end', type=float, default=1.0, help='Final sampling rate.')

    # side-chain parameters
    parser.add_argument('--sample-chi-every', type=int, default=100, help="Frequency (in epochs) of side-chain sampling.")
    parser.add_argument('--use-gt-chi', type=int, default=0, help="Flag to use ground truth side-chain angles, for debugging purposes.")
    
    # genetic parameters
    parser.add_argument('--replication-factor', type=int, default=2, help='Number of replications at each selection step.')
    parser.add_argument('--activate-replication', type=int, default=1, help="Number of epochs to wait before activating the selection/replication.")
    parser.add_argument('--select-best-every', type=int, default=500, help='Frequency (in epochs) of selection/replication (-1 to de-activate).')

    # initialization parameters
    parser.add_argument('--seed', type=int, default=0, help="Random seed.")
    parser.add_argument('--init-gt', type=int, default=0, help="Flag to initialize the model from the deposited PDB, for debugging purposes.")
    parser.add_argument('--std-dev-init', type=float, default=0.0, help="Intensity of Gaussian random noise added on ground truth.")
    parser.add_argument('--eps-init', type=float, default=0.0, help="Size of initial deviation to ground truth in the direction (1, 1, 1).")

    # logging parameters
    parser.add_argument('--log-every', type=int, default=10, help="Frequency (in epochs) for logging.")

    args = parser.parse_args()
    main(args)