Upload 22 files

EncDec-ODE-Gamma/checkpoint.ckpt

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

EncDec-ODE-Gamma/train.yaml

@@ -0,0 +1,168 @@
+model:
+  si:
+    class_path: omg.si.stochastic_interpolants.StochasticInterpolants
+    init_args:
+      stochastic_interpolants:
+        # chemical species
+        - class_path: omg.si.discrete_flow_matching_mask.DiscreteFlowMatchingMask
+          init_args:
+            noise: 0.8465315128078521
+        # fractional coordinates
+        - class_path: omg.si.single_stochastic_interpolant.SingleStochasticInterpolant
+          init_args:
+            interpolant: omg.si.interpolants.PeriodicEncoderDecoderInterpolant
+            gamma: omg.si.gamma.LatentGammaEncoderDecoder
+            epsilon: null
+            differential_equation_type: "ODE"
+            integrator_kwargs:
+              method: "euler"
+            velocity_annealing_factor: 10.274308845621986
+            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.08217016314129522
+            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: 840
+  relative_si_costs:
+    species_loss: 0.2648412544596816
+    pos_loss_b: 0.7267924862588087
+    cell_loss_b: 0.008366259281509825
+  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: mp_20
+      species_distribution:
+        class_path: omg.sampler.distributions.MaskDistribution
+  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: "dng_eval"
+  dataset_name: "mp_20"
+data:
+  train_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/mp_20/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/mp_20/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/mp_20/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_dng_eval"
+        save_top_k: 1
+        monitor: "dng_eval"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_wdist_density"
+        save_top_k: 1
+        monitor: "wdist_density"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_wdist_Nary"
+        save_top_k: 1
+        monitor: "wdist_Nary"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_wdist_CN"
+        save_top_k: 1
+        monitor: "wdist_CN"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_cov_precision"
+        save_top_k: 1
+        monitor: "cov_precision"
+        mode: "max"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_cov_recall"
+        save_top_k: 1
+        monitor: "cov_recall"
+        mode: "max"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_validity"
+        save_top_k: 1
+        monitor: "validity_rate"
+        mode: "max"
+        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.00012021943412654004

EncDec-SDE-Gamma/checkpoint.ckpt

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

EncDec-SDE-Gamma/train.yaml

@@ -0,0 +1,195 @@
+model:
+  si:
+    class_path: omg.si.stochastic_interpolants.StochasticInterpolants
+    init_args:
+      stochastic_interpolants:
+        # chemical species
+        - class_path: omg.si.discrete_flow_matching_mask.DiscreteFlowMatchingMask
+          init_args:
+            noise: 19.77565076697948
+        # fractional coordinates
+        - class_path: omg.si.single_stochastic_interpolant.SingleStochasticInterpolant
+          init_args:
+            interpolant: 
+              class_path: omg.si.interpolants.PeriodicEncoderDecoderInterpolant
+              init_args:
+                switch_time: 0.7261434470495144
+                power: 0.5
+            gamma: 
+              class_path: omg.si.gamma.LatentGammaEncoderDecoder
+              init_args:
+                a: 0.10379253121526635
+                switch_time: 0.7261434470495144
+                power: 0.5
+            epsilon: 
+              class_path: omg.si.epsilon.VanishingEpsilon
+              init_args:
+                c: 7.095366578175936
+                mu: 0.18874541537289413
+                sigma: 0.020535877713041894
+            differential_equation_type: "SDE"
+            integrator_kwargs:
+              method: "euler"
+              dt: 0.0016387520590797067
+            velocity_annealing_factor: 7.868988752162741
+            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: 1.6505152121826552
+            epsilon: null
+            differential_equation_type: "ODE"
+            integrator_kwargs:
+              method: "euler"
+            velocity_annealing_factor: 3.919302532270132
+            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: 610
+  relative_si_costs:
+    species_loss: 0.4341084689667317
+    pos_loss_b: 0.21431903999859292
+    pos_loss_z: 0.1968226417175204
+    cell_loss_b: 0.15474984931715505
+  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: mp_20
+      species_distribution:
+        class_path: omg.sampler.distributions.MaskDistribution
+  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: "dng_eval"
+  dataset_name: "mp_20"
+data:
+  train_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/mp_20/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/mp_20/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/mp_20/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_dng_eval"
+        save_top_k: 1
+        monitor: "dng_eval"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_wdist_density"
+        save_top_k: 1
+        monitor: "wdist_density"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_wdist_Nary"
+        save_top_k: 1
+        monitor: "wdist_Nary"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_wdist_CN"
+        save_top_k: 1
+        monitor: "wdist_CN"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_cov_precision"
+        save_top_k: 1
+        monitor: "cov_precision"
+        mode: "max"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_cov_recall"
+        save_top_k: 1
+        monitor: "cov_recall"
+        mode: "max"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_validity"
+        save_top_k: 1
+        monitor: "validity_rate"
+        mode: "max"
+        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
+  gradient_clip_algorithm: "value"
+  num_sanity_val_steps: 0
+  precision: "32-true"
+  max_epochs: 2000
+  enable_progress_bar: false
+  check_val_every_n_epoch: 100
+optimizer:
+  class_path: torch.optim.AdamW
+  init_args:
+    lr: 3.139610174577985e-05
+    weight_decay: 3.560067412494533e-05
+lr_scheduler:
+  class_path: torch.optim.lr_scheduler.CosineAnnealingLR
+  init_args:
+    T_max: 2000
+    eta_min: 1e-07
+

