Wyckoff Transformer

https://github.com/symmetryadvantage/wyckofftransformer
Admin
Wyckoff Transformer (WyFormer) is a generative model for symmetric crystals that formally conditions...
Tokens:88,334
Snippets:719
Trust Score:5.3
License:Apache-2.0
Update:6 months ago
Show doc for...
Tokens:
Raw
### Prepare Start Token Distribution

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/research_notebooks/loss_study.ipynb

Calculates and adjusts the distribution of start tokens for generation, excluding underrepresented tokens. This ensures a more balanced generation process.

```python
generation_size = 15000 #len(next(iter(tensors["test"].values())))*2
del tensors["test"]

max_start = len(tokenisers[config.model.start_token])
start_counts = torch.bincount(tensors["train"][config.model.start_token], minlength=max_start) +
                torch.bincount(tensors["val"][config.model.start_token], minlength=max_start)
underrepresented = start_counts < 1000
start_counts[underrepresented] = 0
print(f"Excluded {underrepresented.sum()} underrepresented start tokens.")
start_distribution = torch.distributions.Categorical(probs=start_counts.float())
```

--------------------------------

### Setup Statistical Evaluator

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/research_notebooks/MP_20_full_eval.ipynb

Reloads the evaluator module and initializes the StatisticalEvaluator with novel test data.

```python
import wyckoff_transformer.evaluation.statistical_evaluator
import importlib
importlib.reload(wyckoff_transformer.evaluation.statistical_evaluator)
test_unique = filter_by_unique_structure(all_datasets[('split', 'test')].data)
test_novel = novelty_filter.get_novel(test_unique)
test_evaluator = wyckoff_transformer.evaluation.statistical_evaluator.StatisticalEvaluator(test_novel)
```

--------------------------------

### Initialize WyckoffTrainer from Config

Source: https://context7.com/symmetryadvantage/wyckofftransformer/llms.txt

Instantiates the trainer using a configuration dictionary. This approach is suitable for local training and custom model setups.

```python
from pathlib import Path
import torch
from omegaconf import OmegaConf
from wyckoff_transformer.trainer import WyckoffTrainer

# Load configuration
config = OmegaConf.load("yamls/models/NextToken/v6/base_sg.yaml")
config['dataset'] = 'mp_20'

tokeniser_config = OmegaConf.load("yamls/tokenisers/mp_20_sg_multiplicity.yaml")
config['tokeniser'] = tokeniser_config

# Create trainer from config
trainer = WyckoffTrainer.from_config(
    config,
    device=torch.device("cuda"),
    use_cached_tensors=True,
    run_path=Path("runs/my_run"),
    production_training=False,
    no_test=False
)

# Access model components
print(f"Model device: {trainer.device}")
print(f"Cascade order: {trainer.cascade_order}")
print(f"Max sequence length: {trainer.max_sequence_length}")
```

--------------------------------

### Import Core Libraries

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/research_notebooks/prepare_matbench.ipynb

Imports the PyTorch and Pandas libraries. Ensure these libraries are installed in your environment.

```python
import torch
torch.no_grad()
import pandas as pd
```

--------------------------------

### Run a pilot model for next token prediction

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/README.md

Executes a sequence of scripts to cache a dataset, tokenize it, and train a pilot model to verify installation.

```bash
python scripts/cache_a_dataset.py mp_20
python scripts/tokenise_a_dataset.py mp_20 yamls/tokenisers/mp_20_sg_multiplicity.yaml --new-tokenizer
python scripts/train.py yamls/models/NextToken/v6/base_sg.yaml mp_20 cuda --pilot
```

--------------------------------

### Execute CIF Filtering

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/scripts/data_preprocesssing/SymmCD_filter_invalid.ipynb

Example of how to call the `filter_cifs` function with specified input and output paths.

```python
generated_path = Path("..", "..", "generated", "Dropbox")
filter_cifs(generated_path / 'mp_20/SymmCD/crystal_symmcd_mp20.csv.gz',
            generated_path / 'mp_20/SymmCD/crystal_symmcd_mp20_valid.json.gz')
```

--------------------------------

### CatBoost Training Progress Output

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/research_notebooks/sg_predict.ipynb

Example output from a CatBoost model training process, showing learning and testing metrics per iteration. This output helps in monitoring convergence and identifying potential overfitting.

