:> GFN configs

:> search updates
:> analysis updates

Author: InfluenceFunctional

configs/crystal_searches/base.yaml

@@ -1,33 +1,33 @@
 device: cuda
-mol_path: D:\crystal_datasets\nehzor\NEHZOR0_std_conf.pt
-dataset_path: D:\crystal_datasets\nehzor\NEHZOR_structures_std_conf.pt #D:\crystal_datasets\xuldud\xuldud_sg61_zp1.pt
-target_path: D:\crystal_datasets\nehzor\NEHZOR0_std_conf.pt #D:\crystal_datasets\mipcas\MIPCAS.pt
-umbrella_path: D:\crystal_datasets\opt_outputs\nehzor_test_umbrella.pt # null
-target_identifier: NEHZOR #MIPCAS
+mol_path: D:\crystal_datasets\nacjaf\nacjaf_nikos.pt
+dataset_path: null #D:\crystal_datasets\nacjaf\nacjaf_nikos.pt #D:\crystal_datasets\xuldud\xuldud_sg61_zp1.pt
+target_path: D:\crystal_datasets\nacjaf\nacjaf_nikos.pt
+umbrella_path: null #D:\crystal_datasets\opt_outputs\nehzor_test_umbrella.pt # null
+target_identifier: Target_XXII_10581_0242_0005_P1211_02_01_07_1 #MIPCAS
 out_dir: D:\crystal_datasets\opt_outputs
-run_name: NEHZOR_rdf
+run_name: NACJAF_chirality
 save_trajs: false
 
 uma_predictor_path: D:\crystal_datasets\esen_s.pt
 
 init_sample_method: random  # 'random' 'reasonable' 'data'
-init_reduced: true  # initialize in our reduced frame
+init_reduced: false  # initialize in our reduced frame
 
 # if method is 'data' optimize provided pre-built crystals
 # else, from random initial conditions
 mol_seed: 0
-opt_seed: 7
+opt_seed: 0
 sampling_mode: all # 'all' or 'random' or 'ordered'
 mols_to_sample: 2
-num_samples: 10000 # per mol, per space group, per Zp
+num_samples: 1000 # per mol, per space group, per Zp
 
 sgs_to_search: [ 14 ]
 zp_to_search: [ 1 ]
 
 batch_size: 50
 grow_batch_size: false
 
-init_target_cp: wide  # a float, null (use target value), 'wide' (0.0-0.7) or 'std' target packing coefficient for initial structure sampling
+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
 
 # Optimization can have multiple stages by duplicating the opt block below and adjusting params
 # the search algo will run consecutive n consecutive optimizations for n the number of list elements
@@ -53,18 +53,19 @@ init_target_cp: wide  # a float, null (use target value), 'wide' (0.0-0.7) or 's
 #        umbrella_epsilon: Optional[float] = None, # repulsion term for umbrella sampling
 
 opt:
-  - optim_target: 'uma' # lj qlj elj silu ellipsoid classification_score rdf_score rdf_dist latent_dist
-    enforce_reduced: true
-    compression_factor: 1.0
-    cutoff: 10  # can be as low as 6 for SiLU, 10 otherwise
-    init_lr: 0.1
-    convergence_eps: 0.001
+  - 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
+    init_lr: 0.001
+    convergence_eps: 0.000001
     optimizer_func: 'rprop'  # NOTE rprop is by far the fastest and most reliable
-    anneal_lr: true
+    anneal_lr: false
     grad_norm_clip: 0.1
     show_tqdm: true
-    max_num_steps: 250
-    target_packing_coeff: null
+    max_num_steps: 40000
+    rdf_warmup: 100
+    target_packing_coeff: 0.6828
     umbrella: false
     umbrella_sigma: 0.25
     umbrella_epsilon: 40.0

mxtaltools/analysis/crystal_rdf.py

