Upload 22 files

EncDec-ODE-Gamma/checkpoint.ckpt

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

EncDec-ODE-Gamma/train.yaml

@@ -0,0 +1,143 @@
+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.796130965510696
+                power: 1.0
+            gamma: 
+              class_path: omg.si.gamma.LatentGammaEncoderDecoder
+              init_args:
+                a: 0.6557615904788995
+                switch_time: 0.796130965510696
+                power: 1.0
+            epsilon: null
+            differential_equation_type: "ODE"
+            integrator_kwargs:
+              method: "euler"
+            velocity_annealing_factor: 14.941666601494628
+            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.3178550359129071
+            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: 460
+  relative_si_costs:
+    species_loss: 0.0
+    pos_loss_b: 0.8563010628686587
+    cell_loss_b: 0.14369893713134133
+  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: perov_5
+      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"
+  dataset_name: "perov_5"
+data:
+  train_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/perov_5/train.lmdb"
+      niggli: True
+  val_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/perov_5/val.lmdb"
+      niggli: True
+  predict_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/perov_5/test.lmdb"
+      niggli: True
+  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: 6000
+  enable_progress_bar: false
+  check_val_every_n_epoch: 100
+optimizer:
+  class_path: torch.optim.Adam
+  init_args:
+    lr: 7.808103295004345e-05

EncDec-SDE-Gamma/checkpoint.ckpt

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

EncDec-SDE-Gamma/train.yaml

@@ -0,0 +1,153 @@
+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.6055018323069807
+                power: 1.0
+            gamma: 
+              class_path: omg.si.gamma.LatentGammaEncoderDecoder
+              init_args:
+                a: 8.454472851641802
+                switch_time: 0.6055018323069807
+                power: 1.0
+            epsilon: 
+              class_path: omg.si.epsilon.VanishingEpsilon
+              init_args:
+                c: 4.609299406421399
+                mu: 0.2674947568710694
+                sigma: 0.04906444616252471
+            differential_equation_type: "SDE"
+            integrator_kwargs:
+              method: "euler"
+              dt: 0.001074273488484323
+            velocity_annealing_factor: 14.554387706860773
+            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.1539379702797485
+            epsilon: null
+            differential_equation_type: "ODE"
+            integrator_kwargs:
+              method: "euler"
+            velocity_annealing_factor: 0.07461560076268103
+            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: 930
+  relative_si_costs:
+    species_loss: 0.0
+    pos_loss_b: 0.28276509270307465
+    pos_loss_z: 0.7168554318065845
+    cell_loss_b: 0.0003794754903409129
+  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: perov_5
+      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"
+  dataset_name: "perov_5"
+data:
+  train_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/perov_5/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/perov_5/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/perov_5/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: 6000
+  enable_progress_bar: false
+  check_val_every_n_epoch: 100
+optimizer:
+  class_path: torch.optim.Adam
+  init_args:
+    lr: 0.0002837109869864481

Linear-ODE-Gamma/checkpoint.ckpt

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

Linear-ODE-Gamma/train.yaml

@@ -0,0 +1,137 @@
+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.03386737488191369
+            epsilon: null
+            differential_equation_type: "ODE"
+            integrator_kwargs:
+              method: "euler"
+            velocity_annealing_factor: 0.007950108070075533
+            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: 12.194837909618993
+            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.9724021294519893
+    cell_loss_b: 0.0275978705480107
+  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: perov_5
+      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"
+  dataset_name: "perov_5"
+data:
+  train_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/perov_5/train.lmdb"
+      niggli: True
+  val_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/perov_5/val.lmdb"
+      niggli: True
+  predict_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/perov_5/test.lmdb"
+      niggli: True
+  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: 6000
+  enable_progress_bar: false
+  check_val_every_n_epoch: 100
+optimizer:
+  class_path: torch.optim.Adam
+  init_args:
+    lr: 3.6259796277646535e-05

Linear-ODE/checkpoint.ckpt

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

Linear-ODE/train.yaml

