import time import argparse import torch import os import json import matplotlib matplotlib.use('agg') import matplotlib.pyplot as plt import numpy as np from tqdm import tqdm from torch.optim import Adam import torch.nn.functional as F from pathlib import Path from types import SimpleNamespace import numpy as np from torch_geometric.data import Batch from torch_geometric.data import DataLoader from torch.distributions import MultivariateNormal from eval_utils import load_model from cdvae.common.data_utils import get_scaler_from_data_list, build_crystal, build_crystal_graph from visualization.visualize_materials import create_materials, plot_material_single, plot_xrd_single from compute_metrics import Crystal, RecEval, GenEval from torch.optim.lr_scheduler import CosineAnnealingLR, CosineAnnealingWarmRestarts from pymatgen.symmetry.analyzer import SpacegroupAnalyzer from pymatgen.io.cif import CifWriter, CifParser import wandb from PIL import Image AVG_COMPOSITION_ERROR = 'composition error rate' AVG_XRD_MSE = 'Scaled XRD mean squared error' AVG_XRD_L1 = 'Scaled XRD mean absolute error' BEST_XRD_MSE = 'Average best scaled XRD mean squared error among candidates' BEST_XRD_L1 = 'Average best scaled XRD mean absolute error among candidates' AVG_PDF_CORRELATION = "Average Pearson's correlation coefficient between PDFs" BEST_PDF_CORRELATION = "Mean best Pearson's correlation coefficient between PDFs" STD_PDF_CORRELATION = "Std of Pearson's correlation coefficieint between PDFs" STD_BEST_PDF_CORRELATION = "Std of best Pearson's correlation coefficient between PDFs" PDF_CORRELATIONS = "All PDF correlations" AVG_R_FACTOR = 'Average r factor' BEST_R_FACTOR = 'Best r factor' STD_R_FACTOR = 'Std of r factors' STD_BEST_R_FACTOR = 'Std of best r factors' R_FACTORS = "All r factors" MATCH_RATE = 'match_rate' RMS_DIST = 'rms_dist' COMPOSITION_VALIDITY = 'comp_valid' STRUCTURE_VALIDITY = 'struct_valid' VALIDITY = 'valid' NUM_ATOM_ACCURACY = '% materials w/ # atoms pred correctly' PDF_CORRELATION = "pearson's r between PDFs" USE_ALL_SPACEGROUPS = "aggregated stats (all spacegroups)" COUNT = "number of crystals" EPS = 1e-10 # Thanks ChatGPT! # If you want to change the colors of the lines and shades, simply modify in the ax.fill_between() and ax.plot() functions # A list of possible colors can be found at: https://matplotlib.org/stable/gallery/color/named_colors.html def plot_overlaid_graphs(actual, prediction_nn, prediction_simulated, Qs, savepath): fig, ax = plt.subplots() # Plot and fill the area under the first curve ax.fill_between(Qs, actual, color="royalblue", alpha=0.2) ax.plot(Qs, actual, color="blue", alpha=0.6, label="Actual") # Curve line # Plot and fill the area under the second curve ax.fill_between(Qs, prediction_nn, color="mistyrose", alpha=0.2) ax.plot(Qs, prediction_nn, color="red", alpha=0.6, linestyle='dotted', linewidth=2, label="Prediction (NN)") # Dotted curve line with increased linewidth # Plot and fill the area under the second curve ax.fill_between(Qs, prediction_simulated, color="lightgreen", alpha=0.2) ax.plot(Qs, prediction_simulated, color="green", alpha=0.6, linestyle='dashed', linewidth=2, label="Prediction (Simulated)") # Dotted curve line with increased linewidth # Customizing the plot ax.set_title("XRD Patterns") ax.set_xlabel(r'$Q (\mathring A^{-1})$') ax.set_ylabel("Scaled Intensity") # ax.set_xlim(0, 180) # Set x-axis limits ax.set_ylim(0, 1) # Set y-axis limits # ax.set_xticks(np.arange(0, 181, 10)) # ax.set_xticklabels(ax.get_xticks(), rotation=70) # Rotate x-axis labels by 70 degrees ax.set_yticks(np.arange(0, 1.1, 0.1)) # Set horizontal gridlines every 0.1 from 0 to 1 ax.grid(True) # Show gridlines ax.legend() # Display the plot #plt.show() plt.tight_layout() plt.savefig(savepath) plt.savefig(savepath.replace('.png', '.pdf')) plt.close() return def plot_smoothed_vs_sinc(smoothed, sincPattern, noiselessPattern, Qs, savepath): fig, ax = plt.subplots() if not isinstance(smoothed, np.ndarray): smoothed = torch.clone(smoothed).squeeze().detach().cpu().numpy() if not isinstance(sincPattern, np.ndarray): sincPattern = torch.clone(sincPattern).squeeze().detach().cpu().numpy() if not isinstance(noiselessPattern, np.ndarray): noiselessPattern = torch.clone(noiselessPattern).squeeze().detach().cpu().numpy() assert Qs.shape == smoothed.shape # Plot and fill the area under the first curve #ax.fill_between(thetas, smoothed, color="hotpink", alpha=0.1) ax.plot(Qs, smoothed, color="deeppink", alpha=0.4, linestyle='dashed', label="Smoothed") # Plot and fill the area under the second curve #ax.fill_between(thetas, sincPattern, color="purple", alpha=0.2) ax.plot(Qs, sincPattern, color="indigo", alpha=0.6, label="Sinc (Raw Nanomaterial)") # Plot and fill the area under the second curve ax.plot(Qs, noiselessPattern, color="gray", alpha=0.8, label="Noiseless (Ideal Material)") # Customizing the plot ax.set_title("XRD Patterns") ax.set_xlabel(r'$Q (\mathring A^{-1})$') ax.set_ylabel("Scaled Intensity") # ax.set_xlim(0, 180) # Set x-axis limits ax.set_ylim(0, 1) # Set y-axis limits # ax.set_xticks(np.arange(0, 181, 10)) # ax.set_xticklabels(ax.get_xticks(), rotation=70) # Rotate x-axis labels by 70 degrees ax.set_yticks(np.arange(0, 1.1, 0.1)) # Set horizontal gridlines every 0.1 from 0 to 1 ax.grid(True) # Show gridlines ax.legend() # Display the plot #plt.show() plt.tight_layout() plt.savefig(savepath) plt.savefig(savepath.replace('.png', '.pdf')) plt.close() return def point_pdf_query(Qs, signal, r): ret_val = 0 assert np.isclose(np.mean(signal), 1) delta_Q = (Qs[-1] - Qs[0]) / (Qs.shape[0] - 1) assert np.isclose(delta_Q, Qs[1] - Qs[0]) for i in range(len(signal)): q = Qs[i] s_q = signal[i] ret_val += 2 / np.pi * q * (s_q - 1) * np.sin(q * r) * delta_Q return ret_val def overall_pdf(Qs, signal, r_min=0, r_max=25, num_samples=1000): assert Qs.shape == signal.shape signal = signal / np.mean(signal) rs = np.linspace(r_min, r_max, num_samples) the_pdf = list() for r in rs: the_pdf.append(point_pdf_query(Qs=Qs, signal=signal, r=r)) return np.array(rs), np.array(the_pdf) def calc_r_factor(gt_xrd, pred_xrd, Qs): delta_Q = (Qs[-1] - Qs[0]) / (Qs.shape[0] - 1) assert np.isclose(delta_Q, Qs[1] - Qs[0]) assert np.isclose(delta_Q, Qs[-1] - Qs[-2]) if isinstance(gt_xrd, torch.Tensor): gt_xrd = gt_xrd.squeeze().detach().cpu().numpy() if isinstance(pred_xrd, torch.Tensor): pred_xrd = pred_xrd.squeeze().detach().cpu().numpy() assert np.isclose(np.max(gt_xrd), 1) assert np.isclose(np.max(pred_xrd), 1) assert np.min(gt_xrd) >= 0 assert np.min(pred_xrd) >= 0 numerator = np.sum(delta_Q * (gt_xrd - pred_xrd)**2) denominator = np.sum(delta_Q * gt_xrd**2) return numerator / denominator def calc_and_plot_pdf_correlation(args, gt_xrd, pred_xrd, Qs, save_dir): # plot XRD if isinstance(gt_xrd, torch.Tensor): gt_xrd = gt_xrd.squeeze().detach().cpu().numpy() if isinstance(pred_xrd, torch.Tensor): pred_xrd = pred_xrd.squeeze().detach().cpu().numpy() assert gt_xrd.shape == pred_xrd.shape plt.plot(Qs, gt_xrd, alpha=0.8, label='GT XRD (noiseless)') plt.plot(Qs, pred_xrd, alpha=0.8, label='Pred XRD (noiseless)') plt.xlabel(r'$Q (\mathring A^{-1})$') plt.ylabel("Scaled Intensity") plt.title('XRD patterns') plt.grid() plt.legend() plt.savefig(os.path.join(save_dir, 'xrd_comparison.png')) plt.savefig(os.path.join(save_dir, 'xrd_comparison.pdf')) plt.close() # create PDF gt_rs, gt_pdf = overall_pdf(Qs=Qs, signal=gt_xrd, r_min=args.r_min, r_max=args.r_max) pred_rs, pred_pdf = overall_pdf(Qs=Qs, signal=pred_xrd, r_min=args.r_min, r_max=args.r_max) assert np.array_equal(gt_rs, pred_rs) # plot PDF plt.plot(gt_rs, gt_pdf, alpha=0.8, label='GT PDF') plt.plot(pred_rs, pred_pdf, alpha=0.8, label='Pred PDF') plt.xlabel(r'$r (\mathring A)$') plt.ylabel("G(r)") # save PDF torch torch.save(torch.from_numpy(gt_pdf), os.path.join(save_dir, 'gt_pdf.pt')) torch.save(torch.from_numpy(pred_pdf), os.path.join(save_dir, 'pred_pdf.pt')) # calculate correlation coefficient correlation_matrix = np.corrcoef(gt_pdf, pred_pdf) pearson_r = correlation_matrix[0, 1] assert np.isclose(correlation_matrix[0, 1], correlation_matrix[1, 0]) assert np.isclose(correlation_matrix[0, 0], 1) and np.isclose(correlation_matrix[1, 1], 1) # save PDF image plt.title(f"Pair Distribution Function Comparison\n(Pearson's r = {pearson_r:.3f})") plt.grid() plt.legend() plt.savefig(os.path.join(save_dir, 'pdf_comparison.png')) plt.savefig(os.path.join(save_dir, 'pdf_comparison.pdf')) plt.close() # return return pearson_r # Thanks ChatGPT! def resize_image_to_same_width(image, width): """Resize an image to the same width, maintaining the aspect ratio.""" ratio = width / float(image.width) new_height = int(image.height * ratio) return image.resize((width, new_height), Image.LANCZOS) def collate_images(gt_material, gt_xrd, pred_material, pred_xrd, width): gt_material = resize_image_to_same_width(Image.open(gt_material), width) gt_xrd = resize_image_to_same_width(Image.open(gt_xrd), width) pred_material = resize_image_to_same_width(Image.open(pred_material), width) pred_xrd = resize_image_to_same_width(Image.open(pred_xrd), width) assert gt_material.height == pred_material.height assert gt_xrd.height == pred_xrd.height total_height = gt_material.height + gt_xrd.height combined_image = Image.new('RGB', (width * 2, total_height)) combined_image.paste(gt_material, (0, 0)) combined_image.paste(gt_xrd, (0, gt_material.height)) combined_image.paste(pred_material, (width, 0)) combined_image.paste(pred_xrd, (width, pred_material.height)) return combined_image # Thanks ChatGPT! def calculate_accuracy(probabilities, labels): """ Calculate accuracy given the softmax probabilities and true labels. :param probabilities: Softmax probabilities of shape (N, C) where N is the number of samples and C is the number of classes. :param labels: True class labels of shape (N,). :return: Accuracy as a Python float. """ # Step 1: Convert softmax probabilities to predicted class indices _, predicted_classes = torch.max(probabilities, dim=1) # Step 2: Compare with true class labels correct_predictions = (predicted_classes == labels).float() # Convert boolean tensor to float for sum operation # Step 3: Calculate accuracy accuracy = correct_predictions.sum() / labels.size(0) return accuracy.item() # Convert to Python float for readability def optimize_latent_code(args, model, batch, target_noisy_xrd, z_init=None): m = MultivariateNormal(torch.zeros(model.hparams.hidden_dim).cuda(), torch.eye(model.hparams.hidden_dim).cuda()) if z_init is None: assert args.start_from_init is None print('random z') z = torch.randn(args.num_starting_points, model.hparams.hidden_dim, device=model.device) else: print(f'init z from pre-existing: {args.start_from_init}') z = z_init.detach() assert z.shape == (args.num_starting_points, model.hparams.hidden_dim) z.requires_grad = True opt = Adam([z], lr=args.lr) total_gradient_steps = args.num_gradient_steps * (1+2+4) - 1 scheduler = CosineAnnealingWarmRestarts(opt, args.num_gradient_steps, T_mult=2, eta_min=args.min_lr) model.freeze() with tqdm(total=total_gradient_steps, desc="Property opt", unit="steps") as pbar: for i in range(total_gradient_steps): opt.zero_grad() xrd_loss = F.l1_loss(model.fc_property(z), target_noisy_xrd.broadcast_to(z.shape[0], target_noisy_xrd.shape[-1])) if args.l1_loss \ else F.mse_loss(model.fc_property(z), target_noisy_xrd.broadcast_to(z.shape[0], target_noisy_xrd.shape[-1])) prob = m.log_prob(z).mean() # predict the number of atoms, lattice, composition (pred_num_atoms, pred_lengths_and_angles, pred_lengths, pred_angles, pred_composition_per_atom) = model.decode_stats( z, batch.num_atoms, batch.lengths, batch.angles, teacher_forcing=False) num_atom_loss = F.cross_entropy(pred_num_atoms, batch.num_atoms.repeat(args.num_starting_points)) # TODO: they do some weird stuff with composition loss: double check (I think it was inconsequential, but idk) composition_loss = F.cross_entropy(pred_composition_per_atom, (batch.atom_types - 1).repeat(args.num_starting_points)) num_atom_accuracy = calculate_accuracy(pred_num_atoms, batch.num_atoms.repeat(args.num_starting_points)) composition_accuracy = calculate_accuracy(pred_composition_per_atom, (batch.atom_types - 1).repeat(args.num_starting_points)) pbar.set_postfix_str(f"XRD loss: {xrd_loss.item():.3e}; Gaussian log PDF: {prob.item():.3e}; " + f"Num atom loss: {num_atom_loss.item():.3e}; Composition loss: {composition_loss.item():.3e}", refresh=True) # Update the progress bar by one step pbar.update(1) # calculate total loss: minimize XRD loss, maximize latent code probability (min neg prob) total_loss = xrd_loss - args.l2_penalty * prob \ + args.num_atom_lambda * num_atom_loss \ + args.composition_lambda * composition_loss if i % 100 == 0 and i > 0: wandb.log( { "total_loss":total_loss, "lr":scheduler.get_last_lr()[0], "xrd_loss":xrd_loss, "log_prob":prob, "num_atom_loss":num_atom_loss, "composition_loss":composition_loss, "num_atom_accuracy":num_atom_accuracy, "composition_accuracy":composition_accuracy }, step=i) # backprop through total loss total_loss.backward() opt.step() scheduler.step() return z def process_candidates(args, xrd_args, j, curr_gen_crystals_list, all_opt_coords, all_opt_atom_types, opt_generated_xrds, min_loss_indices, curr_material_folder, all_bestPred_crystals, gt_noiseless_xrd, target_noisy_xrd, final_pred_xrds, opt_sinc_only_xrds, noiseless_generated_xrds, curr_gt_crystal, gt_atom_types, gt_material_filepath, gt_xrd_filepath, all_xrd_l1_errors, all_xrd_l2_errors, all_composition_errors, has_correct_num_atoms, all_pdf_correlations, all_r_factors, Qs): candidate_xrd_l1_errors = list() candidate_xrd_l2_errors = list() candidate_match_status = list() candidate_composition_errors = list() candidate_has_correct_num_atoms = list() candidate_pdf_correlations = list() candidate_r_factors = list() print(f'crystal {j} has {len(min_loss_indices)} candidates') best_rms_dist = 1e6 # By default, log best (lowest loss) crystal for metrics best_crystal = Crystal(curr_gen_crystals_list[min_loss_indices[0]]) for i, min_loss_idx in enumerate(min_loss_indices): # for each candidate filename = f'candidate_{i}.png' # construct the corresponding crystal opt_coords = all_opt_coords[min_loss_idx] opt_atom_types = all_opt_atom_types[min_loss_idx] opt_xrd = opt_generated_xrds[min_loss_idx, :].cpu().numpy() curr_pred_crystal = Crystal(curr_gen_crystals_list[min_loss_idx]) curr_candidate_folder = os.path.join(curr_material_folder, 'pred', f'candidate{i}') # save the optimal crystal and its xrd opt_material_folder_cand = os.path.join(curr_candidate_folder, 'visUnitCell') os.makedirs(opt_material_folder_cand, exist_ok=True) pred_material_filepath = plot_material_single(opt_coords, opt_atom_types, opt_material_folder_cand, idx=j, filename=filename) opt_xrd_folder_cand = os.path.join(curr_candidate_folder, 'xrd_opt_gen') os.makedirs(opt_xrd_folder_cand, exist_ok=True) pred_xrd_filepath = plot_xrd_single(xrd_args, opt_xrd, opt_xrd_folder_cand, idx=j, filename=filename, x_axis=Qs, x_label=r'Q $({A^{\circ}}^{-1})$') torch.save(opt_generated_xrds[min_loss_idx, :], os.path.join(opt_xrd_folder_cand, f'candidate_{i}.pt')) pred_opt_xrd_folder_cand = os.path.join(curr_candidate_folder, 'xrd_ml_pred') os.makedirs(pred_opt_xrd_folder_cand, exist_ok=True) pred_opt_xrd_filepath = plot_xrd_single(xrd_args, final_pred_xrds[min_loss_idx].detach().cpu().numpy(), pred_opt_xrd_folder_cand, idx=j, filename=filename, x_axis=Qs, x_label=r'Q $({A^{\circ}}^{-1})$') torch.save(final_pred_xrds[min_loss_idx].detach(), os.path.join(pred_opt_xrd_folder_cand, f'candidate_{i}.pt')) opt_cif_folder_cand = os.path.join(curr_candidate_folder, 'cif') os.makedirs(opt_cif_folder_cand, exist_ok=True) curr_pred_crystal.structure.to(filename=f'{opt_cif_folder_cand}/noSpacegroup_material{j}_candidate{i}.cif', fmt='cif') try: pred_cif_writer = CifWriter(curr_pred_crystal.structure, symprec=0.01) pred_cif_writer.write_file(filename=f'{opt_cif_folder_cand}/material{j}_candidate{i}.cif') except Exception as e: print(e) print(f"Fail on candidate {i}") # Log image log_img = collate_images(gt_material=gt_material_filepath, gt_xrd=gt_xrd_filepath, pred_material=pred_material_filepath, pred_xrd=pred_xrd_filepath, width=600) wandb.log({"prediction": wandb.Image(log_img)}) # metrics assert target_noisy_xrd.squeeze().shape == opt_generated_xrds[min_loss_idx].squeeze().shape the_curr_opt_generated_xrd = opt_generated_xrds[min_loss_idx].to(target_noisy_xrd.device).squeeze() xrd_l1_error = F.l1_loss(target_noisy_xrd.squeeze(), the_curr_opt_generated_xrd).item() xrd_l2_error = F.mse_loss(target_noisy_xrd.squeeze(), the_curr_opt_generated_xrd).item() candidate_xrd_l1_errors.append(xrd_l1_error) candidate_xrd_l2_errors.append(xrd_l2_error) print(f'xrd l1 error: {xrd_l1_error}') print(f'xrd l2 error: {xrd_l2_error}') composition_error = compare_composition(gt_atom_types, opt_atom_types) candidate_composition_errors.append(composition_error) print(f'composition error: {composition_error}') is_num_atoms_correct = compare_num_atoms(gt_atom_types=gt_atom_types, pred_atom_types=opt_atom_types) candidate_has_correct_num_atoms.append(int(is_num_atoms_correct)) print(f'num atoms: {len(gt_atom_types)} (gt) vs {len(opt_atom_types)} (pred)') # Check if this matches curr_match_stats = check_structure_match( gt_structures=[curr_gt_crystal], pred_structures=[curr_pred_crystal]) candidate_match_status.append(curr_match_stats) # Pick crystal with lowest RMS dist as our candidate if curr_match_stats[MATCH_RATE] > 0.5 and curr_match_stats[RMS_DIST] < best_rms_dist: assert int(curr_match_stats[MATCH_RATE]) == 1 best_rms_dist = curr_match_stats[RMS_DIST] best_crystal = curr_pred_crystal plot_overlaid_graphs(actual=target_noisy_xrd.squeeze().detach().cpu().numpy(), prediction_nn=final_pred_xrds[min_loss_idx].detach().cpu().numpy(), prediction_simulated=opt_xrd, Qs=Qs, savepath=f'{opt_xrd_folder_cand}/overlaidXRD.png') # plot smoothed vs sinc: opt plot_smoothed_vs_sinc(smoothed=the_curr_opt_generated_xrd, sincPattern=opt_sinc_only_xrds[min_loss_idx], noiselessPattern=noiseless_generated_xrds[min_loss_idx], Qs=Qs, savepath=os.path.join(opt_xrd_folder_cand, f'sincVsSmoothed.png')) # compare and pair distribution functions opt_pdf_folder_cand = os.path.join(curr_candidate_folder, 'pdf') os.makedirs(opt_pdf_folder_cand, exist_ok=True) pdf_correlation = calc_and_plot_pdf_correlation(args=args, gt_xrd=gt_noiseless_xrd, pred_xrd=noiseless_generated_xrds[min_loss_idx], Qs=Qs, save_dir=opt_pdf_folder_cand) candidate_pdf_correlations.append(pdf_correlation) print(f"pdf correlation: {pdf_correlation}") r_factor = calc_r_factor(gt_xrd=gt_noiseless_xrd, pred_xrd=noiseless_generated_xrds[min_loss_idx], Qs=Qs) candidate_r_factors.append(r_factor) print(f"r factor: {r_factor}") # Log the crystal with lowest RMS dist all_bestPred_crystals.append(best_crystal) curr_material_metrics = { AVG_XRD_MSE: np.mean(candidate_xrd_l2_errors), AVG_XRD_L1: np.mean(candidate_xrd_l1_errors), BEST_XRD_MSE: np.min(candidate_xrd_l2_errors), BEST_XRD_L1: np.min(candidate_xrd_l1_errors), MATCH_RATE: candidate_match_status, AVG_PDF_CORRELATION: np.mean(candidate_pdf_correlations), BEST_PDF_CORRELATION: np.max(candidate_pdf_correlations), AVG_R_FACTOR: np.mean(candidate_r_factors), BEST_R_FACTOR: np.max(candidate_r_factors), PDF_CORRELATIONS: candidate_pdf_correlations, R_FACTORS: candidate_r_factors } metrics_folder = os.path.join(curr_material_folder, 'metrics') os.makedirs(metrics_folder, exist_ok=True) with open(f'{metrics_folder}/material{j}.json', 'w') as fout: json.dump(curr_material_metrics, fout, indent=4) print(json.dumps(curr_material_metrics, indent=4)) all_xrd_l1_errors.append(candidate_xrd_l1_errors) all_xrd_l2_errors.append(candidate_xrd_l2_errors) all_composition_errors.append(candidate_composition_errors) has_correct_num_atoms.append(candidate_has_correct_num_atoms) all_pdf_correlations.append(candidate_pdf_correlations) all_r_factors.append(candidate_r_factors) wandb.finish() return def write_histogram(values, save_folder, title, xlabel, ylabel, standard_range=True): plt.grid() if standard_range: plt.xticks(np.linspace(0, 1, 11)) plt.xlim(0, 1) bins = np.linspace(0, 1, 21) else: plt.xticks(np.linspace(int(np.min(values)), int(np.max(values)), 11)) plt.xlim(int(np.min(values)), int(np.max(values))) bins = np.linspace(int(np.min(values)), int(np.max(values)), 21) plt.hist(values, density=True, cumulative=True, bins=bins) plt.yticks(np.linspace(0, 1, 11)) plt.ylim(0, 1) plt.xlabel(xlabel) plt.ylabel(ylabel) plt.title(title) plt.tight_layout() plt.savefig(os.path.join(save_folder, f'{title}.png')) plt.savefig(os.path.join(save_folder, f'{title}.pdf')) plt.close() return def write_pdf_histogram(pdf_rs, save_folder, title): xlabel = "Pearson's Correlation (r) between Predicted and GT PDFs" ylabel = "Cumulative Density\n(% of Materials at or below r)" write_histogram(values=pdf_rs, save_folder=save_folder, title=title, xlabel=xlabel, ylabel=ylabel) return def write_r_factor_histogram(r_factors, save_folder, title): xlabel = "R-Factor (Residuals Function) between\nPredicted and GT XRDs (Noiseless)" ylabel = "Cumulative Density\n(% of Materials at or below R)" write_histogram(values=r_factors, save_folder=save_folder, title=title, xlabel=xlabel, ylabel=ylabel, standard_range=False) return def create_xrd_args(args): alt_args = SimpleNamespace() alt_args.wave_source = args.wave_source alt_args.num_materials = args.num_starting_points alt_args.xrd_vector_dim = args.xrd_vector_dim alt_args.max_theta = args.max_theta alt_args.min_theta = args.min_theta return alt_args def smooth_xrds(opt_generated_xrds, data_loader): smoothed_xrds = list() sinc_xrds = list() for i in range(opt_generated_xrds.shape[0]): smoothed_xrd, sincOnly, _, _ = data_loader.dataset.augment_xrdStrip(torch.tensor(opt_generated_xrds[i,:]), return_both=True) smoothed_xrds.append(smoothed_xrd) sinc_xrds.append(sincOnly) opt_generated_xrds = torch.stack(smoothed_xrds, dim=0) opt_sinc_xrds = torch.stack(sinc_xrds, dim=0) return opt_generated_xrds, opt_sinc_xrds def plot_filter(filter, Qs, filter_viz_folder, nanomaterial_size): resolution = Qs.shape[0] Q_min = Qs[0] Q_max = Qs[-1] _, ax = plt.subplots() # sim_filter = nanomaterial_size * np.sinc((np.pi * nanomaterial_size * Qs)/(2 * np.pi)) # sim_filter = sim_filter / np.max(sim_filter) filter = filter# / np.max(filter) # plot filter # ax.plot(Qs, sim_filter, alpha=0.5, label='simulated filter') ax.plot(Qs, filter, alpha=0.5) #, label='true filter') ax.set_xlabel(r'Q $({A^{\circ}}^{-1})$') ax.set_ylabel('Filter value') ax.grid(True) # Show gridlines # ax.legend() plt.tight_layout() plt.savefig(f'{filter_viz_folder}/filter_Q.png') plt.savefig(f'{filter_viz_folder}/filter_Q.pdf') plt.close() # plot filter in spatial domain # inverse shift the signal and fourier transform to freq domain F = np.fft.ifft(np.fft.fftshift(filter)) # shift the signal back in freq domain F_shifted = np.fft.ifftshift(F) # calculate frequency bins d = -(resolution - 1) / (2 * resolution * Q_min) spatial_bins = d * np.arange(resolution) spatial_bins_shifted = spatial_bins - d * resolution / 2 # scale and plot _, ax = plt.subplots() ax.plot(spatial_bins_shifted, np.real(F_shifted)) ax.set_xlim(-100, 100) ax.set_xlabel(r'Spatial $(A^{\circ})$') ax.set_ylabel('Amplitude') # ax.set_xticks(np.arange(-5, 5, 1)) # ax.set_xticklabels(ax.get_xticks(), rotation=70) # Rotate x-axis labels by 70 degrees # ax.set_yticks(np.arange(-0.01, 0.05, 0.01)) # Set horizontal gridlines every 0.1 from 0 to 1 ax.grid(True) # Show gridlines plt.tight_layout() plt.savefig(f'{filter_viz_folder}/filter_spatial.png') plt.savefig(f'{filter_viz_folder}/filter_spatial.pdf') plt.close() def create_z_from_init(args, batch, model, cif_path): batch = batch.clone() # do not do the overwriting assert os.path.exists(cif_path), f'{cif_path} does not exist' with open(cif_path, 'r') as fin: cif_str = ''.join(fin.readlines()) crystal = build_crystal(cif_str) frac_coords, atom_types, lengths, angles, edge_indices, to_jimages, num_atoms = \ build_crystal_graph(crystal) batch.frac_coords = torch.Tensor(frac_coords).to(device=batch.frac_coords.device) batch.atom_types = torch.LongTensor(atom_types).to(device=batch.atom_types.device) batch.lengths = torch.Tensor(lengths).view(1, -1).to(device=batch.lengths.device) batch.angles = torch.Tensor(angles).view(1, -1).to(device=batch.angles.device) batch.edge_index = torch.LongTensor(edge_indices.T).contiguous().to(device=batch.edge_index.device) batch.to_jimages = torch.LongTensor(to_jimages).to(device=batch.to_jimages.device) batch.num_atoms = torch.tensor([num_atoms]).to(device=batch.num_atoms.device) batch.num_bonds = torch.tensor([edge_indices.shape[0]]).to(device=batch.num_bonds.device) mu, log_var, z = model.encode(batch) assert mu.shape == (1, model.hparams.hidden_dim), f"actually, mu's shape is: {mu.shape}" assert log_var.shape == (1, model.hparams.hidden_dim), f"actually, log_var's shape is: {log_var.shape}" assert z.shape == (1, model.hparams.hidden_dim), f"actually, z's shape is: {z.shape}" mu = mu.repeat(args.num_starting_points, 1) log_var = log_var.repeat(args.num_starting_points, 1) z = model.reparameterize(mu, log_var) assert z.shape == (args.num_starting_points, model.hparams.hidden_dim) return z # data = Data( # frac_coords=torch.Tensor(frac_coords), # atom_types=torch.LongTensor(atom_types), # lengths=torch.Tensor(lengths).view(1, -1), # angles=torch.Tensor(angles).view(1, -1), # edge_index=torch.LongTensor( # edge_indices.T).contiguous(), # shape (2, num_edges) # to_jimages=torch.LongTensor(to_jimages), # num_atoms=num_atoms, # num_bonds=edge_indices.shape[0], # spacegroup=data_dict['spacegroup.number'], # pretty_formula=data_dict['pretty_formula'], # mpid=data_dict['mp_id'], # num_nodes=num_atoms, # special attribute used for batching in pytorch geometric # y=prop, # raw_sinc=raw_sinc, # raw_sinc_presubsample=raw_sinc_presubsample, # xrd_presubsample=xrd_presubsample, # raw_xrd=torch.tensor(data_dict['rawXRD']) # ) def optimization(args, model, ld_kwargs, data_loader): assert data_loader is not None downsampled_Qs = np.linspace( data_loader.dataset.Qs[0], data_loader.dataset.Qs[-1], data_loader.dataset.n_postsubsample ) # assert filtering matches the configs assert args.xrd_filter == data_loader.dataset.xrd_filter, "XRD filter in config does not match the one in the dataset" base_output_dir = f'{args.output_dir}/{args.label}' os.makedirs(base_output_dir, exist_ok=True) with open(os.path.join(base_output_dir, 'parameters.json'), 'w') as fout: json.dump(vars(args), fout, indent=4) filter_viz_folder = f'{base_output_dir}/filter_viz' os.makedirs(filter_viz_folder, exist_ok=True) # visualize filter and transform if args.xrd_filter == 'sinc' or args.xrd_filter == 'both': Qs_shifted, sinc_filter = data_loader.dataset.Qs_shifted, data_loader.dataset.sinc_filt plot_filter(filter=sinc_filter, Qs=Qs_shifted, filter_viz_folder=filter_viz_folder, nanomaterial_size=data_loader.dataset.nanomaterial_size) all_gt_crystals = list() all_bestPred_crystals = list() all_composition_errors = list() all_xrd_l1_errors = list() all_xrd_l2_errors = list() has_correct_num_atoms = list() all_pdf_correlations = list() all_r_factors = list() spacegroups = list() formula_strs = list() mpids = list() for j, batch in enumerate(data_loader): if j < args.first_idx: continue wandb.init(mode="disabled") # wandb.init(config=args, project='new conditional generation', name=f'crystal {j}', group=args.label) if j == args.num_tested_materials: break batch = batch.to(model.device) spacegroups.append(int(batch.spacegroup[0])) formula_strs.append(batch.pretty_formula[0]) mpids.append(batch.mpid[0]) curr_material_folder = f'{base_output_dir}/material{j}_{mpids[-1]}_{formula_strs[-1]}' os.makedirs(curr_material_folder, exist_ok=True) # get xrd xrd_dim = data_loader.dataset.n_postsubsample target_noisy_xrd = batch.y.reshape(1, xrd_dim) target_sincOnly = batch.raw_sinc.reshape(1, xrd_dim) raw_sinc = batch.raw_sinc.reshape(1, xrd_dim) gt_noiseless_xrd = batch.raw_xrd.reshape(1, xrd_dim) if args.start_from_init: init_material_folder = f'{args.start_from_init}/material{j}_{mpids[-1]}_{formula_strs[-1]}' cif_path = os.path.join(init_material_folder, 'pred', 'candidate0', 'cif', f'noSpacegroup_material{j}_candidate0.cif') z_init = create_z_from_init(args, batch, model, cif_path) else: z_init = None z = optimize_latent_code(args=args, model=model, batch=batch, target_noisy_xrd=target_noisy_xrd, z_init=z_init) # get predicted xrd for all optimized candidates final_pred_xrds = model.fc_property(z).reshape(-1, xrd_dim) # TODO: speed this one up init_num_atoms = batch.num_atoms.repeat(args.num_starting_points) if args.num_atom_lambda > EPS else None init_atom_types = batch.atom_types.repeat(args.num_starting_points) if args.composition_lambda > EPS else None print('know num atoms:', init_num_atoms is not None) print('know atom types:', init_atom_types is not None) crystals = model.langevin_dynamics(z, ld_kwargs, gt_num_atoms=init_num_atoms, gt_atom_types=init_atom_types) if not args.save_traj: crystals = {k: crystals[k] for k in ['frac_coords', 'atom_types', 'num_atoms', 'lengths', 'angles']} else: crystals = {k: crystals[k] for k in ['frac_coords', 'atom_types', 'num_atoms', 'lengths', 'angles', 'all_frac_coords', 'all_atom_types']} xrd_args = create_xrd_args(args) # predictions frac_coords = crystals['frac_coords'] num_atoms = crystals['num_atoms'] atom_types = crystals['atom_types'] lengths = crystals['lengths'] angles = crystals['angles'] all_opt_coords, all_opt_atom_types, opt_generated_xrds, curr_gen_crystals_list = create_materials(xrd_args, frac_coords, num_atoms, atom_types, lengths, angles, create_xrd=True, symprec=0.01) # plot base truth frac_coords = batch.frac_coords num_atoms = batch.num_atoms atom_types = batch.atom_types lengths = batch.lengths angles = batch.angles assert num_atoms.shape[0] == 1 assert frac_coords.shape[0] == atom_types.shape[0] the_coords, atom_types, bt_generated_xrds, singleton_gt_crystal_list = create_materials(xrd_args, frac_coords, num_atoms, atom_types, lengths, angles, create_xrd=True, symprec=0.01) the_coords = np.array(the_coords)[0] atom_types = np.array(atom_types)[0] assert len(singleton_gt_crystal_list) == 1 curr_gt_crystal = Crystal(singleton_gt_crystal_list[0]) all_gt_crystals.append(curr_gt_crystal) # save cif gt_cif_folder = os.path.join(curr_material_folder, 'gt', 'cif') os.makedirs(gt_cif_folder, exist_ok=True) curr_gt_crystal.structure.to(filename=f'{gt_cif_folder}/noSpacegroup_material{j}_{mpids[-1]}_{formula_strs[-1]}.cif', fmt='cif') # TODO: this will sometimes change the # of atoms in the outputted unit cell in the cif file gt_cif_writer = CifWriter(curr_gt_crystal.structure, symprec=0.01) gt_cif_writer.write_file(filename=f'{gt_cif_folder}/material{j}_{mpids[-1]}_{formula_strs[-1]}.cif') # plot image gt_vis_folder = os.path.join(curr_material_folder, 'gt', 'visUnitCell') os.makedirs(gt_vis_folder, exist_ok=True) gt_material_filepath = plot_material_single(the_coords, atom_types, gt_vis_folder, idx=j) # plot xrd gt_xrd_folder = os.path.join(curr_material_folder, 'gt', 'xrd') os.makedirs(gt_xrd_folder, exist_ok=True) gt_xrd_filepath = plot_xrd_single(xrd_args, target_noisy_xrd.squeeze().cpu().numpy(), gt_xrd_folder, idx=j, x_axis=downsampled_Qs, x_label=r'Q $({A^{\circ}}^{-1})$') # save the noisy xrd torch.save(target_noisy_xrd.squeeze().cpu(), os.path.join(gt_xrd_folder, f'material{j}.pt')) # save sinc only xrd torch.save(target_sincOnly.squeeze().cpu(), os.path.join(gt_xrd_folder, f'sincOnly{j}.pt')) # apply smoothing to the XRD patterns noiseless_generated_xrds = np.array([data_loader.dataset.sample(an_xrd) for an_xrd in opt_generated_xrds.tolist()]) opt_generated_xrds, opt_sinc_only_xrds = smooth_xrds(opt_generated_xrds=opt_generated_xrds, data_loader=data_loader) opt_generated_xrds = opt_generated_xrds.to(model.device) opt_sinc_only_xrds = opt_sinc_only_xrds.to(model.device) assert noiseless_generated_xrds.shape == opt_generated_xrds.shape assert noiseless_generated_xrds.shape == opt_sinc_only_xrds.shape # plot smoothed vs sinc: gt plot_smoothed_vs_sinc(smoothed=target_noisy_xrd, sincPattern=raw_sinc, noiselessPattern=gt_noiseless_xrd, Qs=downsampled_Qs, savepath=os.path.join(gt_xrd_folder, f'sincVsSmoothed{j}.png')) # compute loss on desired and generated xrds target = target_noisy_xrd.broadcast_to(bt_generated_xrds.shape[0], xrd_dim).to(model.device) loss = F.l1_loss(opt_generated_xrds.to(model.device), target.to(model.device), reduction='none').mean(dim=-1) if args.l1_loss \ else F.mse_loss(opt_generated_xrds.to(model.device), target.to(model.device), reduction='none').mean(dim=-1) # find the (num_candidates) minimum loss elements min_loss_indices = torch.argsort(loss.reshape(-1))[:args.num_candidates].tolist() if args.save_traj: n_steps = args.n_step_each * len(model.sigmas) for item in crystals: print(item) assert crystals['all_frac_coords'].shape[0] == n_steps, f"{crystals['all_frac_coords'].shape[0]} != {n_steps}" assert crystals['all_atom_types'].shape[0] == n_steps, f"{crystals['all_atom_types'].shape[0]} != {n_steps}" print(f'{n_steps} total steps: save traj') traj_folder = os.path.join(curr_material_folder, 'pred', 'diffusion_vis') os.makedirs(traj_folder, exist_ok=True) for step in range(0, n_steps, args.n_step_each): curr_frac_coords = crystals['all_frac_coords'][step] curr_atom_types = crystals['all_atom_types'][step] assert curr_frac_coords.shape == crystals['frac_coords'].shape, f"{curr_frac_coords.shape} != {crystals['frac_coords'].shape}" assert curr_atom_types.shape == crystals['atom_types'].shape, f"{curr_atom_types.shape} != {crystals['atom_types'].shape}" curr_step_coords, curr_step_atom_types, curr_step_xrds, curr_step_singleton_crystal_list = \ create_materials(xrd_args, curr_frac_coords, crystals['num_atoms'], curr_atom_types, crystals['lengths'], crystals['angles'], create_xrd=True, symprec=0.01) assert len(curr_step_singleton_crystal_list) == args.num_starting_points curr_step_crystal = Crystal(curr_step_singleton_crystal_list[min_loss_indices[0]]) # just save 1 # TODO: save these curr_step_crystal.structure.to(filename=f'{traj_folder}/step{step}_material{j}_candidate{0}_{mpids[-1]}_{formula_strs[-1]}.cif', fmt='cif') process_candidates(args=args, xrd_args=xrd_args, j=j, curr_gen_crystals_list=curr_gen_crystals_list, all_opt_coords=all_opt_coords, all_opt_atom_types=all_opt_atom_types, opt_generated_xrds=opt_generated_xrds, min_loss_indices=min_loss_indices, curr_material_folder=curr_material_folder, all_bestPred_crystals=all_bestPred_crystals, gt_noiseless_xrd=gt_noiseless_xrd, target_noisy_xrd=target_noisy_xrd, final_pred_xrds=final_pred_xrds, opt_sinc_only_xrds=opt_sinc_only_xrds, noiseless_generated_xrds=noiseless_generated_xrds, curr_gt_crystal=curr_gt_crystal, gt_atom_types=atom_types, gt_material_filepath=gt_material_filepath, gt_xrd_filepath=gt_xrd_filepath, all_xrd_l1_errors=all_xrd_l1_errors, all_xrd_l2_errors=all_xrd_l2_errors, all_composition_errors=all_composition_errors, has_correct_num_atoms=has_correct_num_atoms, all_pdf_correlations=all_pdf_correlations, all_r_factors=all_r_factors, Qs=downsampled_Qs) ret_val = dict() for curr_spacegroup in set([USE_ALL_SPACEGROUPS] + spacegroups): if curr_spacegroup != USE_ALL_SPACEGROUPS: curr_count = np.sum(np.array(spacegroups) == curr_spacegroup) if curr_count == 0 or curr_count == len(spacegroups): continue curr_results = calculate_metrics(all_gt_crystals=all_gt_crystals, all_bestPred_crystals=all_bestPred_crystals, all_xrd_l1_errors=all_xrd_l1_errors, all_xrd_l2_errors=all_xrd_l2_errors, all_composition_errors=all_composition_errors, has_correct_num_atoms=has_correct_num_atoms, all_pdf_correlations=all_pdf_correlations, all_r_factors=all_r_factors, spacegroups=spacegroups, desired_spacegroup=curr_spacegroup) ret_val[curr_spacegroup] = curr_results metrics_folder = os.path.join(base_output_dir, 'metrics') os.makedirs(metrics_folder, exist_ok=True) with open(f'{metrics_folder}/aggregate_metrics.json', 'w') as fout: json.dump(ret_val, fout, indent=4) write_pdf_histogram(pdf_rs=np.array(all_pdf_correlations).flatten(), save_folder=metrics_folder, title='All Predicted PDFs vs. Ground Truth') write_pdf_histogram(pdf_rs=np.max(np.array(all_pdf_correlations), axis=1), save_folder=metrics_folder, title='Best PDFs (per Material) vs. Ground Truth') write_r_factor_histogram(r_factors=np.array(all_r_factors).flatten(), save_folder=metrics_folder, title='All Predicted R-Factors vs. Ground Truth') write_r_factor_histogram(r_factors=np.max(np.array(all_r_factors), axis=1), save_folder=metrics_folder, title='Best R-Factors (per Material) vs. Ground Truth') print(json.dumps(ret_val, indent=4)) return ret_val def calculate_metrics(all_gt_crystals, all_bestPred_crystals, all_xrd_l1_errors, all_xrd_l2_errors, all_composition_errors, has_correct_num_atoms, all_pdf_correlations, all_r_factors, spacegroups, desired_spacegroup): # turn into numpy arrays spacegroups = np.array(spacegroups) all_gt_crystals = np.array(all_gt_crystals) all_bestPred_crystals = np.array(all_bestPred_crystals) all_xrd_l1_errors = np.array(all_xrd_l1_errors) all_xrd_l2_errors = np.array(all_xrd_l2_errors) all_composition_errors = np.array(all_composition_errors) has_correct_num_atoms = np.array(has_correct_num_atoms) all_pdf_correlations = np.array(all_pdf_correlations) all_r_factors = np.array(all_r_factors) num_materials_in_spacegroup = len(spacegroups) if desired_spacegroup != USE_ALL_SPACEGROUPS: index_mask = spacegroups == desired_spacegroup assert len(index_mask) == len(spacegroups) assert np.sum(index_mask) > 0 and np.sum(index_mask) < len(spacegroups) all_gt_crystals = all_gt_crystals[index_mask] all_bestPred_crystals = all_bestPred_crystals[index_mask] all_xrd_l1_errors = all_xrd_l1_errors[index_mask] all_xrd_l2_errors = all_xrd_l2_errors[index_mask] all_composition_errors = all_composition_errors[index_mask] has_correct_num_atoms = has_correct_num_atoms[index_mask] all_pdf_correlations = all_pdf_correlations[index_mask] all_r_factors = all_r_factors[index_mask] num_materials_in_spacegroup = np.sum(index_mask) # average xrd errors avg_xrd_mse = np.mean(all_xrd_l2_errors) avg_xrd_l1 = np.mean(all_xrd_l1_errors) avg_pdf_correlation = np.mean(all_pdf_correlations) std_pdf_correlation = np.std(all_pdf_correlations) avg_r_factor = np.mean(all_r_factors) std_r_factor = np.std(all_r_factors) # best of candidate xrd errors best_xrd_mse = np.mean([np.min(list) for list in all_xrd_l2_errors]) best_xrd_l1 = np.mean([np.min(list) for list in all_xrd_l1_errors]) best_pdf_correlation = np.mean(np.max(all_pdf_correlations, axis=1)) std_best_pdf_correlation = np.std(np.max(all_pdf_correlations, axis=1)) best_r_factor = np.mean(np.min(all_r_factors, axis=1)) std_best_r_factor = np.std(np.min(all_r_factors, axis=1)) ret_val = { COUNT: int(num_materials_in_spacegroup), AVG_COMPOSITION_ERROR: np.mean(all_composition_errors), AVG_XRD_MSE: avg_xrd_mse, AVG_XRD_L1: avg_xrd_l1, BEST_XRD_MSE: best_xrd_mse, BEST_XRD_L1: best_xrd_l1, AVG_PDF_CORRELATION: avg_pdf_correlation, BEST_PDF_CORRELATION: best_pdf_correlation, STD_PDF_CORRELATION: std_pdf_correlation, STD_BEST_PDF_CORRELATION: std_best_pdf_correlation, AVG_R_FACTOR: avg_r_factor, STD_R_FACTOR: std_r_factor, BEST_R_FACTOR: best_r_factor, STD_BEST_R_FACTOR: std_best_r_factor } ret_val.update(check_structure_match(gt_structures=all_gt_crystals, pred_structures=all_bestPred_crystals)) ret_val.update(check_validity(gt_structures=all_gt_crystals, pred_structures=all_bestPred_crystals)) ret_val[NUM_ATOM_ACCURACY] = np.sum(has_correct_num_atoms) \ / (has_correct_num_atoms.shape[0] * has_correct_num_atoms.shape[1]) return ret_val def get_elemental_ratios(atom_types): one_hot = np.zeros(119+1) for atom in atom_types: one_hot[atom.Z] += 1 one_hot /= np.sum(one_hot) return one_hot def compare_composition(gt_atom_types, pred_atom_types): gt_atom_types = get_elemental_ratios(gt_atom_types) pred_atom_types = get_elemental_ratios(pred_atom_types) result = np.sum(np.abs(gt_atom_types - pred_atom_types)) / 2 assert result >= -1e-5 and result <= 1 + 1e-5 return result def compare_num_atoms(gt_atom_types, pred_atom_types): if isinstance(gt_atom_types, list): assert not isinstance(gt_atom_types[0], list) else: assert len(gt_atom_types.shape) == 1 if isinstance(pred_atom_types, list): assert not isinstance(pred_atom_types[-1], list) else: assert len(pred_atom_types.shape) == 1 return len(gt_atom_types) == len(pred_atom_types) def check_structure_match(gt_structures, pred_structures): """ Input: lists of Crystal() objects return {'match_rate': match_rate, 'rms_dist': mean_rms_dist} """ structure_evaluator = RecEval(pred_crys=pred_structures, gt_crys=gt_structures, stol=0.3, angle_tol=5, ltol=0.2) match_rate_rms_dict = structure_evaluator.get_match_rate_and_rms() return match_rate_rms_dict def check_validity(gt_structures, pred_structures): """ Input: lists of Crystal() objects return {'comp_valid': comp_valid, 'struct_valid': struct_valid, 'valid': valid} """ validity_checker = GenEval(pred_crys=pred_structures, gt_crys=gt_structures, n_samples=len(gt_structures)) return validity_checker.get_validity() def main(args): # load_data if do reconstruction. model_path = Path(args.model_path) model, test_loader, cfg = load_model( model_path, load_data=True, batch_size=1) ld_kwargs = SimpleNamespace(n_step_each=args.n_step_each, step_lr=args.step_lr, min_sigma=args.min_sigma, save_traj=args.save_traj, disable_bar=args.disable_bar) if torch.cuda.is_available(): model.to('cuda') print('Evaluate model on the property optimization task.') if args.start_from == 'data': loader = test_loader else: loader = None optimization(args=args, model=model, ld_kwargs=ld_kwargs, data_loader=loader) return if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('--model_path', required=True) parser.add_argument('--data_dir', default='data', type=str) parser.add_argument('--n_step_each', default=100, type=int) parser.add_argument('--step_lr', default=1e-4, type=float) parser.add_argument('--save_traj', action='store_true') parser.add_argument('--min_sigma', default=0, type=float) parser.add_argument('--disable_bar', default=False, type=bool) parser.add_argument('--start_from', default='data', type=str) parser.add_argument('--l2_penalty', default=1e-5, type=float) parser.add_argument('--num_atom_lambda', default=1e-3, type=float) parser.add_argument('--lattice_lambda', default=1e-3, type=float) parser.add_argument('--composition_lambda', default=1e-3, type=float) parser.add_argument('--num_starting_points', default=1000, type=int) parser.add_argument('--num_gradient_steps', default=5000, type=int) parser.add_argument('--lr', default=1e-3, type=float) parser.add_argument('--min_lr', default=1e-5, type=float) parser.add_argument('--num_tested_materials', default=10, type=int) parser.add_argument('--l1_loss', action='store_true') parser.add_argument('--label', default='') parser.add_argument('--num_candidates', default=5, type=int) parser.add_argument('--xrd_filter', default='both') parser.add_argument('--output_dir', default='materials_viz', type=str) parser.add_argument('--first_idx', default=0, type=int) parser.add_argument('--r_min', default=0, type=float) parser.add_argument('--r_max', default=30, type=float) parser.add_argument('--wave_source', default='CuKa', type=str) parser.add_argument('--xrd_vector_dim', default=4096, type=int) parser.add_argument('--min_theta', default=0, type=float) parser.add_argument('--max_theta', default=180, type=float) parser.add_argument('--start_from_init', default=None, type=str) args = parser.parse_args() print('starting eval', args) main(args)