Upload MPTS-52 checkpoints and training files

EncDec-ODE-Gamma/checkpoint.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5fd2f9060e6de1382938db662f8d987aa38fa727cf9e96bd8f17571160290149
+size 49644411

EncDec-ODE-Gamma/train.yaml

@@ -0,0 +1,155 @@
+model:
+  si:
+    class_path: omg.si.stochastic_interpolants.StochasticInterpolants
+    init_args:
+      stochastic_interpolants:
+        # chemical species
+        - class_path: omg.si.single_stochastic_interpolant_identity.SingleStochasticInterpolantIdentity
+        # fractional coordinates
+        - class_path: omg.si.single_stochastic_interpolant.SingleStochasticInterpolant
+          init_args:
+            interpolant: 
+              class_path: omg.si.interpolants.PeriodicEncoderDecoderInterpolant
+              init_args:
+                switch_time: 0.6487086666110259
+                power: 1.0
+            gamma: 
+              class_path: omg.si.gamma.LatentGammaEncoderDecoder
+              init_args:
+                a: 1.9883383838119686
+                switch_time: 0.6487086666110259
+                power: 1.0
+            epsilon: null
+            differential_equation_type: "ODE"
+            integrator_kwargs:
+              method: "euler"
+            velocity_annealing_factor: 12.290317841755964
+            correct_center_of_mass_motion: true
+        # lattice vectors
+        - class_path: omg.si.single_stochastic_interpolant.SingleStochasticInterpolant
+          init_args:
+            interpolant: omg.si.interpolants.TrigonometricInterpolant
+            gamma: 
+              class_path: omg.si.gamma.LatentGammaSqrt
+              init_args:
+                a: 0.21935645939922985
+            epsilon: 
+              class_path: omg.si.epsilon.VanishingEpsilon
+              init_args:
+                c: 9.431054439782873
+                mu: 0.21809909486896933
+                sigma: 0.03292165737293197
+            differential_equation_type: "SDE"
+            integrator_kwargs:
+              method: "euler"
+              dt: 0.001218559336848557
+            velocity_annealing_factor: 4.302804708170181
+            correct_center_of_mass_motion: false
+      data_fields:
+        # if the order of the data_fields changes,
+        # the order of the above StochasticInterpolant inputs must also change
+        - "species"
+        - "pos"
+        - "cell"
+      integration_time_steps: 820
+  relative_si_costs:
+    species_loss: 0.0
+    pos_loss_b: 0.689192251322191
+    cell_loss_b: 0.12351464867571432
+    cell_loss_z: 0.18729310000209468
+  sampler:
+    class_path: omg.sampler.sample_from_rng.SampleFromRNG
+    init_args:
+      pos_distribution: null
+      cell_distribution: 
+        class_path: omg.sampler.distributions.InformedLatticeDistribution
+        init_args:
+          dataset_name: mpts_52
+      species_distribution:
+        class_path: omg.sampler.distributions.MirrorData
+  model:
+    class_path: omg.model.model.Model
+    init_args:
+      encoder:
+        class_path: omg.model.encoders.cspnet_full.CSPNetFull
+      head:
+        class_path: omg.model.heads.pass_through.PassThrough
+      time_embedder:
+        class_path: omg.model.model_utils.SinusoidalTimeEmbeddings
+        init_args:
+          dim: 256
+  use_min_perm_dist: False
+  float_32_matmul_precision: "high"
+  validation_mode: "match_rate"
+  number_cpus: 7
+  dataset_name: "mpts_52"
+data:
+  train_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/mpts_52/train.lmdb"
+      niggli: False
+  val_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/mpts_52/val.lmdb"
+      niggli: False
+  predict_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/mpts_52/test.lmdb"
+      niggli: False
+  batch_size: 256
+  num_workers: 4
+  pin_memory: True
+  persistent_workers: True
+trainer:
+  callbacks:
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_loss_total"
+        save_top_k: 1
+        monitor: "val_loss_total"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_match_rate"
+        save_top_k: 1
+        monitor: "match_rate"
+        save_weights_only: true
+        mode: 'max'
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_rmsd"
+        save_top_k: 1
+        monitor: "mean_rmsd"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        save_top_k: -1  # Store every checkpoint after 100 epochs.
+        monitor: "val_loss_total"
+        every_n_epochs: 100
+        save_weights_only: false
+  gradient_clip_val: 0.5
+  num_sanity_val_steps: 0
+  precision: "32-true"
+  max_epochs: 2000
+  enable_progress_bar: true
+  limit_val_batches: 0.5
+  check_val_every_n_epoch: 100
+optimizer:
+  class_path: torch.optim.Adam
+  init_args:
+    lr: 0.00047748599389170053

EncDec-SDE-Gamma/checkpoint.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ecffeaba4cee77a76640b78e1954ee032f58e7e7b2662ed850dfa8ff5123b1fd
+size 49644475

EncDec-SDE-Gamma/train.yaml