```text
0:	learn: 0.1687795	test: 0.1660215	best: 0.1660215 (0)	total: 75.8ms	remaining: 6m 18s
1:	learn: 0.1765183	test: 0.1765291	best: 0.1765291 (1)	total: 143ms	remaining: 5m 57s
2:	learn: 0.1767025	test: 0.1763079	best: 0.1765291 (1)	total: 207ms	remaining: 5m 44s
3:	learn: 0.1831884	test: 0.1829444	best: 0.1829444 (3)	total: 261ms	remaining: 5m 25s
4:	learn: 0.1829673	test: 0.1827232	best: 0.1829444 (3)	total: 313ms	remaining: 5m 12s
5:	learn: 0.1875369	test: 0.1853777	best: 0.1853777 (5)	total: 365ms	remaining: 5m 4s
6:	learn: 0.1876843	test: 0.1859308	best: 0.1859308 (6)	total: 421ms	remaining: 5m
7:	learn: 0.1882739	test: 0.1863732	best: 0.1863732 (7)	total: 474ms	remaining: 4m 56s
8:	learn: 0.1890478	test: 0.1872580	best: 0.1872580 (8)	total: 529ms	remaining: 4m 53s
9:	learn: 0.1894531	test: 0.1881429	best: 0.1881429 (9)	total: 582ms	remaining: 4m 50s
10:	learn: 0.1903007	test: 0.1896914	best: 0.1896914 (10)	total: 637ms	remaining: 4m 48s
11:	learn: 0.1905218	test: 0.1898020	best: 0.1898020 (11)	total: 690ms	remaining: 4m 46s
12:	learn: 0.1911114	test: 0.1899126	best: 0.1899126 (12)	total: 741ms	remaining: 4m 44s
13:	learn: 0.1919959	test: 0.1903550	best: 0.1903550 (13)	total: 795ms	remaining: 4m 43s
14:	learn: 0.1951282	test: 0.1952218	best: 0.1952218 (14)	total: 849ms	remaining: 4m 42s
15:	learn: 0.1991450	test: 0.1990930	best: 0.1990930 (15)	total: 903ms	remaining: 4m 41s
16:	learn: 0.2002874	test: 0.1996461	best: 0.1996461 (16)	total: 957ms	remaining: 4m 40s
17:	learn: 0.2006191	test: 0.1995354	best: 0.1996461 (16)	total: 1.01s	remaining: 4m 40s
18:	learn: 0.2016509	test: 0.2009733	best: 0.2009733 (18)	total: 1.06s	remaining: 4m 38s
19:	learn: 0.2052255	test: 0.2057295	best: 0.2057295 (19)	total: 1.11s	remaining: 4m 37s
20:	learn: 0.2059994	test: 0.2060613	best: 0.2060613 (20)	total: 1.17s	remaining: 4m 37s
21:	learn: 0.2080999	test: 0.2087159	best: 0.2087159 (21)	total: 1.22s	remaining: 4m 36s
22:	learn: 0.2088001	test: 0.2097113	best: 0.2097113 (22)	total: 1.28s	remaining: 4m 36s
23:	learn: 0.2093529	test: 0.2105962	best: 0.2105962 (23)	total: 1.33s	remaining: 4m 36s
24:	learn: 0.2114903	test: 0.2134720	best: 0.2134720 (24)	total: 1.38s	remaining: 4m 35s
25:	learn: 0.2152491	test: 0.2169008	best: 0.2169008 (25)	total: 1.44s	remaining: 4m 34s
26:	learn: 0.2154702	test: 0.2171220	best: 0.2171220 (26)	total: 1.49s	remaining: 4m 33s
27:	learn: 0.2161335	test: 0.2171220	best: 0.2171220 (26)	total: 1.54s	remaining: 4m 33s
28:	learn: 0.2166126	test: 0.2175644	best: 0.2175644 (28)	total: 1.59s	remaining: 4m 33s
29:	learn: 0.2174602	test: 0.2188917	best: 0.2188917 (29)	total: 1.65s	remaining: 4m 32s
30:	learn: 0.2204083	test: 0.2225418	best: 0.2225418 (30)	total: 1.7s	remaining: 4m 32s
31:	learn: 0.2215139	test: 0.2242009	best: 0.2242009 (31)	total: 1.75s	remaining: 4m 31s
32:	learn: 0.2226931	test: 0.2249751	best: 0.2249751 (32)	total: 1.8s	remaining: 4m 31s
33:	learn: 0.2281103	test: 0.2305055	best: 0.2305055 (33)	total: 1.85s	remaining: 4m 30s
34:	learn: 0.2287736	test: 0.2312797	best: 0.2312797 (34)	total: 1.91s	remaining: 4m 30s
35:	learn: 0.2323113	test: 0.2345979	best: 0.2345979 (35)	total: 1.96s	remaining: 4m 29s
36:	learn: 0.2320165	test: 0.2344873	best: 0.2345979 (35)	total: 2.01s	remaining: 4m 29s
37:	learn: 0.2337117	test: 0.2354828	best: 0.2354828 (37)	total: 2.07s	remaining: 4m 29s
38:	learn: 0.2345224	test: 0.2361464	best: 0.2361464 (38)	total: 2.12s	remaining: 4m 29s
39:	learn: 0.2350383	test: 0.2366995	best: 0.2366995 (39)	total: 2.17s	remaining: 4m 29s
40:	learn: 0.2351489	test: 0.2361464	best: 0.2366995 (39)	total: 2.22s	remaining: 4m 28s
41:	learn: 0.2371020	test: 0.2376949	best: 0.2376949 (41)	total: 2.27s	remaining: 4m 28s
42:	learn: 0.2372863	test: 0.2379162	best: 0.2379162 (42)	total: 2.32s	remaining: 4m 27s
43:	learn: 0.2394236	test: 0.2403495	best: 0.2403495 (43)	total: 2.37s	remaining: 4m 27s
44:	learn: 0.2400870	test: 0.2402389	best: 0.2403495 (43)	total: 2.42s	remaining: 4m 26s
45:	learn: 0.2423349	test: 0.2431147	best: 0.2431147 (45)	total: 2.48s	remaining: 4m 26s
46:	learn: 0.2431088	test: 0.2439996	best: 0.2439996 (46)	total: 2.53s	remaining: 4m 26s
47:	learn: 0.2444354	test: 0.2449950	best: 0.2449950 (47)	total: 2.58s	remaining: 4m 26s
48:	learn: 0.2459832	test: 0.2473178	best: 0.2473178 (48)	total: 2.63s	remaining: 4m 25s
49:	learn: 0.2475678	test: 0.2482026	best: 0.2482026 (49)	total: 2.69s	remaining: 4m 25s
50:	learn: 0.2481206	test: 0.2489769	best: 0.2489769 (50)	total: 2.74s	remaining: 4m 25s
51:	learn: 0.2487839	test: 0.2497511	best: 0.2497511 (51)	total: 2.79s	remaining: 4m 25s
52:	learn: 0.2512529	test: 0.2518527	best: 0.2518527 (52)	total: 2.84s	remaining: 4m 25s
53:	learn: 0.2520268	test: 0.2529587	best: 0.2529587 (53)	total: 2.89s	remaining: 4m 25s
54:	learn: 0.2540168	test: 0.2537330	best: 0.2537330 (54)	total: 2.95s	remaining: 4m 24s
```

