auto_resume: true
checkpoint_path: null
data_module:
  _recursive_: true
  _target_: mattergen.common.data.datamodule.CrystDataModule
  average_density: 0.05771451654022283
  batch_size:
    test: 8
    train: 7
    val: 8
  dataset_transforms:
  - _partial_: true
    _target_: mattergen.common.data.dataset_transform.filter_sparse_properties
  max_epochs: 900
  num_workers:
    test: 0
    train: 0
    val: 0
  properties: []
  root_dir: /data/users/sathyae/mclm/external/mattergen/mattergen/../datasets/cache/broader_csp
  test_dataset:
    _target_: mattergen.common.data.dataset.CrystalDataset.from_cache_path
    cache_path: /data/users/sathyae/mclm/external/mattergen/mattergen/../datasets/cache/broader_csp/test
    dataset_transforms:
    - _partial_: true
      _target_: mattergen.common.data.dataset_transform.filter_sparse_properties
    properties: []
    transforms:
    - _partial_: true
      _target_: mattergen.common.data.transform.symmetrize_lattice
    - _partial_: true
      _target_: mattergen.common.data.transform.set_chemical_system_string
    - _partial_: true
      _target_: mattergen.common.data.transform.set_composition_count
  train_dataset:
    _target_: mattergen.common.data.dataset.CrystalDataset.from_cache_path
    cache_path: /data/users/sathyae/mclm/external/mattergen/mattergen/../datasets/cache/broader_csp/train
    dataset_transforms:
    - _partial_: true
      _target_: mattergen.common.data.dataset_transform.filter_sparse_properties
    properties: []
    transforms:
    - _partial_: true
      _target_: mattergen.common.data.transform.symmetrize_lattice
    - _partial_: true
      _target_: mattergen.common.data.transform.set_chemical_system_string
    - _partial_: true
      _target_: mattergen.common.data.transform.set_composition_count
  transforms:
  - _partial_: true
    _target_: mattergen.common.data.transform.symmetrize_lattice
  - _partial_: true
    _target_: mattergen.common.data.transform.set_chemical_system_string
  - _partial_: true
    _target_: mattergen.common.data.transform.set_composition_count
  val_dataset:
    _target_: mattergen.common.data.dataset.CrystalDataset.from_cache_path
    cache_path: /data/users/sathyae/mclm/external/mattergen/mattergen/../datasets/cache/broader_csp/val
    dataset_transforms:
    - _partial_: true
      _target_: mattergen.common.data.dataset_transform.filter_sparse_properties
    properties: []
    transforms:
    - _partial_: true
      _target_: mattergen.common.data.transform.symmetrize_lattice
    - _partial_: true
      _target_: mattergen.common.data.transform.set_chemical_system_string
    - _partial_: true
      _target_: mattergen.common.data.transform.set_composition_count
lightning_module:
  _target_: mattergen.diffusion.lightning_module.DiffusionLightningModule
  diffusion_module:
    _target_: mattergen.diffusion.diffusion_module.DiffusionModule
    corruption:
      _target_: mattergen.diffusion.corruption.multi_corruption.MultiCorruption
      sdes:
        cell:
          _target_: mattergen.common.diffusion.corruption.LatticeVPSDE.from_vpsde_config
          vpsde_config:
            beta_max: 20
            beta_min: 0.1
            limit_density: 0.05771451654022283
            limit_var_scaling_constant: 0.25
        pos:
          _target_: mattergen.common.diffusion.corruption.NumAtomsVarianceAdjustedWrappedVESDE
          limit_info_key: num_atoms
          sigma_max: 5.0
          wrapping_boundary: 1.0
    loss_fn:
      _target_: mattergen.common.loss.MaterialsLoss
      include_atomic_numbers: false
      include_cell: true
      include_pos: true
      reduce: sum
      weights:
        cell: 1.0
        pos: 0.1
    model:
      _target_: mattergen.denoiser.GemNetTDenoiser
      atom_type_diffusion: mask
      denoise_atom_types: true
      gemnet:
        _target_: mattergen.common.gemnet.gemnet.GemNetT
        atom_embedding:
          _target_: mattergen.common.gemnet.layers.embedding_block.AtomEmbedding
          emb_size: 512
          with_mask_type: true
        cutoff: 7.0
        emb_size_atom: 512
        emb_size_edge: 512
        latent_dim: 512
        max_cell_images_per_dim: 5
        max_neighbors: 50
        num_blocks: 4
        num_targets: 1
        otf_graph: true
        regress_stress: true
        scale_file: /data/users/sathyae/mclm/external/mattergen/mattergen/common/gemnet/gemnet-dT.json
      hidden_dim: 512
      property_embeddings: {}
      property_embeddings_adapt: {}
    pre_corruption_fn:
      _target_: mattergen.property_embeddings.SetEmbeddingType
      dropout_fields_iid: false
      p_unconditional: 0.2
  optimizer_partial:
    _partial_: true
    _target_: torch.optim.Adam
    lr: 0.0001
  scheduler_partials:
  - frequency: 1
    interval: epoch
    monitor: loss_train
    scheduler:
      _partial_: true
      _target_: torch.optim.lr_scheduler.ReduceLROnPlateau
      factor: 0.6
      min_lr: 1.0e-06
      patience: 100
    strict: true
load_original: false
params: {}
trainer:
  _target_: pytorch_lightning.Trainer
  accelerator: gpu
  accumulate_grad_batches: 1
  callbacks:
  - _target_: pytorch_lightning.callbacks.LearningRateMonitor
    log_momentum: false
    logging_interval: step
  - _target_: pytorch_lightning.callbacks.ModelCheckpoint
    every_n_epochs: 1
    filename: '{epoch}-{loss_val:.2f}'
    mode: min
    monitor: loss_val
    save_last: true
    save_top_k: 1
    verbose: false
  - _target_: pytorch_lightning.callbacks.TQDMProgressBar
    refresh_rate: 50
  - _target_: mattergen.common.data.callback.SetPropertyScalers
  check_val_every_n_epoch: 5
  devices: 8
  gradient_clip_algorithm: value
  gradient_clip_val: 0.5
  max_epochs: 900
  max_steps: 200000
  num_nodes: 9
  precision: 32
  strategy:
    _target_: pytorch_lightning.strategies.ddp.DDPStrategy
    find_unused_parameters: true