@@ -0,0 +1,143 @@
+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: 0.004755207270677389
+            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.46351945271978645
+                power: 1.0
+            gamma: 
+              class_path: omg.si.gamma.LatentGammaEncoderDecoder
+              init_args:
+                a: 0.8167071952445664
+                switch_time: 0.46351945271978645
+                power: 1.0
+            epsilon: null
+            differential_equation_type: "ODE"
+            integrator_kwargs:
+              method: "euler"
+            velocity_annealing_factor: 13.921408921615031
+            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: 480
+  relative_si_costs:
+    species_loss: 0.0
+    pos_loss_b: 0.9860929911452281
+    cell_loss_b: 0.01390700885477196
+  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: perov_5
+      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"
+  dataset_name: "perov_5"
+data:
+  train_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/perov_5/train.lmdb"
+      niggli: True
+  val_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/perov_5/val.lmdb"
+      niggli: True
+  predict_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/perov_5/test.lmdb"
+      niggli: True
+  batch_size: 1024
+  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: 6000
+  enable_progress_bar: false
+  check_val_every_n_epoch: 100
+optimizer:
+  class_path: torch.optim.Adam
+  init_args:
+    lr: 0.001147361965964576

Linear-SDE-Gamma/checkpoint.ckpt

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

Linear-SDE-Gamma/train.yaml

@@ -0,0 +1,147 @@
+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.027547642683482473
+            epsilon: 
+              class_path: omg.si.epsilon.VanishingEpsilon
+              init_args:
+                c: 8.26092465709134
+                mu: 0.1083235196756059
+                sigma: 0.03686939437589988
+            differential_equation_type: "SDE"
+            integrator_kwargs:
+              method: "euler"
+              dt: 0.001097909756936133
+            velocity_annealing_factor: 8.19603285406944
+            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.012871444447488238
+            epsilon: null
+            differential_equation_type: "ODE"
+            integrator_kwargs:
+              method: "euler"
+            velocity_annealing_factor: 1.4603041330880495
+            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: 910
+  relative_si_costs:
+    species_loss: 0.0
+    pos_loss_b: 0.0023778886849979398
+    pos_loss_z: 0.9924707469401747
+    cell_loss_b: 0.0051513643748272876
+  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: perov_5
+      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"
+  dataset_name: "perov_5"
+data:
+  train_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/perov_5/train.lmdb"
+      niggli: True
+  val_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/perov_5/val.lmdb"
+      niggli: True
+  predict_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/perov_5/test.lmdb"
+      niggli: True
+  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: 6000
+  enable_progress_bar: false
+  check_val_every_n_epoch: 100
+optimizer:
+  class_path: torch.optim.Adam
+  init_args:
+    lr: 1.3254493006246477e-05

Trig-ODE-Gamma/checkpoint.ckpt

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

Trig-ODE-Gamma/train.yaml

@@ -0,0 +1,137 @@
+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.02648500626802044
+            epsilon: null
+            differential_equation_type: "ODE"
+            integrator_kwargs:
+              method: "euler"
+            velocity_annealing_factor: 3.8566932544902413
+            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: 14.219455036917472
+            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: 970
+  relative_si_costs:
+    species_loss: 0.0
+    pos_loss_b: 0.8133671709485343
+    cell_loss_b: 0.1866328290514657
+  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: perov_5
+      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: "perov_5"
+data:
+  train_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/perov_5/train.lmdb"
+      niggli: True
+  val_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/perov_5/val.lmdb"
+      niggli: True
+  predict_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/perov_5/test.lmdb"
+      niggli: True
+  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: 6000
+  enable_progress_bar: false
+  check_val_every_n_epoch: 100
+optimizer:
+  class_path: torch.optim.Adam
+  init_args:
+    lr: 1.620271269284964e-05

Trig-ODE/checkpoint.ckpt

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

Trig-ODE/train.yaml