--------------------------------

### Initialize pyXtal for Structure Generation

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/research_notebooks/diffcsp_reseach.ipynb

Imports the pyXtal library, typically used for generating and analyzing crystal structures. This is a common starting point for crystal structure research.

```python
from pyxtal import pyxtal

print("Generating ideal structures with pyXtal and analyzing them...")
```

--------------------------------

### Fit CatBoost Model

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/research_notebooks/sg_predict.ipynb

Fits a CatBoost model to the training data, with an optional evaluation set for monitoring performance during training. Ensure CatBoost is installed and data is preprocessed.

```python
catboost_model.fit(cb_train, eval_set=cb_val)
```

--------------------------------

### Progress Bar Example

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/research_notebooks/perov_5.ipynb

This snippet displays a progress bar, typically used to indicate the progress of a long-running task. It's often seen during data processing or iteration.

```text
0%|          | 0/2 [00:00<?, ?it/s]
```

--------------------------------

### Get Wyckoff Letter from Site Symmetry Enumeration

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/research_notebooks/enum_study.ipynb

Retrieve the mapping from site symmetry enumeration values to Wyckoff letters for space group 167. This example shows a defaultdict structure.

```python
letter_from_ss_enum[167]
```

--------------------------------

### Initialize Plotting

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/research_notebooks/NeurIPS-2024-eval.ipynb

Prepare matplotlib figure and bins for visualization.

```python
import numpy as np
import matplotlib.pyplot as plt
fig, ax = plt.subplots()
bins = np.arange(-0.5, 3, 0.05)
```

--------------------------------

### Get Site Symmetry Enumeration for Space Group 167

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/research_notebooks/enum_study.ipynb

Retrieve the dictionary mapping Wyckoff site symmetry letters to their enumeration values for space group 167. All values are 0 in this example.

```python
enum_from_ss_letter[167]
```

--------------------------------

### Train Model

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/README.md

Initiates the model training process. Requires the path to the model configuration YAML, the dataset name, and the target device (e.g., 'cpu' or 'cuda'). Model weights and metadata are saved to runs/<run-id> and WanDB. Add the --pilot flag to run for a small number of epochs.

```bash
python scripts/train.py <path-to-model-yaml> <dataset-name> <device>
```

--------------------------------

### Get Site Symmetry for a Wyckoff Position

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/research_notebooks/enum_study.ipynb

Retrieve the site symmetry for a specific Wyckoff position within a space group. The example accesses the first Wyckoff position (index 1) of Group 194.

```python
w=Group(194)[1]
w.get_site_symmetry()
```

--------------------------------

### Get Masked Multiclass Cascade Data

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/research_notebooks/loss_study.ipynb

Retrieves masked data and targets from the dataset for model inference. This function is used to prepare input for sequence generation tasks, filtering based on underrepresented start tokens.

```python
known_seq_len = 4
known_cascade_len = 0
with torch.no_grad():
    start_tokens, masked_data, target = validation_dataset.get_masked_multiclass_cascade_data(
                            known_seq_len, known_cascade_len,
                            target_type=TargetClass.NextToken, multiclass_target=True)
    selected = ~underrepresented[start_tokens]
    start_tokens = start_tokens[selected]
    for i in range(len(masked_data)):
        masked_data[i] = masked_data[i][selected]
    target = target[selected]
    logits = model(start_tokens, masked_data, None, known_cascade_len)
    loss_fn = torch.nn.CrossEntropyLoss(reduction="none")
    loss = loss_fn(logits, target)

    logits_temperated = logits / 1
    loss_temperated = loss_fn(logits_temperated, target)
    element_present = batched_bincount(masked_data[0], dim=1, max_value=num_classes_dict["elements"], dtype=torch.bool)
    element_counts = element_present.sum(dim=1)
    # at max element count
    # ternary + MASK
    at_max_element_count = element_counts == 4
    print(f"Total {target.size(0)}; At max element count: {at_max_element_count.sum()}")
    fixed_logits = logits.clone()

    element_present |= ~at_max_element_count[:, None]
    element_present[:, tokenisers["elements"].stop_token] = True
    fixed_logits[~element_present] -= 1e7
    fixed_loss = loss_fn(fixed_logits, target)
    fig, ax = plt.subplots(1, 1, figsize=(10, 10))
    bins = np.linspace(0, 10, 100)
    ax.hist(loss.cpu().numpy(), bins=bins, alpha=0.5, label=f"Loss, {loss.mean().item():.3f}")
    ax.hist(loss_temperated.cpu().numpy(), bins=bins, alpha=0.5, label=f"Temperated loss, {loss_temperated.mean().item():.3f}")
    ax.hist(fixed_loss.cpu().numpy(), bins=bins, alpha=0.5, label=f"Fixed loss, {fixed_loss.mean().item():.3f}")
    ax.legend()
    ax.set_xlabel("LogLoss")
    ax.set_xlim(0)
```