@@ -475,7 +475,7 @@ def compute_rdf_distmat_parallel(rdf_record, rr, num_cpus, chunk_size=250):
     return rdf_dists
 
 
-def compute_rdf_distance(rdf1, rdf2, rr, n_parallel_rdf2: int = None, return_numpy: bool = False):
+def compute_rdf_distance(rdf1, rdf2, rr, n_parallel_rdf2: int = None, return_numpy: bool = False, channel_weights: torch.tensor=None):
     """
     Compute a distance metric between two radial distribution functions including sub_rdfs where sub_rdfs are e.g., particular interatomic RDFS within a certain sample (elementwise or atomwise modes).
 
@@ -511,7 +511,7 @@ def compute_rdf_distance(rdf1, rdf2, rr, n_parallel_rdf2: int = None, return_num
     else:
         torch_rdf1_f = torch_rdf1
 
-    bin_range = (torch_range[-1] - torch_range[0])
+    #bin_range = (torch_range[-1] - torch_range[0])
     bin_width = torch_range[1] - torch_range[0]
 
     # RDF should measure how much mass needs to move, by how far
@@ -525,7 +525,12 @@ def compute_rdf_distance(rdf1, rdf2, rr, n_parallel_rdf2: int = None, return_num
     # take the raw average over nonzero element pairs
     eps = 1e-12
     active = (torch_rdf1_f.sum(dim=-1) > eps) | (torch_rdf2.sum(dim=-1) > eps)
-    distance = (emd * active).sum(dim=-1) / active.sum(dim=-1).clamp_min(1)  # ignore unused channels
+
+    if channel_weights is not None:
+        w = channel_weights.to(emd.device)[None, :]  # [1, n_channels]
+        distance = (emd * active * w).sum(dim=-1) / (active * w).sum(dim=-1).clamp_min(1e-10)
+    else:
+        distance = (emd * active).sum(dim=-1) / active.sum(dim=-1).clamp_min(1)
 
     # distance = emd.mean(-1)
     #

mxtaltools/constants/space_group_info.py

@@ -179,20 +179,22 @@
 SYM_OPS = {1: [array([[1., 0., 0., 0.],
                       [0., 1., 0., 0.],
                       [0., 0., 1., 0.],
-                      [0., 0., 0., 1.]])], 2: [array([[1., 0., 0., 0.],
-                                                      [0., 1., 0., 0.],
-                                                      [0., 0., 1., 0.],
-                                                      [0., 0., 0., 1.]]), array([[-1., 0., 0., 0.],
-                                                                                 [0., -1., 0., 0.],
-                                                                                 [0., 0., -1., 0.],
-                                                                                 [0., 0., 0., 1.]])],
+                      [0., 0., 0., 1.]])],
+           2: [array([[1., 0., 0., 0.],
+                      [0., 1., 0., 0.],
+                      [0., 0., 1., 0.],
+                      [0., 0., 0., 1.]]), array([[-1., 0., 0., 0.],
+                                                 [0., -1., 0., 0.],
+                                                 [0., 0., -1., 0.],
+                                                 [0., 0., 0., 1.]])],
            3: [array([[1., 0., 0., 0.],
                       [0., 1., 0., 0.],
                       [0., 0., 1., 0.],
                       [0., 0., 0., 1.]]), array([[-1., 0., 0., 0.],
                                                  [0., 1., 0., 0.],
                                                  [0., 0., -1., 0.],
-                                                 [0., 0., 0., 1.]])], 4: [array([[1., 0., 0., 0.],
+                                                 [0., 0., 0., 1.]])],
+           4: [array([[1., 0., 0., 0.],
                                                                                  [0., 1., 0., 0.],
                                                                                  [0., 0., 1., 0.],
                                                                                  [0., 0., 0., 1.]]),
@@ -213,7 +215,8 @@
                array([[-1., 0., 0., 0.5],
                       [0., 1., 0., 0.5],
                       [0., 0., -1., 0.],
-                      [0., 0., 0., 1.]])], 6: [array([[1., 0., 0., 0.],
+                      [0., 0., 0., 1.]])],
+           6: [array([[1., 0., 0., 0.],
                                                       [0., 1., 0., 0.],
                                                       [0., 0., 1., 0.],
                                                       [0., 0., 0., 1.]]), array([[1., 0., 0., 0.],

mxtaltools/crystal_search/crystal_opt_utils.py

@@ -121,6 +121,7 @@ def gradient_descent_optimization(  # todo consolidate kwargs somewhere
         umbrella_sigma: Optional[float] = None,  # bandwidth term for umbrella sampling
         umbrella_epsilon: Optional[float] = None,  # repulsion term for umbrella sampling
         umbrella_record: Optional[list] = None,
+        rdf_warmup: Optional[torch.tensor] = 500,
 ):
     """
     do a local optimization via gradient descent on some score function