Linear-ODE-Gamma/checkpoint.ckpt

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

Linear-ODE-Gamma/train.yaml

@@ -0,0 +1,195 @@
+model:
+  si:
+    class_path: omg.si.stochastic_interpolants.StochasticInterpolants
+    init_args:
+      stochastic_interpolants:
+        # chemical species
+        - class_path: omg.si.discrete_flow_matching_mask.DiscreteFlowMatchingMask
+          init_args:
+            noise: 23.870491382634235
+        # 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: 1.4501684803942854
+            epsilon: null
+            differential_equation_type: "ODE"
+            integrator_kwargs:
+              method: "euler"
+            velocity_annealing_factor: 14.825022083056373
+            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.13766881993242777
+                power: 1.0
+            gamma: 
+              class_path: omg.si.gamma.LatentGammaEncoderDecoder
+              init_args:
+                a: 7.882046441638109
+                switch_time: 0.13766881993242777
+                power: 1.0
+            epsilon: 
+              class_path: omg.si.epsilon.VanishingEpsilon
+              init_args:
+                c: 5.487699104615115
+                mu: 0.2899409657474152
+                sigma: 0.010062500495585096
+            differential_equation_type: "SDE"
+            integrator_kwargs:
+              method: "euler"
+              dt: 0.007736434228718281
+            velocity_annealing_factor: 5.9072140831863305
+            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: 130
+  relative_si_costs:
+    species_loss: 0.22163297905676768
+    pos_loss_b: 0.7682858351574293
+    cell_loss_b: 0.008860420349864338
+    cell_loss_z: 0.001220765435938645
+  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: mp_20
+      species_distribution:
+        class_path: omg.sampler.distributions.MaskDistribution
+  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: "dng_eval"
+  dataset_name: "mp_20"
+data:
+  train_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/mp_20/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/mp_20/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/mp_20/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_dng_eval"
+        save_top_k: 1
+        monitor: "dng_eval"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_wdist_density"
+        save_top_k: 1
+        monitor: "wdist_density"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_wdist_Nary"
+        save_top_k: 1
+        monitor: "wdist_Nary"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_wdist_CN"
+        save_top_k: 1
+        monitor: "wdist_CN"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_cov_precision"
+        save_top_k: 1
+        monitor: "cov_precision"
+        mode: "max"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_cov_recall"
+        save_top_k: 1
+        monitor: "cov_recall"
+        mode: "max"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_validity"
+        save_top_k: 1
+        monitor: "validity_rate"
+        mode: "max"
+        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
+  gradient_clip_algorithm: "value"
+  num_sanity_val_steps: 0
+  precision: "32-true"
+  max_epochs: 2000
+  enable_progress_bar: false
+  check_val_every_n_epoch: 100
+optimizer:
+  class_path: torch.optim.AdamW
+  init_args:
+    lr: 0.00019511523812262233
+    weight_decay: 0.0003813804936812436
+lr_scheduler:
+  class_path: torch.optim.lr_scheduler.CosineAnnealingLR
+  init_args:
+    T_max: 2000
+    eta_min: 1e-07
+

Linear-ODE/checkpoint.ckpt

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

Linear-ODE/train.yaml