@@ -0,0 +1,155 @@
+model:
+  si:
+    class_path: omg.si.stochastic_interpolants.StochasticInterpolants
+    init_args:
+      stochastic_interpolants:
+        # chemical species
+        - class_path: omg.si.single_stochastic_interpolant_identity.SingleStochasticInterpolantIdentity
+        # fractional coordinates
+        - class_path: omg.si.single_stochastic_interpolant.SingleStochasticInterpolant
+          init_args:
+            interpolant: 
+              class_path: omg.si.interpolants.PeriodicEncoderDecoderInterpolant
+              init_args:
+                switch_time: 0.42184997325946555
+                power: 0.5
+            gamma: 
+              class_path: omg.si.gamma.LatentGammaEncoderDecoder
+              init_args:
+                a: 0.03989185248799893
+                switch_time: 0.42184997325946555
+                power: 0.5
+            epsilon: 
+              class_path: omg.si.epsilon.VanishingEpsilon
+              init_args:
+                c: 2.3996529332194574
+                mu: 0.25251095399328916
+                sigma: 0.03759134500470063
+            differential_equation_type: "SDE"
+            integrator_kwargs:
+              method: "euler"
+              dt: 0.0014076164225116372
+            velocity_annealing_factor: 3.7755089557808477
+            correct_center_of_mass_motion: true
+        # lattice vectors
+        - class_path: omg.si.single_stochastic_interpolant.SingleStochasticInterpolant
+          init_args:
+            interpolant: omg.si.interpolants.LinearInterpolant
+            gamma: 
+              class_path: omg.si.gamma.LatentGammaSqrt
+              init_args:
+                a: 4.961271013084809
+            epsilon: null
+            differential_equation_type: "ODE"
+            integrator_kwargs:
+              method: "euler"
+            velocity_annealing_factor: 1.1379701544400436
+            correct_center_of_mass_motion: false
+      data_fields:
+        # if the order of the data_fields changes,
+        # the order of the above StochasticInterpolant inputs must also change
+        - "species"
+        - "pos"
+        - "cell"
+      integration_time_steps: 710
+  relative_si_costs:
+    species_loss: 0.0
+    pos_loss_b: 0.6143090042317803
+    pos_loss_z: 0.3794040725288834
+    cell_loss_b: 0.00628692323933625
+  sampler:
+    class_path: omg.sampler.sample_from_rng.SampleFromRNG
+    init_args:
+      pos_distribution: null
+      cell_distribution: 
+        class_path: omg.sampler.distributions.InformedLatticeDistribution
+        init_args:
+          dataset_name: mpts_52
+      species_distribution:
+        class_path: omg.sampler.distributions.MirrorData
+  model:
+    class_path: omg.model.model.Model
+    init_args:
+      encoder:
+        class_path: omg.model.encoders.cspnet_full.CSPNetFull
+      head:
+        class_path: omg.model.heads.pass_through.PassThrough
+      time_embedder:
+        class_path: omg.model.model_utils.SinusoidalTimeEmbeddings
+        init_args:
+          dim: 256
+  use_min_perm_dist: False
+  float_32_matmul_precision: "high"
+  validation_mode: "match_rate"
+  number_cpus: 7
+  dataset_name: "mpts_52"
+data:
+  train_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/mpts_52/train.lmdb"
+      niggli: False
+  val_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/mpts_52/val.lmdb"
+      niggli: False
+  predict_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/mpts_52/test.lmdb"
+      niggli: False
+  batch_size: 32
+  num_workers: 4
+  pin_memory: True
+  persistent_workers: True
+trainer:
+  callbacks:
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_loss_total"
+        save_top_k: 1
+        monitor: "val_loss_total"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_match_rate"
+        save_top_k: 1
+        monitor: "match_rate"
+        save_weights_only: true
+        mode: 'max'
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_rmsd"
+        save_top_k: 1
+        monitor: "mean_rmsd"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        save_top_k: -1  # Store every checkpoint after 100 epochs.
+        monitor: "val_loss_total"
+        every_n_epochs: 100
+        save_weights_only: false
+  gradient_clip_val: 0.5
+  num_sanity_val_steps: 0
+  precision: "32-true"
+  max_epochs: 2000
+  enable_progress_bar: true
+  limit_val_batches: 0.5
+  check_val_every_n_epoch: 100
+optimizer:
+  class_path: torch.optim.Adam
+  init_args:
+    lr: 0.00018567271191860665

Linear-ODE-Gamma/checkpoint.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4a2de33650021f1e54c181668bfb34de6b845c4b5ce6422abb4f44e85f31ebb5
+size 49644411

Linear-ODE-Gamma/train.yaml

@@ -0,0 +1,142 @@
+model:
+  si:
+    class_path: omg.si.stochastic_interpolants.StochasticInterpolants
+    init_args:
+      stochastic_interpolants:
+        # chemical species
+        - class_path: omg.si.single_stochastic_interpolant_identity.SingleStochasticInterpolantIdentity
+        # fractional coordinates
+        - class_path: omg.si.single_stochastic_interpolant.SingleStochasticInterpolant
+          init_args:
+            interpolant: omg.si.interpolants.PeriodicLinearInterpolant
+            gamma: 
+              class_path: omg.si.gamma.LatentGammaSqrt
+              init_args:
+                a: 0.2575112227566439
+            epsilon: null
+            differential_equation_type: "ODE"
+            integrator_kwargs:
+              method: "euler"
+            velocity_annealing_factor: 7.7611189744870925
+            correct_center_of_mass_motion: true
+        # lattice vectors
+        - class_path: omg.si.single_stochastic_interpolant.SingleStochasticInterpolant
+          init_args:
+            interpolant: omg.si.interpolants.TrigonometricInterpolant
+            gamma: 
+              class_path: omg.si.gamma.LatentGammaSqrt
+              init_args:
+                a: 2.9759856920732597
+            epsilon: null
+            differential_equation_type: "ODE"
+            integrator_kwargs:
+              method: "euler"
+            velocity_annealing_factor: 4.116061496782678
+            correct_center_of_mass_motion: false
+      data_fields:
+        # if the order of the data_fields changes,
+        # the order of the above StochasticInterpolant inputs must also change
+        - "species"
+        - "pos"
+        - "cell"
+      integration_time_steps: 690
+  relative_si_costs:
+    species_loss: 0.0
+    pos_loss_b: 0.9976417941296929
+    cell_loss_b: 0.002358205870307133
+  sampler:
+    class_path: omg.sampler.sample_from_rng.SampleFromRNG
+    init_args:
+      pos_distribution: null
+      cell_distribution: 
+        class_path: omg.sampler.distributions.InformedLatticeDistribution
+        init_args:
+          dataset_name: mpts_52
+      species_distribution:
+        class_path: omg.sampler.distributions.MirrorData
+  model:
+    class_path: omg.model.model.Model
+    init_args:
+      encoder:
+        class_path: omg.model.encoders.cspnet_full.CSPNetFull
+      head:
+        class_path: omg.model.heads.pass_through.PassThrough
+      time_embedder:
+        class_path: omg.model.model_utils.SinusoidalTimeEmbeddings
+        init_args:
+          dim: 256
+  use_min_perm_dist: False
+  float_32_matmul_precision: "high"
+  validation_mode: "match_rate"
+  number_cpus: 7
+  dataset_name: "mpts_52"
+data:
+  train_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/mpts_52/train.lmdb"
+      niggli: False
+  val_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/mpts_52/val.lmdb"
+      niggli: False
+  predict_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/mpts_52/test.lmdb"
+      niggli: False
+  batch_size: 128
+  num_workers: 4
+  pin_memory: True
+  persistent_workers: True
+trainer:
+  callbacks:
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_loss_total"
+        save_top_k: 1
+        monitor: "val_loss_total"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_match_rate"
+        save_top_k: 1
+        monitor: "match_rate"
+        save_weights_only: true
+        mode: 'max'
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_rmsd"
+        save_top_k: 1
+        monitor: "mean_rmsd"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        save_top_k: -1  # Store every checkpoint after 100 epochs.
+        monitor: "val_loss_total"
+        every_n_epochs: 100
+        save_weights_only: false
+  gradient_clip_val: 0.5
+  num_sanity_val_steps: 0
+  precision: "32-true"
+  max_epochs: 2000
+  enable_progress_bar: true
+  limit_val_batches: 0.5
+  check_val_every_n_epoch: 100
+optimizer:
+  class_path: torch.optim.Adam
+  init_args:
+    lr: 4.006249666984122e-05

