[WIP] Multimodal

CHANGELOG.md

@@ -1,6 +1,30 @@
 # Changelog
 All notable changes to this project will be documented in this file.
 
+## WIP [0.11.0]
+### Added
+- Build multi-fidelity model, SevenNet-MF, based on given modality in the yaml
+- Modality support for sevenn_inference, sevenn_get_modal, and SevenNetCalculator
+- [cli] sevenn_cp tool for checkpoint summary, input generation, multi-modal routines
+- Modality append / assign using sevenn_cp
+- Loss weighting for energy, force and stress for corresponding data label
+- Ignore unlabelled data when calculating loss. (e.g. stress data for non-pbc structure)
+- Dict style dataset input for multi-modal and data-weight
+
+### Added (code)
+- sevenn.train.modal_dataset SevenNetMultiModalDataset
+- sevenn.scripts.backward_compatibility.py
+- sevenn.checkpoint.py
+
+### Changed
+- Sort instructions of tensor product in convolution (+ fix flipped w3j coeff of old model)
+- Lazy initialization for `IrrepsLinear` and `SelfConnection*`
+- Checkpoint things using `sevenn/checkpoint.py`
+- e3nn >= 0.5.0, to ensure changed CG coeff later on
+- pandas as dependency
+
+### Fixed
+- More refactor for shift scale things + few bug fixes
 
 ## [0.10.3]
 ### Added

README.md