@@ -0,0 +1,168 @@
+model:
+  si:
+    class_path: omg.si.stochastic_interpolants.StochasticInterpolants
+    init_args:
+      stochastic_interpolants:
+        # chemical species
+        - class_path: omg.si.discrete_flow_matching_mask.DiscreteFlowMatchingMask
+          init_args:
+            noise: 7.080372063368751
+        # 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: 13.620695525269845
+            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: 1.0679662602192312
+            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: 150
+  relative_si_costs:
+    species_loss: 0.021815399034596464
+    pos_loss_b: 0.9775483266595605
+    cell_loss_b: 0.0006362743058429793
+  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: mp_20
+      species_distribution:
+        class_path: omg.sampler.distributions.MaskDistribution
+  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: "dng_eval"
+  dataset_name: "mp_20"
+data:
+  train_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/mp_20/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/mp_20/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/mp_20/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_dng_eval"
+        save_top_k: 1
+        monitor: "dng_eval"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_wdist_density"
+        save_top_k: 1
+        monitor: "wdist_density"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_wdist_Nary"
+        save_top_k: 1
+        monitor: "wdist_Nary"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_wdist_CN"
+        save_top_k: 1
+        monitor: "wdist_CN"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_cov_precision"
+        save_top_k: 1
+        monitor: "cov_precision"
+        mode: "max"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_cov_recall"
+        save_top_k: 1
+        monitor: "cov_recall"
+        mode: "max"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_validity"
+        save_top_k: 1
+        monitor: "validity_rate"
+        mode: "max"
+        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.001736512450391209

Linear-SDE-Gamma/checkpoint.ckpt

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

Linear-SDE-Gamma/train.yaml

@@ -0,0 +1,186 @@
+model:
+  si:
+    class_path: omg.si.stochastic_interpolants.StochasticInterpolants
+    init_args:
+      stochastic_interpolants:
+        # chemical species
+        - class_path: omg.si.discrete_flow_matching_mask.DiscreteFlowMatchingMask
+          init_args:
+            noise: 0.18946955217679085
+        # 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.018159684059653552
+            epsilon: 
+              class_path: omg.si.epsilon.VanishingEpsilon
+              init_args:
+                c: 9.74900863316411
+                mu: 0.17191546490562354
+                sigma: 0.029425925880471573
+            differential_equation_type: "SDE"
+            integrator_kwargs:
+              method: "euler"
+              dt: 0.0014076164225116372
+            velocity_annealing_factor: 6.334345265874859
+            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: 1.0674474901964888
+            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.5918064683979826
+    pos_loss_b: 0.13091891010303253
+    pos_loss_z: 0.27077286248215743
+    cell_loss_b: 0.006501759016827464
+  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: mp_20
+      species_distribution:
+        class_path: omg.sampler.distributions.MaskDistribution
+  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: "dng_eval"
+  dataset_name: "mp_20"
+data:
+  train_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/mp_20/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/mp_20/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/mp_20/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_dng_eval"
+        save_top_k: 1
+        monitor: "dng_eval"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_wdist_density"
+        save_top_k: 1
+        monitor: "wdist_density"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_wdist_Nary"
+        save_top_k: 1
+        monitor: "wdist_Nary"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_wdist_CN"
+        save_top_k: 1
+        monitor: "wdist_CN"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_cov_precision"
+        save_top_k: 1
+        monitor: "cov_precision"
+        mode: "max"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_cov_recall"
+        save_top_k: 1
+        monitor: "cov_recall"
+        mode: "max"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_validity"
+        save_top_k: 1
+        monitor: "validity_rate"
+        mode: "max"
+        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
+  gradient_clip_algorithm: "value"
+  num_sanity_val_steps: 0
+  precision: "32-true"
+  max_epochs: 2000
+  enable_progress_bar: false
+  check_val_every_n_epoch: 100
+optimizer:
+  class_path: torch.optim.AdamW
+  init_args:
+    lr: 0.00019745455354877462
+    weight_decay: 0.0003111161289640361
+lr_scheduler:
+  class_path: torch.optim.lr_scheduler.CosineAnnealingLR
+  init_args:
+    T_max: 2000
+    eta_min: 1e-07
+

Trig-ODE-Gamma/checkpoint.ckpt

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

Trig-ODE-Gamma/train.yaml