@@ -160,7 +161,7 @@ def gradient_descent_optimization(  # todo consolidate kwargs somewhere
     if target_latent is not None:
         target_latent = target_latent.to(init_sample.device)
 
-    if target_rdf is not None:
+    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
@@ -172,6 +173,7 @@ def gradient_descent_optimization(  # todo consolidate kwargs somewhere
         'score_model': score_model,
         'optim_target': optim_target,
         'target_latent': target_latent,
+        'rdf_warmup': rdf_warmup,
     })
 
     aux_config = dict2namespace({
@@ -270,7 +272,7 @@ def gradient_descent_optimization(  # todo consolidate kwargs somewhere
                     else:
                         loss_and_backprop(cluster_batch, crystal_batch, grad_norm_clip,
                                           optimizer, outputs, param_module, records,
-                                          loss_config, aux_config)
+                                          loss_config, aux_config, s_ind)
 
                     if s_ind % 10 == 0:
                         gc.collect()
@@ -344,7 +346,10 @@ def gradient_descent_optimization(  # todo consolidate kwargs somewhere
     
     """
     if optim_target == 'rdf_dist':
-        if torch.amin(records['loss']).log() < -2.5:
+        timesteps = torch.arange(s_ind).repeat(init_crystal_batch.num_graphs, 1).T
+        traj_fig(timesteps, torch.log(records['loss']), names=['time', 'loss'])
+
+        if torch.amin(records['loss'][-1]).log() < -2.5:
             print("Found the crystal!")
             good_ind = torch.argmin(records['loss'][-1]).item()
             sample = crystal_batch.batch_to_list()[good_ind]
@@ -403,8 +408,8 @@ def update_record(crystal_batch, outputs, params_record, records, s_ind):
 
 
 def loss_and_backprop(cluster_batch, crystal_batch, grad_norm_clip, optimizer, outputs, param_module, records,
-                      loss_config, aux_config):
-    loss = compute_loss(cluster_batch, crystal_batch, outputs, loss_config)
+                      loss_config, aux_config, opt_step):
+    loss = compute_loss(cluster_batch, crystal_batch, outputs, loss_config, opt_step)
     loss = compute_auxiliary_loss(cluster_batch, loss, outputs, aux_config)
 
     records['loss'].append(loss.detach().cpu())
@@ -541,7 +546,7 @@ def ema_trajectory(traj: torch.Tensor, alpha: float = 0.1) -> torch.Tensor:
     return numer / denom
 
 
-def compute_loss(cluster_batch, crystal_batch, outputs, config):
+def compute_loss(cluster_batch, crystal_batch, outputs, config, opt_step):
     if config.optim_target.lower() == 'lj':  # todo obviate this with analysis keys
         loss = outputs['lj']
 
@@ -573,8 +578,27 @@ def compute_loss(cluster_batch, crystal_batch, outputs, config):
         loss = outputs['uma']
 
     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
+        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]))
+                                    torch.linspace(0, config.cutoff, config.target_rdf.shape[-1]),
+                                    channel_weights=channel_weights)
 
     elif config.optim_target.lower() == 'latent_dist':
         loss = (config.target_latent - crystal_batch.latent_params()).norm(dim=-1)

mxtaltools/crystal_search/utils.py

@@ -12,7 +12,7 @@
 from mxtaltools.mlip_interfaces.uma_utils import init_uma_crystal_predictor
 
 
-def save_umbrella_record(record, new_latents, path, sigma = 0.2, epsilon = 10):
+def save_umbrella_record(record, new_latents, path, sigma=0.2, epsilon=10):
     if len(record) > 0 and len(new_latents) > 0:
         dists = torch.cdist(new_latents, record)
         repulsion = epsilon * torch.exp(-dists ** 2 / (2 * sigma ** 2)).sum(dim=1).clip(max=10)
@@ -22,6 +22,7 @@ def save_umbrella_record(record, new_latents, path, sigma = 0.2, epsilon = 10):
     torch.save(record, path)
     return record
 
+
 def rdf_clustering(packing_coeff, rdf, rdf_cutoff, rr, samples, vdw, num_cpus=None):
     """cluster samples according to rdf distances"""
     # rdf_dists = compute_rdf_distmat(rdf, rr)
@@ -209,13 +210,14 @@ def get_initial_state(config, crystal_batch, device, batch_idx, target):
             crystal_batch.sample_random_crystal_parameters(
                 target_packing_coeff=target_cp,
                 seed=config.opt_seed + int(batch_idx * 10000))
-        standard_cell = target.compute_standard_cell()
-        crystal_batch.cell_lengths = torch.tensor(standard_cell[0, :3], dtype=torch.float32,
-                                                  device=crystal_batch.device).repeat(crystal_batch.num_graphs, 1)
-        crystal_batch.cell_angles = torch.tensor(standard_cell[0, 3:], dtype=torch.float32,
-                                                 device=crystal_batch.device).repeat(crystal_batch.num_graphs,
-                                                                                     1) * torch.pi / 2 / 90
-        crystal_batch.box_analysis()
+        assert False, "Below is deprecated and probably unnecessary for now"
+        # standard_cell = target.compute_standard_cell()
+        # crystal_batch.cell_lengths = torch.tensor(standard_cell[0, :3], dtype=torch.float32,
+        #                                           device=crystal_batch.device).repeat(crystal_batch.num_graphs, 1)
+        # crystal_batch.cell_angles = torch.tensor(standard_cell[0, 3:], dtype=torch.float32,
+        #                                          device=crystal_batch.device).repeat(crystal_batch.num_graphs,
+        #                                                                              1) * torch.pi / 2 / 90
+        # crystal_batch.box_analysis()
 
     else:
         assert False
@@ -228,7 +230,10 @@ def init_samples_to_optim(config, target=None):
     """
     if config.init_sample_method == 'data':
         samples_to_optim = torch.load(config.dataset_path, weights_only=False)
-        index_block = torch.arange(config.mol_seed * config.num_samples, (config.mol_seed + 1) * config.num_samples)
+        if not isinstance(samples_to_optim, list):
+            samples_to_optim = [samples_to_optim]
+        index_block = [0 for _ in range(
+            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
     else:
@@ -305,7 +310,9 @@ def parse_opt_config(opt_config, config, device, target):
         opt_config['elementwise'] = False
         opt_config['atomwise'] = True
         tbatch = collate_data_list([target])
-        out = tbatch.analyze(['rdf'], cutoff=10, rdf_cutoff=10, elementwise=False, atomwise=True, bins=100)
+        out = tbatch.analyze(['rdf'], cutoff=opt_config['cutoff'],
+                             rdf_cutoff=opt_config['cutoff'],
+                             elementwise=False, atomwise=True, bins=100)
         opt_config['target_rdf'] = out['rdf'][0]
     if opt_config['optim_target'] in ['latent_dist']:
         tbatch = collate_data_list([target])