:> various analyses

:> var boosted training change and configs
> search updates

Author: InfluenceFunctional

configs/crystal_searches/base.yaml

@@ -19,12 +19,12 @@ mol_seed: 0
 opt_seed: 0
 sampling_mode: all # 'all' or 'random' or 'ordered'
 mols_to_sample: 2
-num_samples: 1000 # per mol, per space group, per Zp
+num_samples: 10000 # per mol, per space group, per Zp
 
 sgs_to_search: [ 14 ]
 zp_to_search: [ 1 ]
 
-batch_size: 50
+batch_size: 2000
 grow_batch_size: false
 
 init_target_cp: 0.6828  # a float, null (use target value), 'wide' (0.0-0.7) or 'std' target packing coefficient for initial structure sampling
@@ -56,15 +56,15 @@ opt:
   - optim_target: 'rdf_dist' # lj qlj elj silu ellipsoid classification_score rdf_score rdf_dist latent_dist
     enforce_reduced: false
     compression_factor: 0.0
-    cutoff: 4  # can be as low as 6 for SiLU, 10 otherwise
+    cutoff: 10 # can be as low as 6 for SiLU, 10 otherwise
     init_lr: 0.001
-    convergence_eps: 0.000001
+    convergence_eps: 0.0001
     optimizer_func: 'rprop'  # NOTE rprop is by far the fastest and most reliable
     anneal_lr: false
     grad_norm_clip: 0.1
     show_tqdm: true
-    max_num_steps: 40000
-    rdf_warmup: 100
+    max_num_steps: 500
+    rdf_warmup: null
     target_packing_coeff: 0.6828
     umbrella: false
     umbrella_sigma: 0.25

examples/crystal_search_reporting.py

@@ -195,14 +195,14 @@ def compack_fig(matches, rmsds, write_fig):
         fig.write_image(r'C:\Users\mikem\OneDrive\NYU\CSD\papers\mxt_code\compack_fig.png', width=900, height=900)
 
 
-def batch_compack(best_sample_inds, optimized_samples, reference_cluster_batch): # todo refactor into analysis code
+def batch_compack(best_sample_inds, optimized_samples, reference_cluster_batch, ref_ind: int = 0): # todo refactor into analysis code
     # generate the crystals in ccdc format
     best_crystals_batch = collate_data_list([optimized_samples[ind] for ind in best_sample_inds])
-    best_cluster_batch = best_crystals_batch.mol2cluster().to('cpu')
+    best_cluster_batch = best_crystals_batch.mol2cluster(cutoff=10).to('cpu')
     _ = cluster_batch_to_ccdc_crystals(best_cluster_batch, np.arange(best_cluster_batch.num_graphs))
-    mol = ase_mol_from_crystaldata(reference_cluster_batch, index=0, mode='unit cell')
-    #mol.info['spacegroup'] = Spacegroup(int(best_cluster_batch.sg_ind[0]), setting=1)
-    mol.write('DAFMUV.cif')
+    mol = ase_mol_from_crystaldata(reference_cluster_batch, index=ref_ind, mode='unit cell')
+    mol.info['spacegroup'] = Spacegroup(int(reference_cluster_batch.sg_ind[ref_ind]), setting=1)
+    mol.write('compack_placeholder.cif')
 
     print(f"Running COMPACK on {len(best_sample_inds)} crystals")
     pool = mp.Pool(8)
@@ -225,7 +225,7 @@ def batch_compack(best_sample_inds, optimized_samples, reference_cluster_batch):
 
 
 def single_compack_run(ind):
-    ref_crystal = CrystalReader('DAFMUV.cif')[0]
+    ref_crystal = CrystalReader('compack_placeholder.cif')[0]
     sample_crystal = CrystalReader(f'temp_{ind}.cif')[0]
     similarity_engine = PackingSimilarity()
     similarity_engine.settings.distance_tolerance = 0.4

mxtaltools/common/geometry_utils.py