@@ -0,0 +1,189 @@
+model:
+  si:
+    class_path: omg.si.stochastic_interpolants.StochasticInterpolants
+    init_args:
+      stochastic_interpolants:
+        # chemical species
+        - class_path: omg.si.discrete_flow_matching_mask.DiscreteFlowMatchingMask
+          init_args:
+            noise: 27.249112246908787
+        # 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.02697960692675219
+            epsilon: null
+            differential_equation_type: "ODE"
+            integrator_kwargs:
+              method: "euler"
+            velocity_annealing_factor: 7.788016364580001
+            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.8480953426626178
+            epsilon: 
+              class_path: omg.si.epsilon.VanishingEpsilon
+              init_args:
+                c: 3.673041703474942
+                mu: 0.0706838941982046
+                sigma: 0.01812545867373373
+            differential_equation_type: "SDE"
+            integrator_kwargs:
+              method: "euler"
+              dt: 0.001469808747060597
+            velocity_annealing_factor: 0.2916081322471492
+            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: 680
+  relative_si_costs:
+    species_loss: 0.43055618267791895
+    pos_loss_b: 0.2322254093385872
+    cell_loss_b: 0.003464180396862092
+    cell_loss_z: 0.33375422758663176
+  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: mp_20
+      species_distribution:
+        class_path: omg.sampler.distributions.MaskDistribution
+  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: "dng_eval"
+  dataset_name: "mp_20"
+data:
+  train_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/mp_20/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/mp_20/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/mp_20/test.lmdb"
+      niggli: True
+  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_dng_eval"
+        save_top_k: 1
+        monitor: "dng_eval"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_wdist_density"
+        save_top_k: 1
+        monitor: "wdist_density"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_wdist_Nary"
+        save_top_k: 1
+        monitor: "wdist_Nary"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_wdist_CN"
+        save_top_k: 1
+        monitor: "wdist_CN"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_cov_precision"
+        save_top_k: 1
+        monitor: "cov_precision"
+        mode: "max"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_cov_recall"
+        save_top_k: 1
+        monitor: "cov_recall"
+        mode: "max"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_validity"
+        save_top_k: 1
+        monitor: "validity_rate"
+        mode: "max"
+        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
+  gradient_clip_algorithm: "value"
+  num_sanity_val_steps: 0
+  precision: "32-true"
+  max_epochs: 2000
+  enable_progress_bar: false
+  check_val_every_n_epoch: 100
+optimizer:
+  class_path: torch.optim.AdamW
+  init_args:
+    lr: 0.00014843344531647814
+    weight_decay: 0.00032033785912707614
+lr_scheduler:
+  class_path: torch.optim.lr_scheduler.CosineAnnealingLR
+  init_args:
+    T_max: 2000
+    eta_min: 1e-07
+

Trig-ODE/checkpoint.ckpt

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

Trig-ODE/train.yaml

@@ -0,0 +1,179 @@
+model:
+  si:
+    class_path: omg.si.stochastic_interpolants.StochasticInterpolants
+    init_args:
+      stochastic_interpolants:
+        # chemical species
+        - class_path: omg.si.discrete_flow_matching_mask.DiscreteFlowMatchingMask
+          init_args:
+            noise: 32.687148090341246
+        # 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: 8.59355026270501
+            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: 1.182707731733833
+            epsilon: null
+            differential_equation_type: "ODE"
+            integrator_kwargs:
+              method: "euler"
+            velocity_annealing_factor: 0.308921595643179
+            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: 860
+  relative_si_costs:
+    species_loss: 0.6674861836833045
+    pos_loss_b: 0.33016840408608256
+    cell_loss_b: 0.002345412230612938
+  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: mp_20
+      species_distribution:
+        class_path: omg.sampler.distributions.MaskDistribution
+  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: "dng_eval"
+  dataset_name: "mp_20"
+data:
+  train_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/mp_20/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/mp_20/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/mp_20/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_dng_eval"
+        save_top_k: 1
+        monitor: "dng_eval"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_wdist_density"
+        save_top_k: 1
+        monitor: "wdist_density"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_wdist_Nary"
+        save_top_k: 1
+        monitor: "wdist_Nary"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_wdist_CN"
+        save_top_k: 1
+        monitor: "wdist_CN"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_cov_precision"
+        save_top_k: 1
+        monitor: "cov_precision"
+        mode: "max"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_cov_recall"
+        save_top_k: 1
+        monitor: "cov_recall"
+        mode: "max"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_validity"
+        save_top_k: 1
+        monitor: "validity_rate"
+        mode: "max"
+        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
+  gradient_clip_algorithm: "value"
+  num_sanity_val_steps: 0
+  precision: "32-true"
+  max_epochs: 2000
+  enable_progress_bar: false
+  check_val_every_n_epoch: 100
+optimizer:
+  class_path: torch.optim.AdamW
+  init_args:
+    lr: 0.002704094492670699
+    weight_decay: 0.0006070253248675564
+lr_scheduler:
+  class_path: torch.optim.lr_scheduler.CosineAnnealingLR
+  init_args:
+    T_max: 2000
+    eta_min: 1e-07
+