--------------------------------

### Load Pickled and Gzipped Dataset

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/research_notebooks/mp_2022_plot.ipynb

Loads a dataset from a .pkl.gz file using gzip and pickle. Ensure the file path is correct.

```python
import gzip
import pickle
with gzip.open("cache/matbench_discovery_mp_2022/data.pkl.gz", "rb") as f:
    dataset_pd = pickle.load(f)
```

--------------------------------

### Load all datasets from configuration

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/research_notebooks/cdvae_metrics.ipynb

Iterates through dataset names and loads all associated configurations, including split data, into a dictionary.

```python
all_datasets = {}
for dataset_name in dataset_names:
    all_datasets[dataset_name] = load_all_from_config(
        datasets=list(config_names.values()) + \
            [("split", "train"), ("split", "val"), ("split", "test")],
        dataset_name=dataset_name)
```

--------------------------------

### Load Datasets and Tokenizers

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/research_notebooks/loss_study.ipynb

Loads datasets and tokenizers based on the provided configuration. Ensure the dataset and tokenizer names in the config are correct.

```python
from pathlib import Path
from omegaconf import OmegaConf
import numpy as np
import torch
import matplotlib.pyplot as plt
import wandb

from wyckoff_transformer.tokenization import load_tensors_and_tokenisers
from wyckoff_transformer.generator import WyckoffGenerator
from cascade_transformer.model import CascadeTransformer
from cascade_transformer.dataset import AugmentedCascadeDataset, TargetClass, batched_bincount

device = torch.device("cpu")
run_id = "1okrz6qt"
wandb_run = wandb.Api().run(f"WyckoffTransformer/{run_id}")
config = OmegaConf.create(dict(wandb_run.config))

# The start tokens will be sampled from the train+validation datasets,
# to preserve the sanctity of the test dataset and ex nihilo generation.
tensors, tokenisers = load_tensors_and_tokenisers(config.dataset, config.tokeniser.name)
```

--------------------------------

### Getting Site Symmetry for a Wyckoff Position

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/research_notebooks/enum_study.ipynb

Assigns a Wyckoff position object to the variable 'w' and then calls a method to get its site symmetry. This is likely part of a larger workflow for analyzing crystal structures.

```python
w=Group(15)[4]
w.get_site_symmetry()
```

--------------------------------

### Get primitive standard structure

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/research_notebooks/symmetry_research.ipynb

Retrieves the primitive standard structure representation from the analyzer.

```python
a.get_primitive_standard_structure()
```

--------------------------------

### Initialize evaluation components and process datasets

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/research_notebooks/cdvae_metrics.ipynb

Sets up defaultdict for results, filters warnings, and iterates through datasets. For each dataset, it initializes a NoveltyFilter and a StatisticalEvaluator, then computes metrics with and without novelty filtering for each configured model.

```python
from collections import defaultdict
import warnings
warnings.filterwarnings("ignore", category=UserWarning, module="pymatgen.core.composition",
                        message=r"No Pauling electronegativity for.*")
cdvae_datasets = {"mp_20": "mp20", "perov_5": "perovskite", "carbon_24": "carbon"}
results = defaultdict(lambda: defaultdict(dict))

for dataset_name, these_dataset in all_datasets.items():
    print(f"Processing dataset: {dataset_name}")
    novelty_reference = these_dataset[('split', 'train')].data
    novelty_filter = NoveltyFilter(novelty_reference, reference_index_type="reduced_composition")
    test_evaluator = StatisticalEvaluator(these_dataset[('split', 'test')].data, cdvae_eval_model_name=cdvae_datasets[dataset_name])
    for name, transformations in tqdm(config_names.items()):
        dataset = these_dataset[transformations]
        results[dataset_name]["no_filter"][name] = \
            test_evaluator.compute_cdvae_metrics(dataset.data, novelty_filter=None, sample_size_for_precision=500)
        results[dataset_name]["only_novel"][name] = \
            test_evaluator.compute_cdvae_metrics(
                dataset.data, novelty_filter=novelty_filter, sample_size_for_precision=500, compute_novelty=True)
```

--------------------------------

### Load CHGNet datasets from configuration

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/ICML_eval/DFT_eval.ipynb

Loads all specified CHGNet datasets using the `load_all_from_config` function. It combines configurations from `chgnet_datasets` and an additional specific fix.

```python
chgnet_data = load_all_from_config(datasets=list(chgnet_datasets.values()) + [('WyckoffTransformer', 'CrySPR', 'CHGNet_fix')])
```

--------------------------------

### Get Observation Data

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/ICML_eval/MP_20_paper_eval.ipynb

Generates a binary observation array based on a column value from a table. Used for statistical comparisons.

```python
import numpy as np
def get_observation(name, column, n_structures=1000):
    all_observations = np.zeros(n_structures)
    all_observations[:int(table.at[name, column])] = 1.
    return all_observations
```

--------------------------------

### Load all datasets from configuration

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/ICML_eval/LLM_vs_WT.ipynb

Loads all specified datasets, including training, validation, and test splits, from a configuration file. The 'mp_20' dataset name is used as a base.

```python
all_datasets = load_all_from_config(
    datasets=list(datasets.values()) + [("split", "train"), ("split", "val"), ("split", "test")],
    dataset_name="mp_20")
```

--------------------------------

### Import StatisticalEvaluator and related functions

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/research_notebooks/NeurIPS-2024-eval.ipynb

