diff --git a/scripts/configs/atomic_tensor.yaml b/scripts/configs/atomic_tensor.yaml new file mode 100644 index 0000000..0d4350e --- /dev/null +++ b/scripts/configs/atomic_tensor.yaml @@ -0,0 +1,117 @@ +## Config files for atomic tensor (i.e. a tensor value for each atom) + +seed_everything: 35 +log_level: info + +data: + tensor_target_name: nmr_tensor + atom_selector: atom_selector + tensor_target_formula: ij=ji + root: . + trainset_filename: /Users/mjwen.admin/Documents/Dataset/NMR_tensor/si_nmr_data_small.json + valset_filename: /Users/mjwen.admin/Documents/Dataset/NMR_tensor/si_nmr_data_small.json + testset_filename: /Users/mjwen.admin/Documents/Dataset/NMR_tensor/si_nmr_data_small.json + r_cut: 5.0 + reuse: false + loader_kwargs: + batch_size: 2 + shuffle: true + +model: + ########## + # embedding + ########## + + # atom species embedding + species_embedding_dim: 16 + + # spherical harmonics embedding of edge direction + irreps_edge_sh: 0e + 1o + 2e + + # radial edge distance embedding + radial_basis_type: bessel + num_radial_basis: 8 + radial_basis_start: 0. + radial_basis_end: 5. + + ########## + # message passing conv layers + ########## + num_layers: 3 + + # radial network + invariant_layers: 2 # number of radial layers + invariant_neurons: 32 # number of hidden neurons in radial function + + # Average number of neighbors used for normalization. Options: + # 1. `auto` to determine it automatically, by setting it to average number + # of neighbors of the training set + # 2. float or int provided here. + # 3. `null` to not use it + average_num_neighbors: auto + + # point convolution + conv_layer_irreps: 32x0o+32x0e + 16x1o+16x1e + 4x2o+4x2e + nonlinearity_type: gate + normalization: batch + resnet: true + + ########## + # output + ########## + + # output_format and output_formula should be used together. + # - output_format (can be `irreps` or `cartesian`) determines what the loss + # function will be on (either on the irreps space or the cartesian space). + # - output_formula gives what the cartesian formula of the tensor is. + # For example, ijkl=jikl=klij specifies a forth-rank elasticity tensor. + output_format: irreps + output_formula: ij=ji + + # pooling node feats to graph feats + reduce: mean + +trainer: + max_epochs: 10 # number of maximum training epochs + num_nodes: 1 + accelerator: cpu + devices: 1 + + callbacks: + - class_path: pytorch_lightning.callbacks.ModelCheckpoint + init_args: + monitor: val/score + mode: min + save_top_k: 3 + save_last: true + verbose: false + - class_path: pytorch_lightning.callbacks.EarlyStopping + init_args: + monitor: val/score + mode: min + patience: 150 + min_delta: 0 + verbose: true + - class_path: pytorch_lightning.callbacks.ModelSummary + init_args: + max_depth: -1 + + #logger: + # class_path: pytorch_lightning.loggers.wandb.WandbLogger + # init_args: + # save_dir: matten_logs + # project: matten_proj + +optimizer: + class_path: torch.optim.Adam + init_args: + lr: 0.01 + weight_decay: 0.00001 + +lr_scheduler: + class_path: torch.optim.lr_scheduler.ReduceLROnPlateau + init_args: + mode: min + factor: 0.5 + patience: 50 + verbose: true diff --git a/scripts/configs/materials_tensor.yaml b/scripts/configs/materials_tensor.yaml index f41a54c..921eff0 100644 --- a/scripts/configs/materials_tensor.yaml +++ b/scripts/configs/materials_tensor.yaml @@ -3,6 +3,7 @@ log_level: info data: root: ../datasets/ + tensor_target_name: elastic_tensor_full trainset_filename: example_crystal_elasticity_tensor_n100.json valset_filename: example_crystal_elasticity_tensor_n100.json testset_filename: example_crystal_elasticity_tensor_n100.json diff --git a/scripts/train_atomic_tensor.py b/scripts/train_atomic_tensor.py new file mode 100644 index 0000000..cad7e5a --- /dev/null +++ b/scripts/train_atomic_tensor.py @@ -0,0 +1,81 @@ +"""Script to train the materials tensor model.""" + +from pathlib import Path +from typing import Dict, List, Union + +import yaml +from loguru import logger +from pytorch_lightning import Trainer, seed_everything +from pytorch_lightning.cli import instantiate_class as lit_instantiate_class + +from matten.dataset.structure_scalar_tensor import TensorDataModule +from matten.log import set_logger +from matten.model_factory.task import TensorRegressionTask +from matten.model_factory.tfn_atomic_tensor import AtomicTensorModel + + +def instantiate_class(d: Union[Dict, List]): + args = tuple() # no positional args + if isinstance(d, dict): + return lit_instantiate_class(args, d) + elif isinstance(d, list): + return [lit_instantiate_class(args, x) for x in d] + else: + raise ValueError(f"Cannot instantiate class from {d}") + + +def get_args(path: Path): + """Get the arguments from the config file.""" + with open(path, "r") as f: + config = yaml.safe_load(f) + return config + + +def main(config: Dict): + dm = TensorDataModule(**config["data"]) + dm.prepare_data() + dm.setup() + + model = AtomicTensorModel( + tasks=TensorRegressionTask(name=config["data"]["tensor_target_name"]), + backbone_hparams=config["model"], + dataset_hparams=dm.get_to_model_info(), + optimizer_hparams=config["optimizer"], + lr_scheduler_hparams=config["lr_scheduler"], + ) + + try: + callbacks = instantiate_class(config["trainer"].pop("callbacks")) + lit_logger = instantiate_class(config["trainer"].pop("logger")) + except KeyError: + callbacks = None + lit_logger = None + + trainer = Trainer( + callbacks=callbacks, + logger=lit_logger, + **config["trainer"], + ) + + logger.info("Start training!") + trainer.fit(model, datamodule=dm) + + # test + logger.info("Start testing!") + trainer.test(ckpt_path="best", datamodule=dm) + + # print path of best checkpoint + logger.info(f"Best checkpoint path: {trainer.checkpoint_callback.best_model_path}") + + +if __name__ == "__main__": + config_file = Path(__file__).parent / "configs" / "atomic_tensor.yaml" + config = get_args(config_file) + + seed = config.get("seed_everything", 1) + seed_everything(seed) + + log_level = config.get("log_level", "INFO") + set_logger(log_level) + + main(config) diff --git a/scripts/train_materials_tensor.py b/scripts/train_materials_tensor.py index 75e6bba..e86202d 100644 --- a/scripts/train_materials_tensor.py +++ b/scripts/train_materials_tensor.py @@ -37,7 +37,7 @@ def main(config: Dict): dm.setup() model = ScalarTensorModel( - tasks=TensorRegressionTask(name="elastic_tensor_full"), + tasks=TensorRegressionTask(name=config["data"]["tensor_target_name"]), backbone_hparams=config["model"], dataset_hparams=dm.get_to_model_info(), optimizer_hparams=config["optimizer"], diff --git a/src/matten/dataset/structure_scalar_tensor.py b/src/matten/dataset/structure_scalar_tensor.py index e55f18d..3046bc8 100644 --- a/src/matten/dataset/structure_scalar_tensor.py +++ b/src/matten/dataset/structure_scalar_tensor.py @@ -50,6 +50,8 @@ class TensorDataset(InMemoryDataset): 10.0, 0:1.0}}, then for data points with minLC less than 0, it will have a weight of 10, and for those with minLC larger than 0, it will have a weight of 1. By default `None` means all data points has a weight of 1. + atom_selector: a list of bools to indicate which atoms to use in the structure + to compute the target. If `None`, all atoms are used. global_featurizer: featurizer to compute global features. normalize_global_features: whether to normalize the global feature. atom_featuruzer: featurizer to compute atom features. @@ -79,6 +81,7 @@ def __init__( log_scalar_targets: List[bool] = None, normalize_scalar_targets: List[bool] = None, tensor_target_weight: Dict[str, Dict[str, float]] = None, + atom_selector: List[bool] = None, global_featurizer: None = None, normalize_global_features: bool = False, atom_featurizer: None = None, @@ -97,6 +100,7 @@ def __init__( self.normalize_tensor_target = normalize_tensor_target self.tensor_target_weight = tensor_target_weight + self.atom_selector = atom_selector self.scalar_target_names = ( [] if scalar_target_names is None else scalar_target_names @@ -259,7 +263,7 @@ def _get_crystals(self, df): # TODO, convert to irreps tensor, assuming all input tensor is Cartesian converter = CartesianTensorWrapper(formula=self.tensor_target_formula) df[self.tensor_target_name] = df[self.tensor_target_name].apply( - lambda x: converter.from_cartesian(x).reshape(1, -1) + lambda x: torch.atleast_2d(converter.from_cartesian(x)) ) elif self.tensor_target_format == "cartesian": df[self.tensor_target_name] = df[self.tensor_target_name].apply( @@ -303,6 +307,10 @@ def _get_crystals(self, df): y[self.tensor_target_name] * self.tensor_target_scale ) + # atom selector + if self.atom_selector is not None: + y["atom_selector"] = torch.as_tensor(row[self.atom_selector]) + x = None if self.global_featurizer: # feats @@ -414,6 +422,7 @@ def __init__( tensor_target_scale: float = 1.0, normalize_tensor_target: bool = False, tensor_target_weight: Dict[str, Dict[str, float]] = None, + atom_selector: List[bool] = None, scalar_target_names: List[str] = None, log_scalar_targets: List[bool] = None, normalize_scalar_targets: List[bool] = None, @@ -457,6 +466,7 @@ def __init__( self.tensor_target_scale = tensor_target_scale self.normalize_tensor_target = normalize_tensor_target self.tensor_target_weight = tensor_target_weight + self.atom_selector = atom_selector self.scalar_target_names = scalar_target_names self.log_scalar_targets = log_scalar_targets @@ -563,6 +573,7 @@ def setup(self, stage: Optional[str] = None): log_scalar_targets=self.log_scalar_targets, normalize_scalar_targets=self.normalize_scalar_targets, tensor_target_weight=self.tensor_target_weight, + atom_selector=self.atom_selector, global_featurizer=gf, normalize_global_features=self.normalize_global_features, atom_featurizer=af, @@ -584,6 +595,7 @@ def setup(self, stage: Optional[str] = None): log_scalar_targets=self.log_scalar_targets, normalize_scalar_targets=self.normalize_scalar_targets, tensor_target_weight=self.tensor_target_weight, + atom_selector=self.atom_selector, global_featurizer=gf, normalize_global_features=self.normalize_global_features, atom_featurizer=af, @@ -605,6 +617,7 @@ def setup(self, stage: Optional[str] = None): log_scalar_targets=self.log_scalar_targets, normalize_scalar_targets=self.normalize_scalar_targets, tensor_target_weight=self.tensor_target_weight, + atom_selector=self.atom_selector, global_featurizer=gf, normalize_global_features=self.normalize_global_features, atom_featurizer=af, diff --git a/src/matten/model/model.py b/src/matten/model/model.py index 09c46b7..2da5b84 100644 --- a/src/matten/model/model.py +++ b/src/matten/model/model.py @@ -338,6 +338,11 @@ def shared_step(self, batch, mode: str): # ========== compute predictions ========== preds = self.decode(graphs) + # select atoms + if "atom_selector" in labels: + selector = labels["atom_selector"] + preds = {k: v[selector] for k, v in preds.items()} + # ========== compute losses ========== target_weight = graphs.get("target_weight", None) individual_loss, total_loss = self.compute_loss( @@ -504,6 +509,8 @@ def preprocess_batch(self, batch: DataPoint) -> Tuple[DataPoint, Dict[str, Tenso # task labels labels = {name: graphs.y[name] for name in self.tasks} + if "atom_selector" in graphs.y: + labels["atom_selector"] = graphs.y["atom_selector"] # convert graphs to a dict to use NequIP stuff graphs = graphs.tensor_property_to_dict() diff --git a/src/matten/model_factory/tfn_atomic_tensor.py b/src/matten/model_factory/tfn_atomic_tensor.py new file mode 100644 index 0000000..30e6463 --- /dev/null +++ b/src/matten/model_factory/tfn_atomic_tensor.py @@ -0,0 +1,228 @@ +""" +Make predictions for atomic tensor, i.e. a tensor at each atomic site. + +More general than tfn_scalar_tensor_global_feats.py in terms of the output format, +can be any scalar, vector, and tensor, which needs to be specified by the output format. + +Also, this won't work for the scalar via tensor case as in +tfn_scalar_tensor_global_feats.py. + +Note, this only works for a single target. +""" + +from collections import OrderedDict +from typing import Any, Dict, Optional, Tuple + +import torch +from e3nn.io import CartesianTensor +from torch import Tensor + +from matten.model.model import ModelForPyGData +from matten.model_factory.utils import create_sequential_module +from matten.nn._nequip import SphericalHarmonicEdgeAttrs +from matten.nn.conv import PointConv, PointConvWithActivation +from matten.nn.embedding import EdgeLengthEmbedding, SpeciesEmbedding +from matten.nn.nodewise import NodewiseLinear, NodewiseReduce +from matten.utils import ToCartesian + +OUT_FIELD_NAME = "my_model_output" + + +class AtomicTensorModel(ModelForPyGData): + def init_backbone( + self, + backbone_hparams: Dict[str, Any], + dataset_hparams: Optional[Dict[str, Any]] = None, + ) -> Tuple[torch.nn.Module, Dict]: + backbone = create_model(backbone_hparams, dataset_hparams) + + formula = (backbone_hparams["output_formula"]).lower() + + # # Last linear layer to covert irreps size + # if formula == "scalar": + # # special treatment of scalars + # irreps_out = Irreps("0e") + # else: + # irreps_out = CartesianTensor(formula=formula) + # irreps_in = backbone_hparams["conv_to_output_hidden_irreps_out"] + # extra_layers_dict = { + # "out_layer": Linear(irreps_in=irreps_in, irreps_out=irreps_out) + # } + + if backbone_hparams["output_format"] == "cartesian": + if formula == "scalar": + self.to_cartesian = None + else: + self.to_cartesian = ToCartesian(formula) + else: + self.to_cartesian = None + + # return backbone, extra_layers_dict + return backbone, None + + def decode(self, model_input) -> Dict[str, Tensor]: + out = self.backbone(model_input) + out = out[OUT_FIELD_NAME] + + # convert backbone to final irreps shape + # out = self.extra_layers_dict["out_layer"](out) + + if self.to_cartesian is not None: + out = self.to_cartesian(out) + + names = list(self.tasks.keys()) + assert len(names) == 1, f"only works for 1 target, get{len(names)}" + + task_name = names[0] + preds = {task_name: out} + + return preds + + def transform_prediction( + self, preds: Dict[str, Tensor], task_name: str = "elastic_tensor_full" + ) -> Dict[str, Tensor]: + """ + Transform the normalized prediction back. + """ + + normalizer = self.tasks[task_name].normalizer + + if normalizer is not None: + out = normalizer.inverse(preds[task_name]) + else: + out = preds[task_name] + + return {task_name: out} + + def transform_target( + self, target: Dict[str, Tensor], task_name: str = "elastic_tensor_full" + ) -> Dict[str, Tensor]: + return self.transform_prediction(target, task_name) + + +def create_model(hparams: Dict[str, Any], dataset_hparams): + """ + The actual function to create the model. + """ + use_atom_feats = hparams.get("use_atom_feats", False) + atom_feats_dim = dataset_hparams.get("atom_feats_size", None) + + # ===== input embedding layers ===== + layers = { + "one_hot": ( + SpeciesEmbedding, + { + "embedding_dim": hparams["species_embedding_dim"], + "allowed_species": dataset_hparams["allowed_species"], + "use_atom_feats": use_atom_feats, + "atom_feats_dim": atom_feats_dim, + }, + ), + "spharm_edges": ( + SphericalHarmonicEdgeAttrs, + {"irreps_edge_sh": hparams["irreps_edge_sh"]}, + ), + "radial_basis": ( + EdgeLengthEmbedding, + { + "num_basis": hparams["num_radial_basis"], + "start": hparams["radial_basis_start"], + "end": hparams["radial_basis_end"], + "basis": hparams["radial_basis_type"], + }, + ), + } + + # ===== convnet layers ===== + # insertion preserves order + + num_neigh = hparams["average_num_neighbors"] + if isinstance(num_neigh, str) and num_neigh.lower() == "auto": + num_neigh = dataset_hparams["average_num_neighbors"] + + for i in range(hparams["num_layers"]): + layers[f"layer{i}_convnet"] = ( + PointConvWithActivation, + { + "conv_layer_irreps": hparams["conv_layer_irreps"], + "activation_type": hparams["nonlinearity_type"], + "fc_num_hidden_layers": hparams["invariant_layers"], + "fc_hidden_size": hparams["invariant_neurons"], + "avg_num_neighbors": num_neigh, + "normalization": hparams["normalization"], + }, + ) + + # conv without applying activation + layers["conv_layer_last"] = ( + PointConv, + { + "conv_layer_irreps": hparams["conv_layer_irreps"], + "fc_num_hidden_layers": hparams["invariant_layers"], + "fc_hidden_size": hparams["invariant_neurons"], + "avg_num_neighbors": num_neigh, + }, + ) + + # ===== output head ===== + # + formula = (hparams["output_formula"]).lower() + irreps_out = CartesianTensor(formula=formula) + + layers.update( + { + # last layer of convnet + # -- output block -- + "conv_to_output_hidden": ( + NodewiseLinear, + { + # "irreps_out": hparams["conv_to_output_hidden_irreps_out"], + "irreps_out": irreps_out, + "out_field": OUT_FIELD_NAME, + }, + ) + } + ) + + # # pooling + # layers["output_pooling"] = ( + # NodewiseReduce, + # { + # "field": OUT_FIELD_NAME, + # "out_field": OUT_FIELD_NAME, + # "reduce": hparams["reduce"], + # }, + # ) + + model = create_sequential_module(modules=OrderedDict(layers)) + + return model + + +if __name__ == "__main__": + from matten.log import set_logger + + set_logger("DEBUG") + + hparams = { + "species_embedding_dim": 16, + # "species_embedding_irreps_out": "16x0e", + "conv_layer_irreps": "32x0o + 32x0e + 16x1o + 16x1e", + "irreps_edge_sh": "0e + 1o", + "num_radial_basis": 8, + "radial_basis_start": 0.0, + "radial_basis_end": 4.0, + "num_layers": 3, + "reduce": "sum", + "invariant_layers": 2, + "invariant_neurons": 64, + "average_num_neighbors": None, + "nonlinearity_type": "gate", + "conv_to_output_hidden_irreps_out": "16x0e", + "normalization": "batch", + "output_format": "irreps", + "output_formula": "2x0e+2x2e+4e", + } + + dataset_hyarmas = {"allowed_species": [6, 1, 8]} + create_model(hparams, dataset_hyarmas)