Trig-SDE-Gamma/checkpoint.ckpt

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

Trig-SDE-Gamma/train.yaml

@@ -0,0 +1,189 @@
+model:
+  si:
+    class_path: omg.si.stochastic_interpolants.StochasticInterpolants
+    init_args:
+      stochastic_interpolants:
+        # chemical species
+        - class_path: omg.si.discrete_flow_matching_mask.DiscreteFlowMatchingMask
+          init_args:
+            noise: 13.14607468893319
+        # 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.022769980672795356
+            epsilon: 
+              class_path: omg.si.epsilon.VanishingEpsilon
+              init_args:
+                c: 2.621699079870832
+                mu: 0.15417087293483117
+                sigma: 0.017962649662214652
+            differential_equation_type: "SDE"
+            integrator_kwargs:
+              method: "euler"
+              dt: 0.0013148878933861852
+            velocity_annealing_factor: 12.80156329264574
+            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.31566788838946225
+            epsilon: null
+            differential_equation_type: "ODE"
+            integrator_kwargs:
+              method: "euler"
+            velocity_annealing_factor: 4.36422932474701
+            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: 760
+  relative_si_costs:
+    species_loss: 0.13597807419582586
+    pos_loss_b: 0.6304347545056598
+    pos_loss_z: 0.0753230160674198
+    cell_loss_b: 0.15826415523109444
+  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: mp_20
+      species_distribution:
+        class_path: omg.sampler.distributions.MaskDistribution
+  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: "dng_eval"
+  dataset_name: "mp_20"
+data:
+  train_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/mp_20/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/mp_20/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/mp_20/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_dng_eval"
+        save_top_k: 1
+        monitor: "dng_eval"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_wdist_density"
+        save_top_k: 1
+        monitor: "wdist_density"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_wdist_Nary"
+        save_top_k: 1
+        monitor: "wdist_Nary"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_wdist_CN"
+        save_top_k: 1
+        monitor: "wdist_CN"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_cov_precision"
+        save_top_k: 1
+        monitor: "cov_precision"
+        mode: "max"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_cov_recall"
+        save_top_k: 1
+        monitor: "cov_recall"
+        mode: "max"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_validity"
+        save_top_k: 1
+        monitor: "validity_rate"
+        mode: "max"
+        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
+  gradient_clip_algorithm: "value"
+  num_sanity_val_steps: 0
+  precision: "32-true"
+  max_epochs: 2000
+  enable_progress_bar: false
+  check_val_every_n_epoch: 100
+optimizer:
+  class_path: torch.optim.AdamW
+  init_args:
+    lr: 0.0007969633652411341
+    weight_decay: 1.803908894626558e-05
+lr_scheduler:
+  class_path: torch.optim.lr_scheduler.CosineAnnealingLR
+  init_args:
+    T_max: 2000
+    eta_min: 1e-07
+

VESBD-ODE/checkpoint.ckpt

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

VESBD-ODE/train.yaml