Linear-ODE/checkpoint.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1c189017e4d30b78995435af4e8dc42ca8942f64633d662762928680c536e211
+size 49644411

Linear-ODE/train.yaml

@@ -0,0 +1,134 @@
+model:
+  si:
+    class_path: omg.si.stochastic_interpolants.StochasticInterpolants
+    init_args:
+      stochastic_interpolants:
+        # chemical species
+        - class_path: omg.si.single_stochastic_interpolant_identity.SingleStochasticInterpolantIdentity
+        # fractional coordinates
+        - class_path: omg.si.single_stochastic_interpolant.SingleStochasticInterpolant
+          init_args:
+            interpolant: omg.si.interpolants.PeriodicLinearInterpolant
+            gamma: null
+            epsilon: null
+            differential_equation_type: "ODE"
+            integrator_kwargs:
+              method: "euler"
+            velocity_annealing_factor: 12.752963137656907
+            correct_center_of_mass_motion: true
+        # lattice vectors
+        - class_path: omg.si.single_stochastic_interpolant.SingleStochasticInterpolant
+          init_args:
+            interpolant: omg.si.interpolants.LinearInterpolant
+            gamma: null
+            epsilon: null
+            differential_equation_type: "ODE"
+            integrator_kwargs:
+              method: "euler"
+            velocity_annealing_factor: 0.9964121490291458
+            correct_center_of_mass_motion: false
+      data_fields:
+        # if the order of the data_fields changes,
+        # the order of the above StochasticInterpolant inputs must also change
+        - "species"
+        - "pos"
+        - "cell"
+      integration_time_steps: 100
+  relative_si_costs:
+    species_loss: 0.0
+    pos_loss_b: 0.9983149306572928
+    cell_loss_b: 0.0016850693427072152
+  sampler:
+    class_path: omg.sampler.sample_from_rng.SampleFromRNG
+    init_args:
+      pos_distribution: null
+      cell_distribution:
+        class_path: omg.sampler.distributions.InformedLatticeDistribution
+        init_args:
+          dataset_name: mpts_52
+      species_distribution:
+        class_path: omg.sampler.distributions.MirrorData
+  model:
+    class_path: omg.model.model.Model
+    init_args:
+      encoder:
+        class_path: omg.model.encoders.cspnet_full.CSPNetFull
+      head:
+        class_path: omg.model.heads.pass_through.PassThrough
+      time_embedder:
+        class_path: omg.model.model_utils.SinusoidalTimeEmbeddings
+        init_args:
+          dim: 256
+  use_min_perm_dist: False
+  float_32_matmul_precision: "high"
+  validation_mode: "match_rate"
+  dataset_name: "mpts_52"
+data:
+  train_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/mpts_52/train.lmdb"
+      niggli: False
+  val_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/mpts_52/val.lmdb"
+      niggli: False
+  predict_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/mpts_52/test.lmdb"
+      niggli: False
+  batch_size: 512
+  num_workers: 4
+  pin_memory: True
+  persistent_workers: True
+trainer:
+  callbacks:
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_loss_total"
+        save_top_k: 1
+        monitor: "val_loss_total"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_match_rate"
+        save_top_k: 1
+        monitor: "match_rate"
+        save_weights_only: true
+        mode: 'max'
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_rmsd"
+        save_top_k: 1
+        monitor: "mean_rmsd"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        save_top_k: -1  # Store every checkpoint after 100 epochs.
+        monitor: "val_loss_total"
+        every_n_epochs: 100
+        save_weights_only: false
+  gradient_clip_val: 0.5
+  num_sanity_val_steps: 0
+  precision: "32-true"
+  max_epochs: 10000
+  enable_progress_bar: false
+  check_val_every_n_epoch: 100
+optimizer:
+  class_path: torch.optim.Adam
+  init_args:
+    lr: 0.0005546288717347031

Linear-SDE-Gamma/checkpoint.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a44fb11e6327471ca725a66137dee984529a626330cd137b51d6882ea39996f4
+size 148120276

Linear-SDE-Gamma/train.yaml