Imports necessary functions for statistical evaluation and data handling. Ensure these modules are installed and accessible.

```python
from wyckoff_transformer.evaluation.cdvae_metrics import timed_smact_validity_from_record
```

--------------------------------

### WyckoffTrainer.from_config

Source: https://context7.com/symmetryadvantage/wyckofftransformer/llms.txt

Creates a WyckoffTrainer instance from a configuration dictionary, loading datasets, tokenizers, and initializing the CascadeTransformer model.

```APIDOC
## WyckoffTrainer.from_config

### Description
Creates a WyckoffTrainer instance from a configuration dictionary, loading datasets, tokenizers, and initializing the CascadeTransformer model.

### Method
`WyckoffTrainer.from_config()`

### Parameters
#### Arguments
- **config** (OmegaConf) - Required - Configuration object loaded from YAML.
- **device** (torch.device) - Required - The device to use for training (e.g., `torch.device("cuda")`).
- **use_cached_tensors** (bool) - Optional - Whether to use cached tensors.
- **run_path** (Path) - Optional - Path for the training run.
- **production_training** (bool) - Optional - Flag for production training.
- **no_test** (bool) - Optional - Flag to disable testing.

### Request Example
```python
from pathlib import Path
import torch
from omegaconf import OmegaConf
from wyckoff_transformer.trainer import WyckoffTrainer

# Load configuration
config = OmegaConf.load("yamls/models/NextToken/v6/base_sg.yaml")
config['dataset'] = 'mp_20'

tokeniser_config = OmegaConf.load("yamls/tokenisers/mp_20_sg_multiplicity.yaml")
config['tokeniser'] = tokeniser_config

# Create trainer from config
trainer = WyckoffTrainer.from_config(
    config,
    device=torch.device("cuda"),
    use_cached_tensors=True,
    run_path=Path("runs/my_run"),
    production_training=False,
    no_test=False
)

# Access model components
print(f"Model device: {trainer.device}")
print(f"Cascade order: {trainer.cascade_order}")
print(f"Max sequence length: {trainer.max_sequence_length}")
```
```

--------------------------------

### Initialize Novelty Filter

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/research_notebooks/NeurIPS-2024-eval.ipynb

Set up a novelty filter using training and validation data as reference.

```python
import pandas as pd
from tqdm.auto import tqdm
from wyckoff_transformer.evaluation.novelty import NoveltyFilter, filter_by_unique_structure

novelty_reference = pd.concat([
    mp_20_data[('split', 'train')].data,
    mp_20_data[('split', 'val')].data], axis=0, verify_integrity=True)
novelty_filter = NoveltyFilter(novelty_reference)
```

--------------------------------

### Import necessary libraries

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/research_notebooks/props.ipynb

Imports the 'pickle' and 'gzip' modules for handling serialized data and compressed files, respectively.

```python
import pickle
import gzip
```

--------------------------------

### TtestResult Output Example

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/ICML_eval/DFT_eval.ipynb

This shows the typical output format of a `TtestResult` object from SciPy, including the statistic, p-value, and degrees of freedom.

```text
Result:
TensileResult(statistic=0.19366734265109437, pvalue=0.8464560505477872, df=1998.0)
```

--------------------------------

### Python NameError Example

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/research_notebooks/harmonic_dev.ipynb

Illustrates a NameError in Python, occurring when a variable or function is used before it has been defined or imported. In this case, 'pd' is not defined.

```python
import pandas as pd
import numpy as np

def assign_to_clusters(
    distances: pd.DataFrame):

    remaining_distances = distances.copy().droplevel((0, 1), axis=0)
    assert (remaining_distances.index == np.arange(remaining_distances.shape[0])).all()
    assert (remaining_distances.columns == np.arange(remaining_distances.shape[1])).all()
    mapping = np.empty(distances.shape[0], dtype=int)
    
    while not remaining_distances.empty:
        row, col = np.unravel_index(np.argmin(remaining_distances.values), remaining_distances.shape)
        row_label = remaining_distances.index[row]
        col_label = remaining_distances.columns[col]
        # enum -> cluster
        mapping[row_label] = col_label
        remaining_distances = remaining_distances.drop(row_label, axis=0).drop(col_label, axis=1)
    return pd.Series(mapping, index=distances.index.get_level_values(2))
```

--------------------------------

### Load all datasets from configuration

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/ICML_eval/mpts_52.ipynb

Loads all specified datasets, including training, validation, and test splits, from a configuration file. Ensure the 'mpts_52' dataset name is correctly configured.

```python
all_datasets = load_all_from_config(
    datasets=list(raw_datasets.values()) + \
        [("split", "train"), ("split", "val"), ("split", "test")],
    dataset_name="mpts_52")
```

--------------------------------

### Initialize evaluation table and loop through datasets

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/ICML_eval/LLM_vs_WT.ipynb

Sets up a pandas DataFrame to store evaluation results and iterates through each defined dataset configuration. Calculates and populates metrics for each method.