@@ -0,0 +1,137 @@
+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: 14.9938835509918
+            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.021443243513445315
+            epsilon: null
+            differential_equation_type: "ODE"
+            integrator_kwargs:
+              method: "euler"
+            velocity_annealing_factor: 14.973558717968908
+            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: 880
+  relative_si_costs:
+    species_loss: 0.0
+    pos_loss_b: 0.9983263145571981
+    cell_loss_b: 0.00167368544280187
+  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: perov_5
+      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"
+  dataset_name: "perov_5"
+data:
+  train_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/perov_5/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/perov_5/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/perov_5/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: 6000
+  enable_progress_bar: false
+  check_val_every_n_epoch: 100
+optimizer:
+  class_path: torch.optim.Adam
+  init_args:
+    lr: 4.4871577022001995e-05

Trig-SDE-Gamma/checkpoint.ckpt

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

Trig-SDE-Gamma/train.yaml

@@ -0,0 +1,147 @@
+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.06271372569234963
+            epsilon: 
+              class_path: omg.si.epsilon.VanishingEpsilon
+              init_args:
+                c: 7.478617683255472
+                mu: 0.06295065489868475
+                sigma: 0.03384344419315302
+            differential_equation_type: "SDE"
+            integrator_kwargs:
+              method: "euler"
+              dt: 0.0011101224226877093
+            velocity_annealing_factor: 3.4362588970266796
+            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.0508912424117229
+            epsilon: null
+            differential_equation_type: "ODE"
+            integrator_kwargs:
+              method: "euler"
+            velocity_annealing_factor: 0.03360577590810462
+            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: 900
+  relative_si_costs:
+    species_loss: 0.0
+    pos_loss_b: 0.6868298746007011
+    pos_loss_z: 0.24887105602907683
+    cell_loss_b: 0.064299069370222
+  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: perov_5
+      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"
+  dataset_name: "perov_5"
+data:
+  train_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/perov_5/train.lmdb"
+      niggli: True
+  val_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/perov_5/val.lmdb"
+      niggli: True
+  predict_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/perov_5/test.lmdb"
+      niggli: True
+  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: 6000
+  enable_progress_bar: false
+  check_val_every_n_epoch: 100
+optimizer:
+  class_path: torch.optim.Adam
+  init_args:
+    lr: 7.173658889538975e-05

VESBD-ODE/checkpoint.ckpt

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

VESBD-ODE/train.yaml

@@ -0,0 +1,154 @@
+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.007753186833706728
+                    sigma_max: 0.5165059747015202
+            epsilon: null
+            differential_equation_type: "ODE"
+            integrator_kwargs:
+              method: "euler"
+            velocity_annealing_factor: 0.0030999124784898413
+            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.TrigonometricInterpolant
+            gamma: 
+              class_path: omg.si.gamma.LatentGammaSqrt
+              init_args:
+                a: 0.024482789522429726
+            epsilon: 
+              class_path: omg.si.epsilon.VanishingEpsilon
+              init_args:
+                c: 9.940425570212101
+                mu: 0.24041599621265147
+                sigma: 0.021132860336543085
+            differential_equation_type: "SDE"
+            integrator_kwargs:
+              method: "euler"
+              dt: 0.0026332451961934566
+            velocity_annealing_factor: 14.933642154361792
+            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: 380
+  relative_si_costs:
+    species_loss: 0.0
+    pos_loss_b: 0.979954187812053
+    cell_loss_b: 0.01866918394074503
+    cell_loss_z: 0.0013766282472020075
+  sampler:
+    class_path: omg.sampler.sample_from_rng.SampleFromRNG
+    init_args:
+      pos_distribution: 
+        class_path: omg.sampler.distributions.NormalDistribution
+        init_args:
+          scale: 8.955438982782663
+      cell_distribution: 
+        class_path: omg.sampler.distributions.InformedLatticeDistribution
+        init_args:
+          dataset_name: perov_5
+      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: "perov_5"
+data:
+  train_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/perov_5/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/perov_5/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/perov_5/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: 6000
+  enable_progress_bar: false
+  check_val_every_n_epoch: 100
+optimizer:
+  class_path: torch.optim.Adam
+  init_args:
+    lr: 0.0077762908469486665

VPSBD-ODE/checkpoint.ckpt

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