@@ -0,0 +1,156 @@
+model:
+  si:
+    class_path: omg.si.stochastic_interpolants.StochasticInterpolants
+    init_args:
+      stochastic_interpolants:
+        # chemical species
+        - class_path: omg.si.single_stochastic_interpolant_identity.SingleStochasticInterpolantIdentity
+        # fractional coordinates
+        - class_path: omg.si.single_stochastic_interpolant.SingleStochasticInterpolant
+          init_args:
+            interpolant: omg.si.interpolants.PeriodicLinearInterpolant
+            gamma: 
+              class_path: omg.si.gamma.LatentGammaSqrt
+              init_args:
+                a: 0.06285652866840548
+            epsilon: 
+              class_path: omg.si.epsilon.VanishingEpsilon
+              init_args:
+                c: 6.097168392667226
+                mu: 0.21833859329765842
+                sigma: 0.04985718977712428
+            differential_equation_type: "SDE"
+            integrator_kwargs:
+              method: "euler"
+              dt: 0.0032297736033797264
+            velocity_annealing_factor: 11.58289329358004
+            correct_center_of_mass_motion: true
+        # lattice vectors
+        - class_path: omg.si.single_stochastic_interpolant.SingleStochasticInterpolant
+          init_args:
+            interpolant: omg.si.interpolants.LinearInterpolant
+            gamma: 
+              class_path: omg.si.gamma.LatentGammaSqrt
+              init_args:
+                a: 0.1317493001266121
+            epsilon: 
+              class_path: omg.si.epsilon.VanishingEpsilon
+              init_args:
+                c: 9.612495617660462
+                mu: 0.08389382419092543
+                sigma: 0.033192886798663945
+            differential_equation_type: "SDE"
+            integrator_kwargs:
+              method: "euler"
+              dt: 0.0032297736033797264
+            velocity_annealing_factor: 5.081210983525862
+            correct_center_of_mass_motion: false
+      data_fields:
+        # if the order of the data_fields changes,
+        # the order of the above StochasticInterpolant inputs must also change
+        - "species"
+        - "pos"
+        - "cell"
+      integration_time_steps: 310
+  relative_si_costs:
+    species_loss: 0.0
+    pos_loss_b: 0.007345481151868809
+    pos_loss_z: 0.9153543617007412
+    cell_loss_b: 0.06421063793348068
+    cell_loss_z: 0.013089519213909303
+  sampler:
+    class_path: omg.sampler.sample_from_rng.SampleFromRNG
+    init_args:
+      pos_distribution: null
+      cell_distribution: 
+        class_path: omg.sampler.distributions.InformedLatticeDistribution
+        init_args:
+          dataset_name: mpts_52
+      species_distribution:
+        class_path: omg.sampler.distributions.MirrorData
+  model:
+    class_path: omg.model.model.Model
+    init_args:
+      encoder:
+        class_path: omg.model.encoders.cspnet_full.CSPNetFull
+      head:
+        class_path: omg.model.heads.pass_through.PassThrough
+      time_embedder:
+        class_path: omg.model.model_utils.SinusoidalTimeEmbeddings
+        init_args:
+          dim: 256
+  use_min_perm_dist: False
+  float_32_matmul_precision: "high"
+  validation_mode: "match_rate"
+  number_cpus: 7
+  dataset_name: "mpts_52"
+data:
+  train_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/mpts_52/train.lmdb"
+      niggli: False
+  val_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/mpts_52/val.lmdb"
+      niggli: False
+  predict_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/mpts_52/test.lmdb"
+      niggli: False
+  batch_size: 256
+  num_workers: 4
+  pin_memory: True
+  persistent_workers: True
+trainer:
+  callbacks:
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_loss_total"
+        save_top_k: 1
+        monitor: "val_loss_total"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_match_rate"
+        save_top_k: 1
+        monitor: "match_rate"
+        save_weights_only: true
+        mode: 'max'
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_rmsd"
+        save_top_k: 1
+        monitor: "mean_rmsd"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        save_top_k: -1  # Store every checkpoint after 100 epochs.
+        monitor: "val_loss_total"
+        every_n_epochs: 100
+        save_weights_only: false
+  gradient_clip_val: 0.5
+  num_sanity_val_steps: 0
+  precision: "32-true"
+  max_epochs: 2000
+  enable_progress_bar: true
+  limit_val_batches: 0.5
+  check_val_every_n_epoch: 100
+optimizer:
+  class_path: torch.optim.Adam
+  init_args:
+    lr: 0.0002629870131361822

Trig-ODE-Gamma/checkpoint.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2300d4bfa2684ba81cdaac6277e4fb2205bb8b8c380022e1346169ffd9eda1fa
+size 148107354

Trig-ODE-Gamma/train.yaml