```python
table = pd.DataFrame(
    index=datasets.keys(), columns=[
        "Novelty (%)", "Structural", "Compositional", 
        "Recall", "Precision",
        r"$
ho$", "$E$", "# Elements",
        "S.U.N. (%)",
        "Novel Template (%)", "P1 (%)",
        "Space Group", "S.S.U.N. (%)"])
table.index.name = "Method"
E_hull_threshold = 0.08

for name, transformations in tqdm(datasets.items()):
    dataset = all_datasets[transformations]
    unique = filter_by_unique_structure(dataset.data)
    print(len(unique), len(dataset.data), len(unique) / len(dataset.data))
    novel_template = ~unique.apply(wyckoff_transformer.evaluation.novelty.record_to_anonymous_fingerprint, axis=1).isin(train_w_template_set)
    table.loc[name, "Novel Template (%)"] = 100 * novel_template.mean()
    novel = novelty_filter.get_novel(unique)
    table.loc[name, "Novelty (%)"] = 100 * len(novel) / len(unique)
    if "structural_validity" in novel.columns:
        table.loc[name, "Structural"] = 100 * novel.structural_validity.mean()
    table.loc[name, "Compositional"] = 100 * novel.smact_validity.mean()
    if "cdvae_crystal" in novel.columns:
        cov_metrics = test_evaluator.get_coverage(novel.cdvae_crystal)    
        table.loc[name, "Recall"] = 100 * cov_metrics["cov_recall"]
        table.loc[name, "Precision"] = 100 * cov_metrics["cov_precision"]
        novel = novel[novel.structural_validity]
        table.loc[name, r"$
ho$"] = test_evaluator.get_density_emd(novel)
        table.loc[name, "$E$"] = test_evaluator.get_cdvae_e_emd(novel)
        table.loc[name, "# Elements"] = test_evaluator.get_num_elements_emd(novel)
    table.loc[name, "P1 (%)"] = 100 * (novel.group == 1).mean()
    # table.loc[name, "# DoF"] = test_evaluator.get_dof_emd(novel)
    table.loc[name, "Space Group"] = test_evaluator.get_sg_chi2(novel)
    #try:
    #    table.loc[name, "SG preserved (%)"] = 100 * is_sg_preserved(novel.spacegroup_number, transformations).mean()
    #except KeyError:
    #    pass
    #table.loc[name, "Elements"] = test_evaluator.get_elements_chi2(novel)
    if "corrected_chgnet_ehull" in novel.columns:
        # S.U.N. is measured with respect to the initial structures
        has_ehull = dataset.data.corrected_chgnet_ehull.notna().sum()
        is_sun = (novel.corrected_chgnet_ehull <= E_hull_threshold)
        table.loc[name, "S.U.N. (%)"] = 100 * is_sun.sum() / has_ehull.sum()
        table.loc[name, "S.S.U.N. (%)"] = 100 * (is_sun & (novel.group != 1)).sum() / has_ehull.sum()
```

--------------------------------

### Get Symmetry-Equivalent Site Positions

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/research_notebooks/symmetry_research.ipynb

Calculates all symmetry-equivalent positions for a given site within a specified space group. It uses `apply_symmetry_operation` and `functools.partial`.

```python
import numpy as np
from functools import partial
from typing import List
from pymatgen.core.structure import Structure
from pymatgen.core.lattice import Lattice
from pymatgen.symmetry.analyzer import SpacegroupAnalyzer
from pymatgen.symmetry.groups import SpaceGroup
from pyxtal import pyxtal

def get_site_images(site_frac_coords: np.ndarray, space_group: SpaceGroup) -> List[np.ndarray]:
    """Get all symmetry-equivalent positions for a site under a space group."""
    apply_op = partial(apply_symmetry_operation, site_frac_coords)
    return [apply_op(op) for op in space_group.symmetry_ops]
```

--------------------------------

### Get Rhombohedral Space Groups

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/research_notebooks/symmetry_research.ipynb

Retrieves a list of all space groups that belong to the rhombohedral lattice system. These are identified by their specific integer numbers.

```python
import numpy as np
from functools import partial
from typing import List
from pymatgen.core.structure import Structure
from pymatgen.core.lattice import Lattice
from pymatgen.symmetry.analyzer import SpacegroupAnalyzer
from pymatgen.symmetry.groups import SpaceGroup
from pyxtal import pyxtal

def get_rhombohedral_space_groups() -> List[SpaceGroup]:
    """Get all space groups with rhombohedral lattice system."""
    rhombohedral_numbers = [146, 148, 155, 160, 161, 166, 167]  # Known rhombohedral space groups
    return [SpaceGroup.from_int_number(i) for i in rhombohedral_numbers]

# Perform symmetry analysis
rhombohedral_sgs = get_rhombohedral_space_groups()
print(f"Found {len(rhombohedral_sgs)} rhombohedral space groups: {[sg.int_number for sg in rhombohedral_sgs]}")
```

--------------------------------

### Get Training Set Site Symmetries Count

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/research_notebooks/enum_study.ipynb

Calculates the number of site symmetries in the training dataset. This is useful for understanding the data distribution.

```python
len(tensors['train']['site_symmetries'])
```

--------------------------------

### Define and Load Datasets

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/research_notebooks/perov_5.ipynb

Configuration for dataset selection and loading using the evaluation framework.

```python
raw_datasets = {
    "WyFormer": ("WyckoffTransformer", "CrySPR", "CHGNet_fix"),
    "DiffCSP": ("DiffCSP",)
}
```

```python
all_datasets = load_all_from_config(
    datasets=list(raw_datasets.values()) + \
        [("split", "train"), ("split", "val"), ("split", "test")],
    dataset_name="perov_5")
```

--------------------------------

### Get Value Counts of Wyckoff Positions

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/research_notebooks/mp_2022_plot.ipynb

Displays the frequency of each unique Wyckoff position count in the dataset. This helps understand the distribution of symmetries.

```python
wp_counts.value_counts().so
```

--------------------------------

### Import GeneratedDataset

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/research_notebooks/E_hull_CHGNet.ipynb