@@ -0,0 +1,196 @@
+model:
+  si:
+    class_path: omg.si.stochastic_interpolants.StochasticInterpolants
+    init_args:
+      stochastic_interpolants:
+        # chemical species
+        - class_path: omg.si.discrete_flow_matching_mask.DiscreteFlowMatchingMask
+          init_args:
+            noise: 5.870180115373019
+        # 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.0020828565391521787
+                    sigma_max: 0.8318656965968637
+            epsilon: null
+            differential_equation_type: "ODE"
+            integrator_kwargs:
+              method: "euler"
+            velocity_annealing_factor: 12.718434028622262
+            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.9128957436969677
+            epsilon: 
+              class_path: omg.si.epsilon.VanishingEpsilon
+              init_args:
+                c: 5.79866710636582
+                mu: 0.2500005563766677
+                sigma: 0.020369240775387647
+            differential_equation_type: "SDE"
+            integrator_kwargs:
+              method: "euler"
+              dt: 0.0030334345065057278
+            velocity_annealing_factor: 0.9750213755072251
+            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: 330
+  relative_si_costs:
+    species_loss: 0.09898280162138558
+    pos_loss_b: 0.22090708494461975
+    cell_loss_b: 0.04296279641692716
+    cell_loss_z: 0.6371473170170674
+  sampler:
+    class_path: omg.sampler.sample_from_rng.SampleFromRNG
+    init_args:
+      pos_distribution: 
+        class_path: omg.sampler.distributions.NormalDistribution
+        init_args:
+          scale: 0.45439306223842724
+      cell_distribution: 
+        class_path: omg.sampler.distributions.InformedLatticeDistribution
+        init_args:
+          dataset_name: mp_20
+      species_distribution:
+        class_path: omg.sampler.distributions.MaskDistribution
+  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: "dng_eval"
+  dataset_name: "mp_20"
+data:
+  train_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/mp_20/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/mp_20/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/mp_20/test.lmdb"
+      niggli: True
+  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_dng_eval"
+        save_top_k: 1
+        monitor: "dng_eval"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_wdist_density"
+        save_top_k: 1
+        monitor: "wdist_density"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_wdist_Nary"
+        save_top_k: 1
+        monitor: "wdist_Nary"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_wdist_CN"
+        save_top_k: 1
+        monitor: "wdist_CN"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_cov_precision"
+        save_top_k: 1
+        monitor: "cov_precision"
+        mode: "max"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_cov_recall"
+        save_top_k: 1
+        monitor: "cov_recall"
+        mode: "max"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_validity"
+        save_top_k: 1
+        monitor: "validity_rate"
+        mode: "max"
+        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
+  gradient_clip_algorithm: "value"
+  num_sanity_val_steps: 0
+  precision: "32-true"
+  max_epochs: 2000
+  enable_progress_bar: false
+  check_val_every_n_epoch: 100
+optimizer:
+  class_path: torch.optim.AdamW
+  init_args:
+    lr: 0.0018098696508625563
+    weight_decay: 0.00026498129464991104
+lr_scheduler:
+  class_path: torch.optim.lr_scheduler.CosineAnnealingLR
+  init_args:
+    T_max: 2000
+    eta_min: 1e-07
+

VPSBD-ODE/checkpoint.ckpt

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

VPSBD-ODE/train.yaml

@@ -0,0 +1,182 @@
+model:
+  si:
+    class_path: omg.si.stochastic_interpolants.StochasticInterpolants
+    init_args:
+      stochastic_interpolants:
+        # chemical species
+        - class_path: omg.si.discrete_flow_matching_mask.DiscreteFlowMatchingMask
+          init_args:
+            noise: 20.267191359937392
+        # 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: 2.301841820941901
+            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: 7.796692096273471
+            epsilon: null
+            differential_equation_type: "ODE"
+            integrator_kwargs:
+              method: "euler"
+            velocity_annealing_factor: 2.741014121366449
+            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.5499726353395065
+    pos_loss_b: 0.40529122146887925
+    cell_loss_b: 0.04473614319161423
+  sampler:
+    class_path: omg.sampler.sample_from_rng.SampleFromRNG
+    init_args:
+      pos_distribution: 
+        class_path: omg.sampler.distributions.NormalDistribution
+        init_args:
+          scale: 0.23441087988918383
+      cell_distribution: 
+        class_path: omg.sampler.distributions.InformedLatticeDistribution
+        init_args:
+          dataset_name: mp_20
+      species_distribution:
+        class_path: omg.sampler.distributions.MaskDistribution
+  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: "dng_eval"
+  dataset_name: "mp_20"
+data:
+  train_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/mp_20/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/mp_20/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/mp_20/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_dng_eval"
+        save_top_k: 1
+        monitor: "dng_eval"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_wdist_density"
+        save_top_k: 1
+        monitor: "wdist_density"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_wdist_Nary"
+        save_top_k: 1
+        monitor: "wdist_Nary"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_wdist_CN"
+        save_top_k: 1
+        monitor: "wdist_CN"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_cov_precision"
+        save_top_k: 1
+        monitor: "cov_precision"
+        mode: "max"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_cov_recall"
+        save_top_k: 1
+        monitor: "cov_recall"
+        mode: "max"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_validity"
+        save_top_k: 1
+        monitor: "validity_rate"
+        mode: "max"
+        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
+  gradient_clip_algorithm: "value"
+  num_sanity_val_steps: 0
+  precision: "32-true"
+  max_epochs: 2000
+  enable_progress_bar: false
+  check_val_every_n_epoch: 100
+optimizer:
+  class_path: torch.optim.AdamW
+  init_args:
+    lr: 0.00797735754708741
+    weight_decay: 1.923837446196394e-05
+lr_scheduler:
+  class_path: torch.optim.lr_scheduler.CosineAnnealingLR
+  init_args:
+    T_max: 2000
+    eta_min: 1e-07
+