@@ -15,6 +15,7 @@ SevenNet (Scalable EquiVariance Enabled Neural Network) is a graph neural networ
  - Python [Atomic Simulation Environment (ASE)](https://wiki.fysik.dtu.dk/ase/) calculator support
  - GPU-parallelized molecular dynamics with LAMMPS
  - CUDA-accelerated D3 (van der Waals) dispersion
+ - Multi-fidelity training for combining multiple database with different calculation settings.
 
 ## Pre-trained models
 So far, we have released three pre-trained SevenNet models. Each model has various hyperparameters and training sets, resulting in different accuracy and speed. Please read the descriptions below carefully and choose the model that best suits your purpose.
@@ -27,6 +28,18 @@ Additionally, `keywords` can be called in other parts of SevenNet, such as `seve
 
 ---
 
+### **SevenNet-MF-0 (16Dec2024)**
+> Keywords in ASE: `7net-MF-0` and `SevenNet-MF-0`
+
+The model is trained on PBE (+U) and $\mathrm{r}^{2}$ SCAN database provided in Materials Project.
+It has the same architecture with **SevenNet-0 (11Jul2024)**, except this model contains additional 'fidelity-dependent' parameters utilized for multi-fidelity training.
+However, overhead of calculations regarding fidelity-dependent parameters are negligible, which results in almost the same inference speed with **SevenNet-0 (11Jul2024)**.
+
+Details in using this model as well as choosing level-of-theory for inference can be found in [here](./sevenn/pretrained_potentials/SevenNet_MF_0).
+
+* Training set MAE ($\mathrm{r}^{2}$ SCAN): 10.8 meV/atom (energy), 0.018 eV/Ang. (force), and 0.58 kbar (stress)
+---
+
 ### **SevenNet-l3i5 (12Dec2024)**
 > Keywords in ASE: `7net-l3i5` and `SevenNet-l3i5`
 
@@ -117,7 +130,7 @@ With the `sevenn_preset` command, the input file that sets the training paramete
 sevenn_preset {preset keyword} > input.yaml
 ```
 
-Available preset keywords are: `base`, `fine_tune`, `sevennet-0`, and `sevennet-l3i5`.
+Available preset keywords are: `base`, `fine_tune`, `multi_modal`, `sevennet-0`, and `sevennet-l3i5`.
 Check comments in the preset yaml files for explanations. For fine-tuning, note that most model hyperparameters cannot be modified unless explicitly indicated.
 To reuse a preprocessed training set, you can specify `sevenn_data/${dataset_name}.pt` to the `load_trainset_path:` in the `input.yaml`.
 
@@ -307,3 +320,16 @@ If you use this code, please cite our paper:
 	pages = {4857--4868},
 }
 ```
+
+If you utilize the multi-fidelity feature of this code or the pretrained model SevenNet-MF-0, please cite the following paper:
+```txt
+@article{kim_sevennet_mf_2024,
+	title = {Data-Efficient Multifidelity Training for High-Fidelity Machine Learning Interatomic Potentials},
+	volume = {xx},
+	doi = {10.1021/jacs.4c14455},
+	number = {xx},
+	journal = {J. Am. Chem. Soc.},
+	author = {Kim, Jaesun and Kim, Jisu and Kim, Jaehoon and Lee, Jiho and Park, Yutack and Kang, Youngho and Han, Seungwu},
+	year = {2024},
+	pages = {xx--xx},
+```

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "sevenn"
-version = "0.10.3"
+version = "0.11.0"
 authors = [
     { name = "Yutack Park", email = "[email protected]" },
     { name = "Haekwan Jeon", email = "[email protected]" },
@@ -21,12 +21,13 @@ dependencies = [
     "ase",
     "braceexpand",
     "pyyaml",
-    "e3nn",
+    "e3nn>=0.5.0",
     "tqdm",
     "scikit-learn",
     "torch_geometric>=2.5.0",
     "numpy<2.0",
     "matscipy",
+    "pandas",
 ]
 [project.optional-dependencies]
 test = ["matscipy", "pytest-cov>=5"]
@@ -39,6 +40,7 @@ sevenn_graph_build = "sevenn.main.sevenn_graph_build:main"
 sevenn_inference = "sevenn.main.sevenn_inference:main"
 sevenn_patch_lammps = "sevenn.main.sevenn_patch_lammps:main"
 sevenn_preset = "sevenn.main.sevenn_preset:main"
+sevenn_cp = "sevenn.main.sevenn_cp:main"
 
 [project.urls]
 Homepage = "https://github.com/MDIL-SNU/SevenNet"
@@ -59,6 +61,7 @@ sevenn = [
     "pretrained_potentials/SevenNet_0__11Jul2024/checkpoint_sevennet_0.pth",
     "pretrained_potentials/SevenNet_0__22May2024/checkpoint_sevennet_0.pth",
     "pretrained_potentials/SevenNet_l3i5/checkpoint_l3i5.pth",
+    "pretrained_potentials/SevenNet_MF_0/checkpoint_sevennet_mf_0.pth",
     "py.typed",
 ]
 

setup.cfg

@@ -10,5 +10,5 @@ include_trailing_comma=True
 force_grid_wrap=0
 use_parentheses=True
 line_length=80
-known_third_party=ase,braceexpand,e3nn,numpy,pytest,setuptools,sklearn,torch,torch_geometric,tqdm,yaml
+known_third_party=ase,braceexpand,e3nn,numpy,pandas,pytest,setuptools,sklearn,torch,torch_geometric,tqdm,yaml
 known_first_party=

sevenn/_const.py

@@ -57,13 +57,16 @@
 SEVENNET_l3i5 = (
     f'{_prefix}/SevenNet_l3i5/checkpoint_l3i5.pth'
 )
+SEVENNET_MF_0 = (
+    f'{_prefix}/SevenNet_MF_0/checkpoint_sevennet_mf_0.pth'
+)
 
 
 # to avoid torch script to compile torch_geometry.data
 AtomGraphDataType = Dict[str, torch.Tensor]
 
 
-class LossType(Enum):
+class LossType(Enum):  # only used for train_v1, do not use it afterwards
     ENERGY = 'energy'  # eV or eV/atom
     FORCE = 'force'  # eV/A
     STRESS = 'stress'  # kB
@@ -80,44 +83,48 @@ def error_record_condition(x):
         if v[0] == 'TotalLoss':
             continue
         if v[1] not in SUPPORTING_METRICS:
-            print('w')
             return False
     return True
 
 
 DEFAULT_E3_EQUIVARIANT_MODEL_CONFIG = {
-    KEY.IRREPS_MANUAL: False,
+    KEY.CUTOFF: 4.5,
     KEY.NODE_FEATURE_MULTIPLICITY: 32,
+    KEY.IRREPS_MANUAL: False,
     KEY.LMAX: 1,
     KEY.LMAX_EDGE: -1,  # -1 means lmax_edge = lmax
     KEY.LMAX_NODE: -1,  # -1 means lmax_node = lmax
     KEY.IS_PARITY: True,
+    KEY.NUM_CONVOLUTION: 3,
     KEY.RADIAL_BASIS: {
         KEY.RADIAL_BASIS_NAME: 'bessel',
     },
     KEY.CUTOFF_FUNCTION: {
         KEY.CUTOFF_FUNCTION_NAME: 'poly_cut',
     },
     KEY.ACTIVATION_RADIAL: 'silu',
-    KEY.CUTOFF: 4.5,
-    KEY.CONVOLUTION_WEIGHT_NN_HIDDEN_NEURONS: [64, 64],
-    KEY.NUM_CONVOLUTION: 3,
     KEY.ACTIVATION_SCARLAR: {'e': 'silu', 'o': 'tanh'},
     KEY.ACTIVATION_GATE: {'e': 'silu', 'o': 'tanh'},
+    KEY.CONVOLUTION_WEIGHT_NN_HIDDEN_NEURONS: [64, 64],
     # KEY.AVG_NUM_NEIGH: True,  # deprecated
     # KEY.TRAIN_AVG_NUM_NEIGH: False,  # deprecated
     KEY.CONV_DENOMINATOR: 'avg_num_neigh',
     KEY.TRAIN_DENOMINTAOR: False,
     KEY.TRAIN_SHIFT_SCALE: False,
     # KEY.OPTIMIZE_BY_REDUCE: True,  # deprecated, always True
     KEY.USE_BIAS_IN_LINEAR: False,
+    KEY.USE_MODAL_NODE_EMBEDDING: False,
+    KEY.USE_MODAL_SELF_INTER_INTRO: False,
+    KEY.USE_MODAL_SELF_INTER_OUTRO: False,
+    KEY.USE_MODAL_OUTPUT_BLOCK: False,
     KEY.READOUT_AS_FCN: False,
     # Applied af readout as fcn is True
     KEY.READOUT_FCN_HIDDEN_NEURONS: [30, 30],
     KEY.READOUT_FCN_ACTIVATION: 'relu',
     KEY.SELF_CONNECTION_TYPE: 'nequip',
     KEY.INTERACTION_TYPE: 'nequip',
     KEY._NORMALIZE_SPH: True,
+    KEY.CUEQUIVARIANCE_CONFIG: {},
 }
 
 
@@ -144,13 +151,18 @@ def error_record_condition(x):
     KEY.TRAIN_SHIFT_SCALE: bool,
     KEY.TRAIN_DENOMINTAOR: bool,
     KEY.USE_BIAS_IN_LINEAR: bool,
+    KEY.USE_MODAL_NODE_EMBEDDING: bool,
+    KEY.USE_MODAL_SELF_INTER_INTRO: bool,
+    KEY.USE_MODAL_SELF_INTER_OUTRO: bool,
+    KEY.USE_MODAL_OUTPUT_BLOCK: bool,
     KEY.READOUT_AS_FCN: bool,
     KEY.READOUT_FCN_HIDDEN_NEURONS: list,
     KEY.READOUT_FCN_ACTIVATION: str,
     KEY.ACTIVATION_RADIAL: str,
     KEY.SELF_CONNECTION_TYPE: lambda x: x in IMPLEMENTED_SELF_CONNECTION_TYPE,
     KEY.INTERACTION_TYPE: lambda x: x in IMPLEMENTED_INTERACTION_TYPE,
     KEY._NORMALIZE_SPH: bool,
+    KEY.CUEQUIVARIANCE_CONFIG: dict,
 }
 
 
@@ -179,8 +191,13 @@ def model_defaults(config):
     KEY.COMPUTE_STATISTICS: True,
     KEY.DATASET_TYPE: 'graph',
     # KEY.USE_SPECIES_WISE_SHIFT_SCALE: False,
+    KEY.USE_MODAL_WISE_SHIFT: False,
+    KEY.USE_MODAL_WISE_SCALE: False,
     KEY.SHIFT: 'per_atom_energy_mean',
     KEY.SCALE: 'force_rms',
+    # KEY.DATA_SHUFFLE: True,
+    # KEY.DATA_WEIGHT: False,
+    # KEY.DATA_MODALITY: False,
 }
 
 DATA_CONFIG_CONDITION = {
@@ -197,8 +214,12 @@ def model_defaults(config):
     # KEY.USE_SPECIES_WISE_SHIFT_SCALE: bool,
     KEY.SHIFT: lambda x: type(x) in [float, list] or x in IMPLEMENTED_SHIFT,
     KEY.SCALE: lambda x: type(x) in [float, list] or x in IMPLEMENTED_SCALE,
+    KEY.USE_MODAL_WISE_SHIFT: bool,
+    KEY.USE_MODAL_WISE_SCALE: bool,
+    # KEY.DATA_SHUFFLE: bool,
     KEY.COMPUTE_STATISTICS: bool,
-    KEY.SAVE_DATASET: str,
+    # KEY.DATA_WEIGHT: bool,
+    # KEY.DATA_MODALITY: bool,
 }
 
 
@@ -221,14 +242,16 @@ def data_defaults(config):
     KEY.FORCE_WEIGHT: 0.1,
     KEY.STRESS_WEIGHT: 1e-6,  # SIMPLE-NN default
     KEY.PER_EPOCH: 5,
-    KEY.USE_TESTSET: False,
+    # KEY.USE_TESTSET: False,
     KEY.CONTINUE: {
         KEY.CHECKPOINT: False,
         KEY.RESET_OPTIMIZER: False,
         KEY.RESET_SCHEDULER: False,
         KEY.RESET_EPOCH: False,
         KEY.USE_STATISTIC_VALUES_OF_CHECKPOINT: True,
+        KEY.USE_STATISTIC_VALUES_FOR_CP_MODAL_ONLY: True,
     },
+    # KEY.DEFAULT_MODAL: 'common',
     KEY.CSV_LOG: 'log.csv',
     KEY.NUM_WORKERS: 0,
     KEY.IS_TRAIN_STRESS: True,
@@ -240,6 +263,8 @@ def data_defaults(config):
         ['TotalLoss', 'None'],
     ],
     KEY.BEST_METRIC: 'TotalLoss',
+    KEY.USE_WEIGHT: False,
+    KEY.USE_MODALITY: False,
 }
 
 
@@ -257,12 +282,16 @@ def data_defaults(config):
         KEY.RESET_SCHEDULER: bool,
         KEY.RESET_EPOCH: bool,
         KEY.USE_STATISTIC_VALUES_OF_CHECKPOINT: bool,
+        KEY.USE_STATISTIC_VALUES_FOR_CP_MODAL_ONLY: bool,
     },
+    KEY.DEFAULT_MODAL: str,
     KEY.IS_TRAIN_STRESS: bool,
     KEY.TRAIN_SHUFFLE: bool,
     KEY.ERROR_RECORD: error_record_condition,
     KEY.BEST_METRIC: str,
     KEY.CSV_LOG: str,
+    KEY.USE_MODALITY: bool,
+    KEY.USE_WEIGHT: bool,
 }
 
 

sevenn/_keys.py

@@ -12,8 +12,6 @@
 
 from typing import Final
 
-from torch.jit import CompilationUnit
-
 # see
 # https://github.com/pytorch/pytorch/issues/52312
 # for FYI
@@ -40,6 +38,10 @@
 NODE_FEATURE: Final[str] = 'x'  # (N, ?) PyG
 NODE_FEATURE_GHOST: Final[str] = 'x_ghost'
 NODE_ATTR: Final[str] = 'node_attr'  # (N, N_species) from one_hot
+MODAL_ATTR: Final[str] = (
+    'modal_attr'  # (1, N_modalities) for handling multi-modal
+)
+MODAL_TYPE: Final[str] = 'modal_type'  # (1) one-hot index of modal
 EDGE_ATTR: Final[str] = 'edge_attr'  # (from spherical harmonics)
 EDGE_EMBEDDING: Final[str] = 'edge_embedding'  # (from edge embedding)
 
@@ -69,7 +71,11 @@
 NUM_ATOMS: Final[str] = 'num_atoms'  # int
 NUM_GHOSTS: Final[str] = 'num_ghosts'
 NLOCAL: Final[str] = 'nlocal'  # only for lammps parallel, must be on cpu
-USER_LABEL: Final[str] = 'user_label'  # Deprecated from v0.9.6
+USER_LABEL: Final[str] = 'user_label'
+DATA_WEIGHT: Final[str] = 'data_weight'  # weight for given data
+DATA_MODALITY: Final[str] = (
+    'data_modality'  # modality of given data. e.g. PBE and SCAN
+)
 BATCH: Final[str] = 'batch'
 
 TAG = 'tag'  # replace USER_LABEL
@@ -125,6 +131,9 @@
 RESET_SCHEDULER = 'reset_scheduler'
 RESET_EPOCH = 'reset_epoch'
 USE_STATISTIC_VALUES_OF_CHECKPOINT = 'use_statistic_values_of_checkpoint'
+USE_STATISTIC_VALUES_FOR_CP_MODAL_ONLY = (
+    'use_statistic_values_for_cp_modal_only'
+)
 
 CSV_LOG = 'csv_log'
 
@@ -140,6 +149,10 @@
 DDP_BACKEND = 'ddp_backend'
 PER_EPOCH = 'per_epoch'
 
+USE_WEIGHT = 'use_weight'
+USE_MODALITY = 'use_modality'
+DEFAULT_MODAL = 'default_modal'
+
 
 # ==================================================#
 # ~~~~~~~~ KEY for model configuration ~~~~~~~~~~~ #
@@ -149,9 +162,12 @@
 MODEL_TYPE = '_model_type'
 CUTOFF = 'cutoff'
 CHEMICAL_SPECIES = 'chemical_species'
+MODAL_LIST = 'modal_list'
 CHEMICAL_SPECIES_BY_ATOMIC_NUMBER = '_chemical_species_by_atomic_number'
 NUM_SPECIES = '_number_of_species'
+NUM_MODALITIES = '_number_of_modalities'
 TYPE_MAP = '_type_map'
+MODAL_MAP = '_modal_map'
 
 # ~~ E3 equivariant model build configuration keys ~~ #
 # see model_build default_config for type
@@ -181,6 +197,11 @@
 
 USE_BIAS_IN_LINEAR = 'use_bias_in_linear'
 
+USE_MODAL_NODE_EMBEDDING = 'use_modal_node_embedding'
+USE_MODAL_SELF_INTER_INTRO = 'use_modal_self_inter_intro'
+USE_MODAL_SELF_INTER_OUTRO = 'use_modal_self_inter_outro'
+USE_MODAL_OUTPUT_BLOCK = 'use_modal_output_block'
+
 READOUT_AS_FCN = 'readout_as_fcn'
 READOUT_FCN_HIDDEN_NEURONS = 'readout_fcn_hidden_neurons'
 READOUT_FCN_ACTIVATION = 'readout_fcn_activation'
@@ -191,11 +212,15 @@
 SCALE = 'scale'
 
 USE_SPECIES_WISE_SHIFT_SCALE = 'use_species_wise_shift_scale'
+USE_MODAL_WISE_SHIFT = 'use_modal_wise_shift'
+USE_MODAL_WISE_SCALE = 'use_modal_wise_scale'
 
 TRAIN_SHIFT_SCALE = 'train_shift_scale'
 TRAIN_DENOMINTAOR = 'train_denominator'
 INTERACTION_TYPE = 'interaction_type'
 TRAIN_AVG_NUM_NEIGH = 'train_avg_num_neigh'  # deprecated
 
+CUEQUIVARIANCE_CONFIG = 'cuequivariance_config'
+
 _NORMALIZE_SPH = '_normalize_sph'
 OPTIMIZE_BY_REDUCE = 'optimize_by_reduce'