@@ -0,0 +1,149 @@
+model:
+  si:
+    class_path: omg.si.stochastic_interpolants.StochasticInterpolants
+    init_args:
+      stochastic_interpolants:
+        # chemical species
+        - class_path: omg.si.single_stochastic_interpolant_identity.SingleStochasticInterpolantIdentity
+        # fractional coordinates
+        - class_path: omg.si.single_stochastic_interpolant.SingleStochasticInterpolant
+          init_args:
+            interpolant: omg.si.interpolants.PeriodicTrigonometricInterpolant
+            gamma: 
+              class_path: omg.si.gamma.LatentGammaSqrt
+              init_args:
+                a: 0.03337798944475465
+            epsilon: null
+            differential_equation_type: "ODE"
+            integrator_kwargs:
+              method: "euler"
+            velocity_annealing_factor: 13.545929738762764
+            correct_center_of_mass_motion: true
+        # lattice vectors
+        - class_path: omg.si.single_stochastic_interpolant.SingleStochasticInterpolant
+          init_args:
+            interpolant: omg.si.interpolants.LinearInterpolant
+            gamma: 
+              class_path: omg.si.gamma.LatentGammaSqrt
+              init_args:
+                a: 0.017261010545698854
+            epsilon: 
+              class_path: omg.si.epsilon.VanishingEpsilon
+              init_args:
+                c: 0.8758328635983847
+                mu: 0.29744423858325936
+                sigma: 0.0052236060273636595
+            differential_equation_type: "SDE"
+            integrator_kwargs:
+              method: "euler"
+              dt: 0.0012811297783628106
+            velocity_annealing_factor: 2.380421528846764
+            correct_center_of_mass_motion: false
+      data_fields:
+        # if the order of the data_fields changes,
+        # the order of the above StochasticInterpolant inputs must also change
+        - "species"
+        - "pos"
+        - "cell"
+      integration_time_steps: 780
+  relative_si_costs:
+    species_loss: 0.0
+    pos_loss_b: 0.983015308902659
+    cell_loss_b: 0.01673796318800159
+    cell_loss_z: 0.0002467279093394523
+  sampler:
+    class_path: omg.sampler.sample_from_rng.SampleFromRNG
+    init_args:
+      pos_distribution: null
+      cell_distribution: 
+        class_path: omg.sampler.distributions.InformedLatticeDistribution
+        init_args:
+          dataset_name: mpts_52
+      species_distribution:
+        class_path: omg.sampler.distributions.MirrorData
+  model:
+    class_path: omg.model.model.Model
+    init_args:
+      encoder:
+        class_path: omg.model.encoders.cspnet_full.CSPNetFull
+      head:
+        class_path: omg.model.heads.pass_through.PassThrough
+      time_embedder:
+        class_path: omg.model.model_utils.SinusoidalTimeEmbeddings
+        init_args:
+          dim: 256
+  use_min_perm_dist: False
+  float_32_matmul_precision: "high"
+  validation_mode: "match_rate"
+  number_cpus: 7
+  dataset_name: "mpts_52"
+data:
+  train_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/mpts_52/train.lmdb"
+      niggli: False
+  val_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/mpts_52/val.lmdb"
+      niggli: False
+  predict_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/mpts_52/test.lmdb"
+      niggli: False
+  batch_size: 32
+  num_workers: 4
+  pin_memory: True
+  persistent_workers: True
+trainer:
+  callbacks:
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_loss_total"
+        save_top_k: 1
+        monitor: "val_loss_total"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_match_rate"
+        save_top_k: 1
+        monitor: "match_rate"
+        save_weights_only: true
+        mode: 'max'
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_rmsd"
+        save_top_k: 1
+        monitor: "mean_rmsd"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        save_top_k: -1  # Store every checkpoint after 100 epochs.
+        monitor: "val_loss_total"
+        every_n_epochs: 100
+        save_weights_only: false
+  gradient_clip_val: 0.5
+  num_sanity_val_steps: 0
+  precision: "32-true"
+  max_epochs: 2000
+  enable_progress_bar: true
+  limit_val_batches: 0.5
+  check_val_every_n_epoch: 100
+optimizer:
+  class_path: torch.optim.Adam
+  init_args:
+    lr: 8.341737878937152e-05

Trig-ODE/checkpoint.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4def0d650b785ca9fe3b818229d140d1885d1ff06bc64e75c15fbb0fb553f2e8
+size 148107226

Trig-ODE/train.yaml

@@ -0,0 +1,152 @@
+model:
+  si:
+    class_path: omg.si.stochastic_interpolants.StochasticInterpolants
+    init_args:
+      stochastic_interpolants:
+        # chemical species
+        - class_path: omg.si.single_stochastic_interpolant_identity.SingleStochasticInterpolantIdentity
+        # fractional coordinates
+        - class_path: omg.si.single_stochastic_interpolant.SingleStochasticInterpolant
+          init_args:
+            interpolant: omg.si.interpolants.PeriodicTrigonometricInterpolant
+            gamma: null
+            epsilon: null
+            differential_equation_type: "ODE"
+            integrator_kwargs:
+              method: "euler"
+            velocity_annealing_factor: 12.34532470785473
+            correct_center_of_mass_motion: true
+        # lattice vectors
+        - class_path: omg.si.single_stochastic_interpolant.SingleStochasticInterpolant
+          init_args:
+            interpolant: 
+              class_path: omg.si.interpolants.EncoderDecoderInterpolant
+              init_args:
+                switch_time: 0.4080329374611481
+                power: 0.5
+            gamma: 
+              class_path: omg.si.gamma.LatentGammaEncoderDecoder
+              init_args:
+                a: 5.270616141661882
+                switch_time: 0.4080329374611481
+                power: 0.5
+            epsilon: 
+              class_path: omg.si.epsilon.VanishingEpsilon
+              init_args:
+                c: 4.354817546796119
+                mu: 0.2923928859901851
+                sigma: 0.04742031136770322
+            differential_equation_type: "SDE"
+            integrator_kwargs:
+              method: "euler"
+              dt: 0.005905325524508953
+            velocity_annealing_factor: 3.6141717997883447
+            correct_center_of_mass_motion: false
+      data_fields:
+        # if the order of the data_fields changes,
+        # the order of the above StochasticInterpolant inputs must also change
+        - "species"
+        - "pos"
+        - "cell"
+      integration_time_steps: 170
+  relative_si_costs:
+    species_loss: 0.0
+    pos_loss_b: 0.9967455480681945
+    cell_loss_b: 0.002271914623580616
+    cell_loss_z: 0.0009825373082248405
+  sampler:
+    class_path: omg.sampler.sample_from_rng.SampleFromRNG
+    init_args:
+      pos_distribution: null
+      cell_distribution: 
+        class_path: omg.sampler.distributions.InformedLatticeDistribution
+        init_args:
+          dataset_name: mpts_52
+      species_distribution:
+        class_path: omg.sampler.distributions.MirrorData
+  model:
+    class_path: omg.model.model.Model
+    init_args:
+      encoder:
+        class_path: omg.model.encoders.cspnet_full.CSPNetFull
+      head:
+        class_path: omg.model.heads.pass_through.PassThrough
+      time_embedder:
+        class_path: omg.model.model_utils.SinusoidalTimeEmbeddings
+        init_args:
+          dim: 256
+  use_min_perm_dist: False
+  float_32_matmul_precision: "high"
+  validation_mode: "match_rate"
+  number_cpus: 7
+  dataset_name: "mpts_52"
+data:
+  train_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/mpts_52/train.lmdb"
+      niggli: False
+  val_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/mpts_52/val.lmdb"
+      niggli: False
+  predict_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/mpts_52/test.lmdb"
+      niggli: False
+  batch_size: 32
+  num_workers: 4
+  pin_memory: True
+  persistent_workers: True
+trainer:
+  callbacks:
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_loss_total"
+        save_top_k: 1
+        monitor: "val_loss_total"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_match_rate"
+        save_top_k: 1
+        monitor: "match_rate"
+        save_weights_only: true
+        mode: 'max'
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_rmsd"
+        save_top_k: 1
+        monitor: "mean_rmsd"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        save_top_k: -1  # Store every checkpoint after 100 epochs.
+        monitor: "val_loss_total"
+        every_n_epochs: 100
+        save_weights_only: false
+  gradient_clip_val: 0.5
+  num_sanity_val_steps: 0
+  precision: "32-true"
+  max_epochs: 2000
+  enable_progress_bar: true
+  limit_val_batches: 0.5
+  check_val_every_n_epoch: 100
+optimizer:
+  class_path: torch.optim.Adam
+  init_args:
+    lr: 3.629490873183724e-05

