mxtaltools.crystal_search.crystal_opt_utils

class mxtaltools.crystal_search.crystal_opt_utils.CrystalParams(init_sample, fixed_dims=None)[source]

update_params(new_values, mask=None)[source]: new_values: [batch_size, param_dim] tensor mask: optional boolean mask [batch_size] to only update accepted samples

mxtaltools.crystal_search.crystal_opt_utils.check_convergence(params_record, s_ind, convergence_eps, optimizer, init_lr)[source]

mxtaltools.crystal_search.crystal_opt_utils.compute_auxiliary_loss(cluster_batch, loss, outputs, config)[source]

mxtaltools.crystal_search.crystal_opt_utils.compute_loss(cluster_batch, crystal_batch, outputs, config, opt_step)[source]

mxtaltools.crystal_search.crystal_opt_utils.dict2namespace(data_dict: dict)[source]

Recursively converts a dictionary and its internal dictionaries into an argparse.Namespace

mxtaltools.crystal_search.crystal_opt_utils.ema_trajectory(traj: Tensor, alpha: float = 0.1) → Tensor[source]: Vectorized EMA along time (dim=0). traj: [T, N, D]

mxtaltools.crystal_search.crystal_opt_utils.get_annealing_factor(start_value, stop_value, total_time, step_iters)[source]

mxtaltools.crystal_search.crystal_opt_utils.get_mc_noised_samples(crystal_batch, log_noise_range)[source]

mxtaltools.crystal_search.crystal_opt_utils.gradient_descent_optimization(init_sample: Tensor, init_crystal_batch, optimizer_func: Optimizer | str, init_lr: float | None = 0.1, max_num_steps: int | None = 500, show_tqdm: bool | None = False, convergence_eps: float = 0.001, grad_norm_clip: float | None = 0.1, optim_target: str | None = 'lj', score_model: Module | None = None, target_packing_coeff: Tensor | None = None, do_box_restriction: bool | None = False, cutoff: float | None = None, compression_factor: float | None = 0, enforce_reduced: bool | None = False, supercell_size: int | None = 10, anneal_lr: bool | None = False, predictor: Optional = None, mc_kT: float = 0.5, log_noise_range=[-2, -1], target_rdf: tensor = None, target_latent: tensor = None, elementwise: bool = True, atomwise: bool = False, repulsion: float = 1.0, umbrella: bool | None = False, umbrella_sigma: float | None = None, umbrella_epsilon: float | None = None, umbrella_record: list | None = None, rdf_warmup: tensor | None = 500)[source]: do a local optimization via gradient descent on some score function

mxtaltools.crystal_search.crystal_opt_utils.init_mc_state(compression_factor, cutoff, do_box_restriction, energy_computes, enforce_reduced, init_crystal_batch, mc_kT, optim_target, param_module, score_model, std_margin, supercell_size, target_packing_coeff, uma_predictor, target_rdf, target_latent)[source]

mxtaltools.crystal_search.crystal_opt_utils.init_opt(init_lr, optimizer_func, param_module)[source]

mxtaltools.crystal_search.crystal_opt_utils.inter_overlap_loss(cluster_batch)[source]

mxtaltools.crystal_search.crystal_opt_utils.loss_and_backprop(cluster_batch, crystal_batch, grad_norm_clip, optimizer, outputs, param_module, records, loss_config, aux_config, opt_step)[source]

mxtaltools.crystal_search.crystal_opt_utils.parse_mc_step(T, cluster_batch, compression_factor, crystal_batch, current_energy, current_state, do_box_restriction, enforce_reduced, optim_target, outputs, param_module, records, score_model, target_packing_coeff, target_rdf, target_latent, cutoff)[source]

mxtaltools.crystal_search.crystal_opt_utils.traj_fig(x, y, names=[None, None], yrange=None, xrange=None)[source]

mxtaltools.crystal_search.crystal_opt_utils.update_record(crystal_batch, outputs, params_record, records, s_ind)[source]