Initializes the environment and imports the necessary class for dataset handling.

```python
%matplotlib widget
from wyckoff_transformer.evaluation.generated_dataset import GeneratedDataset
```

--------------------------------

### Get Shape of Harmonic Site Symmetries

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/research_notebooks/harmonic_dev.ipynb

Retrieves the shape of the harmonic site symmetries tensor from the training data. This provides the dimensions of the symmetry data.

```python
tensors['train']['harmonic_site_symmetries'].shape
```

--------------------------------

### Get Dimension of a Tensor

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/research_notebooks/cluster_generation.ipynb

Calculates and returns the number of dimensions (rank) of a PyTorch tensor created from the scalar value 0. A scalar has 0 dimensions.

```python
torch.tensor(0).dim()
```

--------------------------------

### Import required libraries

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/scripts/data_preprocesssing/wbm.ipynb

Initial imports for file path handling, data manipulation, and ASE I/O operations.

```python
from pathlib import Path
import pandas as pd
import ase.io
import ase.io.cif
import matbench_discovery.data
```

--------------------------------

### Load all datasets from configuration

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/ICML_eval/MP_20_biternary.ipynb

Loads all specified datasets, including training, validation, and test splits, using a configuration dictionary. This function is central to accessing the project's data.

```python
all_datasets = load_all_from_config(
    datasets=list(datasets.values()) + [("split", "train"), ("split", "val"), ("split", "test")],
    dataset_name="mp_20_biternary")
```

--------------------------------

### Initialize Group with Space Group Number

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/research_notebooks/enum_study.ipynb

Instantiate a Group object with a given space group number. This is the starting point for querying space group properties.

```python
Group(194)
```

--------------------------------

### Load Datasets

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/research_notebooks/MP_20_full_eval.ipynb

Loads all datasets defined in the configuration.

```python
all_datasets = load_all_from_config()
```

--------------------------------

### Getting site symmetry of a Wyckoff position

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/research_notebooks/enum_study.ipynb

Retrieves the site symmetry of a specific Wyckoff position. This indicates the symmetry elements that leave the Wyckoff position invariant.

```python
w.get_site_symmetry()
```

--------------------------------

### Import necessary libraries

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/research_notebooks/cdvae_metrics.ipynb

Imports pandas, tqdm for progress bars, and specific modules from the wyckoff_transformer library for dataset loading and evaluation.

```python
import pandas as pd
from tqdm.notebook import tqdm
from wyckoff_transformer.evaluation.generated_dataset import load_all_from_config
from wyckoff_transformer.evaluation.novelty import NoveltyFilter
from wyckoff_transformer.evaluation.statistical_evaluator import StatisticalEvaluator
```

--------------------------------

### Get Index Level Values

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/research_notebooks/cluster_generation.ipynb

Retrieves the values of the first level of the index from the database associated with the 'harmonic_cluster' engineer. This is useful for inspecting the structure of indexed data.

```python
engineers["harmonic_cluster"].db.index.get_level_values(0)
```

--------------------------------

### Import Required Libraries

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/ICML_eval/pyxtal_time.ipynb

Imports gzip for decompression, json for parsing, and pyxtal for crystal structure handling.

```python
import gzip
import json
from pyxtal import pyxtal
```

--------------------------------

### Import necessary libraries for WyckoffTransformer

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/ICML_eval/MP_20_paper_eval.ipynb

Imports essential libraries for data manipulation, structure analysis, and evaluation within the WyckoffTransformer project. Ensure these libraries are installed.

```python
from typing import Tuple
import pandas as pd
from tqdm.notebook import tqdm
from pymatgen.core import Structure
from wyckoff_transformer.evaluation.generated_dataset import GeneratedDataset, load_all_from_config
from wyckoff_transformer.evaluation.novelty import NoveltyFilter, filter_by_unique_structure
```

--------------------------------

### Load Datasets from Config

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/research_notebooks/NeurIPS-2024-eval.ipynb

Load datasets using the configuration-based loader.

```python
from wyckoff_transformer.evaluation.generated_dataset import load_all_from_config
mp_20_data = load_all_from_config(
    datasets=list(mp_20_datasets.values()) + \
        [("split", "train"), ("split", "val"), ("split", "test"), ("WyckoffTransformer", "CrySPR", "CHGNet_fix")],
    dataset_name="mp_20")
```

```python
mp_20_biternary_data = load_all_from_config(
    datasets=list(mp_20_biternary_datasets.values()) + \
        [("split", "train"), ("split", "val"), ("split", "test")],
    dataset_name="mp_20_biternary")
```

--------------------------------

### Run CrySPR + CHGNet Generation

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/README.md

This command initiates the CrySPR + CHGNet pipeline for generating 3D structures. It requires the path to the decompressed JSON data, and optionally a model name. The output includes CSV files with energy and structure information, and CIF files for relaxed structures.

```bash
python ./cryspr_pyxtal_chgnet.py 0 -1 ./WyckoffTransformer_mp_20.json model_name
```

--------------------------------

### Get Validation Set Site Symmetries Count

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/research_notebooks/enum_study.ipynb

Calculates the number of site symmetries in the validation dataset. This helps in assessing the model's performance on unseen data.

```python
len(tensors['val']['site_symmetries'])
```

--------------------------------

### Import necessary libraries

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/research_notebooks/generation_investigation.ipynb

Imports the required libraries for the project, including wandb for experiment tracking, omegaconf for configuration management, and the WyckoffTrainer class.

