| import argparse |
| import json |
| import os |
| from pathlib import Path |
|
|
| from model.common.utils.data_classes import MatterGenCheckpointInfo |
| from model.generator import CrystalGenerator |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description="Generate crystals with MatterGen") |
| parser.add_argument("--checkpoint", required=True) |
| parser.add_argument("--output", default="outputs/mattergen") |
| parser.add_argument("--batch-size", type=int, default=1) |
| parser.add_argument("--num-batches", type=int, default=1) |
| parser.add_argument( |
| "--properties", |
| type=json.loads, |
| default=None, |
| help='Condition values as JSON, for example {"dft_mag_density": 0.15}.', |
| ) |
| parser.add_argument("--record-trajectories", action="store_true") |
| args = parser.parse_args() |
| if not os.path.isdir(args.checkpoint): |
| parser.error(f"checkpoint directory does not exist: {args.checkpoint}") |
| checkpoint_info = MatterGenCheckpointInfo( |
| model_path=Path(args.checkpoint).expanduser().resolve(), |
| load_epoch="last", |
| ) |
| generator = CrystalGenerator( |
| checkpoint_info=checkpoint_info, |
| batch_size=args.batch_size, |
| num_batches=args.num_batches, |
| properties_to_condition_on=args.properties, |
| record_trajectories=args.record_trajectories, |
| ) |
| structures = generator.generate( |
| output_dir=Path(args.output).expanduser().resolve() |
| ) |
| print(f"Generated {len(structures)} structures in {args.output}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