@@ -1663,8 +1663,8 @@ def compute_latent_distance(latents1: torch.Tensor,
 
     rot_dists = []
     for zp in range(z_prime):  # this should be replaced with a proper vector distance
-        rmat1 = rotvec2rotmat(lat_sph_rotvec1[...,3 * zp:3 * zp + 3], 'spherical')
-        rmat2 = rotvec2rotmat(lat_sph_rotvec2[...,3 * zp:3 * zp + 3], 'spherical')
+        rmat1 = rotvec2rotmat(lat_sph_rotvec1[..., 3 * zp:3 * zp + 3], 'spherical')
+        rmat2 = rotvec2rotmat(lat_sph_rotvec2[..., 3 * zp:3 * zp + 3], 'spherical')
 
         R_delta = rmat1 @ rmat2.transpose(-1, -2)
 
@@ -1676,10 +1676,11 @@ def compute_latent_distance(latents1: torch.Tensor,
     rot_dists = torch.stack(rot_dists).sum(0)
 
     "overall distance metric"
-    #dists = 0.5 * box_dist + 0.25 * (positions_dist / z_prime / 2.5) + 0.25 * (rot_dists / z_prime / 2)
-    #scales = [2 * sqrt(6), 2*sqrt(3), torch.pi] # maximum variation per dist
-    scales = [1, 0.836, 0.293]# [0.0127, 0.0152, 0.0433]  # empirical std over CSD samples
-    dists = scales[0] * 0.5 * box_dist + scales[1] * 0.25 * (positions_dist / z_prime) + scales[2] * 0.25 * (rot_dists / z_prime)
+    # dists = 0.5 * box_dist + 0.25 * (positions_dist / z_prime / 2.5) + 0.25 * (rot_dists / z_prime / 2)
+    # scales = [2 * sqrt(6), 2*sqrt(3), torch.pi] # maximum variation per dist
+    scales = [1, 0.836, 0.293]  # [0.0127, 0.0152, 0.0433]  # empirical std over CSD samples
+    dists = scales[0] * 0.5 * box_dist + scales[1] * 0.25 * (positions_dist / z_prime) + scales[2] * 0.25 * (
+                rot_dists / z_prime)
 
     return dists
 

mxtaltools/crystal_search/crystal_opt_utils.py

@@ -11,16 +11,13 @@
 from torch_scatter import scatter
 from tqdm import tqdm
 
-
 from mxtaltools.analysis.crystal_rdf import compute_rdf_distance
 from mxtaltools.common.geometry_utils import enforce_crystal_system
 from mxtaltools.common.utils import is_cuda_oom
 from mxtaltools.dataset_utils.utils import collate_data_list
 from mxtaltools.models.utils import enforce_1d_bound, softmax_and_score
 
 
-
-
 def dict2namespace(data_dict: dict):
     """
     Recursively converts a dictionary and its internal dictionaries into an
@@ -131,16 +128,16 @@ def gradient_descent_optimization(  # todo consolidate kwargs somewhere
         # lennard jones need 10 angstroms to nicely converge
         cutoff = 10
 
-    energy_computes = ['lj']
+    energy_computes = ['lj', 'elj']
     min_num_steps = 50
     num_samples = init_crystal_batch.num_graphs
 
     if optim_target.lower() == 'silu':
         energy_computes.append('silu')
     elif optim_target.lower() == 'qlj':
         energy_computes.append('qlj')
-    elif optim_target.lower() == 'elj':
-        energy_computes.append('elj')
+    # elif optim_target.lower() == 'elj': # always do this
+    #     energy_computes.append('elj')
     elif optim_target.lower() == 'ellipsoid':
         energy_computes.append('ellipsoid')
     elif optim_target.lower() == 'reduce':
@@ -161,7 +158,7 @@ def gradient_descent_optimization(  # todo consolidate kwargs somewhere
     if target_latent is not None:
         target_latent = target_latent.to(init_sample.device)
 
-    if False: #target_rdf is not None:  # assumes we already have the box
+    if False:  # target_rdf is not None:  # assumes we already have the box
         fixed_dims = [0, 1, 2, 3, 4, 5]
     else:
         fixed_dims = None
@@ -280,8 +277,8 @@ def gradient_descent_optimization(  # todo consolidate kwargs somewhere
 
                     scheduler1.step()  # shrink
                     s_ind += 1
-                    if s_ind % 50 == 0:
-                        pbar.update(50)
+                    if s_ind % 10 == 0:
+                        pbar.update(10)
                     if s_ind >= min(max_num_steps, max(50, min_num_steps)):
                         converged = check_convergence(params_record, s_ind, convergence_eps,
                                                       optimizer, init_lr)
@@ -580,25 +577,32 @@ def compute_loss(cluster_batch, crystal_batch, outputs, config, opt_step):
     elif config.optim_target.lower() == 'rdf_dist':
         n_channels = config.target_rdf.shape[-2]  # 120
         if config.rdf_warmup is not None:
-            # channel_warmup = config.rdf_warmup
-            # channel_onsets = torch.linspace(0, channel_warmup, n_channels)  # evenly spaced turn-on times
-            # channel_weights = torch.sigmoid((opt_step - channel_onsets) / (channel_warmup / n_channels * 0.5))
-
-            n_waves = 3
-            base_periods = torch.tensor([1.0, 1.6, 2.5]) * config.rdf_warmup
-            channel_idx = torch.arange(n_channels, dtype=torch.float32)
-
-            modulation = torch.zeros(n_channels)
-            for i in range(n_waves):
-                modulation += torch.sin(
-                    2 * torch.pi * opt_step / base_periods[i] + 2 * torch.pi * channel_idx / n_channels * (i + 1))
-            modulation = modulation / n_waves
-            channel_weights = 0.5 + 0.5 * modulation
+            channel_warmup = config.rdf_warmup
+            channel_onsets = torch.linspace(0, channel_warmup, n_channels)  # evenly spaced turn-on times
+            channel_weights = torch.sigmoid((opt_step - channel_onsets) / (channel_warmup / n_channels * 0.5))
+
+            # n_waves = 3
+            # base_periods = torch.tensor([1.0, 1.6, 2.5]) * config.rdf_warmup
+            # channel_idx = torch.arange(n_channels, dtype=torch.float32)
+            #
+            # modulation = torch.zeros(n_channels)
+            # for i in range(n_waves):
+            #     modulation += torch.sin(
+            #         2 * torch.pi * opt_step / base_periods[i] + 2 * torch.pi * channel_idx / n_channels * (i + 1))
+            # modulation = modulation / n_waves
+            # channel_weights = 0.5 + 0.5 * modulation
         else:
             channel_weights = torch.ones(n_channels)
-        loss = compute_rdf_distance(outputs['rdf'][0], config.target_rdf,
-                                    torch.linspace(0, config.cutoff, config.target_rdf.shape[-1]),
-                                    channel_weights=channel_weights)
+        rdf_loss = compute_rdf_distance(outputs['rdf'][0], config.target_rdf,
+                                        torch.linspace(0, config.cutoff, config.target_rdf.shape[-1]),
+                                        channel_weights=channel_weights)
+
+        en_cut = -307  # set equal or higher to the target energy
+        beta = 5
+        lj_en = outputs['elj']
+        lj_loss = F.softplus(beta * (lj_en - en_cut)) / beta
+
+        loss = rdf_loss + lj_loss / 100
 
     elif config.optim_target.lower() == 'latent_dist':
         loss = (config.target_latent - crystal_batch.latent_params()).norm(dim=-1)
@@ -637,7 +641,7 @@ def compute_auxiliary_loss(cluster_batch, loss, outputs, config):
             record = config.umbrella_record.to(cluster_batch.device)
             dists = torch.cdist(latents, record)
             penalty = torch.exp(-dists ** 2 / (2 * config.umbrella_sigma ** 2)).sum(dim=1).clip(max=10)
-            loss = loss + config.umbrella_epsilon *  penalty
+            loss = loss + config.umbrella_epsilon * penalty
 
     return loss
 

mxtaltools/crystal_search/run_search.py

@@ -15,16 +15,7 @@
     recover_opt_state, process_target, save_umbrella_record
 from mxtaltools.dataset_utils.utils import collate_data_list
 
-if __name__ == '__main__':
-    args = parse_args()  # call config with "python run_search.py --config /path/to/config.yaml
-    source_dir = Path(__file__).resolve().parent.parent.parent
-    if args.config is None:
-        config_path = source_dir / 'configs' / 'crystal_searches' / 'base.yaml'
-    else:
-        config_path = Path(args.config)
-
-    config = dict2namespace(load_yaml(config_path))
-
+def crystal_search(config):
     device = config.device
     umbrella_path = config.umbrella_path
 
@@ -102,14 +93,6 @@
                     umbrella_record = torch.load(umbrella_path, weights_only=False)
                     save_umbrella_record(umbrella_record, new_latents, umbrella_path, opt_config['umbrella_sigma'], opt_config['umbrella_epsilon'])
 
-                # cursor = 0
-                # bsz = config.batch_size
-                # while cursor < len(opt_outs):
-                #     batch = collate_data_list(opt_outs[cursor:cursor + bsz])
-                #     en = batch.elj
-                #     print([en.quantile(ii) for ii in torch.linspace(0, 1, 10)])
-                #     cursor += bsz
-
             cursor += config.batch_size
             prev_best_samples = None
             pbar.update(min(config.batch_size, num_samples - cursor))  # safe final update
@@ -143,6 +126,9 @@
             else:
                 raise e
 
+
+    return opt_outs
+
     print(f"Sampling complete! Optimized a total of {len(opt_outs)} crystal samples.")
 
     # batch = collate_data_list(opt_outs)
@@ -151,6 +137,19 @@
     # batch.plot_batch_density_funnel(split_by_sg=True)
 
     aa = 1
+
+if __name__ == '__main__':
+    args = parse_args()  # call config with "python run_search.py --config /path/to/config.yaml
+    source_dir = Path(__file__).resolve().parent.parent.parent
+    if args.config is None:
+        config_path = source_dir / 'configs' / 'crystal_searches' / 'base.yaml'
+    else:
+        config_path = Path(args.config)
+
+    config = dict2namespace(load_yaml(config_path))
+
+    crystal_search(config)
+
 """
 
 

mxtaltools/crystal_search/utils.py

@@ -162,7 +162,7 @@ def coarse_crystal_filter(lj_record, lj_cutoff, packing_coeff_record, packing_cu
 
 def get_initial_state(config, crystal_batch, device, batch_idx, target):
     # sample initial parameters
-    if config.init_sample_method == 'data':
+    if config.init_sample_method == 'data' or config.init_sample_method == 'in_config':
         return crystal_batch
 
     if config.init_target_cp == 'std':
@@ -236,6 +236,9 @@ def init_samples_to_optim(config, target=None):
             config.num_samples)]  # torch.arange(config.mol_seed * config.num_samples, (config.mol_seed + 1) * config.num_samples)
         samples_to_optim = [samples_to_optim[ind] for ind in index_block]
         return samples_to_optim
+    elif config.init_sample_method == 'in_config':
+        samples_to_optim = config.samples_to_optim
+        return samples_to_optim
     else:
         if target is None:
             mol_list = torch.load(config.mol_path, weights_only=False)

mxtaltools/dataset_utils/data_class_methods/crystal_ops.py

@@ -963,25 +963,36 @@ def destandardize_aunit_orientation(self, std_aunit_orientation):
         orientation_means = torch.tensor([[0, 0, torch.pi / 2]], dtype=torch.float32, device=self.device)
         return std_aunit_orientation * orientation_stds + orientation_means
 
-    def _build_feature_labels(self):
+    def _build_feature_labels(self, space):
         lattice_features = ['a', 'b', 'c',
                             r'$\alpha$', r'$\beta$', r'$\gamma$']
         if self.max_z_prime == 1:
             lattice_features.extend([
                 f'u', f'v', f'w',
             ])
-            lattice_features.extend([
-                f'x', f'y', f'z'
-            ])
+            if space == 'latent':
+                lattice_features.extend([
+                    f'θ', f'φ', f'r'
+                ])
+            else:
+                lattice_features.extend([
+                    f'x', f'y', f'z'
+                ])
+
         else:
             for zp in range(self.max_z_prime):
                 lattice_features.extend([
                     f'aunit{zp} u', f'aunit{zp} v', f'aunit{zp} w',
                 ])
             for zp in range(self.max_z_prime):
-                lattice_features.extend([
-                    f'x{zp}', f'y{zp}', f'z{zp}'
-                ])
+                if space == 'latent':
+                    lattice_features.extend([
+                        f'θ{zp}', f'φ{zp}', f'r{zp}'
+                    ])
+                else:
+                    lattice_features.extend([
+                        f'x{zp}', f'y{zp}', f'z{zp}'
+                    ])
         return lattice_features
 
     def _set_cell_ranges(self, space, samples):
@@ -1225,7 +1236,7 @@ def plot_batch_cell_params(self, space='real',
             print("Cell statistics only works for a batch of crystal data objects")
             return None
 
-        lattice_features = self._build_feature_labels()
+        lattice_features = self._build_feature_labels(space=space)
         samples = self._get_samples(space)
         num_dists, dist_names, dists = self._collect_sample_dists(samples, ref_dist, quantiles, split_by_sg,
                                                                   split_by_zp, aux_dists, override_energy)
@@ -1270,6 +1281,10 @@ def plot_batch_cell_params(self, space='real',
         #     showgrid=False, zeroline=False, ticks='outside',
         #     tickwidth=1, mirror=True
         # )
+        fig.update_yaxes(
+            showgrid=False, zeroline=False, showticklabels=False, ticks='',
+            mirror=True
+        )
         if len(dists) > 1:
             fig.update_traces(opacity=0.5)
         if show:
@@ -1290,7 +1305,7 @@ def plot_batch_staircase(self, space='real',
                              ref_dist=None,
                              ):
 
-        labels = self._build_feature_labels()
+        labels = self._build_feature_labels(space=space)
         samples = self._get_samples(space)
         if torch.is_tensor(samples):
             samples = samples.detach().cpu().numpy()