```python
import wandb
from omegaconf import OmegaConf
from wyckoff_transformer.trainer import WyckoffTrainer
```

--------------------------------

### Importing Matbench Discovery modules

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/scripts/data_preprocesssing/mptrj_extract_all.ipynb

Initial imports required for accessing dataset files and energy calculation utilities.

```python
import matbench_discovery.data
import matbench_discovery.energy
```

--------------------------------

### Get Shape of Augmented Harmonic Site Symmetries

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/research_notebooks/harmonic_dev.ipynb

Determines the shape of the augmented harmonic site symmetries tensor for the training data. This indicates the dimensions of the symmetry data.

```python
tensors['train']['harmonic_site_symmetries_augmented'][0][0].shape
```

--------------------------------

### Load Datasets

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/ICML_eval/tokenization_discussion.ipynb

Imports necessary modules and loads training, validation, and test datasets from configuration or cache.

```python
from wyckoff_transformer.evaluation.generated_dataset import load_all_from_config, GeneratedDataset
import pandas as pd
```

```python
train_val_ = load_all_from_config(datasets=[("split", "train"), ("split", "val")])
train_val = pd.concat((d.data for d in train_val_.values()), axis=0, ignore_index=False)
```

```python
test = GeneratedDataset.from_cache(("split", "test")).data
```

```python
wt = GeneratedDataset.from_cache(("WyckoffTransformer", )).data.sample(len(test), random_state=42)
```

--------------------------------

### Get Name of Inverted Series

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/research_notebooks/cluster_generation.ipynb

Applies the `inverse_series` function to the database used by the 'harmonic_cluster' engineer and then retrieves the name of the resulting series. This helps in understanding the transformation applied.

```python
inverse_series(engineers['harmonic_cluster'].db).name
```

--------------------------------

### Define dataset configurations

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/ICML_eval/MP_20_paper_eval.ipynb

Sets up dictionaries to define various datasets and their associated configurations, including raw and processed versions. These configurations are used for loading and processing data.

```python
datasets = {
    "WyckoffTransformer-raw": ("WyckoffTransformer",),
    "WyFormer-harmonic-raw": ("WyckoffTransformer-harmonic",),
    "WyFormer-letters": ("WyckoffTransformer-letters",),
    "WyFormer-letters-DiffCSP++": ("WyckoffTransformer-letters", "DiffCSP++", "CHGNet_fix"),
    "SymmCD": ("SymmCD", "CHGNet_fix"),
    "SymmCD-raw": ("SymmCD",),
    "WyFormer": ("WyckoffTransformer", "CrySPR", "CHGNet_fix_release"),
    "WyFormer-harmonic-DiffCSP++": ("WyckoffTransformer-harmonic", "DiffCSP++", "CHGNet_fix"),
    "WyForDiffCSP++": ("WyckoffTransformer", "DiffCSP++", "CHGNet_fix"),
    "WyLLM-naive-DiffCSP++": ("WyckoffLLM-naive", "DiffCSP++", "CHGNet_fix"),
    "WyLLM-vanilla-DiffCSP++": ("WyckoffLLM-vanilla", "DiffCSP++"),
    "WyLLM-site-symmetry-DiffCSP++": ("WyckoffLLM-site-symmetry", "DiffCSP++"),
    #"WyckoffTransformer-free": ("WyckoffTransformer", "CrySPR", "CHGNet_free"),
    "CrystalFormer": ("CrystalFormer", "CHGNet_fix_release"),
    #"DiffCSP++ raw": ("DiffCSP++",),
    "DiffCSP++": ("DiffCSP++", "CHGNet_fix_release"),
    "DiffCSP": ("DiffCSP", "CHGNet_fix"),
    "FlowMM": ("FlowMM", "CHGNet_fix"),
    "MiAD": ("MiAD", "CHGNet_free"),
    "MatterGen": ("MatterGen", "MatterGen_10k", "CHGNet_fix")
    #"MP-20 train": ("split", "train"),
    #"MP-20 test": ("split", "test"),
}
raw_datasets = {
    "SymmCD": ("SymmCD",),
    "WyFormer": ("WyckoffTransformer", "CrySPR", "CHGNet_fix_release"),
    "WyFormerDiffCSP++": ("WyckoffTransformer", "DiffCSP++"),
    "WyFormer-harmonic-DiffCSP++": ("WyckoffTransformer-harmonic", "DiffCSP++"),
    "WyFormer-letters-DiffCSP++": ("WyckoffTransformer-letters", "DiffCSP++"),
    "WyLLM-DiffCSP++": ("WyckoffLLM-naive", "DiffCSP++"),
    "CrystalFormer": ("CrystalFormer",),
    "DiffCSP++": ("DiffCSP++",),
    "DiffCSP": ("DiffCSP",),
    "FlowMM": ("FlowMM",),
    "MatterGen": ("MatterGen", "MatterGen_10k")
}
```

--------------------------------

### Get Site Symmetry Mapping for Space Group 167

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/research_notebooks/enum_study.ipynb

Retrieve the dictionary mapping Wyckoff site symmetry symbols to their corresponding enumeration values for space group 167.

```python
ss_from_letter[167]
```

--------------------------------

### Get Shape of Unsqueezed Mask Tensor

Source: https://github.com/symmetryadvantage/wyckofftransformer/blob/main/research_notebooks/cluster_generation.ipynb

Adds a dimension to the mask tensor and then prints its shape. This operation is common when preparing tensors for specific layer inputs in neural networks.

```python
mask.unsqueeze(0).shape
```
Tokens:
Raw