Trig-SDE-Gamma/checkpoint.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7fd313b4ce18394628596dce790ba5a6e58bd37a1fe5f7c0c78be6a2e3ab0fb5
+size 148082778

Trig-SDE-Gamma/train.yaml

@@ -0,0 +1,146 @@
+model:
+  si:
+    class_path: omg.si.stochastic_interpolants.StochasticInterpolants
+    init_args:
+      stochastic_interpolants:
+        # chemical species
+        - class_path: omg.si.single_stochastic_interpolant_identity.SingleStochasticInterpolantIdentity
+        # fractional coordinates
+        - class_path: omg.si.single_stochastic_interpolant.SingleStochasticInterpolant
+          init_args:
+            interpolant: omg.si.interpolants.PeriodicTrigonometricInterpolant
+            gamma: 
+              class_path: omg.si.gamma.LatentGammaSqrt
+              init_args:
+                a: 0.049242906264339095
+            epsilon: 
+              class_path: omg.si.epsilon.VanishingEpsilon
+              init_args:
+                c: 9.418703639528207
+                mu: 0.1967838464371502
+                sigma: 0.040028404066547216
+            differential_equation_type: "SDE"
+            integrator_kwargs:
+              method: "euler"
+              dt: 0.0013504737289622426
+            velocity_annealing_factor: 11.483173553510193
+            correct_center_of_mass_motion: true
+        # lattice vectors
+        - class_path: omg.si.single_stochastic_interpolant.SingleStochasticInterpolant
+          init_args:
+            interpolant: omg.si.interpolants.TrigonometricInterpolant
+            gamma: null
+            epsilon: null
+            differential_equation_type: "ODE"
+            integrator_kwargs:
+              method: "euler"
+            velocity_annealing_factor: 0.4337356395028541
+            correct_center_of_mass_motion: false
+      data_fields:
+        # if the order of the data_fields changes,
+        # the order of the above StochasticInterpolant inputs must also change
+        - "species"
+        - "pos"
+        - "cell"
+      integration_time_steps: 740
+  relative_si_costs:
+    species_loss: 0.0
+    pos_loss_b: 0.24677273761024368
+    pos_loss_z: 0.7231540118244248
+    cell_loss_b: 0.030073250565331323
+  sampler:
+    class_path: omg.sampler.sample_from_rng.SampleFromRNG
+    init_args:
+      pos_distribution: null
+      cell_distribution: 
+        class_path: omg.sampler.distributions.InformedLatticeDistribution
+        init_args:
+          dataset_name: mpts_52
+      species_distribution:
+        class_path: omg.sampler.distributions.MirrorData
+  model:
+    class_path: omg.model.model.Model
+    init_args:
+      encoder:
+        class_path: omg.model.encoders.cspnet_full.CSPNetFull
+      head:
+        class_path: omg.model.heads.pass_through.PassThrough
+      time_embedder:
+        class_path: omg.model.model_utils.SinusoidalTimeEmbeddings
+        init_args:
+          dim: 256
+  use_min_perm_dist: True
+  float_32_matmul_precision: "high"
+  validation_mode: "match_rate"
+  number_cpus: 7
+  dataset_name: "mpts_52"
+data:
+  train_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/mpts_52/train.lmdb"
+      niggli: False
+  val_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/mpts_52/val.lmdb"
+      niggli: False
+  predict_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/mpts_52/test.lmdb"
+      niggli: False
+  batch_size: 32
+  num_workers: 4
+  pin_memory: True
+  persistent_workers: True
+trainer:
+  callbacks:
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_loss_total"
+        save_top_k: 1
+        monitor: "val_loss_total"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_match_rate"
+        save_top_k: 1
+        monitor: "match_rate"
+        save_weights_only: true
+        mode: 'max'
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_rmsd"
+        save_top_k: 1
+        monitor: "mean_rmsd"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        save_top_k: -1  # Store every checkpoint after 100 epochs.
+        monitor: "val_loss_total"
+        every_n_epochs: 100
+        save_weights_only: false
+  gradient_clip_val: 0.5
+  num_sanity_val_steps: 0
+  precision: "32-true"
+  max_epochs: 2000
+  enable_progress_bar: true
+  limit_val_batches: 0.5
+  check_val_every_n_epoch: 100
+optimizer:
+  class_path: torch.optim.Adam
+  init_args:
+    lr: 9.320780466656964e-05

VESBD-ODE/checkpoint.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:64cb0f57a91fd0d53d8710f6a6e5e86baefc4ee18a5bf7c0db7e74aad10e89b5
+size 49644411

VESBD-ODE/train.yaml