VPSBD-ODE/train.yaml

@@ -0,0 +1,144 @@
+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: 12.792912174596323
+            correct_center_of_mass_motion: true
+            predict_velocity: true
+        # lattice vectors
+        - class_path: omg.si.single_stochastic_interpolant_os.SingleStochasticInterpolantOS
+          init_args:
+            interpolant: omg.si.interpolants.ScoreBasedDiffusionModelInterpolant
+            epsilon: 
+              class_path: omg.si.epsilon.VanishingEpsilon
+              init_args:
+                c: 8.480198053500128
+                mu: 0.12906782653832816
+                sigma: 0.0485371724887369
+            differential_equation_type: "SDE"
+            integrator_kwargs:
+              method: "euler"
+              dt: 0.007736434228718281
+            velocity_annealing_factor: 2.690266084902449
+            correct_center_of_mass_motion: false
+            predict_velocity: true
+      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: 130
+  relative_si_costs:
+    species_loss: 0.0
+    pos_loss_b: 0.003496793110817246
+    cell_loss_b: 0.01211289827585164
+    cell_loss_z: 0.9843903086133311
+  sampler:
+    class_path: omg.sampler.sample_from_rng.SampleFromRNG
+    init_args:
+      pos_distribution: 
+        class_path: omg.sampler.distributions.NormalDistribution
+        init_args:
+          scale: 0.2775300948889965
+      cell_distribution: 
+        class_path: omg.sampler.distributions.NormalDistribution
+        init_args:
+          scale: 0.6055540534879594
+      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"
+  dataset_name: "perov_5"
+data:
+  train_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/perov_5/train.lmdb"
+      niggli: True
+  val_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/perov_5/val.lmdb"
+      niggli: True
+  predict_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/perov_5/test.lmdb"
+      niggli: True
+  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: 6000
+  enable_progress_bar: false
+  check_val_every_n_epoch: 100
+optimizer:
+  class_path: torch.optim.Adam
+  init_args:
+    lr: 0.0008719662356797908

VPSBD-SDE/checkpoint.ckpt

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

VPSBD-SDE/train.yaml

@@ -0,0 +1,154 @@
+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: 6.705334122560177
+                mu: 0.1759894495124853
+                sigma: 0.02684743624891644
+            differential_equation_type: "SDE"
+            integrator_kwargs:
+              method: "euler"
+              dt: 0.002859598957002163
+            velocity_annealing_factor: 11.540982308844075
+            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.TrigonometricInterpolant
+            gamma: 
+              class_path: omg.si.gamma.LatentGammaSqrt
+              init_args:
+                a: 0.028688550336857962
+            epsilon: 
+              class_path: omg.si.epsilon.VanishingEpsilon
+              init_args:
+                c: 8.616058241205366
+                mu: 0.20683060178158524
+                sigma: 0.010467959402930785
+            differential_equation_type: "SDE"
+            integrator_kwargs:
+              method: "euler"
+              dt: 0.002859598957002163
+            velocity_annealing_factor: 11.528499292207702
+            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: 350
+  relative_si_costs:
+    species_loss: 0.0
+    pos_loss_b: 0.2897890800401683
+    pos_loss_z: 0.3259349777392057
+    cell_loss_b: 0.19601072982998402
+    cell_loss_z: 0.18826521239064184
+  sampler:
+    class_path: omg.sampler.sample_from_rng.SampleFromRNG
+    init_args:
+      pos_distribution: 
+        class_path: omg.sampler.distributions.NormalDistribution
+        init_args:
+          scale: 0.12777034312154512
+      cell_distribution: 
+        class_path: omg.sampler.distributions.InformedLatticeDistribution
+        init_args:
+          dataset_name: perov_5
+      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"
+  dataset_name: "perov_5"
+data:
+  train_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/perov_5/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/perov_5/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/perov_5/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: 6000
+  enable_progress_bar: false
+  check_val_every_n_epoch: 100
+optimizer:
+  class_path: torch.optim.Adam
+  init_args:
+    lr: 3.829398871139748e-05