VPSBD-SDE/checkpoint.ckpt

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

VPSBD-SDE/train.yaml

@@ -0,0 +1,196 @@
+model:
+  si:
+    class_path: omg.si.stochastic_interpolants.StochasticInterpolants
+    init_args:
+      stochastic_interpolants:
+        # chemical species
+        - class_path: omg.si.discrete_flow_matching_mask.DiscreteFlowMatchingMask
+          init_args:
+            noise: 8.517811450005286
+        # 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: 9.268934476283913
+                mu: 0.25243331190144214
+                sigma: 0.04584669320169394
+            differential_equation_type: "SDE"
+            integrator_kwargs:
+              method: "euler"
+              dt: 0.00114844657946378
+            velocity_annealing_factor: 9.059507260865466
+            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: 3.0998874989979193
+            epsilon: 
+              class_path: omg.si.epsilon.VanishingEpsilon
+              init_args:
+                c: 4.885221687138149
+                mu: 0.08252713485380268
+                sigma: 0.010096412051586533
+            differential_equation_type: "SDE"
+            integrator_kwargs:
+              method: "euler"
+              dt: 0.00114844657946378
+            velocity_annealing_factor: 11.76695215049249
+            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: 870
+  relative_si_costs:
+    species_loss: 0.35838215846575894
+    pos_loss_b: 0.5183735506925028
+    pos_loss_z: 0.0007900924652522647
+    cell_loss_b: 0.0044136759736567365
+    cell_loss_z: 0.11804052240282935
+  sampler:
+    class_path: omg.sampler.sample_from_rng.SampleFromRNG
+    init_args:
+      pos_distribution: 
+        class_path: omg.sampler.distributions.NormalDistribution
+        init_args:
+          scale: 7.139671140709246
+      cell_distribution: 
+        class_path: omg.sampler.distributions.InformedLatticeDistribution
+        init_args:
+          dataset_name: mp_20
+      species_distribution:
+        class_path: omg.sampler.distributions.MaskDistribution
+  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: "dng_eval"
+  dataset_name: "mp_20"
+data:
+  train_dataset:
+    class_path: omg.datamodule.dataloader.OMGTorchDataset
+    init_args:
+      dataset:
+        class_path: omg.datamodule.datamodule.DataModule
+        init_args:
+          lmdb_paths:
+           - "data/mp_20/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/mp_20/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/mp_20/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_dng_eval"
+        save_top_k: 1
+        monitor: "dng_eval"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_wdist_density"
+        save_top_k: 1
+        monitor: "wdist_density"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_wdist_Nary"
+        save_top_k: 1
+        monitor: "wdist_Nary"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_wdist_CN"
+        save_top_k: 1
+        monitor: "wdist_CN"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_cov_precision"
+        save_top_k: 1
+        monitor: "cov_precision"
+        mode: "max"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_cov_recall"
+        save_top_k: 1
+        monitor: "cov_recall"
+        mode: "max"
+        save_weights_only: true
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        filename: "best_val_validity"
+        save_top_k: 1
+        monitor: "validity_rate"
+        mode: "max"
+        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
+  gradient_clip_algorithm: "value"
+  num_sanity_val_steps: 0
+  precision: "32-true"
+  max_epochs: 2000
+  enable_progress_bar: false
+  check_val_every_n_epoch: 100
+optimizer:
+  class_path: torch.optim.AdamW
+  init_args:
+    lr: 0.0014461332672089323
+    weight_decay: 0.0007097046414614019
+lr_scheduler:
+  class_path: torch.optim.lr_scheduler.CosineAnnealingLR
+  init_args:
+    T_max: 2000
+    eta_min: 1e-07
+