@@ -0,0 +1,156 @@
+model:
+  si:
+    class_path: omg.si.stochastic_interpolants.StochasticInterpolants
+    init_args:
+      stochastic_interpolants:
+        # chemical species
+        - class_path: omg.si.single_stochastic_interpolant_identity.SingleStochasticInterpolantIdentity
+        # fractional coordinates
+        - class_path: omg.si.single_stochastic_interpolant_os.SingleStochasticInterpolantOS
+          init_args:
+            interpolant: 
+              class_path: omg.si.interpolants.PeriodicScoreBasedDiffusionModelInterpolantVE
+              init_args:
+                sigma:
+                  class_path: omg.si.sigma.GeometricSigma
+                  init_args:
+                    sigma_min: 0.004705415831077799
+                    sigma_max: 0.9967130801483843
+            epsilon: null
+            differential_equation_type: "ODE"
+            integrator_kwargs:
+              method: "euler"
+            velocity_annealing_factor: 8.284579088906593
+            correct_center_of_mass_motion: true
+            predict_velocity: true
+        # lattice vectors
+        - class_path: omg.si.single_stochastic_interpolant.SingleStochasticInterpolant
+          init_args:
+            interpolant: omg.si.interpolants.LinearInterpolant
+            gamma: 
+              class_path: omg.si.gamma.LatentGammaSqrt
+              init_args:
+                a: 0.016616684357970132
+            epsilon: 
+              class_path: omg.si.epsilon.VanishingEpsilon
+              init_args:
+                c: 3.9372558236242052
+                mu: 0.2649556265396099
+                sigma: 0.03578203230805775
+            differential_equation_type: "SDE"
+            integrator_kwargs:
+              method: "euler"
+              dt: 0.0015144158387556672
+            velocity_annealing_factor: 0.42775377056075214
+            correct_center_of_mass_motion: false
+      data_fields:
+        # if the order of the data_fields changes,
+        # the order of the above StochasticInterpolant inputs must also change
+        - "species"
+        - "pos"
+        - "cell"
+      integration_time_steps: 660
+  relative_si_costs:
+    species_loss: 0.0
+    pos_loss_b: 0.9813067351598369
+    cell_loss_b: 0.0005256953168558359
+    cell_loss_z: 0.018167569523307267
+  sampler:
+    class_path: omg.sampler.sample_from_rng.SampleFromRNG
+    init_args:
+      pos_distribution: 
+        class_path: omg.sampler.distributions.NormalDistribution
+        init_args:
+          scale: 9.77149759679434
+      cell_distribution: 
+        class_path: omg.sampler.distributions.InformedLatticeDistribution
+        init_args:
+          dataset_name: mpts_52
+      species_distribution:
+        class_path: omg.sampler.distributions.MirrorData
+  model:
+    class_path: omg.model.model.Model
+    init_args:
+      encoder:
+        class_path: omg.model.encoders.cspnet_full.CSPNetFull
+      head:
+        class_path: omg.model.heads.pass_through.PassThrough
+      time_embedder:
+        class_path: omg.model.model_utils.SinusoidalTimeEmbeddings
+        init_args:
+          dim: 256
+  use_min_perm_dist: False
+  float_32_matmul_precision: "high"
+  validation_mode: "match_rate"
+  number_cpus: 7
+  dataset_name: "mpts_52"
+data:
+  train_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/mpts_52/train.lmdb"
+      niggli: False
+  val_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/mpts_52/val.lmdb"
+      niggli: False
+  predict_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/mpts_52/test.lmdb"
+      niggli: False
+  batch_size: 256
+  num_workers: 4
+  pin_memory: True
+  persistent_workers: True
+trainer:
+  callbacks:
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_loss_total"
+        save_top_k: 1
+        monitor: "val_loss_total"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_match_rate"
+        save_top_k: 1
+        monitor: "match_rate"
+        save_weights_only: true
+        mode: 'max'
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_rmsd"
+        save_top_k: 1
+        monitor: "mean_rmsd"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        save_top_k: -1  # Store every checkpoint after 100 epochs.
+        monitor: "val_loss_total"
+        every_n_epochs: 100
+        save_weights_only: false
+  gradient_clip_val: 0.5
+  num_sanity_val_steps: 0
+  precision: "32-true"
+  max_epochs: 2000
+  enable_progress_bar: true
+  limit_val_batches: 0.5
+  check_val_every_n_epoch: 100
+optimizer:
+  class_path: torch.optim.Adam
+  init_args:
+    lr: 0.000296636127734534

VPSBD-ODE/checkpoint.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f39a43f38e06f45c8185d13fee14d4d14687b6979120958a9640c461bd0e6181
+size 148069280

VPSBD-ODE/train.yaml

@@ -0,0 +1,139 @@
+model:
+  si:
+    class_path: omg.si.stochastic_interpolants.StochasticInterpolants
+    init_args:
+      stochastic_interpolants:
+        # chemical species
+        - class_path: omg.si.single_stochastic_interpolant_identity.SingleStochasticInterpolantIdentity
+        # fractional coordinates
+        - class_path: omg.si.single_stochastic_interpolant_os.SingleStochasticInterpolantOS
+          init_args:
+            interpolant: omg.si.interpolants.PeriodicScoreBasedDiffusionModelInterpolant
+            epsilon: null
+            differential_equation_type: "ODE"
+            integrator_kwargs:
+              method: "euler"
+            velocity_annealing_factor: 6.613808424917352
+            correct_center_of_mass_motion: true
+            predict_velocity: true
+        # lattice vectors
+        - class_path: omg.si.single_stochastic_interpolant.SingleStochasticInterpolant
+          init_args:
+            interpolant: omg.si.interpolants.LinearInterpolant
+            gamma: null
+            epsilon: null
+            differential_equation_type: "ODE"
+            integrator_kwargs:
+              method: "euler"
+            velocity_annealing_factor: 2.447993013544224
+            correct_center_of_mass_motion: false
+      data_fields:
+        # if the order of the data_fields changes,
+        # the order of the above StochasticInterpolant inputs must also change
+        - "species"
+        - "pos"
+        - "cell"
+      integration_time_steps: 890
+  relative_si_costs:
+    species_loss: 0.0
+    pos_loss_b: 0.9597565150933746
+    cell_loss_b: 0.04024348490662539
+  sampler:
+    class_path: omg.sampler.sample_from_rng.SampleFromRNG
+    init_args:
+      pos_distribution: 
+        class_path: omg.sampler.distributions.NormalDistribution
+        init_args:
+          scale: 0.22006712732536396
+      cell_distribution: 
+        class_path: omg.sampler.distributions.InformedLatticeDistribution
+        init_args:
+          dataset_name: mpts_52
+      species_distribution:
+        class_path: omg.sampler.distributions.MirrorData
+  model:
+    class_path: omg.model.model.Model
+    init_args:
+      encoder:
+        class_path: omg.model.encoders.cspnet_full.CSPNetFull
+      head:
+        class_path: omg.model.heads.pass_through.PassThrough
+      time_embedder:
+        class_path: omg.model.model_utils.SinusoidalTimeEmbeddings
+        init_args:
+          dim: 256
+  use_min_perm_dist: True
+  float_32_matmul_precision: "high"
+  validation_mode: "match_rate"
+  number_cpus: 7
+  dataset_name: "mpts_52"
+data:
+  train_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/mpts_52/train.lmdb"
+      niggli: False
+  val_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/mpts_52/val.lmdb"
+      niggli: False
+  predict_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/mpts_52/test.lmdb"
+      niggli: False
+  batch_size: 64
+  num_workers: 4
+  pin_memory: True
+  persistent_workers: True
+trainer:
+  callbacks:
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_loss_total"
+        save_top_k: 1
+        monitor: "val_loss_total"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_match_rate"
+        save_top_k: 1
+        monitor: "match_rate"
+        save_weights_only: true
+        mode: 'max'
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_rmsd"
+        save_top_k: 1
+        monitor: "mean_rmsd"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        save_top_k: -1  # Store every checkpoint after 100 epochs.
+        monitor: "val_loss_total"
+        every_n_epochs: 100
+        save_weights_only: false
+  gradient_clip_val: 0.5
+  num_sanity_val_steps: 0
+  precision: "32-true"
+  max_epochs: 2000
+  enable_progress_bar: true
+  limit_val_batches: 0.5
+  check_val_every_n_epoch: 100
+optimizer:
+  class_path: torch.optim.Adam
+  init_args:
+    lr: 2.519765029616902e-05

VPSBD-SDE/checkpoint.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2d49763d63ce905025b23efa6babcc6ab237cb1684daec33f055b4452dd3825c
+size 49644475

VPSBD-SDE/train.yaml

@@ -0,0 +1,149 @@
+model:
+  si:
+    class_path: omg.si.stochastic_interpolants.StochasticInterpolants
+    init_args:
+      stochastic_interpolants:
+        # chemical species
+        - class_path: omg.si.single_stochastic_interpolant_identity.SingleStochasticInterpolantIdentity
+        # fractional coordinates
+        - class_path: omg.si.single_stochastic_interpolant_os.SingleStochasticInterpolantOS
+          init_args:
+            interpolant: omg.si.interpolants.PeriodicScoreBasedDiffusionModelInterpolant
+            epsilon: 
+              class_path: omg.si.epsilon.VanishingEpsilon
+              init_args:
+                c: 2.4729222108905815
+                mu: 0.17656358406313838
+                sigma: 0.02379822283154629
+            differential_equation_type: "SDE"
+            integrator_kwargs:
+              method: "euler"
+              dt: 0.0016661101253703237
+            velocity_annealing_factor: 6.459028320375323
+            correct_center_of_mass_motion: true
+            predict_velocity: true
+        # lattice vectors
+        - class_path: omg.si.single_stochastic_interpolant.SingleStochasticInterpolant
+          init_args:
+            interpolant: omg.si.interpolants.LinearInterpolant
+            gamma: 
+              class_path: omg.si.gamma.LatentGammaSqrt
+              init_args:
+                a: 3.683542379054881
+            epsilon: null
+            differential_equation_type: "ODE"
+            integrator_kwargs:
+              method: "euler"
+            velocity_annealing_factor: 0.6692350794589719
+            correct_center_of_mass_motion: false
+      data_fields:
+        # if the order of the data_fields changes,
+        # the order of the above StochasticInterpolant inputs must also change
+        - "species"
+        - "pos"
+        - "cell"
+      integration_time_steps: 600
+  relative_si_costs:
+    species_loss: 0.0
+    pos_loss_b: 0.6060249654155797
+    pos_loss_z: 0.3828230559814603
+    cell_loss_b: 0.011151978602959979
+  sampler:
+    class_path: omg.sampler.sample_from_rng.SampleFromRNG
+    init_args:
+      pos_distribution: 
+        class_path: omg.sampler.distributions.NormalDistribution
+        init_args:
+          scale: 2.2937003279036148
+      cell_distribution: 
+        class_path: omg.sampler.distributions.InformedLatticeDistribution
+        init_args:
+          dataset_name: mpts_52
+      species_distribution:
+        class_path: omg.sampler.distributions.MirrorData
+  model:
+    class_path: omg.model.model.Model
+    init_args:
+      encoder:
+        class_path: omg.model.encoders.cspnet_full.CSPNetFull
+      head:
+        class_path: omg.model.heads.pass_through.PassThrough
+      time_embedder:
+        class_path: omg.model.model_utils.SinusoidalTimeEmbeddings
+        init_args:
+          dim: 256
+  use_min_perm_dist: True
+  float_32_matmul_precision: "high"
+  validation_mode: "match_rate"
+  number_cpus: 7
+  dataset_name: "mpts_52"
+data:
+  train_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/mpts_52/train.lmdb"
+      niggli: False
+  val_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/mpts_52/val.lmdb"
+      niggli: False
+  predict_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/mpts_52/test.lmdb"
+      niggli: False
+  batch_size: 64
+  num_workers: 4
+  pin_memory: True
+  persistent_workers: True
+trainer:
+  callbacks:
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_loss_total"
+        save_top_k: 1
+        monitor: "val_loss_total"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_match_rate"
+        save_top_k: 1
+        monitor: "match_rate"
+        save_weights_only: true
+        mode: 'max'
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_rmsd"
+        save_top_k: 1
+        monitor: "mean_rmsd"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        save_top_k: -1  # Store every checkpoint after 100 epochs.
+        monitor: "val_loss_total"
+        every_n_epochs: 100
+        save_weights_only: false
+  gradient_clip_val: 0.5
+  num_sanity_val_steps: 0
+  precision: "32-true"
+  max_epochs: 2000
+  enable_progress_bar: true
+  limit_val_batches: 0.5
+  check_val_every_n_epoch: 100
+optimizer:
+  class_path: torch.optim.Adam
+  init_args:
+    lr: 0.0003030820420973639