diff --git a/ocpmodels/__init__.py b/ocpmodels/__init__.py deleted file mode 100644 index c17674b3f0bcb88e2dd7ad3faeb589434b1356a5..0000000000000000000000000000000000000000 --- a/ocpmodels/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" diff --git a/ocpmodels/checkpointing.py b/ocpmodels/checkpointing.py deleted file mode 100644 index 87bc4f3b3b0a4961cebb1d92da03148130084079..0000000000000000000000000000000000000000 --- a/ocpmodels/checkpointing.py +++ /dev/null @@ -1,328 +0,0 @@ -import torch -import itertools -import warnings -from collections import OrderedDict -from typing import Any, List, Mapping, namedtuple - -# from torch/nn/modules/module.py - -_EXTRA_STATE_KEY_SUFFIX = "_extra_state" - - -class _IncompatibleKeys( - namedtuple("IncompatibleKeys", ["missing_keys", "unexpected_keys"]), -): - def __repr__(self): - if not self.missing_keys and not self.unexpected_keys: - return "" - return super().__repr__() - - __str__ = __repr__ - - -def _load_from_state_dict( - self, - state_dict, - prefix, - local_metadata, - strict, - missing_keys, - unexpected_keys, - error_msgs, -): - r"""Copy parameters and buffers from :attr:`state_dict` into only this module, but not its descendants. - - This is called on every submodule - in :meth:`~torch.nn.Module.load_state_dict`. Metadata saved for this - module in input :attr:`state_dict` is provided as :attr:`local_metadata`. - For state dicts without metadata, :attr:`local_metadata` is empty. - Subclasses can achieve class-specific backward compatible loading using - the version number at `local_metadata.get("version", None)`. - Additionally, :attr:`local_metadata` can also contain the key - `assign_to_params_buffers` that indicates whether keys should be - assigned their corresponding tensor in the state_dict. - - .. note:: - :attr:`state_dict` is not the same object as the input - :attr:`state_dict` to :meth:`~torch.nn.Module.load_state_dict`. So - it can be modified. - - Args: - state_dict (dict): a dict containing parameters and - persistent buffers. - prefix (str): the prefix for parameters and buffers used in this - module - local_metadata (dict): a dict containing the metadata for this module. - See - strict (bool): whether to strictly enforce that the keys in - :attr:`state_dict` with :attr:`prefix` match the names of - parameters and buffers in this module - missing_keys (list of str): if ``strict=True``, add missing keys to - this list - unexpected_keys (list of str): if ``strict=True``, add unexpected - keys to this list - error_msgs (list of str): error messages should be added to this - list, and will be reported together in - :meth:`~torch.nn.Module.load_state_dict` - """ - for hook in self._load_state_dict_pre_hooks.values(): - hook( - state_dict, - prefix, - local_metadata, - strict, - missing_keys, - unexpected_keys, - error_msgs, - ) - - persistent_buffers = { - k: v - for k, v in self._buffers.items() - if k not in self._non_persistent_buffers_set - } - local_name_params = itertools.chain( - self._parameters.items(), persistent_buffers.items() - ) - local_state = {k: v for k, v in local_name_params if v is not None} - assign_to_params_buffers = local_metadata.get("assign_to_params_buffers", False) - use_swap_tensors = torch.__future__.get_swap_module_params_on_conversion() - - for name, param in local_state.items(): - key = prefix + name - if key in state_dict: - input_param = state_dict[key] - if not torch.overrides.is_tensor_like(input_param): - error_msgs.append( - f'While copying the parameter named "{key}", ' - "expected torch.Tensor or Tensor-like object from checkpoint but " - f"received {type(input_param)}" - ) - continue - - # This is used to avoid copying uninitialized parameters into - # non-lazy modules, since they dont have the hook to do the checks - # in such case, it will error when accessing the .shape attribute. - is_param_lazy = torch.nn.parameter.is_lazy(param) - # Backward compatibility: loading 1-dim tensor from 0.3.* to version 0.4+ - if ( - not is_param_lazy - and len(param.shape) == 0 - and len(input_param.shape) == 1 - ): - input_param = input_param[0] - - if not is_param_lazy and input_param.shape != param.shape: - # local shape should match the one in checkpoint - error_msgs.append( - f"size mismatch for {key}: copying a param with shape {input_param.shape} from checkpoint, " - f"the shape in current model is {param.shape}." - ) - continue - - if ( - param.is_meta - and not input_param.is_meta - and not assign_to_params_buffers - ): - warnings.warn( - f"for {key}: copying from a non-meta parameter in the checkpoint to a meta " - "parameter in the current model, which is a no-op. (Did you mean to " - "pass `assign=True` to assign items in the state dictionary to their " - "corresponding key in the module instead of copying them in place?)" - ) - - try: - with torch.no_grad(): - if use_swap_tensors: - new_input_param = param.module_load( - input_param, assign=assign_to_params_buffers - ) - if id(new_input_param) == id(input_param) or id( - new_input_param - ) == id(param): - raise RuntimeError( - "module_load returned one of self or other, please .detach() " - "the result if returning one of the inputs in module_load" - ) - if isinstance(param, torch.nn.Parameter): - if not isinstance(new_input_param, torch.nn.Parameter): - new_input_param = torch.nn.Parameter( - new_input_param, - requires_grad=param.requires_grad, - ) - else: - new_input_param.requires_grad_(param.requires_grad) - torch.utils.swap_tensors(param, new_input_param) - del new_input_param - elif assign_to_params_buffers: - # Shape checks are already done above - if isinstance(param, torch.nn.Parameter): - if not isinstance(input_param, torch.nn.Parameter): - input_param = torch.nn.Parameter( - input_param, requires_grad=param.requires_grad - ) - else: - input_param.requires_grad_(param.requires_grad) - setattr(self, name, input_param) - else: - param.copy_(input_param) - except Exception as ex: - action = "swapping" if use_swap_tensors else "copying" - error_msgs.append( - f'While {action} the parameter named "{key}", ' - f"whose dimensions in the model are {param.size()} and " - f"whose dimensions in the checkpoint are {input_param.size()}, " - f"an exception occurred : {ex.args}." - ) - elif strict: - missing_keys.append(key) - - extra_state_key = prefix + _EXTRA_STATE_KEY_SUFFIX - if ( - getattr(self.__class__, "set_extra_state", torch.nn.Module.set_extra_state) - is not torch.nn.Module.set_extra_state - ): - if extra_state_key in state_dict: - self.set_extra_state(state_dict[extra_state_key]) - elif strict: - missing_keys.append(extra_state_key) - elif strict and (extra_state_key in state_dict): - unexpected_keys.append(extra_state_key) - - if strict: - for key in state_dict.keys(): - if key.startswith(prefix) and key != extra_state_key: - input_name = key[len(prefix) :].split(".", 1) - # Must be Module if it have attributes - if len(input_name) > 1: - if input_name[0] not in self._modules: - unexpected_keys.append(key) - elif input_name[0] not in local_state: - unexpected_keys.append(key) - - -def load_state_dict( - self, - state_dict: Mapping[str, Any], - assign: bool = False, - strict_unexpected_keys: bool = False, - strict_missing_keys: bool = False, -): - r"""Copy parameters and buffers from :attr:`state_dict` into this module and its descendants. - - If :attr:`strict` is ``True``, then - the keys of :attr:`state_dict` must exactly match the keys returned - by this module's :meth:`~torch.nn.Module.state_dict` function. - - .. warning:: - If :attr:`assign` is ``True`` the optimizer must be created after - the call to :attr:`load_state_dict` unless - :func:`~torch.__future__.get_swap_module_params_on_conversion` is ``True``. - - Args: - state_dict (dict): a dict containing parameters and - persistent buffers. - strict (bool, optional): whether to strictly enforce that the keys - in :attr:`state_dict` match the keys returned by this module's - :meth:`~torch.nn.Module.state_dict` function. Default: ``True`` - assign (bool, optional): When ``False``, the properties of the tensors - in the current module are preserved while when ``True``, the - properties of the Tensors in the state dict are preserved. The only - exception is the ``requires_grad`` field of :class:`~torch.nn.Parameter`s - for which the value from the module is preserved. - Default: ``False`` - - Returns: - ``NamedTuple`` with ``missing_keys`` and ``unexpected_keys`` fields: - * **missing_keys** is a list of str containing any keys that are expected - by this module but missing from the provided ``state_dict``. - * **unexpected_keys** is a list of str containing the keys that are not - expected by this module but present in the provided ``state_dict``. - - Note: - If a parameter or buffer is registered as ``None`` and its corresponding key - exists in :attr:`state_dict`, :meth:`load_state_dict` will raise a - ``RuntimeError``. - """ - if not isinstance(state_dict, Mapping): - raise TypeError(f"Expected state_dict to be dict-like, got {type(state_dict)}.") - - missing_keys: List[str] = [] - unexpected_keys: List[str] = [] - error_msgs: List[str] = [] - - # copy state_dict so _load_from_state_dict can modify it - metadata = getattr(state_dict, "_metadata", None) - state_dict = OrderedDict(state_dict) - if metadata is not None: - # mypy isn't aware that "_metadata" exists in state_dict - state_dict._metadata = metadata # type: ignore[attr-defined] - - def load(module, local_state_dict, prefix=""): - local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {}) - if assign: - local_metadata["assign_to_params_buffers"] = assign - module._load_from_state_dict( - local_state_dict, - prefix, - local_metadata, - True, - missing_keys, - unexpected_keys, - error_msgs, - ) - for name, child in module._modules.items(): - if child is not None: - child_prefix = prefix + name + "." - child_state_dict = { - k: v - for k, v in local_state_dict.items() - if k.startswith(child_prefix) - } - load(child, child_state_dict, child_prefix) # noqa: F821 - - # Note that the hook can modify missing_keys and unexpected_keys. - incompatible_keys = _IncompatibleKeys(missing_keys, unexpected_keys) - for hook in module._load_state_dict_post_hooks.values(): - out = hook(module, incompatible_keys) - assert out is None, ( - "Hooks registered with ``register_load_state_dict_post_hook`` are not" - "expected to return new values, if incompatible_keys need to be modified," - "it should be done inplace." - ) - - load(self, state_dict) - del load - - print("\nload_state_dict()") - print(f"Unexpected keys: {len(unexpected_keys)} / {len(state_dict.keys())}") - print(unexpected_keys[:10]) - print(f"Missing keys: {len(missing_keys)} / {len(state_dict.keys())}") - print(missing_keys[:10]) - - if strict_unexpected_keys: - if len(unexpected_keys) > 0: - error_msgs.insert( - 0, - "Unexpected key(s) in state_dict: {}. ".format( - ", ".join(f'"{k}"' for k in unexpected_keys[:20]) - ), - ) - - if strict_missing_keys: - if len(missing_keys) > 0: - error_msgs.insert( - 0, - "Missing key(s) in state_dict: {}. ".format( - ", ".join(f'"{k}"' for k in missing_keys[:20]) - ), - ) - - if len(error_msgs) > 0: - raise RuntimeError( - "Error(s) in loading state_dict for {}:\n\t{}".format( - self.__class__.__name__, "\n\t".join(error_msgs) - ) - ) - return _IncompatibleKeys(missing_keys, unexpected_keys) diff --git a/ocpmodels/common/__init__.py b/ocpmodels/common/__init__.py deleted file mode 100644 index c17674b3f0bcb88e2dd7ad3faeb589434b1356a5..0000000000000000000000000000000000000000 --- a/ocpmodels/common/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" diff --git a/ocpmodels/common/data_parallel.py b/ocpmodels/common/data_parallel.py deleted file mode 100644 index 3bb3afbd040f1c2a8a2a12631fc7fe87cb616042..0000000000000000000000000000000000000000 --- a/ocpmodels/common/data_parallel.py +++ /dev/null @@ -1,268 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import heapq -import logging -from itertools import chain -from pathlib import Path -from typing import List, Literal, Protocol, Union, runtime_checkable - -import numba -import numpy as np -import torch -from torch.utils.data import BatchSampler, DistributedSampler, Sampler - -from ocpmodels.common import distutils, gp_utils -from ocpmodels.datasets import data_list_collater - - -class OCPDataParallel(torch.nn.DataParallel): - def __init__(self, module, output_device, num_gpus): - if num_gpus < 0: - raise ValueError("# GPUs must be positive.") - if num_gpus > torch.cuda.device_count(): - raise ValueError("# GPUs specified larger than available") - - self.src_device = torch.device(output_device) - - self.cpu = False - if num_gpus == 0: - self.cpu = True - elif num_gpus == 1: - device_ids = [self.src_device] - else: - if self.src_device.type == "cuda" and self.src_device.index >= num_gpus: - raise ValueError("Main device must be less than # of GPUs") - device_ids = list(range(num_gpus)) - - if self.cpu: - super(torch.nn.DataParallel, self).__init__() - self.module = module - - else: - super(OCPDataParallel, self).__init__( - module=module, - device_ids=device_ids, - output_device=self.src_device, - ) - - def forward(self, batch_list, **kwargs): - if self.cpu: - return self.module(batch_list[0]) - - if len(self.device_ids) == 1: - return self.module(batch_list[0].to(f"cuda:{self.device_ids[0]}"), **kwargs) - - for t in chain(self.module.parameters(), self.module.buffers()): - if t.device != self.src_device: - raise RuntimeError( - ( - "Module must have its parameters and buffers on device " - "{} but found one of them on device {}." - ).format(self.src_device, t.device) - ) - - inputs = [ - batch.to(f"cuda:{self.device_ids[i]}") for i, batch in enumerate(batch_list) - ] - replicas = self.replicate(self.module, self.device_ids[: len(inputs)]) - outputs = self.parallel_apply(replicas, inputs, kwargs) - return self.gather(outputs, self.output_device) - - -class ParallelCollater: - def __init__(self, num_gpus, otf_graph=False): - self.num_gpus = num_gpus - self.otf_graph = otf_graph - - def __call__(self, data_list): - if self.num_gpus in [0, 1]: # adds cpu-only case - batch = data_list_collater(data_list, otf_graph=self.otf_graph) - return [batch] - - else: - num_devices = min(self.num_gpus, len(data_list)) - - count = torch.tensor([data.num_nodes for data in data_list]) - cumsum = count.cumsum(0) - cumsum = torch.cat([cumsum.new_zeros(1), cumsum], dim=0) - device_id = num_devices * cumsum.to(torch.float) / cumsum[-1].item() - device_id = (device_id[:-1] + device_id[1:]) / 2.0 - device_id = device_id.to(torch.long) - split = device_id.bincount().cumsum(0) - split = torch.cat([split.new_zeros(1), split], dim=0) - split = torch.unique(split, sorted=True) - split = split.tolist() - - return [ - data_list_collater(data_list[split[i] : split[i + 1]]) - for i in range(len(split) - 1) - ] - - -@numba.njit -def balanced_partition(sizes, num_parts): - """ - Greedily partition the given set by always inserting - the largest element into the smallest partition. - """ - sort_idx = np.argsort(-sizes) # Sort in descending order - heap = [] - for idx in sort_idx[:num_parts]: - heap.append((sizes[idx], [idx])) - heapq.heapify(heap) - for idx in sort_idx[num_parts:]: - smallest_part = heapq.heappop(heap) - new_size = smallest_part[0] + sizes[idx] - new_idx = smallest_part[1] + [idx] - heapq.heappush(heap, (new_size, new_idx)) - idx_balanced = [part[1] for part in heap] - return idx_balanced - - -@runtime_checkable -class _HasMetadata(Protocol): - @property - def metadata_path(self) -> Path: ... - - -class BalancedBatchSampler(Sampler): - def _load_dataset(self, dataset, mode: Literal["atoms", "neighbors"]): - errors: List[str] = [] - if not isinstance(dataset, _HasMetadata): - errors.append(f"Dataset {dataset} does not have a metadata_path attribute.") - return None, errors - if not dataset.metadata_path.exists(): - errors.append(f"Metadata file {dataset.metadata_path} does not exist.") - return None, errors - - key = {"atoms": "natoms", "neighbors": "neighbors"}[mode] - sizes = np.load(dataset.metadata_path)[key] - - return sizes, errors - - def __init__( - self, - dataset, - batch_size, - num_replicas, - rank, - device, - mode: Union[str, bool] = "atoms", - shuffle=True, - drop_last=False, - force_balancing=False, - throw_on_error=False, - ): - if mode is True: - mode = "atoms" - - if isinstance(mode, str): - mode = mode.lower() - if mode not in ("atoms", "neighbors"): - raise ValueError( - f"Invalid mode {mode}. Must be one of 'atoms', 'neighbors', or a boolean." - ) - - self.dataset = dataset - self.batch_size = batch_size - self.num_replicas = num_replicas - self.rank = rank - self.device = device - self.mode = mode - self.shuffle = shuffle - self.drop_last = drop_last - - self.single_sampler = DistributedSampler( - self.dataset, - num_replicas=num_replicas, - rank=rank, - shuffle=shuffle, - drop_last=drop_last, - ) - self.batch_sampler = BatchSampler( - self.single_sampler, - batch_size, - drop_last=drop_last, - ) - - self.sizes = None - self.balance_batches = False - - if self.num_replicas <= 1: - logging.info("Batch balancing is disabled for single GPU training.") - return - - if self.mode is False: - logging.info( - "Batch balancing is disabled because `optim.load_balancing` is `False`" - ) - return - - self.sizes, errors = self._load_dataset(dataset, self.mode) - if self.sizes is None: - self.balance_batches = force_balancing - if force_balancing: - errors.append( - "BalancedBatchSampler has to load the data to determine batch sizes, which incurs significant overhead! " - "You can disable balancing by setting `optim.load_balancing` to `False`." - ) - else: - errors.append( - "Batches will not be balanced, which can incur significant overhead!" - ) - else: - self.balance_batches = True - - if errors: - msg = "BalancedBatchSampler: " + " ".join(errors) - if throw_on_error: - raise RuntimeError(msg) - else: - logging.warning(msg) - - def __len__(self): - return len(self.batch_sampler) - - def set_epoch(self, epoch): - self.single_sampler.set_epoch(epoch) - - def __iter__(self): - if not self.balance_batches: - yield from self.batch_sampler - return - - for batch_idx in self.batch_sampler: - if self.sizes is None: - # Unfortunately, we need to load the data to know the image sizes - data_list = [self.dataset[idx] for idx in batch_idx] - - if self.mode == "atoms": - sizes = [data.num_nodes for data in data_list] - elif self.mode == "neighbors": - sizes = [data.edge_index.shape[1] for data in data_list] - else: - raise NotImplementedError( - f"Unknown load balancing mode: {self.mode}" - ) - else: - sizes = [self.sizes[idx] for idx in batch_idx] - - idx_sizes = torch.stack([torch.tensor(batch_idx), torch.tensor(sizes)]) - idx_sizes_all = distutils.all_gather(idx_sizes, device=self.device) - idx_sizes_all = torch.cat(idx_sizes_all, dim=-1).cpu() - if gp_utils.initialized(): - idx_sizes_all = torch.unique(input=idx_sizes_all, dim=1) - idx_all = idx_sizes_all[0] - sizes_all = idx_sizes_all[1] - - local_idx_balanced = balanced_partition( - sizes_all.numpy(), num_parts=self.num_replicas - ) - # Since DistributedSampler pads the last batch - # this should always have an entry for each replica. - yield idx_all[local_idx_balanced[self.rank]] diff --git a/ocpmodels/common/distutils.py b/ocpmodels/common/distutils.py deleted file mode 100644 index eb98f145846dc36149894205b937101faa51fab2..0000000000000000000000000000000000000000 --- a/ocpmodels/common/distutils.py +++ /dev/null @@ -1,157 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import logging -import os -import subprocess - -import torch -import torch.distributed as dist - - -def setup(config): - if config["submit"]: - node_list = os.environ.get("SLURM_STEP_NODELIST") - if node_list is None: - node_list = os.environ.get("SLURM_JOB_NODELIST") - if node_list is not None: - try: - hostnames = subprocess.check_output( - ["scontrol", "show", "hostnames", node_list] - ) - config["init_method"] = "tcp://{host}:{port}".format( - host=hostnames.split()[0].decode("utf-8"), - port=config["distributed_port"], - ) - nnodes = int(os.environ.get("SLURM_NNODES")) - ntasks_per_node = os.environ.get("SLURM_NTASKS_PER_NODE") - if ntasks_per_node is not None: - ntasks_per_node = int(ntasks_per_node) - else: - ntasks = int(os.environ.get("SLURM_NTASKS")) - nnodes = int(os.environ.get("SLURM_NNODES")) - assert ntasks % nnodes == 0 - ntasks_per_node = int(ntasks / nnodes) - if ntasks_per_node == 1: - assert config["world_size"] % nnodes == 0 - gpus_per_node = config["world_size"] // nnodes - node_id = int(os.environ.get("SLURM_NODEID")) - config["rank"] = node_id * gpus_per_node - config["local_rank"] = 0 - else: - assert ntasks_per_node == config["world_size"] // nnodes - config["rank"] = int(os.environ.get("SLURM_PROCID")) - config["local_rank"] = int(os.environ.get("SLURM_LOCALID")) - - logging.info( - f"Init: {config['init_method']}, {config['world_size']}, {config['rank']}" - ) - - # ensures GPU0 does not have extra context/higher peak memory - torch.cuda.set_device(config["local_rank"]) - - dist.init_process_group( - backend=config["distributed_backend"], - init_method=config["init_method"], - world_size=config["world_size"], - rank=config["rank"], - ) - except subprocess.CalledProcessError as e: # scontrol failed - raise e - except FileNotFoundError: # Slurm is not installed - pass - elif config["summit"]: - world_size = int(os.environ["OMPI_COMM_WORLD_SIZE"]) - world_rank = int(os.environ["OMPI_COMM_WORLD_RANK"]) - get_master = ( - "echo $(cat {} | sort | uniq | grep -v batch | grep -v login | head -1)" - ).format(os.environ["LSB_DJOB_HOSTFILE"]) - os.environ["MASTER_ADDR"] = str( - subprocess.check_output(get_master, shell=True) - )[2:-3] - os.environ["MASTER_PORT"] = "23456" - os.environ["WORLD_SIZE"] = os.environ["OMPI_COMM_WORLD_SIZE"] - os.environ["RANK"] = os.environ["OMPI_COMM_WORLD_RANK"] - # NCCL and MPI initialization - dist.init_process_group( - backend="nccl", - rank=world_rank, - world_size=world_size, - init_method="env://", - ) - else: - dist.init_process_group( - backend=config["distributed_backend"], init_method="env://" - ) - # TODO: SLURM - - -def cleanup(): - dist.destroy_process_group() - - -def initialized(): - return dist.is_available() and dist.is_initialized() - - -def get_rank(): - return dist.get_rank() if initialized() else 0 - - -def get_world_size(): - return dist.get_world_size() if initialized() else 1 - - -def is_master(): - return get_rank() == 0 - - -def synchronize(): - if get_world_size() == 1: - return - dist.barrier() - - -def broadcast(tensor, src, group=dist.group.WORLD, async_op=False): - if get_world_size() == 1: - return - dist.broadcast(tensor, src, group, async_op) - - -def all_reduce(data, group=dist.group.WORLD, average=False, device=None): - if get_world_size() == 1: - return data - tensor = data - if not isinstance(data, torch.Tensor): - tensor = torch.tensor(data) - if device is not None: - tensor = tensor.cuda(device) - dist.all_reduce(tensor, group=group) - if average: - tensor /= get_world_size() - if not isinstance(data, torch.Tensor): - result = tensor.cpu().numpy() if tensor.numel() > 1 else tensor.item() - else: - result = tensor - return result - - -def all_gather(data, group=dist.group.WORLD, device=None): - if get_world_size() == 1: - return data - tensor = data - if not isinstance(data, torch.Tensor): - tensor = torch.tensor(data) - if device is not None: - tensor = tensor.cuda(device) - tensor_list = [tensor.new_zeros(tensor.shape) for _ in range(get_world_size())] - dist.all_gather(tensor_list, tensor, group=group) - if not isinstance(data, torch.Tensor): - result = [tensor.cpu().numpy() for tensor in tensor_list] - else: - result = tensor_list - return result diff --git a/ocpmodels/common/flags.py b/ocpmodels/common/flags.py deleted file mode 100644 index 28c4fff889215586b19a9970297354d31a8821fa..0000000000000000000000000000000000000000 --- a/ocpmodels/common/flags.py +++ /dev/null @@ -1,140 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import argparse -from pathlib import Path - - -class Flags: - def __init__(self): - self.parser = argparse.ArgumentParser( - description="Graph Networks for Electrocatalyst Design" - ) - self.add_core_args() - - def get_parser(self): - return self.parser - - def add_core_args(self): - self.parser.add_argument_group("Core Arguments") - self.parser.add_argument( - "--mode", - choices=["train", "predict", "run-relaxations", "validate"], - required=True, - help="Whether to train the model, make predictions, or to run relaxations", - ) - self.parser.add_argument( - "--config-yml", - required=True, - type=Path, - help="Path to a config file listing data, model, optim parameters.", - ) - self.parser.add_argument( - "--identifier", - default="", - type=str, - help="Experiment identifier to append to checkpoint/log/result directory", - ) - self.parser.add_argument( - "--debug", - action="store_true", - help="Whether this is a debugging run or not", - ) - self.parser.add_argument( - "--run-dir", - default="./", - type=str, - help="Directory to store checkpoint/log/result directory", - ) - self.parser.add_argument( - "--print-every", - default=10, - type=int, - help="Log every N iterations (default: 10)", - ) - self.parser.add_argument( - "--seed", default=0, type=int, help="Seed for torch, cuda, numpy" - ) - self.parser.add_argument( - "--amp", action="store_true", help="Use mixed-precision training" - ) - self.parser.add_argument( - "--checkpoint", type=str, help="Model checkpoint to load" - ) - self.parser.add_argument( - "--timestamp-id", - default=None, - type=str, - help="Override time stamp ID. " - "Useful for seamlessly continuing model training in logger.", - ) - # Cluster args - self.parser.add_argument( - "--sweep-yml", - default=None, - type=Path, - help="Path to a config file with parameter sweeps", - ) - self.parser.add_argument( - "--submit", action="store_true", help="Submit job to cluster" - ) - self.parser.add_argument( - "--summit", action="store_true", help="Running on Summit cluster" - ) - self.parser.add_argument( - "--logdir", default="logs", type=Path, help="Where to store logs" - ) - self.parser.add_argument( - "--slurm-partition", - default="ocp", - type=str, - help="Name of partition", - ) - self.parser.add_argument( - "--slurm-mem", default=80, type=int, help="Memory (in gigabytes)" - ) - self.parser.add_argument( - "--slurm-timeout", default=72, type=int, help="Time (in hours)" - ) - self.parser.add_argument( - "--num-gpus", default=1, type=int, help="Number of GPUs to request" - ) - self.parser.add_argument( - "--distributed", action="store_true", help="Run with DDP" - ) - self.parser.add_argument( - "--cpu", action="store_true", help="Run CPU only training" - ) - self.parser.add_argument( - "--num-nodes", - default=1, - type=int, - help="Number of Nodes to request", - ) - self.parser.add_argument( - "--distributed-port", - type=int, - default=13356, - help="Port on master for DDP", - ) - self.parser.add_argument( - "--distributed-backend", - type=str, - default="nccl", - help="Backend for DDP", - ) - self.parser.add_argument("--local_rank", default=0, type=int, help="Local rank") - self.parser.add_argument("--no-ddp", action="store_true", help="Do not use DDP") - self.parser.add_argument( - "--gp-gpus", - type=int, - default=None, - help="Number of GPUs to split the graph over (only for Graph Parallel training)", - ) - - -flags = Flags() diff --git a/ocpmodels/common/gp_utils.py b/ocpmodels/common/gp_utils.py deleted file mode 100644 index cc598d8cfbcd7b9791457a2d902910a939d01968..0000000000000000000000000000000000000000 --- a/ocpmodels/common/gp_utils.py +++ /dev/null @@ -1,290 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import math -from typing import Any, List - -import torch -from torch import distributed as dist - -""" -Functions to support graph parallel training. -This is based on the Megatron-LM implementation: -https://github.com/facebookresearch/fairscale/blob/main/fairscale/nn/model_parallel/initialize.py -""" - -########## INITIALIZATION ########## - -_GRAPH_PARALLEL_GROUP = None -_DATA_PARALLEL_GROUP = None - - -def ensure_div(a, b): - assert a % b == 0 - - -def divide_and_check_no_remainder(a: int, b: int): - ensure_div(a, b) - return a // b - - -def setup_gp(config): - gp_size = config["gp_gpus"] - backend = config["distributed_backend"] - assert torch.distributed.is_initialized() - world_size = torch.distributed.get_world_size() - - gp_size = min(gp_size, world_size) - ensure_div(world_size, gp_size) - dp_size = world_size // gp_size - rank = dist.get_rank() - - if rank == 0: - print("> initializing graph parallel with size {}".format(gp_size)) - print("> initializing ddp with size {}".format(dp_size)) - - groups = torch.arange(world_size).reshape(dp_size, gp_size) - found = [x.item() for x in torch.where(groups == rank)] - - global _DATA_PARALLEL_GROUP - assert _DATA_PARALLEL_GROUP is None, "data parallel group is already initialized" - for j in range(gp_size): - group = dist.new_group(groups[:, j].tolist(), backend=backend) - if j == found[1]: - _DATA_PARALLEL_GROUP = group - global _GRAPH_PARALLEL_GROUP - assert _GRAPH_PARALLEL_GROUP is None, "graph parallel group is already initialized" - for i in range(dp_size): - group = dist.new_group(groups[i, :].tolist(), backend=backend) - if i == found[0]: - _GRAPH_PARALLEL_GROUP = group - - -def cleanup_gp(): - dist.destroy_process_group(_DATA_PARALLEL_GROUP) - dist.destroy_process_group(_GRAPH_PARALLEL_GROUP) - - -def initialized(): - return _GRAPH_PARALLEL_GROUP is not None - - -def get_dp_group(): - return _DATA_PARALLEL_GROUP - - -def get_gp_group(): - return _GRAPH_PARALLEL_GROUP - - -def get_dp_rank(): - return dist.get_rank(group=get_dp_group()) - - -def get_gp_rank(): - return dist.get_rank(group=get_gp_group()) - - -def get_dp_world_size(): - return dist.get_world_size(group=get_dp_group()) - - -def get_gp_world_size(): - return 1 if not initialized() else dist.get_world_size(group=get_gp_group()) - - -########## DIST METHODS ########## - - -def pad_tensor(tensor: torch.Tensor, dim: int = -1, target_size: int = None): - size = tensor.size(dim) - if target_size is None: - world_size = get_gp_world_size() - if size % world_size == 0: - pad_size = 0 - else: - pad_size = world_size - size % world_size - else: - pad_size = target_size - size - if pad_size == 0: - return tensor - pad_shape = list(tensor.shape) - pad_shape[dim] = pad_size - padding = torch.empty(pad_shape, device=tensor.device, dtype=tensor.dtype) - return torch.cat([tensor, padding], dim=dim) - - -def trim_tensor(tensor: torch.Tensor, sizes: torch.Tensor = None, dim: int = 0): - size = tensor.size(dim) - world_size = get_gp_world_size() - if size % world_size == 0: - return tensor, sizes - trim_size = size - size % world_size - if dim == 0: - tensor = tensor[:trim_size] - elif dim == 1: - tensor = tensor[:, :trim_size] - else: - raise ValueError - if sizes is not None: - sizes[-1] = sizes[-1] - size % world_size - return tensor, sizes - - -def _split_tensor( - tensor: torch.Tensor, - num_parts: int, - dim: int = -1, - contiguous_chunks: bool = False, -): - part_size = math.ceil(tensor.size(dim) / num_parts) - tensor_list = torch.split(tensor, part_size, dim=dim) - if contiguous_chunks: - return tuple(chunk.contiguous() for chunk in tensor_list) - return tensor_list - - -def _reduce(ctx: Any, input: torch.Tensor) -> torch.Tensor: - group = get_gp_group() - if ctx: - ctx.mark_dirty(input) - if dist.get_world_size(group) == 1: - return input - dist.all_reduce(input, group=group) - return input - - -def _split(input: torch.Tensor, dim: int = -1) -> torch.Tensor: - group = get_gp_group() - rank = get_gp_rank() - world_size = dist.get_world_size(group=group) - if world_size == 1: - return input - input_list = _split_tensor(input, world_size, dim=dim) - return input_list[rank].contiguous() - - -def _gather(input: torch.Tensor, dim: int = -1) -> torch.Tensor: - group = get_gp_group() - rank = get_gp_rank() - world_size = dist.get_world_size(group=group) - if world_size == 1: - return input - tensor_list = [torch.empty_like(input) for _ in range(world_size)] - tensor_list[rank] = input - dist.all_gather(tensor_list, input, group=group) - return torch.cat(tensor_list, dim=dim).contiguous() - - -def _gather_with_padding(input: torch.Tensor, dim: int = -1): - group = get_gp_group() - rank = get_gp_rank() - world_size = dist.get_world_size(group=group) - if world_size == 1: - return input - - # Gather sizes - size_list = [ - torch.empty(1, device=input.device, dtype=torch.long) for _ in range(world_size) - ] - size = torch.tensor([input.size(dim)], device=input.device, dtype=torch.long) - size_list[rank] = size - dist.all_gather(size_list, size, group=group) - - # Gather the inputs - max_size = max([size.item() for size in size_list]) - input = pad_tensor(input, dim, max_size) - shape = list(input.shape) - shape[dim] = max_size - tensor_list = [ - torch.empty(shape, device=input.device, dtype=input.dtype) - for _ in range(world_size) - ] - tensor_list[rank] = input - dist.all_gather(tensor_list, input, group=group) - - # Trim and cat - if dim == 0: - tensor_list = [tensor[:size] for tensor, size in zip(tensor_list, size_list)] - elif dim == 1: - tensor_list = [tensor[:, :size] for tensor, size in zip(tensor_list, size_list)] - else: - raise ValueError - return torch.cat(tensor_list, dim=dim).contiguous() - - -class CopyToModelParallelRegion(torch.autograd.Function): - @staticmethod - def forward(ctx, input: torch.Tensor): - return input - - @staticmethod - def backward(ctx, grad_output: torch.Tensor): - return _reduce(None, grad_output) - - -class ReduceFromModelParallelRegion(torch.autograd.Function): - @staticmethod - def forward(ctx, input: torch.Tensor): - return _reduce(ctx, input) - - @staticmethod - def backward(ctx, grad_output: torch.Tensor): - world_size = 1 - return grad_output.mul_(world_size) - - -class ScatterToModelParallelRegion(torch.autograd.Function): - @staticmethod - def forward(ctx, input: torch.Tensor, dim: int = -1): - result = _split(input, dim) - ctx.save_for_backward(torch.tensor(dim)) - return result - - @staticmethod - def backward(ctx, grad_output: torch.Tensor): - (dim,) = ctx.saved_tensors - world_size = 1 - return ( - _gather_with_padding(grad_output, dim.item()).div_(world_size), - None, - ) - - -class GatherFromModelParallelRegion(torch.autograd.Function): - @staticmethod - def forward(ctx, input: torch.Tensor, dim: int = -1): - ctx.save_for_backward(torch.tensor(dim)) - result = _gather_with_padding(input, dim) - return result - - @staticmethod - def backward(ctx, grad_output: torch.Tensor): - (dim,) = ctx.saved_tensors - result = _split(grad_output, dim.item()) - world_size = 1 - return result.mul_(world_size), None - - -def copy_to_model_parallel_region(input: torch.Tensor) -> torch.Tensor: - return CopyToModelParallelRegion.apply(input) - - -def reduce_from_model_parallel_region(input: torch.Tensor) -> torch.Tensor: - return ReduceFromModelParallelRegion.apply(input) - - -def scatter_to_model_parallel_region( - input: torch.Tensor, dim: int = -1 -) -> torch.Tensor: - return ScatterToModelParallelRegion.apply(input, dim) - - -def gather_from_model_parallel_region( - input: torch.Tensor, dim: int = -1 -) -> torch.Tensor: - return GatherFromModelParallelRegion.apply(input, dim) diff --git a/ocpmodels/common/hpo_utils.py b/ocpmodels/common/hpo_utils.py deleted file mode 100644 index 239a6f4821897c3159794f6054cdd5326c44e4e6..0000000000000000000000000000000000000000 --- a/ocpmodels/common/hpo_utils.py +++ /dev/null @@ -1,55 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import math - -from ray import tune - - -def tune_reporter( - iters, - train_metrics, - val_metrics, - test_metrics=None, - metric_to_opt="val_loss", - min_max="min", -): - """ - Wrapper function for tune.report() - - Args: - iters(dict): dict with training iteration info (e.g. steps, epochs) - train_metrics(dict): train metrics dict - val_metrics(dict): val metrics dict - test_metrics(dict, optional): test metrics dict, default is None - metric_to_opt(str, optional): str for val metric to optimize, default is val_loss - min_max(str, optional): either "min" or "max", determines whether metric_to_opt is to be minimized or maximized, default is min - - """ - # labels metric dicts - train = label_metric_dict(train_metrics, "train") - val = label_metric_dict(val_metrics, "val") - # this enables tolerance for NaNs assumes val set is used for optimization - if math.isnan(val[metric_to_opt]): - if min_max == "min": - val[metric_to_opt] = 100000.0 - if min_max == "max": - val[metric_to_opt] = 0.0 - if test_metrics: - test = label_metric_dict(test_metrics, "test") - else: - test = {} - # report results to Ray Tune - tune.report(**iters, **train, **val, **test) - - -def label_metric_dict(metric_dict, split): - new_dict = {} - for key in metric_dict: - new_dict["{}_{}".format(split, key)] = metric_dict[key] - metric_dict = new_dict - return metric_dict diff --git a/ocpmodels/common/logger.py b/ocpmodels/common/logger.py deleted file mode 100644 index 8fac6f6e0b47a050999d70729f99624575df7b14..0000000000000000000000000000000000000000 --- a/ocpmodels/common/logger.py +++ /dev/null @@ -1,115 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import logging -from abc import ABC, abstractmethod - -import torch -import wandb -from torch.utils.tensorboard import SummaryWriter - -from ocpmodels.common.registry import registry - - -class Logger(ABC): - """Generic class to interface with various logging modules, e.g. wandb, - tensorboard, etc. - """ - - def __init__(self, config): - self.config = config - - @abstractmethod - def watch(self, model): - """ - Monitor parameters and gradients. - """ - pass - - def log(self, update_dict, step=None, split=""): - """ - Log some values. - """ - assert step is not None - if split != "": - new_dict = {} - for key in update_dict: - new_dict["{}/{}".format(split, key)] = update_dict[key] - update_dict = new_dict - return update_dict - - @abstractmethod - def log_plots(self, plots): - pass - - @abstractmethod - def mark_preempting(self): - pass - - -@registry.register_logger("wandb") -class WandBLogger(Logger): - def __init__(self, config): - super().__init__(config) - project = ( - self.config["logger"].get("project", None) - if isinstance(self.config["logger"], dict) - else None - ) - - wandb.init( - config=self.config, - id=self.config["cmd"]["timestamp_id"], - name=self.config["cmd"]["identifier"], - dir=self.config["cmd"]["logs_dir"], - project=project, - resume="allow", - ) - - def watch(self, model): - wandb.watch(model) - - def log(self, update_dict, step=None, split=""): - update_dict = super().log(update_dict, step, split) - wandb.log(update_dict, step=int(step)) - - def log_plots(self, plots, caption=""): - assert isinstance(plots, list) - plots = [wandb.Image(x, caption=caption) for x in plots] - wandb.log({"data": plots}) - - def mark_preempting(self): - wandb.mark_preempting() - - -@registry.register_logger("tensorboard") -class TensorboardLogger(Logger): - def __init__(self, config): - super().__init__(config) - self.writer = SummaryWriter(self.config["cmd"]["logs_dir"]) - - # TODO: add a model hook for watching gradients. - def watch(self, model): - logging.warning("Model gradient logging to tensorboard not yet supported.") - return False - - def log(self, update_dict, step=None, split=""): - update_dict = super().log(update_dict, step, split) - for key in update_dict: - if torch.is_tensor(update_dict[key]): - self.writer.add_scalar(key, update_dict[key].item(), step) - else: - assert isinstance(update_dict[key], int) or isinstance( - update_dict[key], float - ) - self.writer.add_scalar(key, update_dict[key], step) - - def mark_preempting(self): - pass - - def log_plots(self, plots): - pass diff --git a/ocpmodels/common/registry.py b/ocpmodels/common/registry.py deleted file mode 100644 index 063fbab6f6a3bdfc1575c08747811ee1b4bb82ea..0000000000000000000000000000000000000000 --- a/ocpmodels/common/registry.py +++ /dev/null @@ -1,310 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -# Copyright (c) Facebook, Inc. and its affiliates. -# Borrowed from https://github.com/facebookresearch/pythia/blob/master/pythia/common/registry.py. -""" -Registry is central source of truth. Inspired from Redux's concept of -global store, Registry maintains mappings of various information to unique -keys. Special functions in registry can be used as decorators to register -different kind of classes. - -Import the global registry object using - -``from ocpmodels.common.registry import registry`` - -Various decorators for registry different kind of classes with unique keys - -- Register a model: ``@registry.register_model`` -""" -import importlib - - -def _get_absolute_mapping(name: str): - # in this case, the `name` should be the fully qualified name of the class - # e.g., `ocpmodels.tasks.base_task.BaseTask` - # we can use importlib to get the module (e.g., `ocpmodels.tasks.base_task`) - # and then import the class (e.g., `BaseTask`) - - module_name = ".".join(name.split(".")[:-1]) - class_name = name.split(".")[-1] - - try: - module = importlib.import_module(module_name) - except (ModuleNotFoundError, ValueError) as e: - raise RuntimeError( - f"Could not import module `{module_name}` for import `{name}`" - ) from e - - try: - return getattr(module, class_name) - except AttributeError as e: - raise RuntimeError( - f"Could not import class `{class_name}` from module `{module_name}`" - ) from e - - -class Registry: - r"""Class for registry object which acts as central source of truth.""" - - mapping = { - # Mappings to respective classes. - "task_name_mapping": {}, - "dataset_name_mapping": {}, - "model_name_mapping": {}, - "logger_name_mapping": {}, - "trainer_name_mapping": {}, - "state": {}, - } - - @classmethod - def register_task(cls, name): - r"""Register a new task to registry with key 'name' - Args: - name: Key with which the task will be registered. - Usage:: - from ocpmodels.common.registry import registry - from ocpmodels.tasks import BaseTask - @registry.register_task("train") - class TrainTask(BaseTask): - ... - """ - - def wrap(func): - cls.mapping["task_name_mapping"][name] = func - return func - - return wrap - - @classmethod - def register_dataset(cls, name): - r"""Register a dataset to registry with key 'name' - - Args: - name: Key with which the dataset will be registered. - - Usage:: - - from ocpmodels.common.registry import registry - from ocpmodels.datasets import BaseDataset - - @registry.register_dataset("qm9") - class QM9(BaseDataset): - ... - """ - - def wrap(func): - cls.mapping["dataset_name_mapping"][name] = func - return func - - return wrap - - @classmethod - def register_model(cls, name): - r"""Register a model to registry with key 'name' - - Args: - name: Key with which the model will be registered. - - Usage:: - - from ocpmodels.common.registry import registry - from ocpmodels.modules.layers import CGCNNConv - - @registry.register_model("cgcnn") - class CGCNN(): - ... - """ - - def wrap(func): - cls.mapping["model_name_mapping"][name] = func - return func - - return wrap - - @classmethod - def register_logger(cls, name): - r"""Register a logger to registry with key 'name' - - Args: - name: Key with which the logger will be registered. - - Usage:: - - from ocpmodels.common.registry import registry - - @registry.register_logger("tensorboard") - class WandB(): - ... - """ - - def wrap(func): - from ocpmodels.common.logger import Logger - - assert issubclass(func, Logger), "All loggers must inherit Logger class" - cls.mapping["logger_name_mapping"][name] = func - return func - - return wrap - - @classmethod - def register_trainer(cls, name): - r"""Register a trainer to registry with key 'name' - - Args: - name: Key with which the trainer will be registered. - - Usage:: - - from ocpmodels.common.registry import registry - - @registry.register_trainer("active_discovery") - class ActiveDiscoveryTrainer(): - ... - """ - - def wrap(func): - cls.mapping["trainer_name_mapping"][name] = func - return func - - return wrap - - @classmethod - def register(cls, name, obj): - r"""Register an item to registry with key 'name' - - Args: - name: Key with which the item will be registered. - - Usage:: - - from ocpmodels.common.registry import registry - - registry.register("config", {}) - """ - path = name.split(".") - current = cls.mapping["state"] - - for part in path[:-1]: - if part not in current: - current[part] = {} - current = current[part] - - current[path[-1]] = obj - - @classmethod - def __import_error(cls, name: str, mapping_name: str): - kind = mapping_name[: -len("_name_mapping")] - mapping = cls.mapping.get(mapping_name, {}) - existing_keys = list(mapping.keys()) - - existing_cls_path = ( - mapping.get(existing_keys[-1], None) if existing_keys else None - ) - if existing_cls_path is not None: - existing_cls_path = ( - f"{existing_cls_path.__module__}.{existing_cls_path.__qualname__}" - ) - else: - existing_cls_path = "ocpmodels.trainers.ForcesTrainer" - - existing_keys = [f"'{name}'" for name in existing_keys] - existing_keys = ", ".join(existing_keys[:-1]) + " or " + existing_keys[-1] - existing_keys_str = f" (one of {existing_keys})" if existing_keys else "" - return RuntimeError( - f"Failed to find the {kind} '{name}'. " - f"You may either use a {kind} from the registry{existing_keys_str} " - f"or provide the full import path to the {kind} (e.g., '{existing_cls_path}')." - ) - - @classmethod - def get_class(cls, name: str, mapping_name: str): - existing_mapping = cls.mapping[mapping_name].get(name, None) - if existing_mapping is not None: - return existing_mapping - - # mapping be class path of type `{module_name}.{class_name}` (e.g., `ocpmodels.trainers.ForcesTrainer`) - if name.count(".") < 1: - raise cls.__import_error(name, mapping_name) - - try: - return _get_absolute_mapping(name) - except RuntimeError as e: - raise cls.__import_error(name, mapping_name) from e - - @classmethod - def get_task_class(cls, name): - return cls.get_class(name, "task_name_mapping") - - @classmethod - def get_dataset_class(cls, name): - return cls.get_class(name, "dataset_name_mapping") - - @classmethod - def get_model_class(cls, name): - return cls.get_class(name, "model_name_mapping") - - @classmethod - def get_logger_class(cls, name): - return cls.get_class(name, "logger_name_mapping") - - @classmethod - def get_trainer_class(cls, name): - return cls.get_class(name, "trainer_name_mapping") - - @classmethod - def get(cls, name, default=None, no_warning=False): - r"""Get an item from registry with key 'name' - - Args: - name (string): Key whose value needs to be retreived. - default: If passed and key is not in registry, default value will - be returned with a warning. Default: None - no_warning (bool): If passed as True, warning when key doesn't exist - will not be generated. Useful for cgcnn's - internal operations. Default: False - Usage:: - - from ocpmodels.common.registry import registry - - config = registry.get("config") - """ - original_name = name - name = name.split(".") - value = cls.mapping["state"] - for subname in name: - value = value.get(subname, default) - if value is default: - break - - if ( - "writer" in cls.mapping["state"] - and value == default - and no_warning is False - ): - cls.mapping["state"]["writer"].write( - "Key {} is not present in registry, returning default value " - "of {}".format(original_name, default) - ) - return value - - @classmethod - def unregister(cls, name): - r"""Remove an item from registry with key 'name' - - Args: - name: Key which needs to be removed. - Usage:: - - from ocpmodels.common.registry import registry - - config = registry.unregister("config") - """ - return cls.mapping["state"].pop(name, None) - - -registry = Registry() diff --git a/ocpmodels/common/relaxation/__init__.py b/ocpmodels/common/relaxation/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/ocpmodels/common/relaxation/ase_utils.py b/ocpmodels/common/relaxation/ase_utils.py deleted file mode 100644 index fc391c7293db75fa5bd3bae2fd2beca384bccb4a..0000000000000000000000000000000000000000 --- a/ocpmodels/common/relaxation/ase_utils.py +++ /dev/null @@ -1,236 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. - - - -Utilities to interface OCP models/trainers with the Atomic Simulation -Environment (ASE) -""" - -import copy -import logging -import os -import yaml - -import numpy as np -import torch -from torch_geometric.data import Data as TGData -from torch_geometric.data import Batch - -from ase import Atoms -from ase.calculators.calculator import Calculator -from ase.calculators.singlepoint import SinglePointCalculator as sp -from ase.constraints import FixAtoms - -from ocpmodels.common.registry import registry -from ocpmodels.common.utils import ( - radius_graph_pbc, - setup_imports, - setup_logging, -) -from ocpmodels.datasets import data_list_collater -from ocpmodels.preprocessing import AtomsToGraphs - - -def batch_to_atoms(batch): - n_systems = batch.natoms.shape[0] - natoms = batch.natoms.tolist() - numbers = torch.split(batch.atomic_numbers, natoms) - fixed = torch.split(batch.fixed, natoms) - forces = torch.split(batch.force, natoms) - positions = torch.split(batch.pos, natoms) - tags = torch.split(batch.tags, natoms) - cells = batch.cell - energies = batch.y.tolist() - - atoms_objects = [] - for idx in range(n_systems): - atoms = Atoms( - numbers=numbers[idx].tolist(), - positions=positions[idx].cpu().detach().numpy(), - tags=tags[idx].tolist(), - cell=cells[idx].cpu().detach().numpy(), - constraint=FixAtoms(mask=fixed[idx].tolist()), - pbc=[True, True, True], - ) - calc = sp( - atoms=atoms, - energy=energies[idx], - forces=forces[idx].cpu().detach().numpy(), - ) - atoms.set_calculator(calc) - atoms_objects.append(atoms) - - return atoms_objects - - -class OCPCalculator(Calculator): - implemented_properties = ["energy", "forces"] - - def __init__( - self, - config_yml=None, - checkpoint=None, - trainer=None, - cutoff=6, - max_neighbors=50, - cpu=True, - ): - """ - OCP-ASE Calculator - - Args: - config_yml (str): - Path to yaml config or could be a dictionary. - checkpoint (str): - Path to trained checkpoint. - trainer (str): - OCP trainer to be used. "forces" for S2EF, "energy" for IS2RE. - cutoff (int): - Cutoff radius to be used for data preprocessing. - max_neighbors (int): - Maximum amount of neighbors to store for a given atom. - cpu (bool): - Whether to load and run the model on CPU. Set `False` for GPU. - """ - setup_imports() - setup_logging() - Calculator.__init__(self) - - # Either the config path or the checkpoint path needs to be provided - assert config_yml or checkpoint is not None - - if config_yml is not None: - if isinstance(config_yml, str): - config = yaml.safe_load(open(config_yml, "r")) - - if "includes" in config: - for include in config["includes"]: - # Change the path based on absolute path of config_yml - path = os.path.join(config_yml.split("configs")[0], include) - include_config = yaml.safe_load(open(path, "r")) - config.update(include_config) - else: - config = config_yml - # Only keeps the train data that might have normalizer values - if isinstance(config["dataset"], list): - config["dataset"] = config["dataset"][0] - elif isinstance(config["dataset"], dict): - config["dataset"] = config["dataset"].get("train", None) - else: - # Loads the config from the checkpoint directly (always on CPU). - config = torch.load(checkpoint, map_location=torch.device("cpu"))["config"] - if trainer is not None: # passing the arg overrides everything else - config["trainer"] = trainer - else: - if "trainer" not in config: # older checkpoint - if config["task"]["dataset"] == "trajectory_lmdb": - config["trainer"] = "forces" - elif config["task"]["dataset"] == "single_point_lmdb": - config["trainer"] = "energy" - else: - logging.warning( - "Unable to identify OCP trainer, defaulting to `forces`. Specify the `trainer` argument into OCPCalculator if otherwise." - ) - config["trainer"] = "forces" - - if "model_attributes" in config: - config["model_attributes"]["name"] = config.pop("model") - config["model"] = config["model_attributes"] - - # for checkpoints with relaxation datasets defined, remove to avoid - # unnecesarily trying to load that dataset - if "relax_dataset" in config["task"]: - del config["task"]["relax_dataset"] - - # Calculate the edge indices on the fly - config["model"]["otf_graph"] = True - - # Save config so obj can be transported over network (pkl) - self.config = copy.deepcopy(config) - self.config["checkpoint"] = checkpoint - - if "normalizer" not in config: - del config["dataset"]["src"] - config["normalizer"] = config["dataset"] - - self.trainer = registry.get_trainer_class(config.get("trainer", "energy"))( - task=config["task"], - model=config["model"], - dataset=None, - normalizer=config["normalizer"], - optimizer=config["optim"], - identifier="", - slurm=config.get("slurm", {}), - local_rank=config.get("local_rank", 0), - is_debug=config.get("is_debug", True), - cpu=cpu, - ) - - if checkpoint is not None: - self.load_checkpoint(checkpoint) - - self.a2g = AtomsToGraphs( - max_neigh=max_neighbors, - radius=cutoff, - r_energy=False, - r_forces=False, - r_distances=False, - r_edges=False, - r_pbc=True, - ) - - def load_checkpoint(self, checkpoint_path): - """ - Load existing trained model - - Args: - checkpoint_path: string - Path to trained model - """ - try: - self.trainer.load_checkpoint(checkpoint_path) - except NotImplementedError: - logging.warning("Unable to load checkpoint!") - - def calculate(self, atoms, properties, system_changes): - Calculator.calculate(self, atoms, properties, system_changes) - data_object = self.a2g.convert(atoms) - batch = data_list_collater([data_object], otf_graph=True) - - predictions = self.trainer.predict(batch, per_image=False, disable_tqdm=True) - if self.trainer.name == "s2ef": - self.results["energy"] = predictions["energy"].item() - self.results["forces"] = predictions["forces"].cpu().numpy() - - elif self.trainer.name == "is2re": - self.results["energy"] = predictions["energy"].item() - - -# by Andreas. We do not care about the graph, because connectivity is computed on the fly? -def ase_atoms_to_torch_geometric(atoms): - """ - Convert ASE Atoms object to torch_geometric Data format expected by Equiformer. - - Args: - atoms: ASE Atoms object - - Returns: - Data: torch_geometric Data object with required attributes - """ - positions = atoms.get_positions().astype(np.float32) - atomic_nums = atoms.get_atomic_numbers() - - # Convert to torch tensors - data = TGData( - pos=torch.tensor(positions, dtype=torch.float32), - z=torch.tensor(atomic_nums, dtype=torch.int64), - charges=torch.tensor(atomic_nums, dtype=torch.int64), - natoms=torch.tensor([len(atomic_nums)], dtype=torch.int64), - ) - data = Batch.from_data_list([data]) - - return data diff --git a/ocpmodels/common/relaxation/ml_relaxation.py b/ocpmodels/common/relaxation/ml_relaxation.py deleted file mode 100644 index a33097f411453e8132d5cdd54a47fb9da4a82826..0000000000000000000000000000000000000000 --- a/ocpmodels/common/relaxation/ml_relaxation.py +++ /dev/null @@ -1,90 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import logging -from collections import deque -from pathlib import Path - -import torch -from torch_geometric.data import Batch - -from ocpmodels.common.registry import registry -from ocpmodels.datasets.lmdb_dataset import data_list_collater - -from .optimizers.lbfgs_torch import LBFGS, TorchCalc - - -def ml_relax( - batch, - model, - steps, - fmax, - relax_opt, - save_full_traj, - device="cuda:0", - transform=None, - early_stop_batch=False, -): - """ - Runs ML-based relaxations. - Args: - batch: object - model: object - steps: int - Max number of steps in the structure relaxation. - fmax: float - Structure relaxation terminates when the max force - of the system is no bigger than fmax. - relax_opt: str - Optimizer and corresponding parameters to be used for structure relaxations. - save_full_traj: bool - Whether to save out the full ASE trajectory. If False, only save out initial and final frames. - """ - batches = deque([batch[0]]) - relaxed_batches = [] - while batches: - batch = batches.popleft() - oom = False - ids = batch.sid - calc = TorchCalc(model, transform) - - # Run ML-based relaxation - traj_dir = relax_opt.get("traj_dir", None) - optimizer = LBFGS( - batch, - calc, - maxstep=relax_opt.get("maxstep", 0.04), - memory=relax_opt["memory"], - damping=relax_opt.get("damping", 1.0), - alpha=relax_opt.get("alpha", 70.0), - device=device, - save_full_traj=save_full_traj, - traj_dir=Path(traj_dir) if traj_dir is not None else None, - traj_names=ids, - early_stop_batch=early_stop_batch, - ) - try: - relaxed_batch = optimizer.run(fmax=fmax, steps=steps) - relaxed_batches.append(relaxed_batch) - except RuntimeError as e: - oom = True - torch.cuda.empty_cache() - - if oom: - # move OOM recovery code outside of except clause to allow tensors to be freed. - data_list = batch.to_data_list() - if len(data_list) == 1: - raise e - logging.info( - f"Failed to relax batch with size: {len(data_list)}, splitting into two" - ) - mid = len(data_list) // 2 - batches.appendleft(data_list_collater(data_list[:mid])) - batches.appendleft(data_list_collater(data_list[mid:])) - - relaxed_batch = Batch.from_data_list(relaxed_batches) - return relaxed_batch diff --git a/ocpmodels/common/relaxation/optimizers/__init__.py b/ocpmodels/common/relaxation/optimizers/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/ocpmodels/common/relaxation/optimizers/lbfgs_torch.py b/ocpmodels/common/relaxation/optimizers/lbfgs_torch.py deleted file mode 100644 index 2ec9058510f2f87997dfe5fd9838c069188ba52f..0000000000000000000000000000000000000000 --- a/ocpmodels/common/relaxation/optimizers/lbfgs_torch.py +++ /dev/null @@ -1,229 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import logging -from collections import deque -from pathlib import Path -from typing import Deque, Optional - -import ase -import torch -from torch_geometric.data import Batch -from torch_scatter import scatter - -from ocpmodels.common.relaxation.ase_utils import batch_to_atoms -from ocpmodels.common.utils import radius_graph_pbc - - -class LBFGS: - def __init__( - self, - batch: Batch, - model: "TorchCalc", - maxstep=0.01, - memory=100, - damping=0.25, - alpha=100.0, - force_consistent=None, - device="cuda:0", - save_full_traj=True, - traj_dir: Path = None, - traj_names=None, - early_stop_batch: bool = False, - ): - self.batch = batch - self.model = model - self.maxstep = maxstep - self.memory = memory - self.damping = damping - self.alpha = alpha - self.H0 = 1.0 / self.alpha - self.force_consistent = force_consistent - self.device = device - self.save_full = save_full_traj - self.traj_dir = traj_dir - self.traj_names = traj_names - self.early_stop_batch = early_stop_batch - self.otf_graph = model.model._unwrapped_model.otf_graph - assert not self.traj_dir or ( - traj_dir and len(traj_names) - ), "Trajectory names should be specified to save trajectories" - logging.info("Step Fmax(eV/A)") - - if not self.otf_graph and "edge_index" not in batch: - self.model.update_graph(self.batch) - - def get_energy_and_forces(self, apply_constraint=True): - energy, forces = self.model.get_energy_and_forces(self.batch, apply_constraint) - return energy, forces - - def set_positions(self, update, update_mask): - if not self.early_stop_batch: - update = torch.where(update_mask.unsqueeze(1), update, 0.0) - self.batch.pos += update.to(dtype=torch.float32) - - if not self.otf_graph: - self.model.update_graph(self.batch) - - def check_convergence(self, iteration, forces=None, energy=None): - if forces is None or energy is None: - energy, forces = self.get_energy_and_forces() - forces = forces.to(dtype=torch.float64) - - max_forces_ = scatter( - (forces**2).sum(axis=1).sqrt(), self.batch.batch, reduce="max" - ) - logging.info( - f"{iteration} " + " ".join(f"{x:0.3f}" for x in max_forces_.tolist()) - ) - - # (batch_size) -> (nAtoms) - max_forces = max_forces_[self.batch.batch] - - return max_forces.ge(self.fmax), energy, forces - - def run(self, fmax, steps): - self.fmax = fmax - self.steps = steps - - self.s = deque(maxlen=self.memory) - self.y = deque(maxlen=self.memory) - self.rho = deque(maxlen=self.memory) - self.r0 = self.f0 = None - - self.trajectories = None - if self.traj_dir: - self.traj_dir.mkdir(exist_ok=True, parents=True) - self.trajectories = [ - ase.io.Trajectory(self.traj_dir / f"{name}.traj_tmp", mode="w") - for name in self.traj_names - ] - - iteration = 0 - converged = False - while iteration < steps and not converged: - update_mask, energy, forces = self.check_convergence(iteration) - converged = torch.all(torch.logical_not(update_mask)) - - if self.trajectories is not None: - if ( - self.save_full - or converged - or iteration == steps - 1 - or iteration == 0 - ): - self.write(energy, forces, update_mask) - - if not converged and iteration < steps - 1: - self.step(iteration, forces, update_mask) - - iteration += 1 - - # GPU memory usage as per nvidia-smi seems to gradually build up as - # batches are processed. This releases unoccupied cached memory. - torch.cuda.empty_cache() - - if self.trajectories is not None: - for traj in self.trajectories: - traj.close() - for name in self.traj_names: - traj_fl = Path(self.traj_dir / f"{name}.traj_tmp", mode="w") - traj_fl.rename(traj_fl.with_suffix(".traj")) - - self.batch.y, self.batch.force = self.get_energy_and_forces( - apply_constraint=False - ) - return self.batch - - def step( - self, - iteration: int, - forces: Optional[torch.Tensor], - update_mask: torch.Tensor, - ): - def determine_step(dr): - steplengths = torch.norm(dr, dim=1) - longest_steps = scatter(steplengths, self.batch.batch, reduce="max") - longest_steps = longest_steps[self.batch.batch] - maxstep = longest_steps.new_tensor(self.maxstep) - scale = (longest_steps + 1e-7).reciprocal() * torch.min( - longest_steps, maxstep - ) - dr *= scale.unsqueeze(1) - return dr * self.damping - - if forces is None: - _, forces = self.get_energy_and_forces() - - r = self.batch.pos.clone().to(dtype=torch.float64) - - # Update s, y, rho - if iteration > 0: - s0 = (r - self.r0).flatten() - self.s.append(s0) - - y0 = -(forces - self.f0).flatten() - self.y.append(y0) - - self.rho.append(1.0 / torch.dot(y0, s0)) - - loopmax = min(self.memory, iteration) - alpha = forces.new_empty(loopmax) - q = -forces.flatten() - - for i in range(loopmax - 1, -1, -1): - alpha[i] = self.rho[i] * torch.dot(self.s[i], q) # b - q -= alpha[i] * self.y[i] - - z = self.H0 * q - for i in range(loopmax): - beta = self.rho[i] * torch.dot(self.y[i], z) - z += self.s[i] * (alpha[i] - beta) - - # descent direction - p = -z.reshape((-1, 3)) - dr = determine_step(p) - if torch.abs(dr).max() < 1e-7: - # Same configuration again (maybe a restart): - return - - self.set_positions(dr, update_mask) - - self.r0 = r - self.f0 = forces - - def write(self, energy, forces, update_mask): - self.batch.y, self.batch.force = energy, forces - atoms_objects = batch_to_atoms(self.batch) - update_mask_ = torch.split(update_mask, self.batch.natoms.tolist()) - for atm, traj, mask in zip(atoms_objects, self.trajectories, update_mask_): - if mask[0] or not self.save_full: - traj.write(atm) - - -class TorchCalc: - def __init__(self, model, transform=None): - self.model = model - self.transform = transform - - def get_energy_and_forces(self, atoms, apply_constraint=True): - predictions = self.model.predict(atoms, per_image=False, disable_tqdm=True) - energy = predictions["energy"] - forces = predictions["forces"] - if apply_constraint: - fixed_idx = torch.where(atoms.fixed == 1)[0] - forces[fixed_idx] = 0 - return energy, forces - - def update_graph(self, atoms): - edge_index, cell_offsets, num_neighbors = radius_graph_pbc(atoms, 6, 50) - atoms.edge_index = edge_index - atoms.cell_offsets = cell_offsets - atoms.neighbors = num_neighbors - if self.transform is not None: - atoms = self.transform(atoms) - return atoms diff --git a/ocpmodels/common/transforms.py b/ocpmodels/common/transforms.py deleted file mode 100644 index d364643dc3fddce11176a58bc9a6f28655459241..0000000000000000000000000000000000000000 --- a/ocpmodels/common/transforms.py +++ /dev/null @@ -1,78 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -# Borrowed from https://github.com/rusty1s/pytorch_geometric/blob/master/torch_geometric/transforms/random_rotate.py -# with changes to keep track of the rotation / inverse rotation matrices. - -import math -import numbers -import random - -import torch -import torch_geometric -from torch_geometric.transforms import LinearTransformation - - -class RandomRotate(object): - r"""Rotates node positions around a specific axis by a randomly sampled - factor within a given interval. - - Args: - degrees (tuple or float): Rotation interval from which the rotation - angle is sampled. If `degrees` is a number instead of a - tuple, the interval is given by :math:`[-\mathrm{degrees}, - \mathrm{degrees}]`. - axes (int, optional): The rotation axes. (default: `[0, 1, 2]`) - """ - - def __init__(self, degrees, axes=[0, 1, 2]): - if isinstance(degrees, numbers.Number): - degrees = (-abs(degrees), abs(degrees)) - assert isinstance(degrees, (tuple, list)) and len(degrees) == 2 - self.degrees = degrees - self.axes = axes - - def __call__(self, data): - if data.pos.size(-1) == 2: - degree = math.pi * random.uniform(*self.degrees) / 180.0 - sin, cos = math.sin(degree), math.cos(degree) - matrix = [[cos, sin], [-sin, cos]] - else: - m1, m2, m3 = torch.eye(3), torch.eye(3), torch.eye(3) - if 0 in self.axes: - degree = math.pi * random.uniform(*self.degrees) / 180.0 - sin, cos = math.sin(degree), math.cos(degree) - m1 = torch.tensor([[1, 0, 0], [0, cos, sin], [0, -sin, cos]]) - if 1 in self.axes: - degree = math.pi * random.uniform(*self.degrees) / 180.0 - sin, cos = math.sin(degree), math.cos(degree) - m2 = torch.tensor([[cos, 0, -sin], [0, 1, 0], [sin, 0, cos]]) - if 2 in self.axes: - degree = math.pi * random.uniform(*self.degrees) / 180.0 - sin, cos = math.sin(degree), math.cos(degree) - m3 = torch.tensor([[cos, sin, 0], [-sin, cos, 0], [0, 0, 1]]) - - matrix = torch.mm(torch.mm(m1, m2), m3) - - data_rotated = LinearTransformation(matrix)(data) - if torch_geometric.__version__.startswith("2."): - matrix = matrix.T - - # LinearTransformation only rotates `.pos`; need to rotate `.cell` too. - if hasattr(data_rotated, "cell"): - data_rotated.cell = torch.matmul(data_rotated.cell, matrix) - - return ( - data_rotated, - matrix, - torch.inverse(matrix), - ) - - def __repr__(self): - return "{}({}, axis={})".format( - self.__class__.__name__, self.degrees, self.axis - ) diff --git a/ocpmodels/common/utils.py b/ocpmodels/common/utils.py deleted file mode 100644 index 2966758c7732c1e5ae5595ad16c6d54d400ed87e..0000000000000000000000000000000000000000 --- a/ocpmodels/common/utils.py +++ /dev/null @@ -1,1019 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import ast -import collections -import copy -import importlib -import itertools -import json -import logging -import os -import sys -import time -from argparse import Namespace -from bisect import bisect -from contextlib import contextmanager -from dataclasses import dataclass -from functools import wraps -from itertools import product -from pathlib import Path -from typing import TYPE_CHECKING, Any, Dict, List, Mapping, Optional - -import numpy as np -import torch -import torch.nn as nn -import torch_geometric -import yaml -from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas -from matplotlib.figure import Figure -from torch_geometric.data import Data -from torch_geometric.utils import remove_self_loops -from torch_scatter import scatter, segment_coo, segment_csr - -if TYPE_CHECKING: - from torch.nn.modules.module import _IncompatibleKeys - - -def pyg2_data_transform(data: Data): - """ - if we're on the new pyg (2.0 or later) and if the Data stored is in older format - we need to convert the data to the new format - """ - if torch_geometric.__version__ >= "2.0" and "_store" not in data.__dict__: - return Data(**{k: v for k, v in data.__dict__.items() if v is not None}) - - return data - - -def save_checkpoint( - state, checkpoint_dir="checkpoints/", checkpoint_file="checkpoint.pt" -): - filename = os.path.join(checkpoint_dir, checkpoint_file) - torch.save(state, filename) - return filename - - -class Complete(object): - def __call__(self, data): - device = data.edge_index.device - - row = torch.arange(data.num_nodes, dtype=torch.long, device=device) - col = torch.arange(data.num_nodes, dtype=torch.long, device=device) - - row = row.view(-1, 1).repeat(1, data.num_nodes).view(-1) - col = col.repeat(data.num_nodes) - edge_index = torch.stack([row, col], dim=0) - - edge_attr = None - if data.edge_attr is not None: - idx = data.edge_index[0] * data.num_nodes + data.edge_index[1] - size = list(data.edge_attr.size()) - size[0] = data.num_nodes * data.num_nodes - edge_attr = data.edge_attr.new_zeros(size) - edge_attr[idx] = data.edge_attr - - edge_index, edge_attr = remove_self_loops(edge_index, edge_attr) - data.edge_attr = edge_attr - data.edge_index = edge_index - - return data - - -def warmup_lr_lambda(current_step, optim_config): - """Returns a learning rate multiplier. - Till `warmup_steps`, learning rate linearly increases to `initial_lr`, - and then gets multiplied by `lr_gamma` every time a milestone is crossed. - """ - - # keep this block for older configs that have warmup_epochs instead of warmup_steps - # and lr_milestones are defined in epochs - if ( - any(x < 100 for x in optim_config["lr_milestones"]) - or "warmup_epochs" in optim_config - ): - raise Exception( - "ConfigError: please define lr_milestones in steps not epochs and define warmup_steps instead of warmup_epochs" - ) - - if current_step <= optim_config["warmup_steps"]: - alpha = current_step / float(optim_config["warmup_steps"]) - return optim_config["warmup_factor"] * (1.0 - alpha) + alpha - else: - idx = bisect(optim_config["lr_milestones"], current_step) - return pow(optim_config["lr_gamma"], idx) - - -def print_cuda_usage(): - print("Memory Allocated:", torch.cuda.memory_allocated() / (1024 * 1024)) - print( - "Max Memory Allocated:", - torch.cuda.max_memory_allocated() / (1024 * 1024), - ) - print("Memory Cached:", torch.cuda.memory_cached() / (1024 * 1024)) - print("Max Memory Cached:", torch.cuda.max_memory_cached() / (1024 * 1024)) - - -def conditional_grad(dec): - "Decorator to enable/disable grad depending on whether force/energy predictions are being made" - - # Adapted from https://stackoverflow.com/questions/60907323/accessing-class-property-as-decorator-argument - def decorator(func): - @wraps(func) - def cls_method(self, *args, **kwargs): - f = func - if self.regress_forces and not getattr(self, "direct_forces", 0): - f = dec(func) - return f(self, *args, **kwargs) - - return cls_method - - return decorator - - -def plot_histogram(data, xlabel="", ylabel="", title=""): - assert isinstance(data, list) - - # Preset - fig = Figure(figsize=(5, 4), dpi=150) - canvas = FigureCanvas(fig) - ax = fig.gca() - - # Plot - ax.hist(data, bins=20, rwidth=0.9, zorder=3) - - # Axes - ax.grid(color="0.95", zorder=0) - ax.set_xlabel(xlabel) - ax.set_ylabel(ylabel) - ax.set_title(title) - fig.tight_layout(pad=2) - - # Return numpy array - canvas.draw() - image_from_plot = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8) - image_from_plot = image_from_plot.reshape( - fig.canvas.get_width_height()[::-1] + (3,) - ) - - return image_from_plot - - -# Override the collation method in `pytorch_geometric.data.InMemoryDataset` -def collate(data_list): - keys = data_list[0].keys - data = data_list[0].__class__() - - for key in keys: - data[key] = [] - slices = {key: [0] for key in keys} - - for item, key in product(data_list, keys): - data[key].append(item[key]) - if torch.is_tensor(item[key]): - s = slices[key][-1] + item[key].size(item.__cat_dim__(key, item[key])) - elif isinstance(item[key], int) or isinstance(item[key], float): - s = slices[key][-1] + 1 - else: - raise ValueError("Unsupported attribute type") - slices[key].append(s) - - if hasattr(data_list[0], "__num_nodes__"): - data.__num_nodes__ = [] - for item in data_list: - data.__num_nodes__.append(item.num_nodes) - - for key in keys: - if torch.is_tensor(data_list[0][key]): - data[key] = torch.cat( - data[key], dim=data.__cat_dim__(key, data_list[0][key]) - ) - else: - data[key] = torch.tensor(data[key]) - slices[key] = torch.tensor(slices[key], dtype=torch.long) - - return data, slices - - -def add_edge_distance_to_graph( - batch, - device="cpu", - dmin=0.0, - dmax=6.0, - num_gaussians=50, -): - # Make sure x has positions. - if not all(batch.pos[0][:] == batch.x[0][-3:]): - batch.x = torch.cat([batch.x, batch.pos.float()], dim=1) - # First set computations to be tracked for positions. - batch.x = batch.x.requires_grad_(True) - # Then compute Euclidean distance between edge endpoints. - pdist = torch.nn.PairwiseDistance(p=2.0) - distances = pdist( - batch.x[batch.edge_index[0]][:, -3:], - batch.x[batch.edge_index[1]][:, -3:], - ) - # Expand it using a gaussian basis filter. - gdf_filter = torch.linspace(dmin, dmax, num_gaussians) - var = gdf_filter[1] - gdf_filter[0] - gdf_filter, var = gdf_filter.to(device), var.to(device) - gdf_distances = torch.exp(-((distances.view(-1, 1) - gdf_filter) ** 2) / var**2) - # Reassign edge attributes. - batch.edge_weight = distances - batch.edge_attr = gdf_distances.float() - return batch - - -def _import_local_file(path: Path, *, project_root: Path): - """ - Imports a Python file as a module - - :param path: The path to the file to import - :type path: Path - :param project_root: The root directory of the project (i.e., the "ocp" folder) - :type project_root: Path - """ - - path = path.resolve() - project_root = project_root.resolve() - - module_name = ".".join( - path.absolute().relative_to(project_root.absolute()).with_suffix("").parts - ) - logging.debug(f"Resolved module name of {path} to {module_name}") - importlib.import_module(module_name) - - -def setup_experimental_imports(project_root: Path): - experimental_folder = (project_root / "experimental").resolve() - if not experimental_folder.exists() or not experimental_folder.is_dir(): - return - - experimental_files = [ - f.resolve().absolute() for f in experimental_folder.rglob("*.py") - ] - # Ignore certain directories within experimental - ignore_file = experimental_folder / ".ignore" - if ignore_file.exists(): - with open(ignore_file, "r") as f: - for line in f.read().splitlines(): - for ignored_file in (experimental_folder / line).rglob("*.py"): - experimental_files.remove(ignored_file.resolve().absolute()) - - for f in experimental_files: - _import_local_file(f, project_root=project_root) - - -def _get_project_root(): - """ - Gets the root folder of the project (the "ocp" folder) - :return: The absolute path to the project root. - """ - from ocpmodels.common.registry import registry - - # Automatically load all of the modules, so that - # they register with registry - root_folder = registry.get("ocpmodels_root", no_warning=True) - - if root_folder is not None: - assert isinstance(root_folder, str), "ocpmodels_root must be a string" - root_folder = Path(root_folder).resolve().absolute() - assert root_folder.exists(), f"{root_folder} does not exist" - assert root_folder.is_dir(), f"{root_folder} is not a directory" - else: - root_folder = Path(__file__).resolve().absolute().parent.parent - - # root_folder is the "ocpmodes" folder, so we need to go up one more level - return root_folder.parent - - -# Copied from https://github.com/facebookresearch/mmf/blob/master/mmf/utils/env.py#L89. -def setup_imports(config: Optional[dict] = None): - from ocpmodels.common.registry import registry - - skip_experimental_imports = (config or {}).get("skip_experimental_imports", None) - - # First, check if imports are already setup - has_already_setup = registry.get("imports_setup", no_warning=True) - if has_already_setup: - return - - try: - project_root = _get_project_root() - logging.info(f"Project root: {project_root}") - importlib.import_module("ocpmodels.common.logger") - - import_keys = ["trainers", "datasets", "models", "tasks"] - for key in import_keys: - for f in (project_root / "ocpmodels" / key).rglob("*.py"): - _import_local_file(f, project_root=project_root) - - if not skip_experimental_imports: - setup_experimental_imports(project_root) - finally: - registry.register("imports_setup", True) - - -def dict_set_recursively(dictionary, key_sequence, val): - top_key = key_sequence.pop(0) - if len(key_sequence) == 0: - dictionary[top_key] = val - else: - if top_key not in dictionary: - dictionary[top_key] = {} - dict_set_recursively(dictionary[top_key], key_sequence, val) - - -def parse_value(value): - """ - Parse string as Python literal if possible and fallback to string. - """ - try: - return ast.literal_eval(value) - except (ValueError, SyntaxError): - # Use as string if nothing else worked - return value - - -def create_dict_from_args(args: list, sep: str = "."): - """ - Create a (nested) dictionary from console arguments. - Keys in different dictionary levels are separated by sep. - """ - return_dict = {} - for arg in args: - arg = arg.strip("--") - keys_concat, val = arg.split("=") - val = parse_value(val) - key_sequence = keys_concat.split(sep) - dict_set_recursively(return_dict, key_sequence, val) - return return_dict - - -def load_config(path: str, previous_includes: list = []): - path = Path(path) - if path in previous_includes: - raise ValueError( - f"Cyclic config include detected. {path} included in sequence {previous_includes}." - ) - previous_includes = previous_includes + [path] - - direct_config = yaml.safe_load(open(path, "r")) - - # Load config from included files. - if "includes" in direct_config: - includes = direct_config.pop("includes") - else: - includes = [] - if not isinstance(includes, list): - raise AttributeError( - "Includes must be a list, '{}' provided".format(type(includes)) - ) - - config = {} - duplicates_warning = [] - duplicates_error = [] - - for include in includes: - include_config, inc_dup_warning, inc_dup_error = load_config( - include, previous_includes - ) - duplicates_warning += inc_dup_warning - duplicates_error += inc_dup_error - - # Duplicates between includes causes an error - config, merge_dup_error = merge_dicts(config, include_config) - duplicates_error += merge_dup_error - - # Duplicates between included and main file causes warnings - config, merge_dup_warning = merge_dicts(config, direct_config) - duplicates_warning += merge_dup_warning - - return config, duplicates_warning, duplicates_error - - -def build_config(args, args_override): - config, duplicates_warning, duplicates_error = load_config(args.config_yml) - if len(duplicates_warning) > 0: - logging.warning( - f"Overwritten config parameters from included configs " - f"(non-included parameters take precedence): {duplicates_warning}" - ) - if len(duplicates_error) > 0: - raise ValueError( - f"Conflicting (duplicate) parameters in simultaneously " - f"included configs: {duplicates_error}" - ) - - # Check for overridden parameters. - if args_override != []: - overrides = create_dict_from_args(args_override) - config, _ = merge_dicts(config, overrides) - - # Some other flags. - config["mode"] = args.mode - config["identifier"] = args.identifier - config["timestamp_id"] = args.timestamp_id - config["seed"] = args.seed - config["is_debug"] = args.debug - config["run_dir"] = args.run_dir - config["print_every"] = args.print_every - config["amp"] = args.amp - config["checkpoint"] = args.checkpoint - config["cpu"] = args.cpu - # Submit - config["submit"] = args.submit - config["summit"] = args.summit - # Distributed - config["local_rank"] = args.local_rank - config["distributed_port"] = args.distributed_port - config["world_size"] = args.num_nodes * args.num_gpus - config["distributed_backend"] = args.distributed_backend - config["noddp"] = args.no_ddp - config["gp_gpus"] = args.gp_gpus - - return config - - -def create_grid(base_config, sweep_file): - def _flatten_sweeps(sweeps, root_key="", sep="."): - flat_sweeps = [] - for key, value in sweeps.items(): - new_key = root_key + sep + key if root_key else key - if isinstance(value, collections.MutableMapping): - flat_sweeps.extend(_flatten_sweeps(value, new_key).items()) - else: - flat_sweeps.append((new_key, value)) - return collections.OrderedDict(flat_sweeps) - - def _update_config(config, keys, override_vals, sep="."): - for key, value in zip(keys, override_vals): - key_path = key.split(sep) - child_config = config - for name in key_path[:-1]: - child_config = child_config[name] - child_config[key_path[-1]] = value - return config - - sweeps = yaml.safe_load(open(sweep_file, "r")) - flat_sweeps = _flatten_sweeps(sweeps) - keys = list(flat_sweeps.keys()) - values = list(itertools.product(*flat_sweeps.values())) - - configs = [] - for i, override_vals in enumerate(values): - config = copy.deepcopy(base_config) - config = _update_config(config, keys, override_vals) - config["identifier"] = config["identifier"] + f"_run{i}" - configs.append(config) - return configs - - -def save_experiment_log(args, jobs, configs): - log_file = args.logdir / "exp" / time.strftime("%Y-%m-%d-%I-%M-%S%p.log") - log_file.parent.mkdir(exist_ok=True, parents=True) - with open(log_file, "w") as f: - for job, config in zip(jobs, configs): - print( - json.dumps( - { - "config": config, - "slurm_id": job.job_id, - "timestamp": time.strftime("%I:%M:%S%p %Z %b %d, %Y"), - } - ), - file=f, - ) - return log_file - - -def get_pbc_distances( - pos, - edge_index, - cell, - cell_offsets, - neighbors, - return_offsets=False, - return_distance_vec=False, -): - row, col = edge_index - - distance_vectors = pos[row] - pos[col] - - # correct for pbc - neighbors = neighbors.to(cell.device) - cell = torch.repeat_interleave(cell, neighbors, dim=0) - offsets = cell_offsets.float().view(-1, 1, 3).bmm(cell.float()).view(-1, 3) - distance_vectors += offsets - - # compute distances - distances = distance_vectors.norm(dim=-1) - - # redundancy: remove zero distances - nonzero_idx = torch.arange(len(distances), device=distances.device)[distances != 0] - edge_index = edge_index[:, nonzero_idx] - distances = distances[nonzero_idx] - - out = { - "edge_index": edge_index, - "distances": distances, - } - - if return_distance_vec: - out["distance_vec"] = distance_vectors[nonzero_idx] - - if return_offsets: - out["offsets"] = offsets[nonzero_idx] - - return out - - -def radius_graph_pbc(data, radius, max_num_neighbors_threshold, pbc=[True, True, True]): - device = data.pos.device - batch_size = len(data.natoms) - - if hasattr(data, "pbc"): - data.pbc = torch.atleast_2d(data.pbc) - for i in range(3): - if not torch.any(data.pbc[:, i]).item(): - pbc[i] = False - elif torch.all(data.pbc[:, i]).item(): - pbc[i] = True - else: - raise RuntimeError( - "Different structures in the batch have different PBC configurations. This is not currently supported." - ) - - # position of the atoms - atom_pos = data.pos - - # Before computing the pairwise distances between atoms, first create a list of atom indices to compare for the entire batch - num_atoms_per_image = data.natoms - num_atoms_per_image_sqr = (num_atoms_per_image**2).long() - - # index offset between images - index_offset = torch.cumsum(num_atoms_per_image, dim=0) - num_atoms_per_image - - index_offset_expand = torch.repeat_interleave(index_offset, num_atoms_per_image_sqr) - num_atoms_per_image_expand = torch.repeat_interleave( - num_atoms_per_image, num_atoms_per_image_sqr - ) - - # Compute a tensor containing sequences of numbers that range from 0 to num_atoms_per_image_sqr for each image - # that is used to compute indices for the pairs of atoms. This is a very convoluted way to implement - # the following (but 10x faster since it removes the for loop) - # for batch_idx in range(batch_size): - # batch_count = torch.cat([batch_count, torch.arange(num_atoms_per_image_sqr[batch_idx], device=device)], dim=0) - num_atom_pairs = torch.sum(num_atoms_per_image_sqr) - index_sqr_offset = ( - torch.cumsum(num_atoms_per_image_sqr, dim=0) - num_atoms_per_image_sqr - ) - index_sqr_offset = torch.repeat_interleave( - index_sqr_offset, num_atoms_per_image_sqr - ) - atom_count_sqr = torch.arange(num_atom_pairs, device=device) - index_sqr_offset - - # Compute the indices for the pairs of atoms (using division and mod) - # If the systems get too large this apporach could run into numerical precision issues - index1 = ( - torch.div(atom_count_sqr, num_atoms_per_image_expand, rounding_mode="floor") - ) + index_offset_expand - index2 = (atom_count_sqr % num_atoms_per_image_expand) + index_offset_expand - # Get the positions for each atom - pos1 = torch.index_select(atom_pos, 0, index1) - pos2 = torch.index_select(atom_pos, 0, index2) - - # Calculate required number of unit cells in each direction. - # Smallest distance between planes separated by a1 is - # 1 / ||(a2 x a3) / V||_2, since a2 x a3 is the area of the plane. - # Note that the unit cell volume V = a1 * (a2 x a3) and that - # (a2 x a3) / V is also the reciprocal primitive vector - # (crystallographer's definition). - - cross_a2a3 = torch.cross(data.cell[:, 1], data.cell[:, 2], dim=-1) - cell_vol = torch.sum(data.cell[:, 0] * cross_a2a3, dim=-1, keepdim=True) - - if pbc[0]: - inv_min_dist_a1 = torch.norm(cross_a2a3 / cell_vol, p=2, dim=-1) - rep_a1 = torch.ceil(radius * inv_min_dist_a1) - else: - rep_a1 = data.cell.new_zeros(1) - - if pbc[1]: - cross_a3a1 = torch.cross(data.cell[:, 2], data.cell[:, 0], dim=-1) - inv_min_dist_a2 = torch.norm(cross_a3a1 / cell_vol, p=2, dim=-1) - rep_a2 = torch.ceil(radius * inv_min_dist_a2) - else: - rep_a2 = data.cell.new_zeros(1) - - if pbc[2]: - cross_a1a2 = torch.cross(data.cell[:, 0], data.cell[:, 1], dim=-1) - inv_min_dist_a3 = torch.norm(cross_a1a2 / cell_vol, p=2, dim=-1) - rep_a3 = torch.ceil(radius * inv_min_dist_a3) - else: - rep_a3 = data.cell.new_zeros(1) - - # Take the max over all images for uniformity. This is essentially padding. - # Note that this can significantly increase the number of computed distances - # if the required repetitions are very different between images - # (which they usually are). Changing this to sparse (scatter) operations - # might be worth the effort if this function becomes a bottleneck. - max_rep = [rep_a1.max(), rep_a2.max(), rep_a3.max()] - - # Tensor of unit cells - cells_per_dim = [ - torch.arange(-rep, rep + 1, device=device, dtype=torch.float) for rep in max_rep - ] - unit_cell = torch.cartesian_prod(*cells_per_dim) - num_cells = len(unit_cell) - unit_cell_per_atom = unit_cell.view(1, num_cells, 3).repeat(len(index2), 1, 1) - unit_cell = torch.transpose(unit_cell, 0, 1) - unit_cell_batch = unit_cell.view(1, 3, num_cells).expand(batch_size, -1, -1) - - # Compute the x, y, z positional offsets for each cell in each image - data_cell = torch.transpose(data.cell, 1, 2) - pbc_offsets = torch.bmm(data_cell, unit_cell_batch) - pbc_offsets_per_atom = torch.repeat_interleave( - pbc_offsets, num_atoms_per_image_sqr, dim=0 - ) - - # Expand the positions and indices for the 9 cells - pos1 = pos1.view(-1, 3, 1).expand(-1, -1, num_cells) - pos2 = pos2.view(-1, 3, 1).expand(-1, -1, num_cells) - index1 = index1.view(-1, 1).repeat(1, num_cells).view(-1) - index2 = index2.view(-1, 1).repeat(1, num_cells).view(-1) - # Add the PBC offsets for the second atom - pos2 = pos2 + pbc_offsets_per_atom - - # Compute the squared distance between atoms - atom_distance_sqr = torch.sum((pos1 - pos2) ** 2, dim=1) - atom_distance_sqr = atom_distance_sqr.view(-1) - - # Remove pairs that are too far apart - mask_within_radius = torch.le(atom_distance_sqr, radius * radius) - # Remove pairs with the same atoms (distance = 0.0) - mask_not_same = torch.gt(atom_distance_sqr, 0.0001) - mask = torch.logical_and(mask_within_radius, mask_not_same) - index1 = torch.masked_select(index1, mask) - index2 = torch.masked_select(index2, mask) - unit_cell = torch.masked_select( - unit_cell_per_atom.view(-1, 3), mask.view(-1, 1).expand(-1, 3) - ) - unit_cell = unit_cell.view(-1, 3) - atom_distance_sqr = torch.masked_select(atom_distance_sqr, mask) - - mask_num_neighbors, num_neighbors_image = get_max_neighbors_mask( - natoms=data.natoms, - index=index1, - atom_distance=atom_distance_sqr, - max_num_neighbors_threshold=max_num_neighbors_threshold, - ) - - if not torch.all(mask_num_neighbors): - # Mask out the atoms to ensure each atom has at most max_num_neighbors_threshold neighbors - index1 = torch.masked_select(index1, mask_num_neighbors) - index2 = torch.masked_select(index2, mask_num_neighbors) - unit_cell = torch.masked_select( - unit_cell.view(-1, 3), mask_num_neighbors.view(-1, 1).expand(-1, 3) - ) - unit_cell = unit_cell.view(-1, 3) - - edge_index = torch.stack((index2, index1)) - - return edge_index, unit_cell, num_neighbors_image - - -def get_max_neighbors_mask(natoms, index, atom_distance, max_num_neighbors_threshold): - """ - Give a mask that filters out edges so that each atom has at most - `max_num_neighbors_threshold` neighbors. - Assumes that `index` is sorted. - """ - device = natoms.device - num_atoms = natoms.sum() - - # Get number of neighbors - # segment_coo assumes sorted index - ones = index.new_ones(1).expand_as(index) - num_neighbors = segment_coo(ones, index, dim_size=num_atoms) - max_num_neighbors = num_neighbors.max() - num_neighbors_thresholded = num_neighbors.clamp(max=max_num_neighbors_threshold) - - # Get number of (thresholded) neighbors per image - image_indptr = torch.zeros(natoms.shape[0] + 1, device=device, dtype=torch.long) - image_indptr[1:] = torch.cumsum(natoms, dim=0) - num_neighbors_image = segment_csr(num_neighbors_thresholded, image_indptr) - - # If max_num_neighbors is below the threshold, return early - if ( - max_num_neighbors <= max_num_neighbors_threshold - or max_num_neighbors_threshold <= 0 - ): - mask_num_neighbors = torch.tensor([True], dtype=bool, device=device).expand_as( - index - ) - return mask_num_neighbors, num_neighbors_image - - # Create a tensor of size [num_atoms, max_num_neighbors] to sort the distances of the neighbors. - # Fill with infinity so we can easily remove unused distances later. - distance_sort = torch.full([num_atoms * max_num_neighbors], np.inf, device=device) - - # Create an index map to map distances from atom_distance to distance_sort - # index_sort_map assumes index to be sorted - index_neighbor_offset = torch.cumsum(num_neighbors, dim=0) - num_neighbors - index_neighbor_offset_expand = torch.repeat_interleave( - index_neighbor_offset, num_neighbors - ) - index_sort_map = ( - index * max_num_neighbors - + torch.arange(len(index), device=device) - - index_neighbor_offset_expand - ) - distance_sort.index_copy_(0, index_sort_map, atom_distance) - distance_sort = distance_sort.view(num_atoms, max_num_neighbors) - - # Sort neighboring atoms based on distance - distance_sort, index_sort = torch.sort(distance_sort, dim=1) - # Select the max_num_neighbors_threshold neighbors that are closest - distance_sort = distance_sort[:, :max_num_neighbors_threshold] - index_sort = index_sort[:, :max_num_neighbors_threshold] - - # Offset index_sort so that it indexes into index - index_sort = index_sort + index_neighbor_offset.view(-1, 1).expand( - -1, max_num_neighbors_threshold - ) - # Remove "unused pairs" with infinite distances - mask_finite = torch.isfinite(distance_sort) - index_sort = torch.masked_select(index_sort, mask_finite) - - # At this point index_sort contains the index into index of the - # closest max_num_neighbors_threshold neighbors per atom - # Create a mask to remove all pairs not in index_sort - mask_num_neighbors = torch.zeros(len(index), device=device, dtype=bool) - mask_num_neighbors.index_fill_(0, index_sort, True) - - return mask_num_neighbors, num_neighbors_image - - -def get_pruned_edge_idx(edge_index, num_atoms=None, max_neigh=1e9): - assert num_atoms is not None - - # removes neighbors > max_neigh - # assumes neighbors are sorted in increasing distance - _nonmax_idx = [] - for i in range(num_atoms): - idx_i = torch.arange(len(edge_index[1]))[(edge_index[1] == i)][:max_neigh] - _nonmax_idx.append(idx_i) - _nonmax_idx = torch.cat(_nonmax_idx) - - return _nonmax_idx - - -def merge_dicts(dict1: dict, dict2: dict): - """Recursively merge two dictionaries. - Values in dict2 override values in dict1. If dict1 and dict2 contain a dictionary as a - value, this will call itself recursively to merge these dictionaries. - This does not modify the input dictionaries (creates an internal copy). - Additionally returns a list of detected duplicates. - Adapted from https://github.com/TUM-DAML/seml/blob/master/seml/utils.py - - Parameters - ---------- - dict1: dict - First dict. - dict2: dict - Second dict. Values in dict2 will override values from dict1 in case they share the same key. - - Returns - ------- - return_dict: dict - Merged dictionaries. - """ - if not isinstance(dict1, dict): - raise ValueError(f"Expecting dict1 to be dict, found {type(dict1)}.") - if not isinstance(dict2, dict): - raise ValueError(f"Expecting dict2 to be dict, found {type(dict2)}.") - - return_dict = copy.deepcopy(dict1) - duplicates = [] - - for k, v in dict2.items(): - if k not in dict1: - return_dict[k] = v - else: - if isinstance(v, dict) and isinstance(dict1[k], dict): - return_dict[k], duplicates_k = merge_dicts(dict1[k], dict2[k]) - duplicates += [f"{k}.{dup}" for dup in duplicates_k] - else: - return_dict[k] = dict2[k] - duplicates.append(k) - - return return_dict, duplicates - - -class SeverityLevelBetween(logging.Filter): - def __init__(self, min_level, max_level): - super().__init__() - self.min_level = min_level - self.max_level = max_level - - def filter(self, record): - return self.min_level <= record.levelno < self.max_level - - -def setup_logging(): - root = logging.getLogger() - - # Perform setup only if logging has not been configured - if not root.hasHandlers(): - root.setLevel(logging.INFO) - - log_formatter = logging.Formatter( - "%(asctime)s (%(levelname)s): %(message)s", - datefmt="%Y-%m-%d %H:%M:%S", - ) - - # Send INFO to stdout - handler_out = logging.StreamHandler(sys.stdout) - handler_out.addFilter(SeverityLevelBetween(logging.INFO, logging.WARNING)) - handler_out.setFormatter(log_formatter) - root.addHandler(handler_out) - - # Send WARNING (and higher) to stderr - handler_err = logging.StreamHandler(sys.stderr) - handler_err.setLevel(logging.WARNING) - handler_err.setFormatter(log_formatter) - root.addHandler(handler_err) - - -def compute_neighbors(data, edge_index): - # Get number of neighbors - # segment_coo assumes sorted index - ones = edge_index[1].new_ones(1).expand_as(edge_index[1]) - num_neighbors = segment_coo(ones, edge_index[1], dim_size=data.natoms.sum()) - - # Get number of neighbors per image - image_indptr = torch.zeros( - data.natoms.shape[0] + 1, device=data.pos.device, dtype=torch.long - ) - image_indptr[1:] = torch.cumsum(data.natoms, dim=0) - neighbors = segment_csr(num_neighbors, image_indptr) - return neighbors - - -def check_traj_files(batch, traj_dir): - if traj_dir is None: - return False - traj_dir = Path(traj_dir) - traj_files = [traj_dir / f"{id}.traj" for id in batch[0].sid.tolist()] - return all(fl.exists() for fl in traj_files) - - -@contextmanager -def new_trainer_context(*, config: Dict[str, Any], args: Namespace): - from ocpmodels.common import distutils, gp_utils - from ocpmodels.common.registry import registry - - if TYPE_CHECKING: - from ocpmodels.tasks.task import BaseTask - from ocpmodels.trainers import BaseTrainer - - @dataclass - class _TrainingContext: - config: Dict[str, Any] - task: "BaseTask" - trainer: "BaseTrainer" - - setup_logging() - original_config = config - config = copy.deepcopy(original_config) - - if args.distributed: - distutils.setup(config) - if config["gp_gpus"] is not None: - gp_utils.setup_gp(config) - try: - setup_imports(config) - trainer_cls = registry.get_trainer_class(config.get("trainer", "energy")) - assert trainer_cls is not None, "Trainer not found" - trainer = trainer_cls( - task=config["task"], - model=config["model"], - dataset=config["dataset"], - optimizer=config["optim"], - identifier=config["identifier"], - timestamp_id=config.get("timestamp_id", None), - run_dir=config.get("run_dir", "./"), - is_debug=config.get("is_debug", False), - print_every=config.get("print_every", 10), - seed=config.get("seed", 0), - logger=config.get("logger", "tensorboard"), - local_rank=config["local_rank"], - amp=config.get("amp", False), - cpu=config.get("cpu", False), - slurm=config.get("slurm", {}), - noddp=config.get("noddp", False), - ) - - task_cls = registry.get_task_class(config["mode"]) - assert task_cls is not None, "Task not found" - task = task_cls(config) - start_time = time.time() - ctx = _TrainingContext(config=original_config, task=task, trainer=trainer) - yield ctx - distutils.synchronize() - if distutils.is_master(): - logging.info(f"Total time taken: {time.time() - start_time}") - finally: - if args.distributed: - distutils.cleanup() - - -def _resolve_scale_factor_submodule(model: nn.Module, name: str): - from ocpmodels.modules.scaling.scale_factor import ScaleFactor - - try: - scale = model.get_submodule(name) - if not isinstance(scale, ScaleFactor): - return None - return scale - except AttributeError: - return None - - -def _report_incompat_keys( - model: nn.Module, - keys: "_IncompatibleKeys", - strict: bool = False, -): - # filter out the missing scale factor keys for the new scaling factor module - missing_keys: List[str] = [] - for full_key_name in keys.missing_keys: - parent_module_name, _ = full_key_name.rsplit(".", 1) - scale_factor = _resolve_scale_factor_submodule(model, parent_module_name) - if scale_factor is not None: - continue - missing_keys.append(full_key_name) - - # filter out unexpected scale factor keys that remain from the old scaling modules - unexpected_keys: List[str] = [] - for full_key_name in keys.unexpected_keys: - parent_module_name, _ = full_key_name.rsplit(".", 1) - scale_factor = _resolve_scale_factor_submodule(model, parent_module_name) - if scale_factor is not None: - continue - unexpected_keys.append(full_key_name) - - error_msgs = [] - if len(unexpected_keys) > 0: - error_msgs.insert( - 0, - "Unexpected key(s) in state_dict: {}. ".format( - ", ".join('"{}"'.format(k) for k in unexpected_keys) - ), - ) - if len(missing_keys) > 0: - error_msgs.insert( - 0, - "Missing key(s) in state_dict: {}. ".format( - ", ".join('"{}"'.format(k) for k in missing_keys) - ), - ) - - if len(error_msgs) > 0: - error_msg = "Error(s) in loading state_dict for {}:\n\t{}".format( - model.__class__.__name__, "\n\t".join(error_msgs) - ) - if strict: - raise RuntimeError(error_msg) - else: - logging.warning(error_msg) - - return missing_keys, unexpected_keys - - -def load_state_dict( - module: nn.Module, - state_dict: Mapping[str, torch.Tensor], - strict: bool = True, -): - incompat_keys = module.load_state_dict(state_dict, strict=False) # type: ignore - return _report_incompat_keys(module, incompat_keys, strict=strict) - - -def scatter_det(*args, **kwargs): - from ocpmodels.common.registry import registry - - if registry.get("set_deterministic_scatter", no_warning=True): - torch.use_deterministic_algorithms(mode=True) - - out = scatter(*args, **kwargs) - - if registry.get("set_deterministic_scatter", no_warning=True): - torch.use_deterministic_algorithms(mode=False) - - return out diff --git a/ocpmodels/datasets/__init__.py b/ocpmodels/datasets/__init__.py deleted file mode 100644 index 9ed38d832f755ed2a3c3697ef224891fa0405c58..0000000000000000000000000000000000000000 --- a/ocpmodels/datasets/__init__.py +++ /dev/null @@ -1,12 +0,0 @@ -# Copyright (c) Facebook, Inc. and its affiliates. -# -# This source code is licensed under the MIT license found in the -# LICENSE file in the root directory of this source tree. - -from .lmdb_dataset import ( - LmdbDataset, - SinglePointLmdbDataset, - TrajectoryLmdbDataset, - data_list_collater, -) -from .oc22_lmdb_dataset import OC22LmdbDataset diff --git a/ocpmodels/datasets/embeddings/__init__.py b/ocpmodels/datasets/embeddings/__init__.py deleted file mode 100644 index 1c411633507834ebf441912e41c4aa9c0df844f1..0000000000000000000000000000000000000000 --- a/ocpmodels/datasets/embeddings/__init__.py +++ /dev/null @@ -1,11 +0,0 @@ -__all__ = [ - "ATOMIC_RADII", - "KHOT_EMBEDDINGS", - "CONTINUOUS_EMBEDDINGS", - "QMOF_KHOT_EMBEDDINGS", -] - -from .atomic_radii import ATOMIC_RADII -from .continuous_embeddings import CONTINUOUS_EMBEDDINGS -from .khot_embeddings import KHOT_EMBEDDINGS -from .qmof_khot_embeddings import QMOF_KHOT_EMBEDDINGS diff --git a/ocpmodels/datasets/embeddings/atomic_radii.py b/ocpmodels/datasets/embeddings/atomic_radii.py deleted file mode 100644 index c9a2a31ebed5f703d0ecbfd699dde427cbb26ef1..0000000000000000000000000000000000000000 --- a/ocpmodels/datasets/embeddings/atomic_radii.py +++ /dev/null @@ -1,109 +0,0 @@ -""" -Atomic radii in picometers - -NaN stored for unavailable parameters. -""" - -ATOMIC_RADII = { - 0: float("NaN"), - 1: 25.0, - 2: 120.0, - 3: 145.0, - 4: 105.0, - 5: 85.0, - 6: 70.0, - 7: 65.0, - 8: 60.0, - 9: 50.0, - 10: 160.0, - 11: 180.0, - 12: 150.0, - 13: 125.0, - 14: 110.0, - 15: 100.0, - 16: 100.0, - 17: 100.0, - 18: 71.0, - 19: 220.0, - 20: 180.0, - 21: 160.0, - 22: 140.0, - 23: 135.0, - 24: 140.0, - 25: 140.0, - 26: 140.0, - 27: 135.0, - 28: 135.0, - 29: 135.0, - 30: 135.0, - 31: 130.0, - 32: 125.0, - 33: 115.0, - 34: 115.0, - 35: 115.0, - 36: float("NaN"), - 37: 235.0, - 38: 200.0, - 39: 180.0, - 40: 155.0, - 41: 145.0, - 42: 145.0, - 43: 135.0, - 44: 130.0, - 45: 135.0, - 46: 140.0, - 47: 160.0, - 48: 155.0, - 49: 155.0, - 50: 145.0, - 51: 145.0, - 52: 140.0, - 53: 140.0, - 54: float("NaN"), - 55: 260.0, - 56: 215.0, - 57: 195.0, - 58: 185.0, - 59: 185.0, - 60: 185.0, - 61: 185.0, - 62: 185.0, - 63: 185.0, - 64: 180.0, - 65: 175.0, - 66: 175.0, - 67: 175.0, - 68: 175.0, - 69: 175.0, - 70: 175.0, - 71: 175.0, - 72: 155.0, - 73: 145.0, - 74: 135.0, - 75: 135.0, - 76: 130.0, - 77: 135.0, - 78: 135.0, - 79: 135.0, - 80: 150.0, - 81: 190.0, - 82: 180.0, - 83: 160.0, - 84: 190.0, - 85: float("NaN"), - 86: float("NaN"), - 87: float("NaN"), - 88: 215.0, - 89: 195.0, - 90: 180.0, - 91: 180.0, - 92: 175.0, - 93: 175.0, - 94: 175.0, - 95: 175.0, - 96: float("NaN"), - 97: float("NaN"), - 98: float("NaN"), - 99: float("NaN"), - 100: float("NaN"), -} diff --git a/ocpmodels/datasets/embeddings/continuous_embeddings.py b/ocpmodels/datasets/embeddings/continuous_embeddings.py deleted file mode 100644 index a15d350cb7bbfd255e33285534a169d17f3629af..0000000000000000000000000000000000000000 --- a/ocpmodels/datasets/embeddings/continuous_embeddings.py +++ /dev/null @@ -1,1120 +0,0 @@ -""" -CGCNN-like embeddings using continuous values instead of original k-hot. - -Properties: - Group number - Period number - Electronegativity - Covalent radius - Valence electrons - First ionization energy - Electron affinity - Block - Atomic Volume - -NaN stored for unavaialable parameters. -""" - -CONTINUOUS_EMBEDDINGS = { - 0: [ - float("NaN"), - float("NaN"), - float("NaN"), - float("NaN"), - float("NaN"), - float("NaN"), - float("NaN"), - float("NaN"), - float("NaN"), - ], - 1: [ - 1.0, - 1.0, - 2.1877708435058594, - 31.0, - 1.0, - 13.598434448242188, - 0.754194974899292, - 1.0, - 14.100000381469727, - ], - 2: [ - 18.0, - 1.0, - 1.0, - 28.0, - 2.0, - 24.587387084960938, - -19.700000762939453, - 1.0, - 31.799999237060547, - ], - 3: [ - 1.0, - 2.0, - 0.04886792600154877, - 128.0, - 1.0, - 5.391714572906494, - 0.6180490255355835, - 1.0, - 13.100000381469727, - ], - 4: [ - 2.0, - 2.0, - 0.1268472671508789, - 96.0, - 2.0, - 9.322698593139648, - -2.4000000953674316, - 1.0, - 5.0, - ], - 5: [ - 13.0, - 2.0, - 0.25462737679481506, - 84.0, - 3.0, - 8.298019409179688, - 0.27972298860549927, - 2.0, - 4.599999904632568, - ], - 6: [ - 14.0, - 2.0, - 0.42752504348754883, - 73.0, - 4.0, - 11.260295867919922, - 1.2621190547943115, - 2.0, - 5.300000190734863, - ], - 7: [ - 15.0, - 2.0, - 0.5774819254875183, - 71.0, - 5.0, - 14.534130096435547, - -1.399999976158142, - 2.0, - 17.299999237060547, - ], - 8: [ - 16.0, - 2.0, - 0.9416494369506836, - 66.0, - 6.0, - 13.618054389953613, - 1.461113452911377, - 2.0, - 14.0, - ], - 9: [ - 17.0, - 2.0, - 1.017681360244751, - 57.0, - 7.0, - 17.422819137573242, - 3.4011898040771484, - 2.0, - 17.100000381469727, - ], - 10: [ - 18.0, - 2.0, - 1.0, - 58.0, - 8.0, - 21.56454086303711, - -3.0, - 2.0, - 16.799999237060547, - ], - 11: [ - 1.0, - 3.0, - 0.09459763765335083, - 166.0, - 1.0, - 5.1390767097473145, - 0.5479260087013245, - 1.0, - 23.700000762939453, - ], - 12: [ - 2.0, - 3.0, - 0.15242105722427368, - 141.0, - 2.0, - 7.64623498916626, - -3.0, - 1.0, - 14.0, - ], - 13: [ - 13.0, - 3.0, - 0.2360926866531372, - 121.0, - 3.0, - 5.9857683181762695, - 0.43283000588417053, - 2.0, - 10.0, - ], - 14: [ - 14.0, - 3.0, - 0.3468157947063446, - 111.0, - 4.0, - 8.15168285369873, - 1.3895211219787598, - 2.0, - 12.100000381469727, - ], - 15: [ - 15.0, - 3.0, - 0.45102688670158386, - 107.0, - 5.0, - 10.486685752868652, - 0.7466070055961609, - 2.0, - 17.0, - ], - 16: [ - 16.0, - 3.0, - 0.6397251486778259, - 105.0, - 6.0, - 10.360010147094727, - 2.077104091644287, - 2.0, - 15.5, - ], - 17: [ - 17.0, - 3.0, - 0.8123772740364075, - 102.0, - 7.0, - 12.967630386352539, - 3.612725019454956, - 2.0, - 18.700000762939453, - ], - 18: [ - 18.0, - 3.0, - 1.0, - 106.0, - 8.0, - 15.759611129760742, - -11.5, - 2.0, - 24.200000762939453, - ], - 19: [ - 1.0, - 4.0, - 0.12183826416730881, - 203.0, - 1.0, - 4.340663433074951, - 0.5014700293540955, - 1.0, - 45.29999923706055, - ], - 20: [ - 2.0, - 4.0, - 0.1901577115058899, - 176.0, - 2.0, - 6.113155364990234, - 0.024550000205636024, - 1.0, - 29.899999618530273, - ], - 21: [ - 3.0, - 4.0, - 0.3038673996925354, - 170.0, - 3.0, - 6.561490058898926, - 0.18799999356269836, - 3.0, - 15.0, - ], - 22: [ - 4.0, - 4.0, - 0.4055461883544922, - 160.0, - 4.0, - 6.828120231628418, - 0.07900000363588333, - 3.0, - 10.600000381469727, - ], - 23: [ - 5.0, - 4.0, - 0.4388898015022278, - 153.0, - 5.0, - 6.746187210083008, - 0.5249999761581421, - 3.0, - 8.350000381469727, - ], - 24: [ - 6.0, - 4.0, - 0.6017723083496094, - 139.0, - 6.0, - 6.766510009765625, - 0.6660000085830688, - 3.0, - 7.230000019073486, - ], - 25: [ - 7.0, - 4.0, - 0.6707264184951782, - 150.0, - 7.0, - 7.434018135070801, - -3.0, - 3.0, - 7.389999866485596, - ], - 26: [ - 8.0, - 4.0, - 0.748727023601532, - 142.0, - 8.0, - 7.902467727661133, - 0.1509999930858612, - 3.0, - 7.099999904632568, - ], - 27: [ - 9.0, - 4.0, - 0.8832423686981201, - 138.0, - 9.0, - 7.881010055541992, - 0.6622564792633057, - 3.0, - 6.699999809265137, - ], - 28: [ - 10.0, - 4.0, - 0.9377039670944214, - 124.0, - 10.0, - 7.639876842498779, - 1.156000018119812, - 3.0, - 6.599999904632568, - ], - 29: [ - 11.0, - 4.0, - 0.9175541996955872, - 132.0, - 11.0, - 7.726379871368408, - 1.2350000143051147, - 3.0, - 7.099999904632568, - ], - 30: [ - 12.0, - 4.0, - 0.8100876808166504, - 122.0, - 12.0, - 9.39419937133789, - -3.0, - 3.0, - 9.199999809265137, - ], - 31: [ - 13.0, - 4.0, - 0.7205410003662109, - 122.0, - 3.0, - 5.999301910400391, - 0.4300000071525574, - 2.0, - 11.800000190734863, - ], - 32: [ - 14.0, - 4.0, - 0.8001470565795898, - 120.0, - 4.0, - 7.899435043334961, - 1.2327120304107666, - 2.0, - 13.600000381469727, - ], - 33: [ - 15.0, - 4.0, - 0.825337290763855, - 119.0, - 5.0, - 9.788999557495117, - 0.8040000200271606, - 2.0, - 13.100000381469727, - ], - 34: [ - 16.0, - 4.0, - 0.9659121036529541, - 120.0, - 6.0, - 9.752391815185547, - 2.020669937133789, - 2.0, - 16.5, - ], - 35: [ - 17.0, - 4.0, - 1.0490256547927856, - 120.0, - 7.0, - 11.813810348510742, - 3.3635880947113037, - 2.0, - 23.5, - ], - 36: [ - 18.0, - 4.0, - 1.0, - 116.0, - 8.0, - 13.999605178833008, - -3.0, - 2.0, - 32.20000076293945, - ], - 37: [ - 1.0, - 5.0, - 0.1764136552810669, - 220.0, - 1.0, - 4.177127838134766, - 0.4859200119972229, - 1.0, - 55.900001525878906, - ], - 38: [ - 2.0, - 5.0, - 0.26317858695983887, - 195.0, - 2.0, - 5.694867134094238, - 0.04800000041723251, - 1.0, - 33.70000076293945, - ], - 39: [ - 3.0, - 5.0, - 0.39239412546157837, - 190.0, - 3.0, - 6.217259883880615, - 0.3070000112056732, - 3.0, - 19.799999237060547, - ], - 40: [ - 4.0, - 5.0, - 0.4744466543197632, - 175.0, - 4.0, - 6.633900165557861, - 0.4259999990463257, - 3.0, - 14.100000381469727, - ], - 41: [ - 5.0, - 5.0, - 0.5561695098876953, - 164.0, - 5.0, - 6.75885009765625, - 0.9174060225486755, - 3.0, - 10.800000190734863, - ], - 42: [ - 6.0, - 5.0, - 0.6852949857711792, - 154.0, - 6.0, - 7.092430114746094, - 0.7480000257492065, - 3.0, - 9.399999618530273, - ], - 43: [ - 7.0, - 5.0, - 0.8753613233566284, - 147.0, - 7.0, - 7.119380950927734, - 0.550000011920929, - 3.0, - 8.5, - ], - 44: [ - 8.0, - 5.0, - 0.9579373002052307, - 146.0, - 8.0, - 7.360499858856201, - 1.0499999523162842, - 3.0, - 8.300000190734863, - ], - 45: [ - 9.0, - 5.0, - 0.9761914610862732, - 142.0, - 9.0, - 7.458899974822998, - 1.1369999647140503, - 3.0, - 8.300000190734863, - ], - 46: [ - 10.0, - 5.0, - 1.1242631673812866, - 139.0, - 12.0, - 8.336859703063965, - 0.5619999766349792, - 3.0, - 8.899999618530273, - ], - 47: [ - 11.0, - 5.0, - 0.9437955021858215, - 145.0, - 11.0, - 7.576233863830566, - 1.3020000457763672, - 3.0, - 10.300000190734863, - ], - 48: [ - 12.0, - 5.0, - 0.8015620112419128, - 144.0, - 12.0, - 8.99382209777832, - -3.0, - 3.0, - 13.100000381469727, - ], - 49: [ - 13.0, - 5.0, - 0.7172747254371643, - 142.0, - 3.0, - 5.786355018615723, - 0.30000001192092896, - 2.0, - 15.699999809265137, - ], - 50: [ - 14.0, - 5.0, - 0.7622796893119812, - 139.0, - 4.0, - 7.343916893005371, - 1.1120669841766357, - 2.0, - 16.299999237060547, - ], - 51: [ - 15.0, - 5.0, - 0.7762722373008728, - 139.0, - 5.0, - 8.608388900756836, - 1.0460000038146973, - 2.0, - 18.399999618530273, - ], - 52: [ - 16.0, - 5.0, - 0.8622506260871887, - 138.0, - 6.0, - 9.009659767150879, - 1.9708759784698486, - 2.0, - 20.5, - ], - 53: [ - 17.0, - 5.0, - 0.9386428594589233, - 139.0, - 7.0, - 10.45125961303711, - 3.0590367317199707, - 2.0, - 25.700000762939453, - ], - 54: [ - 18.0, - 5.0, - 1.0, - 140.0, - 8.0, - 12.129842758178711, - -0.0560000017285347, - 2.0, - 42.900001525878906, - ], - 55: [ - 1.0, - 6.0, - 0.18145304918289185, - 244.0, - 1.0, - 3.8939056396484375, - 0.47162601351737976, - 1.0, - 70.0, - ], - 56: [ - 2.0, - 6.0, - 0.3032951354980469, - 215.0, - 2.0, - 5.211664199829102, - 0.14462000131607056, - 1.0, - 39.0, - ], - 57: [ - 3.0, - 6.0, - 0.39465051889419556, - 207.0, - 3.0, - 5.576900005340576, - 0.4699999988079071, - 3.0, - 22.5, - ], - 58: [ - 4.0, - 6.0, - 0.5356179475784302, - 204.0, - 2.0, - 5.538599967956543, - 0.6499999761581421, - 4.0, - 21.0, - ], - 59: [ - 5.0, - 6.0, - 0.4288040101528168, - 203.0, - 2.0, - 5.4730000495910645, - 0.9620000123977661, - 4.0, - 20.799999237060547, - ], - 60: [ - 6.0, - 6.0, - 0.44721803069114685, - 201.0, - 2.0, - 5.525000095367432, - 1.9160000085830688, - 4.0, - 20.600000381469727, - ], - 61: [ - 7.0, - 6.0, - 0.4585537314414978, - 199.0, - 2.0, - 5.581999778747559, - -3.0, - 4.0, - 20.229999542236328, - ], - 62: [ - 8.0, - 6.0, - 0.47021451592445374, - 198.0, - 2.0, - 5.643710136413574, - -3.0, - 4.0, - 19.899999618530273, - ], - 63: [ - 9.0, - 6.0, - 0.5085079669952393, - 198.0, - 2.0, - 5.670384883880615, - 0.8640000224113464, - 4.0, - 28.899999618530273, - ], - 64: [ - 10.0, - 6.0, - 0.5033860206604004, - 196.0, - 2.0, - 6.149796009063721, - -3.0, - 4.0, - 19.899999618530273, - ], - 65: [ - 11.0, - 6.0, - 0.5163695216178894, - 194.0, - 2.0, - 5.863800048828125, - 1.1649999618530273, - 4.0, - 19.200000762939453, - ], - 66: [ - 12.0, - 6.0, - 0.5297338366508484, - 192.0, - 2.0, - 5.939050197601318, - 0.35199999809265137, - 4.0, - 19.0, - ], - 67: [ - 13.0, - 6.0, - 0.5434919595718384, - 192.0, - 2.0, - 6.021500110626221, - -3.0, - 4.0, - 18.700000762939453, - ], - 68: [ - 14.0, - 6.0, - 0.5576573014259338, - 189.0, - 2.0, - 6.107699871063232, - -3.0, - 4.0, - 18.399999618530273, - ], - 69: [ - 15.0, - 6.0, - 0.5722439289093018, - 190.0, - 2.0, - 6.184309959411621, - 1.0290000438690186, - 4.0, - 18.100000381469727, - ], - 70: [ - 16.0, - 6.0, - 0.517667829990387, - 187.0, - 2.0, - 6.254159927368164, - -0.019999999552965164, - 4.0, - 24.799999237060547, - ], - 71: [ - 17.0, - 6.0, - 0.6027398109436035, - 187.0, - 2.0, - 5.425870895385742, - 0.3400000035762787, - 4.0, - 17.799999237060547, - ], - 72: [ - 4.0, - 6.0, - 0.7352124452590942, - 175.0, - 4.0, - 6.825069904327393, - 0.014000000432133675, - 3.0, - 13.600000381469727, - ], - 73: [ - 5.0, - 6.0, - 0.8358832001686096, - 170.0, - 5.0, - 7.549570083618164, - 0.32199999690055847, - 3.0, - 10.899999618530273, - ], - 74: [ - 6.0, - 6.0, - 1.0192831754684448, - 162.0, - 6.0, - 7.864029884338379, - 0.8162599802017212, - 3.0, - 9.529999732971191, - ], - 75: [ - 7.0, - 6.0, - 1.1745918989181519, - 151.0, - 7.0, - 7.83351993560791, - 0.15000000596046448, - 3.0, - 8.850000381469727, - ], - 76: [ - 8.0, - 6.0, - 1.2392759323120117, - 144.0, - 8.0, - 8.43822956085205, - 1.100000023841858, - 3.0, - 8.430000305175781, - ], - 77: [ - 9.0, - 6.0, - 1.4759982824325562, - 141.0, - 9.0, - 8.967020034790039, - 1.5637999773025513, - 3.0, - 8.539999961853027, - ], - 78: [ - 10.0, - 6.0, - 1.4510095119476318, - 136.0, - 10.0, - 8.958829879760742, - 2.128000020980835, - 3.0, - 9.100000381469727, - ], - 79: [ - 11.0, - 6.0, - 1.4267007112503052, - 136.0, - 11.0, - 9.225552558898926, - 2.3086299896240234, - 3.0, - 10.199999809265137, - ], - 80: [ - 12.0, - 6.0, - 1.1647894382476807, - 132.0, - 12.0, - 10.437503814697266, - -3.0, - 3.0, - 14.800000190734863, - ], - 81: [ - 13.0, - 6.0, - 0.924509584903717, - 145.0, - 3.0, - 6.1082868576049805, - 0.37700000405311584, - 2.0, - 17.200000762939453, - ], - 82: [ - 14.0, - 6.0, - 0.9313225746154785, - 146.0, - 4.0, - 7.416679382324219, - 0.3567431569099426, - 2.0, - 18.299999237060547, - ], - 83: [ - 15.0, - 6.0, - 0.8136501312255859, - 148.0, - 5.0, - 7.285515785217285, - 0.9423620104789734, - 2.0, - 21.299999237060547, - ], - 84: [ - 16.0, - 6.0, - 0.9256306886672974, - 140.0, - 6.0, - 8.413999557495117, - 1.899999976158142, - 2.0, - 22.700000762939453, - ], - 85: [ - float("NaN"), - float("NaN"), - float("NaN"), - float("NaN"), - float("NaN"), - float("NaN"), - float("NaN"), - float("NaN"), - float("NaN"), - ], - 86: [18.0, 6.0, 1.0, 150.0, 8.0, 10.748499870300293, -3.0, 2.0, 50.5], - 87: [ - float("NaN"), - float("NaN"), - float("NaN"), - float("NaN"), - float("NaN"), - float("NaN"), - float("NaN"), - float("NaN"), - float("NaN"), - ], - 88: [ - 2.0, - 7.0, - 0.3596253991127014, - 221.0, - 2.0, - 5.27842378616333, - 0.10000000149011612, - 1.0, - 45.0, - ], - 89: [ - 3.0, - 7.0, - 0.4583164155483246, - 215.0, - 3.0, - 5.380226135253906, - 0.3499999940395355, - 3.0, - 22.540000915527344, - ], - 90: [ - 4.0, - 7.0, - 0.5557018518447876, - 206.0, - 2.0, - 6.306700229644775, - -3.0, - 4.0, - 19.799999237060547, - ], - 91: [ - 5.0, - 7.0, - 0.623065710067749, - 200.0, - 2.0, - 5.889999866485596, - -3.0, - 4.0, - 15.0, - ], - 92: [ - 6.0, - 7.0, - 0.6181179881095886, - 196.0, - 2.0, - 6.194049835205078, - -3.0, - 4.0, - 12.5, - ], - 93: [ - 7.0, - 7.0, - 0.6132539510726929, - 190.0, - 2.0, - 6.265500068664551, - -3.0, - 4.0, - 21.100000381469727, - ], - 94: [ - 8.0, - 7.0, - 0.6084716320037842, - 187.0, - 2.0, - 6.0258002281188965, - -3.0, - 4.0, - 12.289999961853027, - ], - 95: [ - 9.0, - 7.0, - 0.6834156513214111, - 180.0, - 2.0, - 5.973800182342529, - -3.0, - 4.0, - 20.799999237060547, - ], - 96: [ - 10.0, - 7.0, - 0.6900094747543335, - 169.0, - 2.0, - 5.991399765014648, - -3.0, - 4.0, - 18.280000686645508, - ], - 97: [ - float("NaN"), - float("NaN"), - float("NaN"), - float("NaN"), - float("NaN"), - float("NaN"), - float("NaN"), - float("NaN"), - float("NaN"), - ], - 98: [ - float("NaN"), - float("NaN"), - float("NaN"), - float("NaN"), - float("NaN"), - float("NaN"), - float("NaN"), - float("NaN"), - float("NaN"), - ], - 99: [ - float("NaN"), - float("NaN"), - float("NaN"), - float("NaN"), - float("NaN"), - float("NaN"), - float("NaN"), - float("NaN"), - float("NaN"), - ], - 100: [ - float("NaN"), - float("NaN"), - float("NaN"), - float("NaN"), - float("NaN"), - float("NaN"), - float("NaN"), - float("NaN"), - float("NaN"), - ], -} diff --git a/ocpmodels/datasets/embeddings/khot_embeddings.py b/ocpmodels/datasets/embeddings/khot_embeddings.py deleted file mode 100644 index 5f55a460f8f1c2b7209cfd4360d9bdc9eec5be23..0000000000000000000000000000000000000000 --- a/ocpmodels/datasets/embeddings/khot_embeddings.py +++ /dev/null @@ -1,9412 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. - - -Original CGCNN k-hot elemental embeddings. -""" - -KHOT_EMBEDDINGS = { - 1: [ - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - ], - 2: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - ], - 3: [ - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - ], - 4: [ - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - ], - 5: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - ], - 6: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - ], - 7: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - ], - 8: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - ], - 9: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - ], - 10: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - ], - 11: [ - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - ], - 12: [ - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - ], - 13: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - ], - 14: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - ], - 15: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - ], - 16: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - ], - 17: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - ], - 18: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - ], - 19: [ - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - ], - 20: [ - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - ], - 21: [ - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - ], - 22: [ - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - ], - 23: [ - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - ], - 24: [ - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - ], - 25: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - ], - 26: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - ], - 27: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - ], - 28: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - ], - 29: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - ], - 30: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - ], - 31: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - ], - 32: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - ], - 33: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - ], - 34: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - ], - 35: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - ], - 36: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - ], - 37: [ - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - ], - 38: [ - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - ], - 39: [ - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - ], - 40: [ - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - ], - 41: [ - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - ], - 42: [ - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - ], - 43: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - ], - 44: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - ], - 45: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - ], - 46: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - ], - 47: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - ], - 48: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - ], - 49: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - ], - 50: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - ], - 51: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - ], - 52: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - ], - 53: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - ], - 54: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - ], - 55: [ - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - ], - 56: [ - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - ], - 57: [ - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - ], - 58: [ - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - ], - 59: [ - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - ], - 60: [ - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - ], - 61: [ - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - ], - 62: [ - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - ], - 63: [ - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - ], - 64: [ - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - ], - 65: [ - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - ], - 66: [ - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - ], - 67: [ - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - ], - 68: [ - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - ], - 69: [ - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - ], - 70: [ - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - ], - 71: [ - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - ], - 72: [ - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - ], - 73: [ - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - ], - 74: [ - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - ], - 75: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - ], - 76: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - ], - 77: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - ], - 78: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - ], - 79: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - ], - 80: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - ], - 81: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - ], - 82: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - ], - 83: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - ], - 84: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - ], - 85: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - ], - 86: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - ], - 87: [ - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - ], - 88: [ - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - ], - 89: [ - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - ], - 90: [ - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - ], - 91: [ - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - ], - 92: [ - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - ], - 93: [ - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - ], - 94: [ - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - ], - 95: [ - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - ], - 96: [ - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - ], - 97: [ - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - ], - 98: [ - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - ], - 99: [ - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - ], - 100: [ - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - ], -} diff --git a/ocpmodels/datasets/embeddings/qmof_khot_embeddings.py b/ocpmodels/datasets/embeddings/qmof_khot_embeddings.py deleted file mode 100644 index 82692a81308b86af48c4125521e60db20188c201..0000000000000000000000000000000000000000 --- a/ocpmodels/datasets/embeddings/qmof_khot_embeddings.py +++ /dev/null @@ -1,7636 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. - - -k-hot elemental embeddings from QMOF, motivated by the following Github Issue threads: -https://github.com/txie-93/cgcnn/issues/2 -https://github.com/arosen93/QMOF/issues/18 -""" - -QMOF_KHOT_EMBEDDINGS = { - 1: [ - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - ], - 2: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 1, - 0, - 0, - 0, - ], - 3: [ - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - ], - 4: [ - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - ], - 5: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - ], - 6: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - ], - 7: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - ], - 8: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - ], - 9: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 0, - ], - 10: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 1, - 0, - 0, - ], - 11: [ - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - ], - 12: [ - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - ], - 13: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - ], - 14: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - ], - 15: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - ], - 16: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - ], - 17: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - ], - 18: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 0, - ], - 19: [ - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - ], - 20: [ - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - ], - 21: [ - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - ], - 22: [ - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - ], - 23: [ - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - ], - 24: [ - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - ], - 25: [ - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - ], - 26: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - ], - 27: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - ], - 28: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - ], - 29: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - ], - 30: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - ], - 31: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - ], - 32: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - ], - 33: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - ], - 34: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - ], - 35: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - ], - 36: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - ], - 37: [ - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - ], - 38: [ - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - ], - 39: [ - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - ], - 40: [ - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - ], - 41: [ - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - ], - 42: [ - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - ], - 43: [ - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - ], - 44: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - ], - 45: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - ], - 46: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - ], - 47: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - ], - 48: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - ], - 49: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - ], - 50: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - ], - 51: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - ], - 52: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - ], - 53: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - ], - 54: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - ], - 55: [ - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - ], - 56: [ - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - ], - 57: [ - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - ], - 58: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - ], - 59: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - ], - 60: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - ], - 61: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - ], - 62: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - ], - 63: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - ], - 64: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - ], - 65: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - ], - 66: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - ], - 67: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - ], - 68: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - ], - 69: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - ], - 70: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - ], - 71: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - ], - 72: [ - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - ], - 73: [ - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - ], - 74: [ - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - ], - 75: [ - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - ], - 76: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - ], - 77: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - ], - 78: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - ], - 79: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - ], - 80: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - ], - 81: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - ], - 82: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - ], - 83: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - ], - 84: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - ], - 85: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - ], - 86: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - ], - 87: [ - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - ], - 88: [ - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - ], - 89: [ - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - ], - 90: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - ], - 91: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - ], - 92: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - ], - 93: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - ], - 94: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - ], - 95: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - ], - 96: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - ], - 97: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - ], - 98: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - ], - 99: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - ], - 100: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - ], - 101: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - ], - 102: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - ], - 103: [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - ], -} diff --git a/ocpmodels/datasets/lmdb_dataset.py b/ocpmodels/datasets/lmdb_dataset.py deleted file mode 100644 index ea8029ce7dd22a1605c6aa90a2a5a21e92eafa4d..0000000000000000000000000000000000000000 --- a/ocpmodels/datasets/lmdb_dataset.py +++ /dev/null @@ -1,178 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import bisect -import logging -import math -import pickle -import random -import warnings -from pathlib import Path - -import lmdb -import numpy as np -import torch -from torch.utils.data import Dataset -from torch_geometric.data import Batch - -from ocpmodels.common import distutils -from ocpmodels.common.registry import registry -from ocpmodels.common.utils import pyg2_data_transform - - -@registry.register_dataset("lmdb") -@registry.register_dataset("single_point_lmdb") -@registry.register_dataset("trajectory_lmdb") -class LmdbDataset(Dataset): - r"""Dataset class to load from LMDB files containing relaxation - trajectories or single point computations. - - Useful for Structure to Energy & Force (S2EF), Initial State to - Relaxed State (IS2RS), and Initial State to Relaxed Energy (IS2RE) tasks. - - Args: - config (dict): Dataset configuration - transform (callable, optional): Data transform function. - (default: :obj:`None`) - """ - - def __init__(self, config, transform=None): - super(LmdbDataset, self).__init__() - self.config = config - - self.path = Path(self.config["src"]) - if not self.path.is_file(): - db_paths = sorted(self.path.glob("*.lmdb")) - assert len(db_paths) > 0, f"No LMDBs found in '{self.path}'" - - self.metadata_path = self.path / "metadata.npz" - - self._keys, self.envs = [], [] - for db_path in db_paths: - self.envs.append(self.connect_db(db_path)) - length = pickle.loads( - self.envs[-1].begin().get("length".encode("ascii")) - ) - self._keys.append(list(range(length))) - - keylens = [len(k) for k in self._keys] - self._keylen_cumulative = np.cumsum(keylens).tolist() - self.num_samples = sum(keylens) - else: - self.metadata_path = self.path.parent / "metadata.npz" - self.env = self.connect_db(self.path) - self._keys = [ - f"{j}".encode("ascii") for j in range(self.env.stat()["entries"]) - ] - self.num_samples = len(self._keys) - - # If specified, limit dataset to only a portion of the entire dataset - # total_shards: defines total chunks to partition dataset - # shard: defines dataset shard to make visible - self.sharded = False - if "shard" in self.config and "total_shards" in self.config: - self.sharded = True - self.indices = range(self.num_samples) - # split all available indices into 'total_shards' bins - self.shards = np.array_split( - self.indices, self.config.get("total_shards", 1) - ) - # limit each process to see a subset of data based off defined shard - self.available_indices = self.shards[self.config.get("shard", 0)] - self.num_samples = len(self.available_indices) - - self.transform = transform - - def __len__(self): - return self.num_samples - - def __getitem__(self, idx): - # if sharding, remap idx to appropriate idx of the sharded set - if self.sharded: - idx = self.available_indices[idx] - if not self.path.is_file(): - # Figure out which db this should be indexed from. - db_idx = bisect.bisect(self._keylen_cumulative, idx) - # Extract index of element within that db. - el_idx = idx - if db_idx != 0: - el_idx = idx - self._keylen_cumulative[db_idx - 1] - assert el_idx >= 0 - - # Return features. - datapoint_pickled = ( - self.envs[db_idx] - .begin() - .get(f"{self._keys[db_idx][el_idx]}".encode("ascii")) - ) - data_object = pyg2_data_transform(pickle.loads(datapoint_pickled)) - data_object.id = f"{db_idx}_{el_idx}" - else: - datapoint_pickled = self.env.begin().get(self._keys[idx]) - data_object = pyg2_data_transform(pickle.loads(datapoint_pickled)) - - if self.transform is not None: - data_object = self.transform(data_object) - - return data_object - - def connect_db(self, lmdb_path=None): - env = lmdb.open( - str(lmdb_path), - subdir=False, - readonly=True, - lock=False, - readahead=False, - meminit=False, - max_readers=1, - ) - return env - - def close_db(self): - if not self.path.is_file(): - for env in self.envs: - env.close() - else: - self.env.close() - - -class SinglePointLmdbDataset(LmdbDataset): - def __init__(self, config, transform=None): - super(SinglePointLmdbDataset, self).__init__(config, transform) - warnings.warn( - "SinglePointLmdbDataset is deprecated and will be removed in the future." - "Please use 'LmdbDataset' instead.", - stacklevel=3, - ) - - -class TrajectoryLmdbDataset(LmdbDataset): - def __init__(self, config, transform=None): - super(TrajectoryLmdbDataset, self).__init__(config, transform) - warnings.warn( - "TrajectoryLmdbDataset is deprecated and will be removed in the future." - "Please use 'LmdbDataset' instead.", - stacklevel=3, - ) - - -def data_list_collater(data_list, otf_graph=False): - batch = Batch.from_data_list(data_list) - - if not otf_graph: - try: - n_neighbors = [] - for i, data in enumerate(data_list): - n_index = data.edge_index[1, :] - n_neighbors.append(n_index.shape[0]) - batch.neighbors = torch.tensor(n_neighbors) - except (NotImplementedError, TypeError): - logging.warning( - "LMDB does not contain edge index information, set otf_graph=True" - ) - - return batch diff --git a/ocpmodels/datasets/oc22_lmdb_dataset.py b/ocpmodels/datasets/oc22_lmdb_dataset.py deleted file mode 100644 index c8ac186f3068fa6dd66ec0174d435d4c2b5fc18b..0000000000000000000000000000000000000000 --- a/ocpmodels/datasets/oc22_lmdb_dataset.py +++ /dev/null @@ -1,203 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import bisect -import logging -import math -import pickle -import random -import warnings -from pathlib import Path - -import lmdb -import numpy as np -import torch -from torch.utils.data import Dataset -from torch_geometric.data import Batch - -from ocpmodels.common import distutils -from ocpmodels.common.registry import registry -from ocpmodels.common.utils import pyg2_data_transform - - -@registry.register_dataset("oc22_lmdb") -class OC22LmdbDataset(Dataset): - r"""Dataset class to load from LMDB files containing relaxation - trajectories or single point computations. - - Useful for Structure to Energy & Force (S2EF), Initial State to - Relaxed State (IS2RS), and Initial State to Relaxed Energy (IS2RE) tasks. - - Args: - config (dict): Dataset configuration - transform (callable, optional): Data transform function. - (default: :obj:`None`) - """ - - def __init__(self, config, transform=None): - super(OC22LmdbDataset, self).__init__() - self.config = config - - self.path = Path(self.config["src"]) - self.data2train = self.config.get("data2train", "all") - if not self.path.is_file(): - db_paths = sorted(self.path.glob("*.lmdb")) - assert len(db_paths) > 0, f"No LMDBs found in '{self.path}'" - - self.metadata_path = self.path / "metadata.npz" - - self._keys, self.envs = [], [] - for db_path in db_paths: - self.envs.append(self.connect_db(db_path)) - try: - length = pickle.loads( - self.envs[-1].begin().get("length".encode("ascii")) - ) - except TypeError: - length = self.envs[-1].stat()["entries"] - self._keys.append(list(range(length))) - - keylens = [len(k) for k in self._keys] - self._keylen_cumulative = np.cumsum(keylens).tolist() - self.num_samples = sum(keylens) - - if self.data2train != "all": - txt_paths = sorted(self.path.glob("*.txt")) - index = 0 - self.indices = [] - for txt_path in txt_paths: - lines = open(txt_path).read().splitlines() - for line in lines: - if self.data2train == "adslabs": - if "clean" not in line: - self.indices.append(index) - if self.data2train == "slabs": - if "clean" in line: - self.indices.append(index) - index += 1 - self.num_samples = len(self.indices) - else: - self.metadata_path = self.path.parent / "metadata.npz" - self.env = self.connect_db(self.path) - self._keys = [ - f"{j}".encode("ascii") for j in range(self.env.stat()["entries"]) - ] - self.num_samples = len(self._keys) - - self.transform = transform - self.lin_ref = self.oc20_ref = False - self.train_total = self.config.get("total_energy", False) - # only needed for oc20 datasets, oc22 is total by default - if self.train_total: - self.oc20_ref = pickle.load(open(config["oc20_ref"], "rb")) - if self.config.get("lin_ref", False): - coeff = np.load(self.config["lin_ref"], allow_pickle=True)["coeff"] - self.lin_ref = torch.nn.Parameter(torch.tensor(coeff), requires_grad=False) - self.subsample = self.config.get("subsample", False) - - def __len__(self): - if self.subsample: - return min(self.subsample, self.num_samples) - return self.num_samples - - def __getitem__(self, idx): - if self.data2train != "all": - idx = self.indices[idx] - if not self.path.is_file(): - # Figure out which db this should be indexed from. - db_idx = bisect.bisect(self._keylen_cumulative, idx) - # Extract index of element within that db. - el_idx = idx - if db_idx != 0: - el_idx = idx - self._keylen_cumulative[db_idx - 1] - assert el_idx >= 0 - - # Return features. - datapoint_pickled = ( - self.envs[db_idx] - .begin() - .get(f"{self._keys[db_idx][el_idx]}".encode("ascii")) - ) - data_object = pyg2_data_transform(pickle.loads(datapoint_pickled)) - data_object.id = f"{db_idx}_{el_idx}" - else: - datapoint_pickled = self.env.begin().get(self._keys[idx]) - data_object = pyg2_data_transform(pickle.loads(datapoint_pickled)) - - if self.transform is not None: - data_object = self.transform(data_object) - # make types consistent - sid = data_object.sid - if isinstance(sid, torch.Tensor): - sid = sid.item() - data_object.sid = sid - if "fid" in data_object: - fid = data_object.fid - if isinstance(fid, torch.Tensor): - fid = fid.item() - data_object.fid = fid - - if hasattr(data_object, "y_relaxed"): - attr = "y_relaxed" - elif hasattr(data_object, "y"): - attr = "y" - # if targets are not available, test data is being used - else: - return data_object - - # convert s2ef energies to raw energies - if attr == "y": - # OC20 data - if "oc22" not in data_object and self.train_total: - randomid = f"random{sid}" - data_object[attr] += self.oc20_ref[randomid] - data_object.nads = 1 - data_object.oc22 = 0 - - # convert is2re energies to raw energies - else: - if "oc22" not in data_object and self.train_total: - randomid = f"random{sid}" - data_object[attr] += self.oc20_ref[randomid] - del data_object.force - del data_object.y_init - data_object.nads = 1 - data_object.oc22 = 0 - - if self.lin_ref is not False: - lin_energy = sum(self.lin_ref[data_object.atomic_numbers.long()]) - data_object[attr] -= lin_energy - - # to jointly train on oc22+oc20, need to delete these oc20-only attributes - # ensure otf_graph=1 in your model configuration - if "edge_index" in data_object: - del data_object.edge_index - if "cell_offsets" in data_object: - del data_object.cell_offsets - if "distances" in data_object: - del data_object.distances - - return data_object - - def connect_db(self, lmdb_path=None): - env = lmdb.open( - str(lmdb_path), - subdir=False, - readonly=True, - lock=False, - readahead=False, - meminit=False, - max_readers=1, - ) - return env - - def close_db(self): - if not self.path.is_file(): - for env in self.envs: - env.close() - else: - self.env.close() diff --git a/ocpmodels/ff_lmdb.py b/ocpmodels/ff_lmdb.py deleted file mode 100644 index c95413ed8612b11f7147df8df05391e7d8f2ace9..0000000000000000000000000000000000000000 --- a/ocpmodels/ff_lmdb.py +++ /dev/null @@ -1,215 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import bisect -import pickle -from pathlib import Path - -import lmdb -import numpy as np -from torch.utils.data import Dataset - -# from torch_geometric.data import Batch - - -class LmdbDataset(Dataset): - r"""Dataset class to load from LMDB files containing relaxation - trajectories or single point computations. - - Useful for Structure to Energy & Force (S2EF), Initial State to - Relaxed State (IS2RS), and Initial State to Relaxed Energy (IS2RE) tasks. - - Args: - config (dict): Dataset configuration - transform (callable, optional): Data transform function. - (default: :obj:`None`) - """ - - def __init__(self, src, transform=None, **kwargs): - super(LmdbDataset, self).__init__() - - self.path = Path(src) - if not self.path.is_file(): - db_paths = sorted(self.path.glob("*.lmdb")) - assert len(db_paths) > 0, f"No LMDBs found in '{self.path}'" - - self.metadata_path = self.path / "metadata.npz" - - self._keys, self.envs = [], [] - for db_path in db_paths: - self.envs.append(self.connect_db(db_path)) - length = pickle.loads( - self.envs[-1].begin().get("length".encode("ascii")) - ) - self._keys.append(list(range(length))) - - keylens = [len(k) for k in self._keys] - self._keylen_cumulative = np.cumsum(keylens).tolist() - self.num_samples = sum(keylens) - else: - self.metadata_path = self.path.parent / "metadata.npz" - self.env = self.connect_db(self.path) - try: - # Try to get the stored length value first - self.num_samples = pickle.loads( - self.env.begin().get("length".encode("ascii")) - ) - except (TypeError, KeyError): - # Fallback to entries count if length key doesn't exist - self.num_samples = self.env.stat()["entries"] - - self._keys = [f"{j}".encode("ascii") for j in range(self.num_samples)] - - self.transform = transform - - def __len__(self): - return self.num_samples - - def __getitem__(self, idx): - if idx >= self.num_samples: - raise IndexError( - f"Index {idx} out of range for dataset with {self.num_samples} samples" - ) - - if not self.path.is_file(): - # Figure out which db this should be indexed from. - db_idx = bisect.bisect(self._keylen_cumulative, idx) - # Extract index of element within that db. - el_idx = idx - if db_idx != 0: - el_idx = idx - self._keylen_cumulative[db_idx - 1] - assert el_idx >= 0 - - # Return features. - datapoint_pickled = ( - self.envs[db_idx] - .begin() - .get(f"{self._keys[db_idx][el_idx]}".encode("ascii")) - ) - data_object = pickle.loads(datapoint_pickled) - data_object.id = f"{db_idx}_{el_idx}" - else: - datapoint_pickled = self.env.begin().get(self._keys[idx]) - if datapoint_pickled is None: - raise KeyError(f"No data found for index {idx}") - data_object = pickle.loads(datapoint_pickled) - - if self.transform is not None: - data_object = self.transform(data_object) - - return data_object - - def connect_db(self, lmdb_path=None): - env = lmdb.open( - str(lmdb_path), - subdir=False, - readonly=True, - lock=False, - readahead=False, - meminit=False, - max_readers=1, - map_size=1099511627776 * 2, - ) - return env - - def close_db(self): - if not self.path.is_file(): - for env in self.envs: - env.close() - else: - self.env.close() - - -def remove_hessian_transform(data): - # Remove 'hessian' if present as attribute - if hasattr(data, "hessian"): - delattr(data, "hessian") - # Remove 'hessian' if present as key (for dict-like) - # if isinstance(data, dict) and 'hessian' in data: - # del data['hessian'] - return data - - -def fix_hessian_eigen_transform(data): - """Fixes errors in old versions of the code. - Only needed for legacy compatibility. - You can probably remove this function. - - Sizes of tensors must match except in dimension 0. - - Hessian eigen information is stored as: - hessian_eigenvalues: torch.Size([2]) - hessian_eigenvectors: torch.Size([2, N*3]) - - Instead save as: - hessian_eigenvalue_1: torch.Size([1]) - hessian_eigenvalue_2: torch.Size([1]) - hessian_eigenvector_1: torch.Size([N, 3]) - hessian_eigenvector_2: torch.Size([N, 3]) - """ - # Check if hessian eigenvalue/eigenvector data exists - if hasattr(data, "hessian_eigenvalues") and hasattr(data, "hessian_eigenvectors"): - eigenvalues = data.hessian_eigenvalues - eigenvectors = data.hessian_eigenvectors - - # Split eigenvalues into separate attributes - data.hessian_eigenvalue_1 = eigenvalues[0:1] # Keep as [1] tensor - data.hessian_eigenvalue_2 = eigenvalues[1:2] # Keep as [1] tensor - - # Reshape and split eigenvectors from [2, N*3] to [N, 3] format - n_atoms = len(data.pos) # Get number of atoms from positions - eigenvector_1 = eigenvectors[0].reshape(n_atoms, 3) - eigenvector_2 = eigenvectors[1].reshape(n_atoms, 3) - - data.hessian_eigenvector_1 = eigenvector_1 - data.hessian_eigenvector_2 = eigenvector_2 - - # Remove original attributes - delattr(data, "hessian_eigenvalues") - delattr(data, "hessian_eigenvectors") - - # Remove batch artifacts manually - for key in ["batch", "ptr"]: - if hasattr(data, key): - delattr(data, key) - - return data - - -if __name__ == "__main__": - import os - - dataset_dir = os.path.expanduser( - "~/.cache/kagglehub/datasets/yunhonghan/hessian-dataset-for-optimizing-reactive-mliphorm/versions/5/" - ) - dataset_files = [ - "ts1x-val.lmdb", - "ts1x_hess_train_big.lmdb", - "RGD1.lmdb", - ] - lmdb_path = os.path.join(dataset_dir, dataset_files[0]) - lmdb_dataset = LmdbDataset(lmdb_path) - print("length of lmdb_dataset:", len(lmdb_dataset)) - print("first element of lmdb_dataset:", lmdb_dataset[0]) - print("first element of lmdb_dataset.pos:", lmdb_dataset[0].pos) - print("first element of lmdb_dataset.ae:", lmdb_dataset[0].ae) - first_elem = lmdb_dataset[0] - print("") - print("hasattr(first_elem, 'hessian'):", hasattr(first_elem, "hessian")) - print("'hessian' in first_elem:", "hessian" in first_elem) - - # Test with transform that removes hessian - lmdb_dataset_no_hessian = LmdbDataset(lmdb_path, transform=remove_hessian_transform) - first_elem = lmdb_dataset_no_hessian[0] - print("") - print("hasattr(first_elem, 'hessian'):", hasattr(first_elem, "hessian")) - print("'hessian' in first_elem:", "hessian" in first_elem) - - for fname in dataset_files: - path = os.path.join(dataset_dir, fname) - ds = LmdbDataset(path) - print(f"Size of {fname}: {len(ds)}") diff --git a/ocpmodels/models/__init__.py b/ocpmodels/models/__init__.py deleted file mode 100644 index 54695f27f15e90e7d5e0eac0a5560e4a1c18e64d..0000000000000000000000000000000000000000 --- a/ocpmodels/models/__init__.py +++ /dev/null @@ -1,31 +0,0 @@ -# Copyright (c) Facebook, Inc. and its affiliates. -# -# This source code is licensed under the MIT license found in the -# LICENSE file in the root directory of this source tree. - -from .base import BaseModel -from .cgcnn import CGCNN - -# These models use torch_sparse -try: - from .dimenet import DimeNetWrap as DimeNet -except ImportError as e: - print("Will not be able to use DimeNet. Error: ", e) - -try: - from .dimenet_plus_plus import DimeNetPlusPlusWrap as DimeNetPlusPlus -except ImportError as e: - print("Will not be able to use DimeNetPlusPlus. Error: ", e) - -try: - from .gemnet.gemnet import GemNetT - from .gemnet_gp.gemnet import GraphParallelGemNetT as GraphParallelGemNetT - from .gemnet_oc.gemnet_oc import GemNetOC -except ImportError as e: - print("Will not be able to use GemNet. Error: ", e) - -from .forcenet import ForceNet -from .painn.painn import PaiNN -from .schnet import SchNetWrap as SchNet -from .scn.scn import SphericalChannelNetwork -from .spinconv import spinconv diff --git a/ocpmodels/models/base.py b/ocpmodels/models/base.py deleted file mode 100644 index ba1cccd89bd53b15e594fb03b637bef01c3fa559..0000000000000000000000000000000000000000 --- a/ocpmodels/models/base.py +++ /dev/null @@ -1,109 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import logging - -import torch -import torch.nn as nn -from torch_geometric.nn import radius_graph - -from ocpmodels.common.utils import ( - compute_neighbors, - conditional_grad, - get_pbc_distances, - radius_graph_pbc, -) - - -class BaseModel(nn.Module): - def __init__(self, num_atoms=None, bond_feat_dim=None, num_targets=None): - super(BaseModel, self).__init__() - self.num_atoms = num_atoms - self.bond_feat_dim = bond_feat_dim - self.num_targets = num_targets - - def forward(self, data): - raise NotImplementedError - - def generate_graph( - self, - data, - cutoff=None, - max_neighbors=None, - use_pbc=None, - otf_graph=None, - ): - cutoff = cutoff or self.cutoff - max_neighbors = max_neighbors or self.max_neighbors - use_pbc = use_pbc or self.use_pbc - otf_graph = otf_graph or self.otf_graph - - if not otf_graph: - try: - edge_index = data.edge_index - - if use_pbc: - cell_offsets = data.cell_offsets - neighbors = data.neighbors - - except AttributeError: - logging.warning( - "Turning otf_graph=True as required attributes not present in data object" - ) - otf_graph = True - - if use_pbc: - if otf_graph: - edge_index, cell_offsets, neighbors = radius_graph_pbc( - data, cutoff, max_neighbors - ) - - out = get_pbc_distances( - data.pos, - edge_index, - data.cell, - cell_offsets, - neighbors, - return_offsets=True, - return_distance_vec=True, - ) - - edge_index = out["edge_index"] - edge_dist = out["distances"] - cell_offset_distances = out["offsets"] - distance_vec = out["distance_vec"] - else: - if otf_graph: - edge_index = radius_graph( - data.pos, - r=cutoff, - batch=data.batch, - max_num_neighbors=max_neighbors, - ) - - j, i = edge_index - distance_vec = data.pos[j] - data.pos[i] - - edge_dist = distance_vec.norm(dim=-1) - cell_offsets = torch.zeros(edge_index.shape[1], 3, device=data.pos.device) - cell_offset_distances = torch.zeros_like( - cell_offsets, device=data.pos.device - ) - neighbors = compute_neighbors(data, edge_index) - - return ( - edge_index, - edge_dist, - distance_vec, - cell_offsets, - cell_offset_distances, - neighbors, - ) - - @property - def num_params(self): - return sum(p.numel() for p in self.parameters()) diff --git a/ocpmodels/models/cgcnn.py b/ocpmodels/models/cgcnn.py deleted file mode 100644 index 46606ddf1050666e8ea2a87d59a26dd0c2217392..0000000000000000000000000000000000000000 --- a/ocpmodels/models/cgcnn.py +++ /dev/null @@ -1,232 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import torch -import torch.nn as nn -from torch_geometric.nn import MessagePassing, global_mean_pool, radius_graph -from torch_geometric.nn.models.schnet import GaussianSmearing - -from ocpmodels.common.registry import registry -from ocpmodels.common.utils import ( - conditional_grad, - get_pbc_distances, - radius_graph_pbc, -) -from ocpmodels.datasets.embeddings import KHOT_EMBEDDINGS, QMOF_KHOT_EMBEDDINGS -from ocpmodels.models.base import BaseModel - - -@registry.register_model("cgcnn") -class CGCNN(BaseModel): - r"""Implementation of the Crystal Graph CNN model from the - `"Crystal Graph Convolutional Neural Networks for an Accurate - and Interpretable Prediction of Material Properties" - `_ paper. - - Args: - num_atoms (int): Number of atoms. - bond_feat_dim (int): Dimension of bond features. - num_targets (int): Number of targets to predict. - use_pbc (bool, optional): If set to :obj:`True`, account for periodic boundary conditions. - (default: :obj:`True`) - regress_forces (bool, optional): If set to :obj:`True`, predict forces by differentiating - energy with respect to positions. - (default: :obj:`True`) - atom_embedding_size (int, optional): Size of atom embeddings. - (default: :obj:`64`) - num_graph_conv_layers (int, optional): Number of graph convolutional layers. - (default: :obj:`6`) - fc_feat_size (int, optional): Size of fully connected layers. - (default: :obj:`128`) - num_fc_layers (int, optional): Number of fully connected layers. - (default: :obj:`4`) - otf_graph (bool, optional): If set to :obj:`True`, compute graph edges on the fly. - (default: :obj:`False`) - cutoff (float, optional): Cutoff distance for interatomic interactions. - (default: :obj:`10.0`) - num_gaussians (int, optional): Number of Gaussians used for smearing. - (default: :obj:`50.0`) - """ - - def __init__( - self, - num_atoms, - bond_feat_dim, - num_targets, - use_pbc=True, - regress_forces=True, - atom_embedding_size=64, - num_graph_conv_layers=6, - fc_feat_size=128, - num_fc_layers=4, - otf_graph=False, - cutoff=6.0, - num_gaussians=50, - embeddings="khot", - ): - super(CGCNN, self).__init__(num_atoms, bond_feat_dim, num_targets) - self.regress_forces = regress_forces - self.use_pbc = use_pbc - self.cutoff = cutoff - self.otf_graph = otf_graph - self.max_neighbors = 50 - # Get CGCNN atom embeddings - if embeddings == "khot": - embeddings = KHOT_EMBEDDINGS - elif embeddings == "qmof": - embeddings = QMOF_KHOT_EMBEDDINGS - else: - raise ValueError( - 'embedding mnust be either "khot" for original CGCNN K-hot elemental embeddings or "qmof" for QMOF K-hot elemental embeddings' - ) - self.embedding = torch.zeros(100, len(embeddings[1])) - for i in range(100): - self.embedding[i] = torch.tensor(embeddings[i + 1]) - self.embedding_fc = nn.Linear(len(embeddings[1]), atom_embedding_size) - - self.convs = nn.ModuleList( - [ - CGCNNConv( - node_dim=atom_embedding_size, - edge_dim=bond_feat_dim, - cutoff=cutoff, - ) - for _ in range(num_graph_conv_layers) - ] - ) - - self.conv_to_fc = nn.Sequential( - nn.Linear(atom_embedding_size, fc_feat_size), nn.Softplus() - ) - - if num_fc_layers > 1: - layers = [] - for _ in range(num_fc_layers - 1): - layers.append(nn.Linear(fc_feat_size, fc_feat_size)) - layers.append(nn.Softplus()) - self.fcs = nn.Sequential(*layers) - self.fc_out = nn.Linear(fc_feat_size, self.num_targets) - - self.cutoff = cutoff - self.distance_expansion = GaussianSmearing(0.0, cutoff, num_gaussians) - - @conditional_grad(torch.enable_grad()) - def _forward(self, data): - # Get node features - if self.embedding.device != data.atomic_numbers.device: - self.embedding = self.embedding.to(data.atomic_numbers.device) - data.x = self.embedding[data.atomic_numbers.long() - 1] - - ( - edge_index, - distances, - distance_vec, - cell_offsets, - _, # cell offset distances - neighbors, - ) = self.generate_graph(data) - - data.edge_index = edge_index - data.edge_attr = self.distance_expansion(distances) - # Forward pass through the network - mol_feats = self._convolve(data) - mol_feats = self.conv_to_fc(mol_feats) - if hasattr(self, "fcs"): - mol_feats = self.fcs(mol_feats) - - energy = self.fc_out(mol_feats) - return energy - - def forward(self, data): - if self.regress_forces: - data.pos.requires_grad_(True) - energy = self._forward(data) - - if self.regress_forces: - forces = -1 * ( - torch.autograd.grad( - energy, - data.pos, - grad_outputs=torch.ones_like(energy), - create_graph=True, - )[0] - ) - return energy, forces - else: - return energy - - def _convolve(self, data): - """ - Returns the output of the convolution layers before they are passed - into the dense layers. - """ - node_feats = self.embedding_fc(data.x) - for f in self.convs: - node_feats = f(node_feats, data.edge_index, data.edge_attr) - mol_feats = global_mean_pool(node_feats, data.batch) - return mol_feats - - -class CGCNNConv(MessagePassing): - """Implements the message passing layer from - `"Crystal Graph Convolutional Neural Networks for an - Accurate and Interpretable Prediction of Material Properties" - `. - """ - - def __init__(self, node_dim, edge_dim, cutoff=6.0, **kwargs): - super(CGCNNConv, self).__init__(aggr="add") - self.node_feat_size = node_dim - self.edge_feat_size = edge_dim - self.cutoff = cutoff - - self.lin1 = nn.Linear( - 2 * self.node_feat_size + self.edge_feat_size, - 2 * self.node_feat_size, - ) - self.bn1 = nn.BatchNorm1d(2 * self.node_feat_size) - self.ln1 = nn.LayerNorm(self.node_feat_size) - - self.reset_parameters() - - def reset_parameters(self): - torch.nn.init.xavier_uniform_(self.lin1.weight) - - self.lin1.bias.data.fill_(0) - - self.bn1.reset_parameters() - self.ln1.reset_parameters() - - def forward(self, x, edge_index, edge_attr): - """ - Arguments: - x has shape [num_nodes, node_feat_size] - edge_index has shape [2, num_edges] - edge_attr is [num_edges, edge_feat_size] - """ - out = self.propagate( - edge_index, x=x, edge_attr=edge_attr, size=(x.size(0), x.size(0)) - ) - out = nn.Softplus()(self.ln1(out) + x) - return out - - def message(self, x_i, x_j, edge_attr): - """ - Arguments: - x_i has shape [num_edges, node_feat_size] - x_j has shape [num_edges, node_feat_size] - edge_attr has shape [num_edges, edge_feat_size] - - Returns: - tensor of shape [num_edges, node_feat_size] - """ - z = self.lin1(torch.cat([x_i, x_j, edge_attr], dim=1)) - z = self.bn1(z) - z1, z2 = z.chunk(2, dim=1) - z1 = nn.Sigmoid()(z1) - z2 = nn.Softplus()(z2) - return z1 * z2 diff --git a/ocpmodels/models/dimenet.py b/ocpmodels/models/dimenet.py deleted file mode 100644 index ece70fbc16ce689fcde87596e8059459aa08a53d..0000000000000000000000000000000000000000 --- a/ocpmodels/models/dimenet.py +++ /dev/null @@ -1,230 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import torch -from torch import nn -from torch_geometric.nn import DimeNet, radius_graph -from torch_scatter import scatter -from torch_sparse import SparseTensor - -from ocpmodels.common.registry import registry -from ocpmodels.common.utils import ( - conditional_grad, - get_pbc_distances, - radius_graph_pbc, -) -from ocpmodels.models.base import BaseModel - - -@registry.register_model("dimenet") -class DimeNetWrap(DimeNet, BaseModel): - r"""Wrapper around the directional message passing neural network (DimeNet) from the - `"Directional Message Passing for Molecular Graphs" - `_ paper. - - DimeNet transforms messages based on the angle between them in a - rotation-equivariant fashion. - - Args: - num_atoms (int): Unused argument - bond_feat_dim (int): Unused argument - num_targets (int): Number of targets to predict. - use_pbc (bool, optional): If set to :obj:`True`, account for periodic boundary conditions. - (default: :obj:`True`) - regress_forces (bool, optional): If set to :obj:`True`, predict forces by differentiating - energy with respect to positions. - (default: :obj:`True`) - hidden_channels (int, optional): Number of hidden channels. - (default: :obj:`128`) - num_blocks (int, optional): Number of building blocks. - (default: :obj:`6`) - num_bilinear (int, optional): Size of the bilinear layer tensor. - (default: :obj:`8`) - num_spherical (int, optional): Number of spherical harmonics. - (default: :obj:`7`) - num_radial (int, optional): Number of radial basis functions. - (default: :obj:`6`) - otf_graph (bool, optional): If set to :obj:`True`, compute graph edges on the fly. - (default: :obj:`False`) - cutoff (float, optional): Cutoff distance for interatomic interactions. - (default: :obj:`10.0`) - envelope_exponent (int, optional): Shape of the smooth cutoff. - (default: :obj:`5`) - num_before_skip: (int, optional): Number of residual layers in the - interaction blocks before the skip connection. (default: :obj:`1`) - num_after_skip: (int, optional): Number of residual layers in the - interaction blocks after the skip connection. (default: :obj:`2`) - num_output_layers: (int, optional): Number of linear layers for the - output blocks. (default: :obj:`3`) - max_angles_per_image (int, optional): The maximum number of angles used - per image. This can be used to reduce memory usage at the cost of - model performance. (default: :obj:`1e6`) - """ - - def __init__( - self, - num_atoms, - bond_feat_dim, # not used - num_targets, - use_pbc=True, - regress_forces=True, - hidden_channels=128, - num_blocks=6, - num_bilinear=8, - num_spherical=7, - num_radial=6, - otf_graph=False, - cutoff=10.0, - envelope_exponent=5, - num_before_skip=1, - num_after_skip=2, - num_output_layers=3, - max_angles_per_image=int(1e6), - ): - self.num_targets = num_targets - self.regress_forces = regress_forces - self.use_pbc = use_pbc - self.cutoff = cutoff - self.otf_graph = otf_graph - self.max_angles_per_image = max_angles_per_image - self.max_neighbors = 50 - - super(DimeNetWrap, self).__init__( - hidden_channels=hidden_channels, - out_channels=num_targets, - num_blocks=num_blocks, - num_bilinear=num_bilinear, - num_spherical=num_spherical, - num_radial=num_radial, - cutoff=cutoff, - envelope_exponent=envelope_exponent, - num_before_skip=num_before_skip, - num_after_skip=num_after_skip, - num_output_layers=num_output_layers, - ) - - def triplets(self, edge_index, cell_offsets, num_nodes): - row, col = edge_index # j->i - - value = torch.arange(row.size(0), device=row.device) - adj_t = SparseTensor( - row=col, col=row, value=value, sparse_sizes=(num_nodes, num_nodes) - ) - adj_t_row = adj_t[row] - num_triplets = adj_t_row.set_value(None).sum(dim=1).to(torch.long) - - # Node indices (k->j->i) for triplets. - idx_i = col.repeat_interleave(num_triplets) - idx_j = row.repeat_interleave(num_triplets) - idx_k = adj_t_row.storage.col() - - # Edge indices (k->j, j->i) for triplets. - idx_kj = adj_t_row.storage.value() - idx_ji = adj_t_row.storage.row() - - # Remove self-loop triplets d->b->d - # Check atom as well as cell offset - cell_offset_kji = cell_offsets[idx_kj] + cell_offsets[idx_ji] - mask = (idx_i != idx_k) | torch.any(cell_offset_kji != 0, dim=-1) - - idx_i, idx_j, idx_k = idx_i[mask], idx_j[mask], idx_k[mask] - idx_kj, idx_ji = idx_kj[mask], idx_ji[mask] - - return col, row, idx_i, idx_j, idx_k, idx_kj, idx_ji - - @conditional_grad(torch.enable_grad()) - def _forward(self, data): - pos = data.pos - batch = data.batch - ( - edge_index, - dist, - _, - cell_offsets, - offsets, - neighbors, - ) = self.generate_graph(data) - - data.edge_index = edge_index - data.cell_offsets = cell_offsets - data.neighbors = neighbors - j, i = edge_index - - _, _, idx_i, idx_j, idx_k, idx_kj, idx_ji = self.triplets( - edge_index, - data.cell_offsets, - num_nodes=data.atomic_numbers.size(0), - ) - - # Cap no. of triplets during training. - if self.training: - sub_ix = torch.randperm(idx_i.size(0))[ - : self.max_angles_per_image * data.natoms.size(0) - ] - idx_i, idx_j, idx_k = ( - idx_i[sub_ix], - idx_j[sub_ix], - idx_k[sub_ix], - ) - idx_kj, idx_ji = idx_kj[sub_ix], idx_ji[sub_ix] - - # Calculate angles. - pos_i = pos[idx_i].detach() - pos_j = pos[idx_j].detach() - if self.use_pbc: - pos_ji, pos_kj = ( - pos[idx_j].detach() - pos_i + offsets[idx_ji], - pos[idx_k].detach() - pos_j + offsets[idx_kj], - ) - else: - pos_ji, pos_kj = ( - pos[idx_j].detach() - pos_i, - pos[idx_k].detach() - pos_j, - ) - - a = (pos_ji * pos_kj).sum(dim=-1) - b = torch.cross(pos_ji, pos_kj).norm(dim=-1) - angle = torch.atan2(b, a) - - rbf = self.rbf(dist) - sbf = self.sbf(dist, angle, idx_kj) - - # Embedding block. - x = self.emb(data.atomic_numbers.long(), rbf, i, j) - P = self.output_blocks[0](x, rbf, i, num_nodes=pos.size(0)) - - # Interaction blocks. - for interaction_block, output_block in zip( - self.interaction_blocks, self.output_blocks[1:] - ): - x = interaction_block(x, rbf, sbf, idx_kj, idx_ji) - P += output_block(x, rbf, i, num_nodes=pos.size(0)) - - energy = P.sum(dim=0) if batch is None else scatter(P, batch, dim=0) - return energy - - def forward(self, data): - if self.regress_forces: - data.pos.requires_grad_(True) - energy = self._forward(data) - - if self.regress_forces: - forces = -1 * ( - torch.autograd.grad( - energy, - data.pos, - grad_outputs=torch.ones_like(energy), - create_graph=True, - )[0] - ) - return energy, forces - else: - return energy - - @property - def num_params(self): - return sum(p.numel() for p in self.parameters()) diff --git a/ocpmodels/models/dimenet_plus_plus.py b/ocpmodels/models/dimenet_plus_plus.py deleted file mode 100644 index 414a6ea662f8a32b5184ecbc15973a0cfe01e52d..0000000000000000000000000000000000000000 --- a/ocpmodels/models/dimenet_plus_plus.py +++ /dev/null @@ -1,465 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. - ---- - -This code borrows heavily from the DimeNet implementation as part of -pytorch-geometric: https://github.com/rusty1s/pytorch_geometric. License: - ---- - -Copyright (c) 2020 Matthias Fey - -Permission is hereby granted, free of charge, to any person obtaining a copy -of this software and associated documentation files (the "Software"), to deal -in the Software without restriction, including without limitation the rights -to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -copies of the Software, and to permit persons to whom the Software is -furnished to do so, subject to the following conditions: - -The above copyright notice and this permission notice shall be included in -all copies or substantial portions of the Software. - -THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN -THE SOFTWARE. -""" - -import torch -from torch import nn -from torch_geometric.nn import radius_graph -from torch_geometric.nn.inits import glorot_orthogonal - -from torch_geometric.nn.models.dimenet import ( - BesselBasisLayer, - EmbeddingBlock, - Envelope, - ResidualLayer, - SphericalBasisLayer, -) -from torch_geometric.nn.resolver import activation_resolver -from torch_scatter import scatter -from torch_sparse import SparseTensor - -from ocpmodels.common.registry import registry -from ocpmodels.common.utils import ( - conditional_grad, - get_pbc_distances, - radius_graph_pbc, -) -from ocpmodels.models.base import BaseModel - -try: - import sympy as sym -except ImportError: - sym = None - - -class InteractionPPBlock(torch.nn.Module): - def __init__( - self, - hidden_channels, - int_emb_size, - basis_emb_size, - num_spherical, - num_radial, - num_before_skip, - num_after_skip, - act="silu", - ): - act = activation_resolver(act) - super(InteractionPPBlock, self).__init__() - self.act = act - - # Transformations of Bessel and spherical basis representations. - self.lin_rbf1 = nn.Linear(num_radial, basis_emb_size, bias=False) - self.lin_rbf2 = nn.Linear(basis_emb_size, hidden_channels, bias=False) - self.lin_sbf1 = nn.Linear( - num_spherical * num_radial, basis_emb_size, bias=False - ) - self.lin_sbf2 = nn.Linear(basis_emb_size, int_emb_size, bias=False) - - # Dense transformations of input messages. - self.lin_kj = nn.Linear(hidden_channels, hidden_channels) - self.lin_ji = nn.Linear(hidden_channels, hidden_channels) - - # Embedding projections for interaction triplets. - self.lin_down = nn.Linear(hidden_channels, int_emb_size, bias=False) - self.lin_up = nn.Linear(int_emb_size, hidden_channels, bias=False) - - # Residual layers before and after skip connection. - self.layers_before_skip = torch.nn.ModuleList( - [ResidualLayer(hidden_channels, act) for _ in range(num_before_skip)] - ) - self.lin = nn.Linear(hidden_channels, hidden_channels) - self.layers_after_skip = torch.nn.ModuleList( - [ResidualLayer(hidden_channels, act) for _ in range(num_after_skip)] - ) - - self.reset_parameters() - - def reset_parameters(self): - glorot_orthogonal(self.lin_rbf1.weight, scale=2.0) - glorot_orthogonal(self.lin_rbf2.weight, scale=2.0) - glorot_orthogonal(self.lin_sbf1.weight, scale=2.0) - glorot_orthogonal(self.lin_sbf2.weight, scale=2.0) - - glorot_orthogonal(self.lin_kj.weight, scale=2.0) - self.lin_kj.bias.data.fill_(0) - glorot_orthogonal(self.lin_ji.weight, scale=2.0) - self.lin_ji.bias.data.fill_(0) - - glorot_orthogonal(self.lin_down.weight, scale=2.0) - glorot_orthogonal(self.lin_up.weight, scale=2.0) - - for res_layer in self.layers_before_skip: - res_layer.reset_parameters() - glorot_orthogonal(self.lin.weight, scale=2.0) - self.lin.bias.data.fill_(0) - for res_layer in self.layers_after_skip: - res_layer.reset_parameters() - - def forward(self, x, rbf, sbf, idx_kj, idx_ji): - # Initial transformations. - x_ji = self.act(self.lin_ji(x)) - x_kj = self.act(self.lin_kj(x)) - - # Transformation via Bessel basis. - rbf = self.lin_rbf1(rbf) - rbf = self.lin_rbf2(rbf) - x_kj = x_kj * rbf - - # Down-project embeddings and generate interaction triplet embeddings. - x_kj = self.act(self.lin_down(x_kj)) - - # Transform via 2D spherical basis. - sbf = self.lin_sbf1(sbf) - sbf = self.lin_sbf2(sbf) - x_kj = x_kj[idx_kj] * sbf - - # Aggregate interactions and up-project embeddings. - x_kj = scatter(x_kj, idx_ji, dim=0, dim_size=x.size(0)) - x_kj = self.act(self.lin_up(x_kj)) - - h = x_ji + x_kj - for layer in self.layers_before_skip: - h = layer(h) - h = self.act(self.lin(h)) + x - for layer in self.layers_after_skip: - h = layer(h) - - return h - - -class OutputPPBlock(torch.nn.Module): - def __init__( - self, - num_radial, - hidden_channels, - out_emb_channels, - out_channels, - num_layers, - act="silu", - ): - act = activation_resolver(act) - super(OutputPPBlock, self).__init__() - self.act = act - - self.lin_rbf = nn.Linear(num_radial, hidden_channels, bias=False) - self.lin_up = nn.Linear(hidden_channels, out_emb_channels, bias=True) - self.lins = torch.nn.ModuleList() - for _ in range(num_layers): - self.lins.append(nn.Linear(out_emb_channels, out_emb_channels)) - self.lin = nn.Linear(out_emb_channels, out_channels, bias=False) - - self.reset_parameters() - - def reset_parameters(self): - glorot_orthogonal(self.lin_rbf.weight, scale=2.0) - glorot_orthogonal(self.lin_up.weight, scale=2.0) - for lin in self.lins: - glorot_orthogonal(lin.weight, scale=2.0) - lin.bias.data.fill_(0) - self.lin.weight.data.fill_(0) - - def forward(self, x, rbf, i, num_nodes=None): - x = self.lin_rbf(rbf) * x - x = scatter(x, i, dim=0, dim_size=num_nodes) - x = self.lin_up(x) - for lin in self.lins: - x = self.act(lin(x)) - return self.lin(x) - - -class DimeNetPlusPlus(torch.nn.Module): - r"""DimeNet++ implementation based on https://github.com/klicperajo/dimenet. - - Args: - hidden_channels (int): Hidden embedding size. - out_channels (int): Size of each output sample. - num_blocks (int): Number of building blocks. - int_emb_size (int): Embedding size used for interaction triplets - basis_emb_size (int): Embedding size used in the basis transformation - out_emb_channels(int): Embedding size used for atoms in the output block - num_spherical (int): Number of spherical harmonics. - num_radial (int): Number of radial basis functions. - cutoff: (float, optional): Cutoff distance for interatomic - interactions. (default: :obj:`5.0`) - envelope_exponent (int, optional): Shape of the smooth cutoff. - (default: :obj:`5`) - num_before_skip: (int, optional): Number of residual layers in the - interaction blocks before the skip connection. (default: :obj:`1`) - num_after_skip: (int, optional): Number of residual layers in the - interaction blocks after the skip connection. (default: :obj:`2`) - num_output_layers: (int, optional): Number of linear layers for the - output blocks. (default: :obj:`3`) - act: (function, optional): The activation funtion. - (default: :obj:`silu`) - """ - - url = "https://github.com/klicperajo/dimenet/raw/master/pretrained" - - def __init__( - self, - hidden_channels, - out_channels, - num_blocks, - int_emb_size, - basis_emb_size, - out_emb_channels, - num_spherical, - num_radial, - cutoff=5.0, - envelope_exponent=5, - num_before_skip=1, - num_after_skip=2, - num_output_layers=3, - act="silu", - ): - - act = activation_resolver(act) - - super(DimeNetPlusPlus, self).__init__() - - self.cutoff = cutoff - - if sym is None: - raise ImportError("Package `sympy` could not be found.") - - self.num_blocks = num_blocks - - self.rbf = BesselBasisLayer(num_radial, cutoff, envelope_exponent) - self.sbf = SphericalBasisLayer( - num_spherical, num_radial, cutoff, envelope_exponent - ) - - self.emb = EmbeddingBlock(num_radial, hidden_channels, act) - - self.output_blocks = torch.nn.ModuleList( - [ - OutputPPBlock( - num_radial, - hidden_channels, - out_emb_channels, - out_channels, - num_output_layers, - act, - ) - for _ in range(num_blocks + 1) - ] - ) - - self.interaction_blocks = torch.nn.ModuleList( - [ - InteractionPPBlock( - hidden_channels, - int_emb_size, - basis_emb_size, - num_spherical, - num_radial, - num_before_skip, - num_after_skip, - act, - ) - for _ in range(num_blocks) - ] - ) - - self.reset_parameters() - - def reset_parameters(self): - self.rbf.reset_parameters() - self.emb.reset_parameters() - for out in self.output_blocks: - out.reset_parameters() - for interaction in self.interaction_blocks: - interaction.reset_parameters() - - def triplets(self, edge_index, cell_offsets, num_nodes): - row, col = edge_index # j->i - - value = torch.arange(row.size(0), device=row.device) - adj_t = SparseTensor( - row=col, col=row, value=value, sparse_sizes=(num_nodes, num_nodes) - ) - adj_t_row = adj_t[row] - num_triplets = adj_t_row.set_value(None).sum(dim=1).to(torch.long) - - # Node indices (k->j->i) for triplets. - idx_i = col.repeat_interleave(num_triplets) - idx_j = row.repeat_interleave(num_triplets) - idx_k = adj_t_row.storage.col() - - # Edge indices (k->j, j->i) for triplets. - idx_kj = adj_t_row.storage.value() - idx_ji = adj_t_row.storage.row() - - # Remove self-loop triplets d->b->d - # Check atom as well as cell offset - cell_offset_kji = cell_offsets[idx_kj] + cell_offsets[idx_ji] - mask = (idx_i != idx_k) | torch.any(cell_offset_kji != 0, dim=-1) - - idx_i, idx_j, idx_k = idx_i[mask], idx_j[mask], idx_k[mask] - idx_kj, idx_ji = idx_kj[mask], idx_ji[mask] - - return col, row, idx_i, idx_j, idx_k, idx_kj, idx_ji - - def forward(self, z, pos, batch=None): - """ """ - raise NotImplementedError - - -@registry.register_model("dimenetplusplus") -class DimeNetPlusPlusWrap(DimeNetPlusPlus, BaseModel): - def __init__( - self, - num_atoms, - bond_feat_dim, # not used - num_targets, - use_pbc=True, - regress_forces=True, - hidden_channels=128, - num_blocks=4, - int_emb_size=64, - basis_emb_size=8, - out_emb_channels=256, - num_spherical=7, - num_radial=6, - otf_graph=False, - cutoff=10.0, - envelope_exponent=5, - num_before_skip=1, - num_after_skip=2, - num_output_layers=3, - ): - self.num_targets = num_targets - self.regress_forces = regress_forces - self.use_pbc = use_pbc - self.cutoff = cutoff - self.otf_graph = otf_graph - self.max_neighbors = 50 - - super(DimeNetPlusPlusWrap, self).__init__( - hidden_channels=hidden_channels, - out_channels=num_targets, - num_blocks=num_blocks, - int_emb_size=int_emb_size, - basis_emb_size=basis_emb_size, - out_emb_channels=out_emb_channels, - num_spherical=num_spherical, - num_radial=num_radial, - cutoff=cutoff, - envelope_exponent=envelope_exponent, - num_before_skip=num_before_skip, - num_after_skip=num_after_skip, - num_output_layers=num_output_layers, - ) - - @conditional_grad(torch.enable_grad()) - def _forward(self, data): - pos = data.pos - batch = data.batch - ( - edge_index, - dist, - _, - cell_offsets, - offsets, - neighbors, - ) = self.generate_graph(data) - - data.edge_index = edge_index - data.cell_offsets = cell_offsets - data.neighbors = neighbors - j, i = edge_index - - _, _, idx_i, idx_j, idx_k, idx_kj, idx_ji = self.triplets( - edge_index, - data.cell_offsets, - num_nodes=data.atomic_numbers.size(0), - ) - - # Calculate angles. - pos_i = pos[idx_i].detach() - pos_j = pos[idx_j].detach() - if self.use_pbc: - pos_ji, pos_kj = ( - pos[idx_j].detach() - pos_i + offsets[idx_ji], - pos[idx_k].detach() - pos_j + offsets[idx_kj], - ) - else: - pos_ji, pos_kj = ( - pos[idx_j].detach() - pos_i, - pos[idx_k].detach() - pos_j, - ) - - a = (pos_ji * pos_kj).sum(dim=-1) - b = torch.cross(pos_ji, pos_kj).norm(dim=-1) - angle = torch.atan2(b, a) - - rbf = self.rbf(dist) - sbf = self.sbf(dist, angle, idx_kj) - - # Embedding block. - x = self.emb(data.atomic_numbers.long(), rbf, i, j) - P = self.output_blocks[0](x, rbf, i, num_nodes=pos.size(0)) - - # Interaction blocks. - for interaction_block, output_block in zip( - self.interaction_blocks, self.output_blocks[1:] - ): - x = interaction_block(x, rbf, sbf, idx_kj, idx_ji) - P += output_block(x, rbf, i, num_nodes=pos.size(0)) - - energy = P.sum(dim=0) if batch is None else scatter(P, batch, dim=0) - - return energy - - def forward(self, data): - if self.regress_forces: - data.pos.requires_grad_(True) - energy = self._forward(data) - - if self.regress_forces: - forces = -1 * ( - torch.autograd.grad( - energy, - data.pos, - grad_outputs=torch.ones_like(energy), - create_graph=True, - )[0] - ) - return energy, forces - else: - return energy - - @property - def num_params(self): - return sum(p.numel() for p in self.parameters()) diff --git a/ocpmodels/models/escn/Jd.pt b/ocpmodels/models/escn/Jd.pt deleted file mode 100644 index 9bf1f88c6fc9186c7df212787ad6ee53c96907b8..0000000000000000000000000000000000000000 --- a/ocpmodels/models/escn/Jd.pt +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:b4059c45be246dcb6c49c545670b65c56550eb0c2e7a9c92b4b50a92d370dbe2 -size 21697 diff --git a/ocpmodels/models/escn/escn.py b/ocpmodels/models/escn/escn.py deleted file mode 100644 index 7107f7574788c8e401bc49122be40dba46f6779f..0000000000000000000000000000000000000000 --- a/ocpmodels/models/escn/escn.py +++ /dev/null @@ -1,958 +0,0 @@ -""" -Copyright (c) Meta, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import logging -import time - -import numpy as np -import torch -import torch.nn as nn -from pyexpat.model import XML_CQUANT_OPT - -from ocpmodels.common.registry import registry -from ocpmodels.common.utils import conditional_grad -from ocpmodels.models.base import BaseModel -from ocpmodels.models.escn.so3 import ( - CoefficientMapping, - SO3_Embedding, - SO3_Grid, - SO3_Rotation, -) -from ocpmodels.models.scn.sampling import CalcSpherePoints -from ocpmodels.models.scn.smearing import ( - GaussianSmearing, - LinearSigmoidSmearing, - SigmoidSmearing, - SiLUSmearing, -) - -try: - from e3nn import o3 -except ImportError: - pass - - -@registry.register_model("escn") -class eSCN(BaseModel): - """Equivariant Spherical Channel Network - Paper: Reducing SO(3) Convolutions to SO(2) for Efficient Equivariant GNNs - - - Args: - use_pbc (bool): Use periodic boundary conditions - regress_forces (bool): Compute forces - otf_graph (bool): Compute graph On The Fly (OTF) - max_neighbors (int): Maximum number of neighbors per atom - cutoff (float): Maximum distance between nieghboring atoms in Angstroms - max_num_elements (int): Maximum atomic number - - num_layers (int): Number of layers in the GNN - lmax_list (int): List of maximum degree of the spherical harmonics (1 to 10) - mmax_list (int): List of maximum order of the spherical harmonics (0 to lmax) - sphere_channels (int): Number of spherical channels (one set per resolution) - hidden_channels (int): Number of hidden units in message passing - num_sphere_samples (int): Number of samples used to approximate the integration of the sphere in the output blocks - edge_channels (int): Number of channels for the edge invariant features - distance_function ("gaussian", "sigmoid", "linearsigmoid", "silu"): Basis function used for distances - basis_width_scalar (float): Width of distance basis function - distance_resolution (float): Distance between distance basis functions in Angstroms - show_timing_info (bool): Show timing and memory info - """ - - def __init__( - self, - num_atoms, # not used - bond_feat_dim, # not used - num_targets, # not used - use_pbc=True, - regress_forces=True, - otf_graph=False, - max_neighbors=40, - cutoff=8.0, - max_num_elements=90, - num_layers=8, - lmax_list=[6], - mmax_list=[2], - sphere_channels=128, - hidden_channels=256, - edge_channels=128, - use_grid=True, - num_sphere_samples=128, - distance_function="gaussian", - basis_width_scalar=1.0, - distance_resolution=0.02, - show_timing_info=True, - ): - super().__init__() - - import sys - - if "e3nn" not in sys.modules: - logging.error("You need to install the e3nn library to use the SCN model") - raise ImportError - - self.regress_forces = regress_forces - self.use_pbc = use_pbc - self.cutoff = cutoff - self.otf_graph = otf_graph - self.show_timing_info = show_timing_info - self.max_num_elements = max_num_elements - self.hidden_channels = hidden_channels - self.num_layers = num_layers - self.num_atoms = 0 - self.num_sphere_samples = num_sphere_samples - self.sphere_channels = sphere_channels - self.max_neighbors = max_neighbors - self.edge_channels = edge_channels - self.distance_resolution = distance_resolution - self.grad_forces = False - self.lmax_list = lmax_list - self.mmax_list = mmax_list - self.num_resolutions = len(self.lmax_list) - self.sphere_channels_all = self.num_resolutions * self.sphere_channels - self.basis_width_scalar = basis_width_scalar - self.distance_function = distance_function - self.device = torch.cuda.current_device() - - # variables used for display purposes - self.counter = 0 - - # non-linear activation function used throughout the network - self.act = nn.SiLU() - - # Weights for message initialization - self.sphere_embedding = nn.Embedding( - self.max_num_elements, self.sphere_channels_all - ) - - # Initialize the function used to measure the distances between atoms - assert self.distance_function in [ - "gaussian", - "sigmoid", - "linearsigmoid", - "silu", - ] - self.num_gaussians = int(cutoff / self.distance_resolution) - if self.distance_function == "gaussian": - self.distance_expansion = GaussianSmearing( - 0.0, - cutoff, - self.num_gaussians, - basis_width_scalar, - ) - if self.distance_function == "sigmoid": - self.distance_expansion = SigmoidSmearing( - 0.0, - cutoff, - self.num_gaussians, - basis_width_scalar, - ) - if self.distance_function == "linearsigmoid": - self.distance_expansion = LinearSigmoidSmearing( - 0.0, - cutoff, - self.num_gaussians, - basis_width_scalar, - ) - if self.distance_function == "silu": - self.distance_expansion = SiLUSmearing( - 0.0, - cutoff, - self.num_gaussians, - basis_width_scalar, - ) - - # Initialize the transformations between spherical and grid representations - self.SO3_grid = nn.ModuleList() - for l in range(max(self.lmax_list) + 1): - SO3_m_grid = nn.ModuleList() - for m in range(max(self.lmax_list) + 1): - SO3_m_grid.append(SO3_Grid(l, m)) - - self.SO3_grid.append(SO3_m_grid) - - # Initialize the blocks for each layer of the GNN - self.layer_blocks = nn.ModuleList() - for i in range(self.num_layers): - block = LayerBlock( - i, - self.sphere_channels, - self.hidden_channels, - self.edge_channels, - self.lmax_list, - self.mmax_list, - self.distance_expansion, - self.max_num_elements, - self.SO3_grid, - self.act, - ) - self.layer_blocks.append(block) - - # Output blocks for energy and forces - self.energy_block = EnergyBlock( - self.sphere_channels_all, self.num_sphere_samples, self.act - ) - if self.regress_forces: - self.force_block = ForceBlock( - self.sphere_channels_all, self.num_sphere_samples, self.act - ) - - # Create a roughly evenly distributed point sampling of the sphere for the output blocks - self.sphere_points = CalcSpherePoints( - self.num_sphere_samples, self.device - ).detach() - - # For each spherical point, compute the spherical harmonic coefficient weights - self.sphharm_weights = [] - for i in range(self.num_resolutions): - self.sphharm_weights.append( - o3.spherical_harmonics( - torch.arange(0, self.lmax_list[i] + 1).tolist(), - self.sphere_points, - False, - ).detach() - ) - - @conditional_grad(torch.enable_grad()) - def forward(self, data): - self.batch_size = len(data.natoms) - self.dtype = data.pos.dtype - - start_time = time.time() - atomic_numbers = data.atomic_numbers.long() - num_atoms = len(atomic_numbers) - pos = data.pos - - ( - edge_index, - edge_distance, - edge_distance_vec, - cell_offsets, - _, # cell offset distances - neighbors, - ) = self.generate_graph(data) - - ############################################################### - # Initialize data structures - ############################################################### - - # Compute 3x3 rotation matrix per edge - edge_rot_mat = self._init_edge_rot_mat(data, edge_index, edge_distance_vec) - - # Initialize the WignerD matrices and other values for spherical harmonic calculations - self.SO3_edge_rot = nn.ModuleList() - for i in range(self.num_resolutions): - self.SO3_edge_rot.append(SO3_Rotation(edge_rot_mat, self.lmax_list[i])) - - ############################################################### - # Initialize node embeddings - ############################################################### - - # Init per node representations using an atomic number based embedding - offset = 0 - x = SO3_Embedding( - num_atoms, - self.lmax_list, - self.sphere_channels, - self.device, - self.dtype, - ) - - offset_res = 0 - offset = 0 - # Initialize the l=0,m=0 coefficients for each resolution - for i in range(self.num_resolutions): - x.embedding[:, offset_res, :] = self.sphere_embedding(atomic_numbers)[ - :, offset : offset + self.sphere_channels - ] - offset = offset + self.sphere_channels - offset_res = offset_res + int((self.lmax_list[i] + 1) ** 2) - - # This can be expensive to compute (not implemented efficiently), so only do it once and pass it along to each layer - mappingReduced = CoefficientMapping(self.lmax_list, self.mmax_list, self.device) - - ############################################################### - # Update spherical node embeddings - ############################################################### - - for i in range(self.num_layers): - if i > 0: - x_message = self.layer_blocks[i]( - x, - atomic_numbers, - edge_distance, - edge_index, - self.SO3_edge_rot, - mappingReduced, - ) - - # Residual layer for all layers past the first - x.embedding = x.embedding + x_message.embedding - - else: - # No residual for the first layer - x = self.layer_blocks[i]( - x, - atomic_numbers, - edge_distance, - edge_index, - self.SO3_edge_rot, - mappingReduced, - ) - - # Sample the spherical channels (node embeddings) at evenly distributed points on the sphere. - # These values are fed into the output blocks. - x_pt = torch.tensor([], device=self.device) - offset = 0 - # Compute the embedding values at every sampled point on the sphere - for i in range(self.num_resolutions): - num_coefficients = int((x.lmax_list[i] + 1) ** 2) - x_pt = torch.cat( - [ - x_pt, - torch.einsum( - "abc, pb->apc", - x.embedding[:, offset : offset + num_coefficients], - self.sphharm_weights[i], - ).contiguous(), - ], - dim=2, - ) - offset = offset + num_coefficients - - x_pt = x_pt.view(-1, self.sphere_channels_all) - - ############################################################### - # Energy estimation - ############################################################### - node_energy = self.energy_block(x_pt) - energy = torch.zeros(len(data.natoms), device=pos.device) - energy.index_add_(0, data.batch, node_energy.view(-1)) - # Scale energy to help balance numerical precision w.r.t. forces - energy = energy * 0.001 - - ############################################################### - # Force estimation - ############################################################### - if self.regress_forces: - forces = self.force_block(x_pt, self.sphere_points) - - if self.show_timing_info is True: - torch.cuda.synchronize() - print( - "{} Time: {}\tMemory: {}\t{}".format( - self.counter, - time.time() - start_time, - len(data.pos), - torch.cuda.max_memory_allocated() / 1000000, - ) - ) - - self.counter = self.counter + 1 - - if not self.regress_forces: - return energy - else: - return energy, forces - - # Initialize the edge rotation matrics - def _init_edge_rot_mat(self, data, edge_index, edge_distance_vec): - edge_vec_0 = edge_distance_vec - edge_vec_0_distance = torch.sqrt(torch.sum(edge_vec_0**2, dim=1)) - - # Make sure the atoms are far enough apart - if torch.min(edge_vec_0_distance) < 0.0001: - print( - "Error edge_vec_0_distance: {}".format(torch.min(edge_vec_0_distance)) - ) - (minval, minidx) = torch.min(edge_vec_0_distance, 0) - print( - "Error edge_vec_0_distance: {} {} {} {} {}".format( - minidx, - edge_index[0, minidx], - edge_index[1, minidx], - data.pos[edge_index[0, minidx]], - data.pos[edge_index[1, minidx]], - ) - ) - - norm_x = edge_vec_0 / (edge_vec_0_distance.view(-1, 1)) - - edge_vec_2 = torch.rand_like(edge_vec_0) - 0.5 - edge_vec_2 = edge_vec_2 / ( - torch.sqrt(torch.sum(edge_vec_2**2, dim=1)).view(-1, 1) - ) - # Create two rotated copys of the random vectors in case the random vector is aligned with norm_x - # With two 90 degree rotated vectors, at least one should not be aligned with norm_x - edge_vec_2b = edge_vec_2.clone() - edge_vec_2b[:, 0] = -edge_vec_2[:, 1] - edge_vec_2b[:, 1] = edge_vec_2[:, 0] - edge_vec_2c = edge_vec_2.clone() - edge_vec_2c[:, 1] = -edge_vec_2[:, 2] - edge_vec_2c[:, 2] = edge_vec_2[:, 1] - vec_dot_b = torch.abs(torch.sum(edge_vec_2b * norm_x, dim=1)).view(-1, 1) - vec_dot_c = torch.abs(torch.sum(edge_vec_2c * norm_x, dim=1)).view(-1, 1) - - vec_dot = torch.abs(torch.sum(edge_vec_2 * norm_x, dim=1)).view(-1, 1) - edge_vec_2 = torch.where(torch.gt(vec_dot, vec_dot_b), edge_vec_2b, edge_vec_2) - vec_dot = torch.abs(torch.sum(edge_vec_2 * norm_x, dim=1)).view(-1, 1) - edge_vec_2 = torch.where(torch.gt(vec_dot, vec_dot_c), edge_vec_2c, edge_vec_2) - - vec_dot = torch.abs(torch.sum(edge_vec_2 * norm_x, dim=1)) - # Check the vectors aren't aligned - assert torch.max(vec_dot) < 0.99 - - norm_z = torch.cross(norm_x, edge_vec_2, dim=1) - norm_z = norm_z / (torch.sqrt(torch.sum(norm_z**2, dim=1, keepdim=True))) - norm_z = norm_z / (torch.sqrt(torch.sum(norm_z**2, dim=1)).view(-1, 1)) - norm_y = torch.cross(norm_x, norm_z, dim=1) - norm_y = norm_y / (torch.sqrt(torch.sum(norm_y**2, dim=1, keepdim=True))) - - # Construct the 3D rotation matrix - norm_x = norm_x.view(-1, 3, 1) - norm_y = -norm_y.view(-1, 3, 1) - norm_z = norm_z.view(-1, 3, 1) - - edge_rot_mat_inv = torch.cat([norm_z, norm_x, norm_y], dim=2) - edge_rot_mat = torch.transpose(edge_rot_mat_inv, 1, 2) - - return edge_rot_mat.detach() - - @property - def num_params(self): - return sum(p.numel() for p in self.parameters()) - - -class LayerBlock(torch.nn.Module): - """ - Layer block: Perform one layer (message passing and aggregation) of the GNN - - Args: - layer_idx (int): Layer number - sphere_channels (int): Number of spherical channels - hidden_channels (int): Number of hidden channels used during the SO(2) conv - edge_channels (int): Size of invariant edge embedding - lmax_list (list:int): List of degrees (l) for each resolution - mmax_list (list:int): List of orders (m) for each resolution - distance_expansion (func): Function used to compute distance embedding - max_num_elements (int): Maximum number of atomic numbers - SO3_grid (SO3_grid): Class used to convert from grid the spherical harmonic representations - act (function): Non-linear activation function - """ - - def __init__( - self, - layer_idx, - sphere_channels, - hidden_channels, - edge_channels, - lmax_list, - mmax_list, - distance_expansion, - max_num_elements, - SO3_grid, - act, - ): - super(LayerBlock, self).__init__() - self.layer_idx = layer_idx - self.act = act - self.lmax_list = lmax_list - self.mmax_list = mmax_list - self.num_resolutions = len(lmax_list) - self.sphere_channels = sphere_channels - self.sphere_channels_all = self.num_resolutions * self.sphere_channels - self.SO3_grid = SO3_grid - - # Message block - self.message_block = MessageBlock( - self.layer_idx, - self.sphere_channels, - hidden_channels, - edge_channels, - self.lmax_list, - self.mmax_list, - distance_expansion, - max_num_elements, - self.SO3_grid, - self.act, - ) - - # Non-linear point-wise comvolution for the aggregated messages - self.fc1_sphere = nn.Linear( - 2 * self.sphere_channels_all, self.sphere_channels_all, bias=False - ) - - self.fc2_sphere = nn.Linear( - self.sphere_channels_all, self.sphere_channels_all, bias=False - ) - - self.fc3_sphere = nn.Linear( - self.sphere_channels_all, self.sphere_channels_all, bias=False - ) - - def forward( - self, - x, - atomic_numbers, - edge_distance, - edge_index, - SO3_edge_rot, - mappingReduced, - ): - - # Compute messages by performing message block - x_message = self.message_block( - x, - atomic_numbers, - edge_distance, - edge_index, - SO3_edge_rot, - mappingReduced, - ) - - # Compute point-wise spherical non-linearity on aggregated messages - max_lmax = max(self.lmax_list) - - # Project to grid - x_grid_message = x_message.to_grid(self.SO3_grid, lmax=max_lmax) - x_grid = x.to_grid(self.SO3_grid, lmax=max_lmax) - x_grid = torch.cat([x_grid, x_grid_message], dim=3) - - # Perform point-wise convolution - x_grid = self.act(self.fc1_sphere(x_grid)) - x_grid = self.act(self.fc2_sphere(x_grid)) - x_grid = self.fc3_sphere(x_grid) - - # Project back to spherical harmonic coefficients - x_message._from_grid(x_grid, self.SO3_grid, lmax=max_lmax) - - # Return aggregated messages - return x_message - - -class MessageBlock(torch.nn.Module): - """ - Message block: Perform message passing - - Args: - layer_idx (int): Layer number - sphere_channels (int): Number of spherical channels - hidden_channels (int): Number of hidden channels used during the SO(2) conv - edge_channels (int): Size of invariant edge embedding - lmax_list (list:int): List of degrees (l) for each resolution - mmax_list (list:int): List of orders (m) for each resolution - distance_expansion (func): Function used to compute distance embedding - max_num_elements (int): Maximum number of atomic numbers - SO3_grid (SO3_grid): Class used to convert from grid the spherical harmonic representations - act (function): Non-linear activation function - """ - - def __init__( - self, - layer_idx, - sphere_channels, - hidden_channels, - edge_channels, - lmax_list, - mmax_list, - distance_expansion, - max_num_elements, - SO3_grid, - act, - ): - super(MessageBlock, self).__init__() - self.layer_idx = layer_idx - self.act = act - self.hidden_channels = hidden_channels - self.sphere_channels = sphere_channels - self.SO3_grid = SO3_grid - self.num_resolutions = len(lmax_list) - self.lmax_list = lmax_list - self.mmax_list = mmax_list - self.edge_channels = edge_channels - - # Create edge scalar (invariant to rotations) features - self.edge_block = EdgeBlock( - self.edge_channels, - distance_expansion, - max_num_elements, - self.act, - ) - - # Create SO(2) convolution blocks - self.so2_block_source = SO2Block( - self.sphere_channels, - self.hidden_channels, - self.edge_channels, - self.lmax_list, - self.mmax_list, - self.act, - ) - self.so2_block_target = SO2Block( - self.sphere_channels, - self.hidden_channels, - self.edge_channels, - self.lmax_list, - self.mmax_list, - self.act, - ) - - def forward( - self, - x, - atomic_numbers, - edge_distance, - edge_index, - SO3_edge_rot, - mappingReduced, - ): - ############################################################### - # Compute messages - ############################################################### - - # Compute edge scalar features (invariant to rotations) - # Uses atomic numbers and edge distance as inputs - x_edge = self.edge_block( - edge_distance, - atomic_numbers[edge_index[0]], # Source atom atomic number - atomic_numbers[edge_index[1]], # Target atom atomic number - ) - - # Copy embeddings for each edge's source and target nodes - x_source = x.clone() - x_target = x.clone() - x_source._expand_edge(edge_index[0, :]) - x_target._expand_edge(edge_index[1, :]) - - # Rotate the irreps to align with the edge - x_source._rotate(SO3_edge_rot, self.lmax_list, self.mmax_list) - x_target._rotate(SO3_edge_rot, self.lmax_list, self.mmax_list) - - # Compute messages - x_source = self.so2_block_source(x_source, x_edge, mappingReduced) - x_target = self.so2_block_target(x_target, x_edge, mappingReduced) - - # Add together the source and target results - x_target.embedding = x_source.embedding + x_target.embedding - - # Point-wise spherical non-linearity - x_target._grid_act(self.SO3_grid, self.act, mappingReduced) - - # Rotate back the irreps - x_target._rotate_inv(SO3_edge_rot, mappingReduced) - - # Compute the sum of the incoming neighboring messages for each target node - x_target._reduce_edge(edge_index[1], len(x.embedding)) - - return x_target - - -class SO2Block(torch.nn.Module): - """ - SO(2) Block: Perform SO(2) convolutions for all m (orders) - - Args: - sphere_channels (int): Number of spherical channels - hidden_channels (int): Number of hidden channels used during the SO(2) conv - edge_channels (int): Size of invariant edge embedding - lmax_list (list:int): List of degrees (l) for each resolution - mmax_list (list:int): List of orders (m) for each resolution - act (function): Non-linear activation function - """ - - def __init__( - self, - sphere_channels, - hidden_channels, - edge_channels, - lmax_list, - mmax_list, - act, - ): - super(SO2Block, self).__init__() - self.sphere_channels = sphere_channels - self.hidden_channels = hidden_channels - self.lmax_list = lmax_list - self.mmax_list = mmax_list - self.num_resolutions = len(lmax_list) - self.act = act - - num_channels_m0 = 0 - for i in range(self.num_resolutions): - num_coefficents = self.lmax_list[i] + 1 - num_channels_m0 = num_channels_m0 + num_coefficents * self.sphere_channels - - # SO(2) convolution for m=0 - self.fc1_dist0 = nn.Linear(edge_channels, self.hidden_channels) - self.fc1_m0 = nn.Linear(num_channels_m0, self.hidden_channels, bias=False) - self.fc2_m0 = nn.Linear(self.hidden_channels, num_channels_m0, bias=False) - - # SO(2) convolution for non-zero m - self.so2_conv = nn.ModuleList() - for m in range(1, max(self.mmax_list) + 1): - so2_conv = SO2Conv( - m, - self.sphere_channels, - self.hidden_channels, - edge_channels, - self.lmax_list, - self.mmax_list, - self.act, - ) - self.so2_conv.append(so2_conv) - - def forward( - self, - x, - x_edge, - mappingReduced, - ): - - num_edges = len(x_edge) - - # Reshape the spherical harmonics based on m (order) - x._m_primary(mappingReduced) - - # Compute m=0 coefficients separately since they only have real values (no imaginary) - - # Compute edge scalar features for m=0 - x_edge_0 = self.act(self.fc1_dist0(x_edge)) - - x_0 = x.embedding[:, 0 : mappingReduced.m_size[0]].contiguous() - x_0 = x_0.view(num_edges, -1) - - x_0 = self.fc1_m0(x_0) - x_0 = x_0 * x_edge_0 - x_0 = self.fc2_m0(x_0) - x_0 = x_0.view(num_edges, -1, x.num_channels) - - # Update the m=0 coefficients - x.embedding[:, 0 : mappingReduced.m_size[0]] = x_0 - - # Compute the values for the m > 0 coefficients - offset = mappingReduced.m_size[0] - for m in range(1, max(self.mmax_list) + 1): - # Get the m order coefficients - x_m = x.embedding[ - :, offset : offset + 2 * mappingReduced.m_size[m] - ].contiguous() - x_m = x_m.view(num_edges, 2, -1) - # Perform SO(2) convolution - x_m = self.so2_conv[m - 1](x_m, x_edge) - x_m = x_m.view(num_edges, -1, x.num_channels) - x.embedding[:, offset : offset + 2 * mappingReduced.m_size[m]] = x_m - - offset = offset + 2 * mappingReduced.m_size[m] - - # Reshape the spherical harmonics based on l (degree) - x._l_primary(mappingReduced) - - return x - - -class SO2Conv(torch.nn.Module): - """ - SO(2) Conv: Perform an SO(2) convolution - - Args: - m (int): Order of the spherical harmonic coefficients - sphere_channels (int): Number of spherical channels - hidden_channels (int): Number of hidden channels used during the SO(2) conv - edge_channels (int): Size of invariant edge embedding - lmax_list (list:int): List of degrees (l) for each resolution - mmax_list (list:int): List of orders (m) for each resolution - act (function): Non-linear activation function - """ - - def __init__( - self, - m, - sphere_channels, - hidden_channels, - edge_channels, - lmax_list, - mmax_list, - act, - ): - super(SO2Conv, self).__init__() - self.hidden_channels = hidden_channels - self.lmax_list = lmax_list - self.mmax_list = mmax_list - self.sphere_channels = sphere_channels - self.num_resolutions = len(self.lmax_list) - self.m = m - self.act = act - - num_channels = 0 - for i in range(self.num_resolutions): - num_coefficents = 0 - if self.mmax_list[i] >= m: - num_coefficents = self.lmax_list[i] - m + 1 - - num_channels = num_channels + num_coefficents * self.sphere_channels - - assert num_channels > 0 - - # Embedding function of the distance - self.fc1_dist = nn.Linear(edge_channels, 2 * self.hidden_channels) - - # Real weights of SO(2) convolution - self.fc1_r = nn.Linear(num_channels, self.hidden_channels, bias=False) - self.fc2_r = nn.Linear(self.hidden_channels, num_channels, bias=False) - - # Imaginary weights of SO(2) convolution - self.fc1_i = nn.Linear(num_channels, self.hidden_channels, bias=False) - self.fc2_i = nn.Linear(self.hidden_channels, num_channels, bias=False) - - def forward(self, x_m, x_edge): - # Compute edge scalar features - x_edge = self.act(self.fc1_dist(x_edge)) - x_edge = x_edge.view(-1, 2, self.hidden_channels) - - # Perform the complex weight multiplication - x_r = self.fc1_r(x_m) - x_r = x_r * x_edge[:, 0:1, :] - x_r = self.fc2_r(x_r) - - x_i = self.fc1_i(x_m) - x_i = x_i * x_edge[:, 1:2, :] - x_i = self.fc2_i(x_i) - - x_m_r = x_r[:, 0] - x_i[:, 1] - x_m_i = x_r[:, 1] + x_i[:, 0] - - return torch.stack((x_m_r, x_m_i), dim=1).contiguous() - - -class EdgeBlock(torch.nn.Module): - """ - Edge Block: Compute invariant edge representation from edge diatances and atomic numbers - - Args: - edge_channels (int): Size of invariant edge embedding - distance_expansion (func): Function used to compute distance embedding - max_num_elements (int): Maximum number of atomic numbers - act (function): Non-linear activation function - """ - - def __init__( - self, - edge_channels, - distance_expansion, - max_num_elements, - act, - ): - super(EdgeBlock, self).__init__() - self.in_channels = distance_expansion.num_output - self.distance_expansion = distance_expansion - self.act = act - self.edge_channels = edge_channels - self.max_num_elements = max_num_elements - - # Embedding function of the distance - self.fc1_dist = nn.Linear(self.in_channels, self.edge_channels) - - # Embedding function of the atomic numbers - self.source_embedding = nn.Embedding(self.max_num_elements, self.edge_channels) - self.target_embedding = nn.Embedding(self.max_num_elements, self.edge_channels) - nn.init.uniform_(self.source_embedding.weight.data, -0.001, 0.001) - nn.init.uniform_(self.target_embedding.weight.data, -0.001, 0.001) - - # Embedding function of the edge - self.fc1_edge_attr = nn.Linear( - self.edge_channels, - self.edge_channels, - ) - - def forward(self, edge_distance, source_element, target_element): - - # Compute distance embedding - x_dist = self.distance_expansion(edge_distance) - x_dist = self.fc1_dist(x_dist) - - # Compute atomic number embeddings - source_embedding = self.source_embedding(source_element) - target_embedding = self.target_embedding(target_element) - - # Compute invariant edge embedding - x_edge = self.act(source_embedding + target_embedding + x_dist) - x_edge = self.act(self.fc1_edge_attr(x_edge)) - - return x_edge - - -class EnergyBlock(torch.nn.Module): - """ - Energy Block: Output block computing the energy - - Args: - num_channels (int): Number of channels - num_sphere_samples (int): Number of samples used to approximate the integral on the sphere - act (function): Non-linear activation function - """ - - def __init__( - self, - num_channels, - num_sphere_samples, - act, - ): - super(EnergyBlock, self).__init__() - self.num_channels = num_channels - self.num_sphere_samples = num_sphere_samples - self.act = act - - self.fc1 = nn.Linear(self.num_channels, self.num_channels) - self.fc2 = nn.Linear(self.num_channels, self.num_channels) - self.fc3 = nn.Linear(self.num_channels, 1, bias=False) - - def forward(self, x_pt): - # x_pt are the values of the channels sampled at different points on the sphere - x_pt = self.act(self.fc1(x_pt)) - x_pt = self.act(self.fc2(x_pt)) - x_pt = self.fc3(x_pt) - x_pt = x_pt.view(-1, self.num_sphere_samples, 1) - node_energy = torch.sum(x_pt, dim=1) / self.num_sphere_samples - - return node_energy - - -class ForceBlock(torch.nn.Module): - """ - Force Block: Output block computing the per atom forces - - Args: - num_channels (int): Number of channels - num_sphere_samples (int): Number of samples used to approximate the integral on the sphere - act (function): Non-linear activation function - """ - - def __init__( - self, - num_channels, - num_sphere_samples, - act, - ): - super(ForceBlock, self).__init__() - self.num_channels = num_channels - self.num_sphere_samples = num_sphere_samples - self.act = act - - self.fc1 = nn.Linear(self.num_channels, self.num_channels) - self.fc2 = nn.Linear(self.num_channels, self.num_channels) - self.fc3 = nn.Linear(self.num_channels, 1, bias=False) - - def forward(self, x_pt, sphere_points): - # x_pt are the values of the channels sampled at different points on the sphere - x_pt = self.act(self.fc1(x_pt)) - x_pt = self.act(self.fc2(x_pt)) - x_pt = self.fc3(x_pt) - x_pt = x_pt.view(-1, self.num_sphere_samples, 1) - forces = x_pt * sphere_points.view(1, self.num_sphere_samples, 3) - forces = torch.sum(forces, dim=1) / self.num_sphere_samples - - return forces diff --git a/ocpmodels/models/escn/so3.py b/ocpmodels/models/escn/so3.py deleted file mode 100644 index f1d1cee7404341220e0e0eaeaadac0d9562fb6ed..0000000000000000000000000000000000000000 --- a/ocpmodels/models/escn/so3.py +++ /dev/null @@ -1,528 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import os - -import torch -import torch.nn as nn - -try: - from e3nn import o3 - from e3nn.o3 import FromS2Grid, ToS2Grid -except ImportError: - pass - -# Borrowed from e3nn @ 0.4.0: -# https://github.com/e3nn/e3nn/blob/0.4.0/e3nn/o3/_wigner.py#L10 -# _Jd is a list of tensors of shape (2l+1, 2l+1) -_Jd = torch.load(os.path.join(os.path.dirname(__file__), "Jd.pt")) - - -class CoefficientMapping: - """ - Helper functions for coefficients used to reshape l<-->m and to get coefficients of specific degree or order - - Args: - lmax_list (list:int): List of maximum degree of the spherical harmonics - mmax_list (list:int): List of maximum order of the spherical harmonics - device: Device of the output - """ - - def __init__( - self, - lmax_list, - mmax_list, - device, - ): - super().__init__() - - self.lmax_list = lmax_list - self.mmax_list = mmax_list - self.num_resolutions = len(lmax_list) - self.device = device - - # Compute the degree (l) and order (m) for each - # entry of the embedding - - self.l_harmonic = torch.tensor([], device=self.device).long() - self.m_harmonic = torch.tensor([], device=self.device).long() - self.m_complex = torch.tensor([], device=self.device).long() - - self.res_size = torch.zeros([self.num_resolutions], device=self.device).long() - offset = 0 - for i in range(self.num_resolutions): - for l in range(0, self.lmax_list[i] + 1): - mmax = min(self.mmax_list[i], l) - m = torch.arange(-mmax, mmax + 1, device=self.device).long() - self.m_complex = torch.cat([self.m_complex, m], dim=0) - self.m_harmonic = torch.cat( - [self.m_harmonic, torch.abs(m).long()], dim=0 - ) - self.l_harmonic = torch.cat([self.l_harmonic, m.fill_(l).long()], dim=0) - self.res_size[i] = len(self.l_harmonic) - offset - offset = len(self.l_harmonic) - - num_coefficients = len(self.l_harmonic) - self.to_m = torch.zeros( - [num_coefficients, num_coefficients], device=self.device - ) - self.m_size = torch.zeros([max(self.mmax_list) + 1], device=self.device).long() - - # The following is implemented poorly - very slow. It only gets called - # a few times so haven't optimized. - offset = 0 - for m in range(max(self.mmax_list) + 1): - idx_r, idx_i = self.complex_idx(m) - - for idx_out, idx_in in enumerate(idx_r): - self.to_m[idx_out + offset, idx_in] = 1.0 - offset = offset + len(idx_r) - self.m_size[m] = int(len(idx_r)) - - for idx_out, idx_in in enumerate(idx_i): - self.to_m[idx_out + offset, idx_in] = 1.0 - offset = offset + len(idx_i) - - self.to_m = self.to_m.detach() - - # Return mask containing coefficients of order m (real and imaginary parts) - def complex_idx(self, m, lmax=-1): - if lmax == -1: - lmax = max(self.lmax_list) - - indices = torch.arange(len(self.l_harmonic), device=self.device) - # Real part - mask_r = torch.bitwise_and(self.l_harmonic.le(lmax), self.m_complex.eq(m)) - mask_idx_r = torch.masked_select(indices, mask_r) - - mask_idx_i = torch.tensor([], device=self.device).long() - # Imaginary part - if m != 0: - mask_i = torch.bitwise_and(self.l_harmonic.le(lmax), self.m_complex.eq(-m)) - mask_idx_i = torch.masked_select(indices, mask_i) - - return mask_idx_r, mask_idx_i - - # Return mask containing coefficients less than or equal to degree (l) and order (m) - def coefficient_idx(self, lmax, mmax): - mask = torch.bitwise_and(self.l_harmonic.le(lmax), self.m_harmonic.le(mmax)) - indices = torch.arange(len(mask), device=self.device) - - return torch.masked_select(indices, mask) - - -class SO3_Embedding(torch.nn.Module): - """ - Helper functions for irreps embedding - - Args: - length (int): Batch size - lmax_list (list:int): List of maximum degree of the spherical harmonics - num_channels (int): Number of channels - device: Device of the output - dtype: type of the output tensors - """ - - def __init__( - self, - length, - lmax_list, - num_channels, - device, - dtype, - ): - super().__init__() - self.num_channels = num_channels - self.device = device - self.dtype = dtype - self.num_resolutions = len(lmax_list) - - self.num_coefficients = 0 - for i in range(self.num_resolutions): - self.num_coefficients = self.num_coefficients + int((lmax_list[i] + 1) ** 2) - - embedding = torch.zeros( - length, - self.num_coefficients, - self.num_channels, - device=self.device, - dtype=self.dtype, - ) - - self.set_embedding(embedding) - self.set_lmax_mmax(lmax_list, lmax_list.copy()) - - # Clone an embedding of irreps - def clone(self): - clone = SO3_Embedding( - 0, - self.lmax_list.copy(), - self.num_channels, - self.device, - self.dtype, - ) - - clone.set_embedding(self.embedding.clone()) - - return clone - - # Initialize an embedding of irreps - def set_embedding(self, embedding): - self.length = len(embedding) - self.embedding = embedding - - # Set the maximum order to be the maximum degree - def set_lmax_mmax(self, lmax_list, mmax_list): - self.lmax_list = lmax_list - self.mmax_list = mmax_list - - # Expand the node embeddings to the number of edges - def _expand_edge(self, edge_index): - embedding = self.embedding[edge_index] - self.set_embedding(embedding) - - # Initialize an embedding of irreps of a neighborhood - def expand_edge(self, edge_index): - x_expand = SO3_Embedding( - 0, - self.lmax_list.copy(), - self.num_channels, - self.device, - self.dtype, - ) - x_expand.set_embedding(self.embedding[edge_index]) - return x_expand - - # Compute the sum of the embeddings of the neighborhood - def _reduce_edge(self, edge_index, num_nodes): - new_embedding = torch.zeros( - num_nodes, - self.num_coefficients, - self.num_channels, - device=self.embedding.device, - dtype=self.embedding.dtype, - ) - new_embedding.index_add_(0, edge_index, self.embedding) - self.set_embedding(new_embedding) - - # Reshape the embedding l-->m - def _m_primary(self, mapping): - self.embedding = torch.einsum("nac,ba->nbc", self.embedding, mapping.to_m) - - # Reshape the embedding m-->l - def _l_primary(self, mapping): - self.embedding = torch.einsum("nac,ab->nbc", self.embedding, mapping.to_m) - - # Rotate the embedding - def _rotate(self, SO3_rotation, lmax_list, mmax_list): - embedding_rotate = torch.tensor([], device=self.device, dtype=self.dtype) - - offset = 0 - for i in range(self.num_resolutions): - num_coefficients = int((self.lmax_list[i] + 1) ** 2) - embedding_i = self.embedding[:, offset : offset + num_coefficients] - embedding_rotate = torch.cat( - [ - embedding_rotate, - SO3_rotation[i].rotate(embedding_i, lmax_list[i], mmax_list[i]), - ], - dim=1, - ) - offset = offset + num_coefficients - - self.embedding = embedding_rotate - self.set_lmax_mmax(lmax_list.copy(), mmax_list.copy()) - - # Rotate the embedding by the inverse of the rotation matrix - def _rotate_inv(self, SO3_rotation, mappingReduced): - embedding_rotate = torch.tensor([], device=self.device, dtype=self.dtype) - - offset = 0 - for i in range(self.num_resolutions): - num_coefficients = mappingReduced.res_size[i] - embedding_i = self.embedding[:, offset : offset + num_coefficients] - embedding_rotate = torch.cat( - [ - embedding_rotate, - SO3_rotation[i].rotate_inv( - embedding_i, self.lmax_list[i], self.mmax_list[i] - ), - ], - dim=1, - ) - offset = offset + num_coefficients - - self.embedding = embedding_rotate - - # Assume mmax = lmax when rotating back - for i in range(self.num_resolutions): - self.mmax_list[i] = int(self.lmax_list[i]) - - self.set_lmax_mmax(self.lmax_list, self.mmax_list) - - # Compute point-wise spherical non-linearity - def _grid_act(self, SO3_grid, act, mappingReduced): - offset = 0 - for i in range(self.num_resolutions): - - num_coefficients = mappingReduced.res_size[i] - - x_res = self.embedding[:, offset : offset + num_coefficients].contiguous() - to_grid_mat = SO3_grid[self.lmax_list[i]][ - self.mmax_list[i] - ].get_to_grid_mat(self.device) - from_grid_mat = SO3_grid[self.lmax_list[i]][ - self.mmax_list[i] - ].get_from_grid_mat(self.device) - - x_grid = torch.einsum("bai,zic->zbac", to_grid_mat, x_res) - x_grid = act(x_grid) - x_res = torch.einsum("bai,zbac->zic", from_grid_mat, x_grid) - - self.embedding[:, offset : offset + num_coefficients] = x_res - offset = offset + num_coefficients - - # Compute a sample of the grid - def to_grid(self, SO3_grid, lmax=-1): - if lmax == -1: - lmax = max(self.lmax_list) - - to_grid_mat_lmax = SO3_grid[lmax][lmax].get_to_grid_mat(self.device) - grid_mapping = SO3_grid[lmax][lmax].mapping - - offset = 0 - x_grid = torch.tensor([], device=self.device) - - for i in range(self.num_resolutions): - num_coefficients = int((self.lmax_list[i] + 1) ** 2) - x_res = self.embedding[:, offset : offset + num_coefficients].contiguous() - to_grid_mat = to_grid_mat_lmax[ - :, - :, - grid_mapping.coefficient_idx(self.lmax_list[i], self.lmax_list[i]), - ] - x_grid = torch.cat( - [x_grid, torch.einsum("bai,zic->zbac", to_grid_mat, x_res)], - dim=3, - ) - offset = offset + num_coefficients - - return x_grid - - # Compute irreps from grid representation - def _from_grid(self, x_grid, SO3_grid, lmax=-1): - if lmax == -1: - lmax = max(self.lmax_list) - - from_grid_mat_lmax = SO3_grid[lmax][lmax].get_from_grid_mat(self.device) - grid_mapping = SO3_grid[lmax][lmax].mapping - - offset = 0 - offset_channel = 0 - for i in range(self.num_resolutions): - from_grid_mat = from_grid_mat_lmax[ - :, - :, - grid_mapping.coefficient_idx(self.lmax_list[i], self.lmax_list[i]), - ] - x_res = torch.einsum( - "bai,zbac->zic", - from_grid_mat, - x_grid[ - :, - :, - :, - offset_channel : offset_channel + self.num_channels, - ], - ) - num_coefficients = int((self.lmax_list[i] + 1) ** 2) - self.embedding[:, offset : offset + num_coefficients] = x_res - offset = offset + num_coefficients - offset_channel = offset_channel + self.num_channels - - -class SO3_Rotation(torch.nn.Module): - """ - Helper functions for Wigner-D rotations - - Args: - rot_mat3x3 (tensor): Rotation matrix - lmax_list (list:int): List of maximum degree of the spherical harmonics - """ - - def __init__( - self, - rot_mat3x3, - lmax, - ): - super().__init__() - self.device = rot_mat3x3.device - self.dtype = rot_mat3x3.dtype - - length = len(rot_mat3x3) - - self.wigner = self.RotationToWignerDMatrix(rot_mat3x3, 0, lmax) - self.wigner_inv = torch.transpose(self.wigner, 1, 2).contiguous() - - self.wigner = self.wigner.detach() - self.wigner_inv = self.wigner_inv.detach() - - self.set_lmax(lmax) - - # Initialize coefficients for reshape l<-->m - def set_lmax(self, lmax): - self.lmax = lmax - self.mapping = CoefficientMapping([self.lmax], [self.lmax], self.device) - - # Rotate the embedding - def rotate(self, embedding, out_lmax, out_mmax): - out_mask = self.mapping.coefficient_idx(out_lmax, out_mmax) - wigner = self.wigner[:, out_mask, :] - return torch.bmm(wigner, embedding) - - # Rotate the embedding by the inverse of the rotation matrix - def rotate_inv(self, embedding, in_lmax, in_mmax): - in_mask = self.mapping.coefficient_idx(in_lmax, in_mmax) - wigner_inv = self.wigner_inv[:, :, in_mask] - - return torch.bmm(wigner_inv, embedding) - - # Compute Wigner matrices from rotation matrix - def RotationToWignerDMatrix(self, edge_rot_mat, start_lmax, end_lmax): - x = edge_rot_mat @ edge_rot_mat.new_tensor([0.0, 1.0, 0.0]) - alpha, beta = o3.xyz_to_angles(x) - R = ( - o3.angles_to_matrix(alpha, beta, torch.zeros_like(alpha)).transpose(-1, -2) - @ edge_rot_mat - ) - gamma = torch.atan2(R[..., 0, 2], R[..., 0, 0]) - - size = (end_lmax + 1) ** 2 - (start_lmax) ** 2 - wigner = torch.zeros(len(alpha), size, size, device=self.device) - start = 0 - for lmax in range(start_lmax, end_lmax + 1): - block = self.wigner_D(lmax, alpha, beta, gamma) - end = start + block.size()[1] - wigner[:, start:end, start:end] = block - start = end - - return wigner.detach() - - # Borrowed from e3nn @ 0.4.0: - # https://github.com/e3nn/e3nn/blob/0.4.0/e3nn/o3/_wigner.py#L37 - # - # In 0.5.0, e3nn shifted to torch.matrix_exp which is significantly slower: - # https://github.com/e3nn/e3nn/blob/0.5.0/e3nn/o3/_wigner.py#L92 - def wigner_D(self, l, alpha, beta, gamma): - if not l < len(_Jd): - raise NotImplementedError( - f"wigner D maximum l implemented is {len(_Jd) - 1}, send us an email to ask for more" - ) - - alpha, beta, gamma = torch.broadcast_tensors(alpha, beta, gamma) - J = _Jd[l].to(dtype=alpha.dtype, device=alpha.device) - Xa = self._z_rot_mat(alpha, l) - Xb = self._z_rot_mat(beta, l) - Xc = self._z_rot_mat(gamma, l) - return Xa @ J @ Xb @ J @ Xc - - def _z_rot_mat(self, angle, l): - shape, device, dtype = angle.shape, angle.device, angle.dtype - M = angle.new_zeros((*shape, 2 * l + 1, 2 * l + 1)) - inds = torch.arange(0, 2 * l + 1, 1, device=device) - reversed_inds = torch.arange(2 * l, -1, -1, device=device) - frequencies = torch.arange(l, -l - 1, -1, dtype=dtype, device=device) - M[..., inds, reversed_inds] = torch.sin(frequencies * angle[..., None]) - M[..., inds, inds] = torch.cos(frequencies * angle[..., None]) - return M - - -class SO3_Grid(torch.nn.Module): - """ - Helper functions for grid representation of the irreps - - Args: - lmax (int): Maximum degree of the spherical harmonics - mmax (int): Maximum order of the spherical harmonics - """ - - def __init__( - self, - lmax, - mmax, - ): - super().__init__() - self.lmax = lmax - self.mmax = mmax - self.lat_resolution = 2 * (self.lmax + 1) - if lmax == mmax: - self.long_resolution = 2 * (self.mmax + 1) + 1 - else: - self.long_resolution = 2 * (self.mmax) + 1 - - self.initialized = False - - def _initialize(self, device): - if self.initialized is True: - return - self.mapping = CoefficientMapping([self.lmax], [self.lmax], device) - - to_grid = ToS2Grid( - self.lmax, - (self.lat_resolution, self.long_resolution), - normalization="integral", - device=device, - ) - - self.to_grid_mat = torch.einsum( - "mbi,am->bai", to_grid.shb, to_grid.sha - ).detach() - self.to_grid_mat = self.to_grid_mat[ - :, :, self.mapping.coefficient_idx(self.lmax, self.mmax) - ] - - from_grid = FromS2Grid( - (self.lat_resolution, self.long_resolution), - self.lmax, - normalization="integral", - device=device, - ) - - self.from_grid_mat = torch.einsum( - "am,mbi->bai", from_grid.sha, from_grid.shb - ).detach() - self.from_grid_mat = self.from_grid_mat[ - :, :, self.mapping.coefficient_idx(self.lmax, self.mmax) - ] - - self.initialized = True - - # Compute matrices to transform irreps to grid - def get_to_grid_mat(self, device): - self._initialize(device) - return self.to_grid_mat - - # Compute matrices to transform grid to irreps - def get_from_grid_mat(self, device): - self._initialize(device) - return self.from_grid_mat - - # Compute grid from irreps representation - def to_grid(self, embedding, lmax, mmax): - self._initialize(embedding.device) - to_grid_mat = self.to_grid_mat[:, :, self.mapping.coefficient_idx(lmax, mmax)] - grid = torch.einsum("bai,zic->zbac", to_grid_mat, embedding) - return grid - - # Compute irreps from grid representation - def from_grid(self, grid, lmax, mmax): - self._initialize(grid.device) - from_grid_mat = self.from_grid_mat[ - :, :, self.mapping.coefficient_idx(lmax, mmax) - ] - embedding = torch.einsum("bai,zbac->zic", from_grid_mat, grid) - return embedding diff --git a/ocpmodels/models/forcenet.py b/ocpmodels/models/forcenet.py deleted file mode 100644 index 17abec790b09bbf71f063eb8c5094df9fdd61bfd..0000000000000000000000000000000000000000 --- a/ocpmodels/models/forcenet.py +++ /dev/null @@ -1,501 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import os -from math import pi as PI - -import numpy as np -import torch -import torch.nn as nn -from torch_geometric.nn import MessagePassing -from torch_scatter import scatter - -from ocpmodels.common.registry import registry -from ocpmodels.common.utils import get_pbc_distances, radius_graph_pbc -from ocpmodels.datasets.embeddings import ATOMIC_RADII, CONTINUOUS_EMBEDDINGS -from ocpmodels.models.base import BaseModel -from ocpmodels.models.utils.activations import Act -from ocpmodels.models.utils.basis import Basis, SphericalSmearing - - -class FNDecoder(nn.Module): - def __init__(self, decoder_type, decoder_activation_str, output_dim): - super(FNDecoder, self).__init__() - self.decoder_type = decoder_type - self.decoder_activation = Act(decoder_activation_str) - self.output_dim = output_dim - - if self.decoder_type == "linear": - self.decoder = nn.Sequential(nn.Linear(self.output_dim, 3)) - elif self.decoder_type == "mlp": - self.decoder = nn.Sequential( - nn.Linear(self.output_dim, self.output_dim), - nn.BatchNorm1d(self.output_dim), - self.decoder_activation, - nn.Linear(self.output_dim, 3), - ) - else: - raise ValueError(f"Undefined force decoder: {self.decoder_type}") - - self.reset_parameters() - - def reset_parameters(self): - for m in self.decoder: - if isinstance(m, nn.Linear): - nn.init.xavier_uniform_(m.weight) - m.bias.data.fill_(0) - - def forward(self, x): - return self.decoder(x) - - -class InteractionBlock(MessagePassing): - def __init__( - self, - hidden_channels, - mlp_basis_dim, - basis_type, - depth_mlp_edge=2, - depth_mlp_trans=1, - activation_str="ssp", - ablation="none", - ): - super(InteractionBlock, self).__init__(aggr="add") - - self.activation = Act(activation_str) - self.ablation = ablation - self.basis_type = basis_type - - # basis function assumes input is in the range of [-1,1] - if self.basis_type != "rawcat": - self.lin_basis = torch.nn.Linear(mlp_basis_dim, hidden_channels) - - if self.ablation == "nocond": - # the edge filter only depends on edge_attr - in_features = ( - mlp_basis_dim if self.basis_type == "rawcat" else hidden_channels - ) - else: - # edge filter depends on edge_attr and current node embedding - in_features = ( - mlp_basis_dim + 2 * hidden_channels - if self.basis_type == "rawcat" - else 3 * hidden_channels - ) - - if depth_mlp_edge > 0: - mlp_edge = [torch.nn.Linear(in_features, hidden_channels)] - for i in range(depth_mlp_edge): - mlp_edge.append(self.activation) - mlp_edge.append(torch.nn.Linear(hidden_channels, hidden_channels)) - else: - ## need batch normalization afterwards. Otherwise training is unstable. - mlp_edge = [ - torch.nn.Linear(in_features, hidden_channels), - torch.nn.BatchNorm1d(hidden_channels), - ] - self.mlp_edge = torch.nn.Sequential(*mlp_edge) - - if not self.ablation == "nofilter": - self.lin = torch.nn.Linear(hidden_channels, hidden_channels) - - if depth_mlp_trans > 0: - mlp_trans = [torch.nn.Linear(hidden_channels, hidden_channels)] - for i in range(depth_mlp_trans): - mlp_trans.append(torch.nn.BatchNorm1d(hidden_channels)) - mlp_trans.append(self.activation) - mlp_trans.append(torch.nn.Linear(hidden_channels, hidden_channels)) - else: - # need batch normalization afterwards. Otherwise, becomes NaN - mlp_trans = [ - torch.nn.Linear(hidden_channels, hidden_channels), - torch.nn.BatchNorm1d(hidden_channels), - ] - - self.mlp_trans = torch.nn.Sequential(*mlp_trans) - - if not self.ablation == "noself": - self.center_W = torch.nn.Parameter(torch.Tensor(1, hidden_channels)) - - self.reset_parameters() - - def reset_parameters(self): - if self.basis_type != "rawcat": - torch.nn.init.xavier_uniform_(self.lin_basis.weight) - self.lin_basis.bias.data.fill_(0) - - for m in self.mlp_trans: - if isinstance(m, torch.nn.Linear): - torch.nn.init.xavier_uniform_(m.weight) - m.bias.data.fill_(0) - - for m in self.mlp_edge: - if isinstance(m, torch.nn.Linear): - torch.nn.init.xavier_uniform_(m.weight) - m.bias.data.fill_(0) - - if not self.ablation == "nofilter": - torch.nn.init.xavier_uniform_(self.lin.weight) - self.lin.bias.data.fill_(0) - - if not self.ablation == "noself": - torch.nn.init.xavier_uniform_(self.center_W) - - def forward(self, x, edge_index, edge_attr, edge_weight): - if self.basis_type != "rawcat": - edge_emb = self.lin_basis(edge_attr) - else: - # for rawcat, we directly use the raw feature - edge_emb = edge_attr - - if self.ablation == "nocond": - emb = edge_emb - else: - emb = torch.cat([edge_emb, x[edge_index[0]], x[edge_index[1]]], dim=1) - - W = self.mlp_edge(emb) * edge_weight.view(-1, 1) - if self.ablation == "nofilter": - x = self.propagate(edge_index, x=x, W=W) + self.center_W - else: - x = self.lin(x) - if self.ablation == "noself": - x = self.propagate(edge_index, x=x, W=W) - else: - x = self.propagate(edge_index, x=x, W=W) + self.center_W * x - x = self.mlp_trans(x) - - return x - - def message(self, x_j, W): - if self.ablation == "nofilter": - return W - else: - return x_j * W - - -# flake8: noqa: C901 -@registry.register_model("forcenet") -class ForceNet(BaseModel): - r"""Implementation of ForceNet architecture. - - Args: - num_atoms (int): Unused argument - bond_feat_dim (int): Unused argument - num_targets (int): Unused argumebt - hidden_channels (int, optional): Number of hidden channels. - (default: :obj:`512`) - num_iteractions (int, optional): Number of interaction blocks. - (default: :obj:`5`) - cutoff (float, optional): Cutoff distance for interatomic interactions. - (default: :obj:`6.0`) - feat (str, optional): Input features to be used - (default: :obj:`full`) - num_freqs (int, optional): Number of frequencies for basis function. - (default: :obj:`50`) - max_n (int, optional): Maximum order of spherical harmonics. - (default: :obj:`6`) - basis (str, optional): Basis function to be used. - (default: :obj:`full`) - depth_mlp_edge (int, optional): Depth of MLP for edges in interaction blocks. - (default: :obj:`2`) - depth_mlp_node (int, optional): Depth of MLP for nodes in interaction blocks. - (default: :obj:`1`) - activation_str (str, optional): Activation function used post linear layer in all message passing MLPs. - (default: :obj:`swish`) - ablation (str, optional): Type of ablation to be performed. - (default: :obj:`none`) - decoder_hidden_channels (int, optional): Number of hidden channels in the decoder. - (default: :obj:`512`) - decoder_type (str, optional): Type of decoder: linear or MLP. - (default: :obj:`mlp`) - decoder_activation_str (str, optional): Activation function used post linear layer in decoder. - (default: :obj:`swish`) - training (bool, optional): If set to :obj:`True`, specify training phase. - (default: :obj:`True`) - otf_graph (bool, optional): If set to :obj:`True`, compute graph edges on the fly. - (default: :obj:`False`) - """ - - def __init__( - self, - num_atoms, # not used - bond_feat_dim, # not used - num_targets, # not used - hidden_channels=512, - num_interactions=5, - cutoff=6.0, - feat="full", - num_freqs=50, - max_n=3, - basis="sphallmul", - depth_mlp_edge=2, - depth_mlp_node=1, - activation_str="swish", - ablation="none", - decoder_hidden_channels=512, - decoder_type="mlp", - decoder_activation_str="swish", - training=True, - otf_graph=False, - use_pbc=True, - ): - - super(ForceNet, self).__init__() - self.training = training - self.ablation = ablation - if self.ablation not in [ - "none", - "nofilter", - "nocond", - "nodistlist", - "onlydist", - "nodelinear", - "edgelinear", - "noself", - ]: - raise ValueError(f"Unknown ablation called {ablation}.") - - """ - Descriptions of ablations: - - none: base ForceNet model - - nofilter: no element-wise filter parameterization in message modeling - - nocond: convolutional filter is only conditioned on edge features, not node embeddings - - nodistlist: no atomic radius information in edge features - - onlydist: edge features only contains distance information. Orientation information is ommited. - - nodelinear: node update MLP function is replaced with linear function followed by batch normalization - - edgelinear: edge MLP transformation function is replaced with linear function followed by batch normalization. - - noself: no self edge of m_t. - """ - - self.otf_graph = otf_graph - self.cutoff = cutoff - self.output_dim = decoder_hidden_channels - self.feat = feat - self.num_freqs = num_freqs - self.num_layers = num_interactions - self.max_n = max_n - self.activation_str = activation_str - self.use_pbc = use_pbc - self.max_neighbors = 50 - - if self.ablation == "edgelinear": - depth_mlp_edge = 0 - - if self.ablation == "nodelinear": - depth_mlp_node = 0 - - # read atom map and atom radii - atom_map = torch.zeros(101, 9) - for i in range(101): - atom_map[i] = torch.tensor(CONTINUOUS_EMBEDDINGS[i]) - - atom_radii = torch.zeros(101) - for i in range(101): - atom_radii[i] = ATOMIC_RADII[i] - atom_radii = atom_radii / 100 - - self.atom_radii = nn.Parameter(atom_radii, requires_grad=False) - self.basis_type = basis - - self.pbc_apply_sph_harm = "sph" in self.basis_type - self.pbc_sph_option = None - - # for spherical harmonics for PBC - if "sphall" in self.basis_type: - self.pbc_sph_option = "all" - elif "sphsine" in self.basis_type: - self.pbc_sph_option = "sine" - elif "sphcosine" in self.basis_type: - self.pbc_sph_option = "cosine" - - self.pbc_sph = None - if self.pbc_apply_sph_harm: - self.pbc_sph = SphericalSmearing( - max_n=self.max_n, option=self.pbc_sph_option - ) - - # self.feat can be "simple" or "full" - if self.feat == "simple": - self.embedding = nn.Embedding(100, hidden_channels) - - # set up dummy atom_map that only contains atomic_number information - atom_map = torch.linspace(0, 1, 101).view(-1, 1).repeat(1, 9) - self.atom_map = nn.Parameter(atom_map, requires_grad=False) - - elif self.feat == "full": - # Normalize along each dimaension - atom_map[0] = np.nan - atom_map_notnan = atom_map[atom_map[:, 0] == atom_map[:, 0]] - atom_map_min = torch.min(atom_map_notnan, dim=0)[0] - atom_map_max = torch.max(atom_map_notnan, dim=0)[0] - atom_map_gap = atom_map_max - atom_map_min - - ## squash to [0,1] - atom_map = (atom_map - atom_map_min.view(1, -1)) / atom_map_gap.view(1, -1) - - self.atom_map = torch.nn.Parameter(atom_map, requires_grad=False) - - in_features = 9 - # first apply basis function and then linear function - if "sph" in self.basis_type: - # spherical basis is only meaningful for edge feature, so use powersine instead - node_basis_type = "powersine" - else: - node_basis_type = self.basis_type - basis = Basis( - in_features, - num_freqs=num_freqs, - basis_type=node_basis_type, - act=self.activation_str, - ) - self.embedding = torch.nn.Sequential( - basis, torch.nn.Linear(basis.out_dim, hidden_channels) - ) - - else: - raise ValueError("Undefined feature type for atom") - - # process basis function for edge feature - if self.ablation == "nodistlist": - # do not consider additional distance edge features - # normalized (x,y,z) + distance - in_feature = 4 - elif self.ablation == "onlydist": - # only consider distance-based edge features - # ignore normalized (x,y,z) - in_feature = 4 - - # if basis_type is spherical harmonics, then reduce to powersine - if "sph" in self.basis_type: - logging.info( - "Under onlydist ablation, spherical basis is reduced to powersine basis." - ) - self.basis_type = "powersine" - self.pbc_sph = None - - else: - in_feature = 7 - self.basis_fun = Basis( - in_feature, - num_freqs, - self.basis_type, - self.activation_str, - sph=self.pbc_sph, - ) - - # process interaction blocks - self.interactions = torch.nn.ModuleList() - for _ in range(num_interactions): - block = InteractionBlock( - hidden_channels, - self.basis_fun.out_dim, - self.basis_type, - depth_mlp_edge=depth_mlp_edge, - depth_mlp_trans=depth_mlp_node, - activation_str=self.activation_str, - ablation=ablation, - ) - self.interactions.append(block) - - self.lin = torch.nn.Linear(hidden_channels, self.output_dim) - self.activation = Act(activation_str) - - # ForceNet decoder - self.decoder = FNDecoder(decoder_type, decoder_activation_str, self.output_dim) - - # Projection layer for energy prediction - self.energy_mlp = nn.Linear(self.output_dim, 1) - - def forward(self, data): - z = data.atomic_numbers.long() - - pos = data.pos - batch = data.batch - - if self.feat == "simple": - h = self.embedding(z) - elif self.feat == "full": - h = self.embedding(self.atom_map[z]) - else: - raise RuntimeError("Undefined feature type for atom") - - ( - edge_index, - edge_dist, - edge_vec, - cell_offsets, - _, # cell offset distances - neighbors, - ) = self.generate_graph(data) - - data.edge_index = edge_index - data.cell_offsets = cell_offsets - data.neighbors = neighbors - - if self.pbc_apply_sph_harm: - edge_vec_normalized = edge_vec / edge_dist.view(-1, 1) - edge_attr_sph = self.pbc_sph(edge_vec_normalized) - - # calculate the edge weight according to the dist - edge_weight = torch.cos(0.5 * edge_dist * PI / self.cutoff) - - # normalized edge vectors - edge_vec_normalized = edge_vec / edge_dist.view(-1, 1) - - # edge distance, taking the atom_radii into account - # each element lies in [0,1] - edge_dist_list = ( - torch.stack( - [ - edge_dist, - edge_dist - self.atom_radii[z[edge_index[0]]], - edge_dist - self.atom_radii[z[edge_index[1]]], - edge_dist - - self.atom_radii[z[edge_index[0]]] - - self.atom_radii[z[edge_index[1]]], - ] - ).transpose(0, 1) - / self.cutoff - ) - - if self.ablation == "nodistlist": - edge_dist_list = edge_dist_list[:, 0].view(-1, 1) - - # make sure distance is positive - edge_dist_list[edge_dist_list < 1e-3] = 1e-3 - - # squash to [0,1] for gaussian basis - if self.basis_type == "gauss": - edge_vec_normalized = (edge_vec_normalized + 1) / 2.0 - - # process raw_edge_attributes to generate edge_attributes - if self.ablation == "onlydist": - raw_edge_attr = edge_dist_list - else: - raw_edge_attr = torch.cat([edge_vec_normalized, edge_dist_list], dim=1) - - if "sph" in self.basis_type: - edge_attr = self.basis_fun(raw_edge_attr, edge_attr_sph) - else: - edge_attr = self.basis_fun(raw_edge_attr) - - # pass edge_attributes through interaction blocks - for i, interaction in enumerate(self.interactions): - h = h + interaction(h, edge_index, edge_attr, edge_weight) - - h = self.lin(h) - h = self.activation(h) - - out = scatter(h, batch, dim=0, reduce="add") - - force = self.decoder(h) - energy = self.energy_mlp(out) - return energy, force - - @property - def num_params(self): - return sum(p.numel() for p in self.parameters()) diff --git a/ocpmodels/models/gemnet/gemnet.py b/ocpmodels/models/gemnet/gemnet.py deleted file mode 100644 index a26e8b2b5df634c96b40babddf2d93617a3e737a..0000000000000000000000000000000000000000 --- a/ocpmodels/models/gemnet/gemnet.py +++ /dev/null @@ -1,601 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -from typing import Optional - -import numpy as np -import torch -from torch_geometric.nn import radius_graph -from torch_scatter import scatter -from torch_sparse import SparseTensor - -from ocpmodels.common.registry import registry -from ocpmodels.common.utils import ( - compute_neighbors, - conditional_grad, - get_pbc_distances, - radius_graph_pbc, -) -from ocpmodels.models.base import BaseModel -from ocpmodels.modules.scaling.compat import load_scales_compat - -from .layers.atom_update_block import OutputBlock -from .layers.base_layers import Dense -from .layers.efficient import EfficientInteractionDownProjection -from .layers.embedding_block import AtomEmbedding, EdgeEmbedding -from .layers.interaction_block import InteractionBlockTripletsOnly -from .layers.radial_basis import RadialBasis -from .layers.spherical_basis import CircularBasisLayer -from .utils import ( - inner_product_normalized, - mask_neighbors, - ragged_range, - repeat_blocks, -) - - -@registry.register_model("gemnet_t") -class GemNetT(BaseModel): - """ - GemNet-T, triplets-only variant of GemNet - - Parameters - ---------- - num_atoms (int): Unused argument - bond_feat_dim (int): Unused argument - num_targets: int - Number of prediction targets. - - num_spherical: int - Controls maximum frequency. - num_radial: int - Controls maximum frequency. - num_blocks: int - Number of building blocks to be stacked. - - emb_size_atom: int - Embedding size of the atoms. - emb_size_edge: int - Embedding size of the edges. - emb_size_trip: int - (Down-projected) Embedding size in the triplet message passing block. - emb_size_rbf: int - Embedding size of the radial basis transformation. - emb_size_cbf: int - Embedding size of the circular basis transformation (one angle). - emb_size_bil_trip: int - Embedding size of the edge embeddings in the triplet-based message passing block after the bilinear layer. - - num_before_skip: int - Number of residual blocks before the first skip connection. - num_after_skip: int - Number of residual blocks after the first skip connection. - num_concat: int - Number of residual blocks after the concatenation. - num_atom: int - Number of residual blocks in the atom embedding blocks. - - regress_forces: bool - Whether to predict forces. Default: True - direct_forces: bool - If True predict forces based on aggregation of interatomic directions. - If False predict forces based on negative gradient of energy potential. - - cutoff: float - Embedding cutoff for interactomic directions in Angstrom. - rbf: dict - Name and hyperparameters of the radial basis function. - envelope: dict - Name and hyperparameters of the envelope function. - cbf: dict - Name and hyperparameters of the cosine basis function. - extensive: bool - Whether the output should be extensive (proportional to the number of atoms) - output_init: str - Initialization method for the final dense layer. - activation: str - Name of the activation function. - scale_file: str - Path to the json file containing the scaling factors. - """ - - def __init__( - self, - num_atoms: Optional[int], - bond_feat_dim: int, - num_targets: int, - num_spherical: int, - num_radial: int, - num_blocks: int, - emb_size_atom: int, - emb_size_edge: int, - emb_size_trip: int, - emb_size_rbf: int, - emb_size_cbf: int, - emb_size_bil_trip: int, - num_before_skip: int, - num_after_skip: int, - num_concat: int, - num_atom: int, - regress_forces: bool = True, - direct_forces: bool = False, - cutoff: float = 6.0, - max_neighbors: int = 50, - rbf: dict = {"name": "gaussian"}, - envelope: dict = {"name": "polynomial", "exponent": 5}, - cbf: dict = {"name": "spherical_harmonics"}, - extensive: bool = True, - otf_graph: bool = False, - use_pbc: bool = True, - output_init: str = "HeOrthogonal", - activation: str = "swish", - num_elements: int = 83, - scale_file: Optional[str] = None, - ): - super().__init__() - self.num_targets = num_targets - assert num_blocks > 0 - self.num_blocks = num_blocks - self.extensive = extensive - - self.cutoff = cutoff - assert self.cutoff <= 6 or otf_graph - - self.max_neighbors = max_neighbors - assert self.max_neighbors == 50 or otf_graph - - self.regress_forces = regress_forces - self.otf_graph = otf_graph - self.use_pbc = use_pbc - - # GemNet variants - self.direct_forces = direct_forces - - ### ---------------------------------- Basis Functions ---------------------------------- ### - self.radial_basis = RadialBasis( - num_radial=num_radial, - cutoff=cutoff, - rbf=rbf, - envelope=envelope, - ) - - radial_basis_cbf3 = RadialBasis( - num_radial=num_radial, - cutoff=cutoff, - rbf=rbf, - envelope=envelope, - ) - self.cbf_basis3 = CircularBasisLayer( - num_spherical, - radial_basis=radial_basis_cbf3, - cbf=cbf, - efficient=True, - ) - ### ------------------------------------------------------------------------------------- ### - - ### ------------------------------- Share Down Projections ------------------------------ ### - # Share down projection across all interaction blocks - self.mlp_rbf3 = Dense( - num_radial, - emb_size_rbf, - activation=None, - bias=False, - ) - self.mlp_cbf3 = EfficientInteractionDownProjection( - num_spherical, num_radial, emb_size_cbf - ) - - # Share the dense Layer of the atom embedding block accross the interaction blocks - self.mlp_rbf_h = Dense( - num_radial, - emb_size_rbf, - activation=None, - bias=False, - ) - self.mlp_rbf_out = Dense( - num_radial, - emb_size_rbf, - activation=None, - bias=False, - ) - ### ------------------------------------------------------------------------------------- ### - - # Embedding block - self.atom_emb = AtomEmbedding(emb_size_atom, num_elements) - self.edge_emb = EdgeEmbedding( - emb_size_atom, num_radial, emb_size_edge, activation=activation - ) - - out_blocks = [] - int_blocks = [] - - # Interaction Blocks - interaction_block = InteractionBlockTripletsOnly # GemNet-(d)T - for i in range(num_blocks): - int_blocks.append( - interaction_block( - emb_size_atom=emb_size_atom, - emb_size_edge=emb_size_edge, - emb_size_trip=emb_size_trip, - emb_size_rbf=emb_size_rbf, - emb_size_cbf=emb_size_cbf, - emb_size_bil_trip=emb_size_bil_trip, - num_before_skip=num_before_skip, - num_after_skip=num_after_skip, - num_concat=num_concat, - num_atom=num_atom, - activation=activation, - name=f"IntBlock_{i+1}", - ) - ) - - for i in range(num_blocks + 1): - out_blocks.append( - OutputBlock( - emb_size_atom=emb_size_atom, - emb_size_edge=emb_size_edge, - emb_size_rbf=emb_size_rbf, - nHidden=num_atom, - num_targets=num_targets, - activation=activation, - output_init=output_init, - direct_forces=direct_forces, - name=f"OutBlock_{i}", - ) - ) - - self.out_blocks = torch.nn.ModuleList(out_blocks) - self.int_blocks = torch.nn.ModuleList(int_blocks) - - self.shared_parameters = [ - (self.mlp_rbf3.linear.weight, self.num_blocks), - (self.mlp_cbf3.weight, self.num_blocks), - (self.mlp_rbf_h.linear.weight, self.num_blocks), - (self.mlp_rbf_out.linear.weight, self.num_blocks + 1), - ] - - load_scales_compat(self, scale_file) - - def get_triplets(self, edge_index, num_atoms): - """ - Get all b->a for each edge c->a. - It is possible that b=c, as long as the edges are distinct. - - Returns - ------- - id3_ba: torch.Tensor, shape (num_triplets,) - Indices of input edge b->a of each triplet b->a<-c - id3_ca: torch.Tensor, shape (num_triplets,) - Indices of output edge c->a of each triplet b->a<-c - id3_ragged_idx: torch.Tensor, shape (num_triplets,) - Indices enumerating the copies of id3_ca for creating a padded matrix - """ - idx_s, idx_t = edge_index # c->a (source=c, target=a) - - value = torch.arange(idx_s.size(0), device=idx_s.device, dtype=idx_s.dtype) - # Possibly contains multiple copies of the same edge (for periodic interactions) - adj = SparseTensor( - row=idx_t, - col=idx_s, - value=value, - sparse_sizes=(num_atoms, num_atoms), - ) - adj_edges = adj[idx_t] - - # Edge indices (b->a, c->a) for triplets. - id3_ba = adj_edges.storage.value() - id3_ca = adj_edges.storage.row() - - # Remove self-loop triplets - # Compare edge indices, not atom indices to correctly handle periodic interactions - mask = id3_ba != id3_ca - id3_ba = id3_ba[mask] - id3_ca = id3_ca[mask] - - # Get indices to reshape the neighbor indices b->a into a dense matrix. - # id3_ca has to be sorted for this to work. - num_triplets = torch.bincount(id3_ca, minlength=idx_s.size(0)) - id3_ragged_idx = ragged_range(num_triplets) - - return id3_ba, id3_ca, id3_ragged_idx - - def select_symmetric_edges(self, tensor, mask, reorder_idx, inverse_neg): - # Mask out counter-edges - tensor_directed = tensor[mask] - # Concatenate counter-edges after normal edges - sign = 1 - 2 * inverse_neg - tensor_cat = torch.cat([tensor_directed, sign * tensor_directed]) - # Reorder everything so the edges of every image are consecutive - tensor_ordered = tensor_cat[reorder_idx] - return tensor_ordered - - def reorder_symmetric_edges( - self, edge_index, cell_offsets, neighbors, edge_dist, edge_vector - ): - """ - Reorder edges to make finding counter-directional edges easier. - - Some edges are only present in one direction in the data, - since every atom has a maximum number of neighbors. Since we only use i->j - edges here, we lose some j->i edges and add others by - making it symmetric. - We could fix this by merging edge_index with its counter-edges, - including the cell_offsets, and then running torch.unique. - But this does not seem worth it. - """ - - # Generate mask - mask_sep_atoms = edge_index[0] < edge_index[1] - # Distinguish edges between the same (periodic) atom by ordering the cells - cell_earlier = ( - (cell_offsets[:, 0] < 0) - | ((cell_offsets[:, 0] == 0) & (cell_offsets[:, 1] < 0)) - | ( - (cell_offsets[:, 0] == 0) - & (cell_offsets[:, 1] == 0) - & (cell_offsets[:, 2] < 0) - ) - ) - mask_same_atoms = edge_index[0] == edge_index[1] - mask_same_atoms &= cell_earlier - mask = mask_sep_atoms | mask_same_atoms - - # Mask out counter-edges - edge_index_new = edge_index[mask[None, :].expand(2, -1)].view(2, -1) - - # Concatenate counter-edges after normal edges - edge_index_cat = torch.cat( - [ - edge_index_new, - torch.stack([edge_index_new[1], edge_index_new[0]], dim=0), - ], - dim=1, - ) - - # Count remaining edges per image - batch_edge = torch.repeat_interleave( - torch.arange(neighbors.size(0), device=edge_index.device), - neighbors, - ) - batch_edge = batch_edge[mask] - neighbors_new = 2 * torch.bincount(batch_edge, minlength=neighbors.size(0)) - - # Create indexing array - edge_reorder_idx = repeat_blocks( - neighbors_new // 2, - repeats=2, - continuous_indexing=True, - repeat_inc=edge_index_new.size(1), - ) - - # Reorder everything so the edges of every image are consecutive - edge_index_new = edge_index_cat[:, edge_reorder_idx] - cell_offsets_new = self.select_symmetric_edges( - cell_offsets, mask, edge_reorder_idx, True - ) - edge_dist_new = self.select_symmetric_edges( - edge_dist, mask, edge_reorder_idx, False - ) - edge_vector_new = self.select_symmetric_edges( - edge_vector, mask, edge_reorder_idx, True - ) - - return ( - edge_index_new, - cell_offsets_new, - neighbors_new, - edge_dist_new, - edge_vector_new, - ) - - def select_edges( - self, - data, - edge_index, - cell_offsets, - neighbors, - edge_dist, - edge_vector, - cutoff=None, - ): - if cutoff is not None: - edge_mask = edge_dist <= cutoff - - edge_index = edge_index[:, edge_mask] - cell_offsets = cell_offsets[edge_mask] - neighbors = mask_neighbors(neighbors, edge_mask) - edge_dist = edge_dist[edge_mask] - edge_vector = edge_vector[edge_mask] - - empty_image = neighbors == 0 - if torch.any(empty_image): - raise ValueError( - f"An image has no neighbors: id={data.id[empty_image]}, " - f"sid={data.sid[empty_image]}, fid={data.fid[empty_image]}" - ) - return edge_index, cell_offsets, neighbors, edge_dist, edge_vector - - def generate_interaction_graph(self, data): - num_atoms = data.atomic_numbers.size(0) - - ( - edge_index, - D_st, - distance_vec, - cell_offsets, - _, # cell offset distances - neighbors, - ) = self.generate_graph(data) - # These vectors actually point in the opposite direction. - # But we want to use col as idx_t for efficient aggregation. - V_st = -distance_vec / D_st[:, None] - - # Mask interaction edges if required - if self.otf_graph or np.isclose(self.cutoff, 6): - select_cutoff = None - else: - select_cutoff = self.cutoff - ( - edge_index, - cell_offsets, - neighbors, - D_st, - V_st, - ) = self.select_edges( - data=data, - edge_index=edge_index, - cell_offsets=cell_offsets, - neighbors=neighbors, - edge_dist=D_st, - edge_vector=V_st, - cutoff=select_cutoff, - ) - - ( - edge_index, - cell_offsets, - neighbors, - D_st, - V_st, - ) = self.reorder_symmetric_edges( - edge_index, cell_offsets, neighbors, D_st, V_st - ) - - # Indices for swapping c->a and a->c (for symmetric MP) - block_sizes = neighbors // 2 - id_swap = repeat_blocks( - block_sizes, - repeats=2, - continuous_indexing=False, - start_idx=block_sizes[0], - block_inc=block_sizes[:-1] + block_sizes[1:], - repeat_inc=-block_sizes, - ) - - id3_ba, id3_ca, id3_ragged_idx = self.get_triplets( - edge_index, num_atoms=num_atoms - ) - - return ( - edge_index, - neighbors, - D_st, - V_st, - id_swap, - id3_ba, - id3_ca, - id3_ragged_idx, - ) - - @conditional_grad(torch.enable_grad()) - def forward(self, data): - pos = data.pos - batch = data.batch - atomic_numbers = data.atomic_numbers.long() - - if self.regress_forces and not self.direct_forces: - pos.requires_grad_(True) - - ( - edge_index, - neighbors, - D_st, - V_st, - id_swap, - id3_ba, - id3_ca, - id3_ragged_idx, - ) = self.generate_interaction_graph(data) - idx_s, idx_t = edge_index - - # Calculate triplet angles - cosφ_cab = inner_product_normalized(V_st[id3_ca], V_st[id3_ba]) - rad_cbf3, cbf3 = self.cbf_basis3(D_st, cosφ_cab, id3_ca) - - rbf = self.radial_basis(D_st) - - # Embedding block - h = self.atom_emb(atomic_numbers) - # (nAtoms, emb_size_atom) - m = self.edge_emb(h, rbf, idx_s, idx_t) # (nEdges, emb_size_edge) - - rbf3 = self.mlp_rbf3(rbf) - cbf3 = self.mlp_cbf3(rad_cbf3, cbf3, id3_ca, id3_ragged_idx) - - rbf_h = self.mlp_rbf_h(rbf) - rbf_out = self.mlp_rbf_out(rbf) - - E_t, F_st = self.out_blocks[0](h, m, rbf_out, idx_t) - # (nAtoms, num_targets), (nEdges, num_targets) - - for i in range(self.num_blocks): - # Interaction block - h, m = self.int_blocks[i]( - h=h, - m=m, - rbf3=rbf3, - cbf3=cbf3, - id3_ragged_idx=id3_ragged_idx, - id_swap=id_swap, - id3_ba=id3_ba, - id3_ca=id3_ca, - rbf_h=rbf_h, - idx_s=idx_s, - idx_t=idx_t, - ) # (nAtoms, emb_size_atom), (nEdges, emb_size_edge) - - E, F = self.out_blocks[i + 1](h, m, rbf_out, idx_t) - # (nAtoms, num_targets), (nEdges, num_targets) - F_st += F - E_t += E - - nMolecules = torch.max(batch) + 1 - if self.extensive: - E_t = scatter( - E_t, batch, dim=0, dim_size=nMolecules, reduce="add" - ) # (nMolecules, num_targets) - else: - E_t = scatter( - E_t, batch, dim=0, dim_size=nMolecules, reduce="mean" - ) # (nMolecules, num_targets) - - if self.regress_forces: - if self.direct_forces: - # map forces in edge directions - F_st_vec = F_st[:, :, None] * V_st[:, None, :] - # (nEdges, num_targets, 3) - F_t = scatter( - F_st_vec, - idx_t, - dim=0, - dim_size=data.atomic_numbers.size(0), - reduce="add", - ) # (nAtoms, num_targets, 3) - F_t = F_t.squeeze(1) # (nAtoms, 3) - else: - if self.num_targets > 1: - forces = [] - for i in range(self.num_targets): - # maybe this can be solved differently - forces += [ - -torch.autograd.grad( - E_t[:, i].sum(), pos, create_graph=True - )[0] - ] - F_t = torch.stack(forces, dim=1) - # (nAtoms, num_targets, 3) - else: - F_t = -torch.autograd.grad(E_t.sum(), pos, create_graph=True)[0] - # (nAtoms, 3) - - return E_t, F_t # (nMolecules, num_targets), (nAtoms, 3) - else: - return E_t - - @property - def num_params(self): - return sum(p.numel() for p in self.parameters()) diff --git a/ocpmodels/models/gemnet/initializers.py b/ocpmodels/models/gemnet/initializers.py deleted file mode 100644 index 39a07ce6a5daadb7dc9ce351199ae24b152b6a1d..0000000000000000000000000000000000000000 --- a/ocpmodels/models/gemnet/initializers.py +++ /dev/null @@ -1,47 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import torch - - -def _standardize(kernel): - """ - Makes sure that N*Var(W) = 1 and E[W] = 0 - """ - eps = 1e-6 - - if len(kernel.shape) == 3: - axis = [0, 1] # last dimension is output dimension - else: - axis = 1 - - var, mean = torch.var_mean(kernel, dim=axis, unbiased=True, keepdim=True) - kernel = (kernel - mean) / (var + eps) ** 0.5 - return kernel - - -def he_orthogonal_init(tensor): - """ - Generate a weight matrix with variance according to He (Kaiming) initialization. - Based on a random (semi-)orthogonal matrix neural networks - are expected to learn better when features are decorrelated - (stated by eg. "Reducing overfitting in deep networks by decorrelating representations", - "Dropout: a simple way to prevent neural networks from overfitting", - "Exact solutions to the nonlinear dynamics of learning in deep linear neural networks") - """ - tensor = torch.nn.init.orthogonal_(tensor) - - if len(tensor.shape) == 3: - fan_in = tensor.shape[:-1].numel() - else: - fan_in = tensor.shape[1] - - with torch.no_grad(): - tensor.data = _standardize(tensor.data) - tensor.data *= (1 / fan_in) ** 0.5 - - return tensor diff --git a/ocpmodels/models/gemnet/layers/atom_update_block.py b/ocpmodels/models/gemnet/layers/atom_update_block.py deleted file mode 100644 index 596a2a69b991da80d471ed238f85f250bc5b9d92..0000000000000000000000000000000000000000 --- a/ocpmodels/models/gemnet/layers/atom_update_block.py +++ /dev/null @@ -1,200 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import torch -from torch_scatter import scatter - -from ocpmodels.modules.scaling import ScaleFactor - -from ..initializers import he_orthogonal_init -from .base_layers import Dense, ResidualLayer - - -class AtomUpdateBlock(torch.nn.Module): - """ - Aggregate the message embeddings of the atoms - - Parameters - ---------- - emb_size_atom: int - Embedding size of the atoms. - emb_size_atom: int - Embedding size of the edges. - nHidden: int - Number of residual blocks. - activation: callable/str - Name of the activation function to use in the dense layers. - """ - - def __init__( - self, - emb_size_atom: int, - emb_size_edge: int, - emb_size_rbf: int, - nHidden: int, - activation=None, - name: str = "atom_update", - ): - super().__init__() - self.name = name - - self.dense_rbf = Dense(emb_size_rbf, emb_size_edge, activation=None, bias=False) - self.scale_sum = ScaleFactor(name + "_sum") - - self.layers = self.get_mlp(emb_size_edge, emb_size_atom, nHidden, activation) - - def get_mlp(self, units_in, units, nHidden, activation): - dense1 = Dense(units_in, units, activation=activation, bias=False) - mlp = [dense1] - res = [ - ResidualLayer(units, nLayers=2, activation=activation) - for i in range(nHidden) - ] - mlp += res - return torch.nn.ModuleList(mlp) - - def forward(self, h, m, rbf, id_j): - """ - Returns - ------- - h: torch.Tensor, shape=(nAtoms, emb_size_atom) - Atom embedding. - """ - nAtoms = h.shape[0] - - mlp_rbf = self.dense_rbf(rbf) # (nEdges, emb_size_edge) - x = m * mlp_rbf - - x2 = scatter(x, id_j, dim=0, dim_size=nAtoms, reduce="sum") - # (nAtoms, emb_size_edge) - x = self.scale_sum(x2, ref=m) - - for layer in self.layers: - x = layer(x) # (nAtoms, emb_size_atom) - - return x - - -class OutputBlock(AtomUpdateBlock): - """ - Combines the atom update block and subsequent final dense layer. - - Parameters - ---------- - emb_size_atom: int - Embedding size of the atoms. - emb_size_atom: int - Embedding size of the edges. - nHidden: int - Number of residual blocks. - num_targets: int - Number of targets. - activation: str - Name of the activation function to use in the dense layers except for the final dense layer. - direct_forces: bool - If true directly predict forces without taking the gradient of the energy potential. - output_init: int - Kernel initializer of the final dense layer. - """ - - def __init__( - self, - emb_size_atom: int, - emb_size_edge: int, - emb_size_rbf: int, - nHidden: int, - num_targets: int, - activation=None, - direct_forces=True, - output_init="HeOrthogonal", - name: str = "output", - **kwargs, - ): - - super().__init__( - name=name, - emb_size_atom=emb_size_atom, - emb_size_edge=emb_size_edge, - emb_size_rbf=emb_size_rbf, - nHidden=nHidden, - activation=activation, - **kwargs, - ) - - assert isinstance(output_init, str) - self.output_init = output_init.lower() - self.direct_forces = direct_forces - - self.seq_energy = self.layers # inherited from parent class - self.out_energy = Dense(emb_size_atom, num_targets, bias=False, activation=None) - - if self.direct_forces: - self.scale_rbf_F = ScaleFactor(name + "_had") - self.seq_forces = self.get_mlp( - emb_size_edge, emb_size_edge, nHidden, activation - ) - self.out_forces = Dense( - emb_size_edge, num_targets, bias=False, activation=None - ) - self.dense_rbf_F = Dense( - emb_size_rbf, emb_size_edge, activation=None, bias=False - ) - - self.reset_parameters() - - def reset_parameters(self): - if self.output_init == "heorthogonal": - self.out_energy.reset_parameters(he_orthogonal_init) - if self.direct_forces: - self.out_forces.reset_parameters(he_orthogonal_init) - elif self.output_init == "zeros": - self.out_energy.reset_parameters(torch.nn.init.zeros_) - if self.direct_forces: - self.out_forces.reset_parameters(torch.nn.init.zeros_) - else: - raise UserWarning(f"Unknown output_init: {self.output_init}") - - def forward(self, h, m, rbf, id_j): - """ - Returns - ------- - (E, F): tuple - - E: torch.Tensor, shape=(nAtoms, num_targets) - - F: torch.Tensor, shape=(nEdges, num_targets) - Energy and force prediction - """ - nAtoms = h.shape[0] - - # -------------------------------------- Energy Prediction -------------------------------------- # - rbf_emb_E = self.dense_rbf(rbf) # (nEdges, emb_size_edge) - x = m * rbf_emb_E - - x_E = scatter(x, id_j, dim=0, dim_size=nAtoms, reduce="sum") - # (nAtoms, emb_size_edge) - x_E = self.scale_sum(x_E, ref=m) - - for layer in self.seq_energy: - x_E = layer(x_E) # (nAtoms, emb_size_atom) - - x_E = self.out_energy(x_E) # (nAtoms, num_targets) - - # --------------------------------------- Force Prediction -------------------------------------- # - if self.direct_forces: - x_F = m - for i, layer in enumerate(self.seq_forces): - x_F = layer(x_F) # (nEdges, emb_size_edge) - - rbf_emb_F = self.dense_rbf_F(rbf) # (nEdges, emb_size_edge) - x_F_rbf = x_F * rbf_emb_F - x_F = self.scale_rbf_F(x_F_rbf, ref=x_F) - - x_F = self.out_forces(x_F) # (nEdges, num_targets) - else: - x_F = 0 - # ----------------------------------------------------------------------------------------------- # - - return x_E, x_F diff --git a/ocpmodels/models/gemnet/layers/base_layers.py b/ocpmodels/models/gemnet/layers/base_layers.py deleted file mode 100644 index 948bb5607b6816e51c4f945affe824b47999eee1..0000000000000000000000000000000000000000 --- a/ocpmodels/models/gemnet/layers/base_layers.py +++ /dev/null @@ -1,106 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import math - -import torch - -from ..initializers import he_orthogonal_init - - -class Dense(torch.nn.Module): - """ - Combines dense layer with scaling for swish activation. - - Parameters - ---------- - units: int - Output embedding size. - activation: str - Name of the activation function to use. - bias: bool - True if use bias. - """ - - def __init__(self, in_features, out_features, bias=False, activation=None): - super().__init__() - - self.linear = torch.nn.Linear(in_features, out_features, bias=bias) - self.reset_parameters() - - if isinstance(activation, str): - activation = activation.lower() - if activation in ["swish", "silu"]: - self._activation = ScaledSiLU() - elif activation == "siqu": - self._activation = SiQU() - elif activation is None: - self._activation = torch.nn.Identity() - else: - raise NotImplementedError( - "Activation function not implemented for GemNet (yet)." - ) - - def reset_parameters(self, initializer=he_orthogonal_init): - initializer(self.linear.weight) - if self.linear.bias is not None: - self.linear.bias.data.fill_(0) - - def forward(self, x): - x = self.linear(x) - x = self._activation(x) - return x - - -class ScaledSiLU(torch.nn.Module): - def __init__(self): - super().__init__() - self.scale_factor = 1 / 0.6 - self._activation = torch.nn.SiLU() - - def forward(self, x): - return self._activation(x) * self.scale_factor - - -class SiQU(torch.nn.Module): - def __init__(self): - super().__init__() - self._activation = torch.nn.SiLU() - - def forward(self, x): - return x * self._activation(x) - - -class ResidualLayer(torch.nn.Module): - """ - Residual block with output scaled by 1/sqrt(2). - - Parameters - ---------- - units: int - Output embedding size. - nLayers: int - Number of dense layers. - layer_kwargs: str - Keyword arguments for initializing the layers. - """ - - def __init__(self, units: int, nLayers: int = 2, layer=Dense, **layer_kwargs): - super().__init__() - self.dense_mlp = torch.nn.Sequential( - *[ - layer(in_features=units, out_features=units, bias=False, **layer_kwargs) - for _ in range(nLayers) - ] - ) - self.inv_sqrt_2 = 1 / math.sqrt(2) - - def forward(self, input): - x = self.dense_mlp(input) - x = input + x - x = x * self.inv_sqrt_2 - return x diff --git a/ocpmodels/models/gemnet/layers/basis_utils.py b/ocpmodels/models/gemnet/layers/basis_utils.py deleted file mode 100644 index 987bdd630ff7a89af37fe48996ee71665fc59d05..0000000000000000000000000000000000000000 --- a/ocpmodels/models/gemnet/layers/basis_utils.py +++ /dev/null @@ -1,280 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import numpy as np -import sympy as sym -from scipy import special as sp -from scipy.optimize import brentq - - -def Jn(r, n): - """ - numerical spherical bessel functions of order n - """ - return sp.spherical_jn(n, r) - - -def Jn_zeros(n, k): - """ - Compute the first k zeros of the spherical bessel functions up to order n (excluded) - """ - zerosj = np.zeros((n, k), dtype="float32") - zerosj[0] = np.arange(1, k + 1) * np.pi - points = np.arange(1, k + n) * np.pi - racines = np.zeros(k + n - 1, dtype="float32") - for i in range(1, n): - for j in range(k + n - 1 - i): - foo = brentq(Jn, points[j], points[j + 1], (i,)) - racines[j] = foo - points = racines - zerosj[i][:k] = racines[:k] - - return zerosj - - -def spherical_bessel_formulas(n): - """ - Computes the sympy formulas for the spherical bessel functions up to order n (excluded) - """ - x = sym.symbols("x") - # j_i = (-x)^i * (1/x * d/dx)^î * sin(x)/x - j = [sym.sin(x) / x] # j_0 - a = sym.sin(x) / x - for i in range(1, n): - b = sym.diff(a, x) / x - j += [sym.simplify(b * (-x) ** i)] - a = sym.simplify(b) - return j - - -def bessel_basis(n, k): - """ - Compute the sympy formulas for the normalized and rescaled spherical bessel functions up to - order n (excluded) and maximum frequency k (excluded). - - Returns: - bess_basis: list - Bessel basis formulas taking in a single argument x. - Has length n where each element has length k. -> In total n*k many. - """ - zeros = Jn_zeros(n, k) - normalizer = [] - for order in range(n): - normalizer_tmp = [] - for i in range(k): - normalizer_tmp += [0.5 * Jn(zeros[order, i], order + 1) ** 2] - normalizer_tmp = ( - 1 / np.array(normalizer_tmp) ** 0.5 - ) # sqrt(2/(j_l+1)**2) , sqrt(1/c**3) not taken into account yet - normalizer += [normalizer_tmp] - - f = spherical_bessel_formulas(n) - x = sym.symbols("x") - bess_basis = [] - for order in range(n): - bess_basis_tmp = [] - for i in range(k): - bess_basis_tmp += [ - sym.simplify( - normalizer[order][i] * f[order].subs(x, zeros[order, i] * x) - ) - ] - bess_basis += [bess_basis_tmp] - return bess_basis - - -def sph_harm_prefactor(l_degree, m_order): - """Computes the constant pre-factor for the spherical harmonic of degree l and order m. - - Parameters - ---------- - l_degree: int - Degree of the spherical harmonic. l >= 0 - m_order: int - Order of the spherical harmonic. -l <= m <= l - - Returns - ------- - factor: float - - """ - # sqrt((2*l+1)/4*pi * (l-m)!/(l+m)! ) - return ( - (2 * l_degree + 1) - / (4 * np.pi) - * np.math.factorial(l_degree - abs(m_order)) - / np.math.factorial(l_degree + abs(m_order)) - ) ** 0.5 - - -def associated_legendre_polynomials(L_maxdegree, zero_m_only=True, pos_m_only=True): - """Computes string formulas of the associated legendre polynomials up to degree L (excluded). - - Parameters - ---------- - L_maxdegree: int - Degree up to which to calculate the associated legendre polynomials (degree L is excluded). - zero_m_only: bool - If True only calculate the polynomials for the polynomials where m=0. - pos_m_only: bool - If True only calculate the polynomials for the polynomials where m>=0. Overwritten by zero_m_only. - - Returns - ------- - polynomials: list - Contains the sympy functions of the polynomials (in total L many if zero_m_only is True else L^2 many). - """ - # calculations from http://web.cmb.usc.edu/people/alber/Software/tomominer/docs/cpp/group__legendre__polynomials.html - z = sym.symbols("z") - P_l_m = [ - [0] * (2 * l_degree + 1) for l_degree in range(L_maxdegree) - ] # for order l: -l <= m <= l - - P_l_m[0][0] = 1 - if L_maxdegree > 0: - if zero_m_only: - # m = 0 - P_l_m[1][0] = z - for l_degree in range(2, L_maxdegree): - P_l_m[l_degree][0] = sym.simplify( - ( - (2 * l_degree - 1) * z * P_l_m[l_degree - 1][0] - - (l_degree - 1) * P_l_m[l_degree - 2][0] - ) - / l_degree - ) - return P_l_m - else: - # for m >= 0 - for l_degree in range(1, L_maxdegree): - P_l_m[l_degree][l_degree] = sym.simplify( - (1 - 2 * l_degree) - * (1 - z**2) ** 0.5 - * P_l_m[l_degree - 1][l_degree - 1] - ) # P_00, P_11, P_22, P_33 - - for m_order in range(0, L_maxdegree - 1): - P_l_m[m_order + 1][m_order] = sym.simplify( - (2 * m_order + 1) * z * P_l_m[m_order][m_order] - ) # P_10, P_21, P_32, P_43 - - for l_degree in range(2, L_maxdegree): - for m_order in range(l_degree - 1): # P_20, P_30, P_31 - P_l_m[l_degree][m_order] = sym.simplify( - ( - (2 * l_degree - 1) * z * P_l_m[l_degree - 1][m_order] - - (l_degree + m_order - 1) * P_l_m[l_degree - 2][m_order] - ) - / (l_degree - m_order) - ) - - if not pos_m_only: - # for m < 0: P_l(-m) = (-1)^m * (l-m)!/(l+m)! * P_lm - for l_degree in range(1, L_maxdegree): - for m_order in range(1, l_degree + 1): # P_1(-1), P_2(-1) P_2(-2) - P_l_m[l_degree][-m_order] = sym.simplify( - (-1) ** m_order - * np.math.factorial(l_degree - m_order) - / np.math.factorial(l_degree + m_order) - * P_l_m[l_degree][m_order] - ) - - return P_l_m - - -def real_sph_harm(L_maxdegree, use_theta, use_phi=True, zero_m_only=True): - """ - Computes formula strings of the the real part of the spherical harmonics up to degree L (excluded). - Variables are either spherical coordinates phi and theta (or cartesian coordinates x,y,z) on the UNIT SPHERE. - - Parameters - ---------- - L_maxdegree: int - Degree up to which to calculate the spherical harmonics (degree L is excluded). - use_theta: bool - - True: Expects the input of the formula strings to contain theta. - - False: Expects the input of the formula strings to contain z. - use_phi: bool - - True: Expects the input of the formula strings to contain phi. - - False: Expects the input of the formula strings to contain x and y. - Does nothing if zero_m_only is True - zero_m_only: bool - If True only calculate the harmonics where m=0. - - Returns - ------- - Y_lm_real: list - Computes formula strings of the the real part of the spherical harmonics up - to degree L (where degree L is not excluded). - In total L^2 many sph harm exist up to degree L (excluded). However, if zero_m_only only is True then - the total count is reduced to be only L many. - """ - z = sym.symbols("z") - P_l_m = associated_legendre_polynomials(L_maxdegree, zero_m_only) - if zero_m_only: - # for all m != 0: Y_lm = 0 - Y_l_m = [[0] for l_degree in range(L_maxdegree)] - else: - Y_l_m = [ - [0] * (2 * l_degree + 1) for l_degree in range(L_maxdegree) - ] # for order l: -l <= m <= l - - # convert expressions to spherical coordiantes - if use_theta: - # replace z by cos(theta) - theta = sym.symbols("theta") - for l_degree in range(L_maxdegree): - for m_order in range(len(P_l_m[l_degree])): - if not isinstance(P_l_m[l_degree][m_order], int): - P_l_m[l_degree][m_order] = P_l_m[l_degree][m_order].subs( - z, sym.cos(theta) - ) - - ## calculate Y_lm - # Y_lm = N * P_lm(cos(theta)) * exp(i*m*phi) - # { sqrt(2) * (-1)^m * N * P_l|m| * sin(|m|*phi) if m < 0 - # Y_lm_real = { Y_lm if m = 0 - # { sqrt(2) * (-1)^m * N * P_lm * cos(m*phi) if m > 0 - - for l_degree in range(L_maxdegree): - Y_l_m[l_degree][0] = sym.simplify( - sph_harm_prefactor(l_degree, 0) * P_l_m[l_degree][0] - ) # Y_l0 - - if not zero_m_only: - phi = sym.symbols("phi") - for l_degree in range(1, L_maxdegree): - # m > 0 - for m_order in range(1, l_degree + 1): - Y_l_m[l_degree][m_order] = sym.simplify( - 2**0.5 - * (-1) ** m_order - * sph_harm_prefactor(l_degree, m_order) - * P_l_m[l_degree][m_order] - * sym.cos(m_order * phi) - ) - # m < 0 - for m_order in range(1, l_degree + 1): - Y_l_m[l_degree][-m_order] = sym.simplify( - 2**0.5 - * (-1) ** m_order - * sph_harm_prefactor(l_degree, -m_order) - * P_l_m[l_degree][m_order] - * sym.sin(m_order * phi) - ) - - # convert expressions to cartesian coordinates - if not use_phi: - # replace phi by atan2(y,x) - x = sym.symbols("x") - y = sym.symbols("y") - for l_degree in range(L_maxdegree): - for m_order in range(len(Y_l_m[l_degree])): - Y_l_m[l_degree][m_order] = sym.simplify( - Y_l_m[l_degree][m_order].subs(phi, sym.atan2(y, x)) - ) - return Y_l_m diff --git a/ocpmodels/models/gemnet/layers/efficient.py b/ocpmodels/models/gemnet/layers/efficient.py deleted file mode 100644 index 3c76de26d11344f01b9c2dd7317a5aadd05d25f7..0000000000000000000000000000000000000000 --- a/ocpmodels/models/gemnet/layers/efficient.py +++ /dev/null @@ -1,171 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import torch - -from ..initializers import he_orthogonal_init - - -class EfficientInteractionDownProjection(torch.nn.Module): - """ - Down projection in the efficient reformulation. - - Parameters - ---------- - emb_size_interm: int - Intermediate embedding size (down-projection size). - kernel_initializer: callable - Initializer of the weight matrix. - """ - - def __init__( - self, - num_spherical: int, - num_radial: int, - emb_size_interm: int, - ): - super().__init__() - - self.num_spherical = num_spherical - self.num_radial = num_radial - self.emb_size_interm = emb_size_interm - - self.reset_parameters() - - def reset_parameters(self): - self.weight = torch.nn.Parameter( - torch.empty((self.num_spherical, self.num_radial, self.emb_size_interm)), - requires_grad=True, - ) - he_orthogonal_init(self.weight) - - def forward(self, rbf, sph, id_ca, id_ragged_idx): - """ - - Arguments - --------- - rbf: torch.Tensor, shape=(1, nEdges, num_radial) - sph: torch.Tensor, shape=(nEdges, Kmax, num_spherical) - id_ca - id_ragged_idx - - Returns - ------- - rbf_W1: torch.Tensor, shape=(nEdges, emb_size_interm, num_spherical) - sph: torch.Tensor, shape=(nEdges, Kmax, num_spherical) - Kmax = maximum number of neighbors of the edges - """ - num_edges = rbf.shape[1] - - # MatMul: mul + sum over num_radial - rbf_W1 = torch.matmul(rbf, self.weight) - # (num_spherical, nEdges , emb_size_interm) - rbf_W1 = rbf_W1.permute(1, 2, 0) - # (nEdges, emb_size_interm, num_spherical) - - # Zero padded dense matrix - # maximum number of neighbors, catch empty id_ca with maximum - if sph.shape[0] == 0: - Kmax = 0 - else: - Kmax = torch.max( - torch.max(id_ragged_idx + 1), - torch.tensor(0).to(id_ragged_idx.device), - ) - - sph2 = sph.new_zeros(num_edges, Kmax, self.num_spherical) - sph2[id_ca, id_ragged_idx] = sph - - sph2 = torch.transpose(sph2, 1, 2) - # (nEdges, num_spherical/emb_size_interm, Kmax) - - return rbf_W1, sph2 - - -class EfficientInteractionBilinear(torch.nn.Module): - """ - Efficient reformulation of the bilinear layer and subsequent summation. - - Parameters - ---------- - units_out: int - Embedding output size of the bilinear layer. - kernel_initializer: callable - Initializer of the weight matrix. - """ - - def __init__( - self, - emb_size: int, - emb_size_interm: int, - units_out: int, - ): - super().__init__() - self.emb_size = emb_size - self.emb_size_interm = emb_size_interm - self.units_out = units_out - - self.reset_parameters() - - def reset_parameters(self): - self.weight = torch.nn.Parameter( - torch.empty( - (self.emb_size, self.emb_size_interm, self.units_out), - requires_grad=True, - ) - ) - he_orthogonal_init(self.weight) - - def forward( - self, - basis, - m, - id_reduce, - id_ragged_idx, - ): - """ - - Arguments - --------- - basis - m: quadruplets: m = m_db , triplets: m = m_ba - id_reduce - id_ragged_idx - - Returns - ------- - m_ca: torch.Tensor, shape=(nEdges, units_out) - Edge embeddings. - """ - # num_spherical is actually num_spherical**2 for quadruplets - (rbf_W1, sph) = basis - # (nEdges, emb_size_interm, num_spherical), (nEdges, num_spherical, Kmax) - nEdges = rbf_W1.shape[0] - - # Create (zero-padded) dense matrix of the neighboring edge embeddings. - Kmax = torch.max( - torch.max(id_ragged_idx) + 1, - torch.tensor(0).to(id_ragged_idx.device), - ) - # maximum number of neighbors, catch empty id_reduce_ji with maximum - m2 = m.new_zeros(nEdges, Kmax, self.emb_size) - m2[id_reduce, id_ragged_idx] = m - # (num_quadruplets or num_triplets, emb_size) -> (nEdges, Kmax, emb_size) - - sum_k = torch.matmul(sph, m2) # (nEdges, num_spherical, emb_size) - - # MatMul: mul + sum over num_spherical - rbf_W1_sum_k = torch.matmul(rbf_W1, sum_k) - # (nEdges, emb_size_interm, emb_size) - - # Bilinear: Sum over emb_size_interm and emb_size - m_ca = torch.matmul(rbf_W1_sum_k.permute(2, 0, 1), self.weight) - # (emb_size, nEdges, units_out) - m_ca = torch.sum(m_ca, dim=0) - # (nEdges, units_out) - - return m_ca diff --git a/ocpmodels/models/gemnet/layers/embedding_block.py b/ocpmodels/models/gemnet/layers/embedding_block.py deleted file mode 100644 index 58f877f69e59582ef038280e6c21d0e6b6809263..0000000000000000000000000000000000000000 --- a/ocpmodels/models/gemnet/layers/embedding_block.py +++ /dev/null @@ -1,93 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import numpy as np -import torch - -from .base_layers import Dense - - -class AtomEmbedding(torch.nn.Module): - """ - Initial atom embeddings based on the atom type - - Parameters - ---------- - emb_size: int - Atom embeddings size - """ - - def __init__(self, emb_size, num_elements): - super().__init__() - self.emb_size = emb_size - - self.embeddings = torch.nn.Embedding(num_elements, emb_size) - # init by uniform distribution - torch.nn.init.uniform_(self.embeddings.weight, a=-np.sqrt(3), b=np.sqrt(3)) - - def forward(self, Z): - """ - Returns - ------- - h: torch.Tensor, shape=(nAtoms, emb_size) - Atom embeddings. - """ - h = self.embeddings(Z - 1) # -1 because Z.min()=1 (==Hydrogen) - return h - - -class EdgeEmbedding(torch.nn.Module): - """ - Edge embedding based on the concatenation of atom embeddings and subsequent dense layer. - - Parameters - ---------- - emb_size: int - Embedding size after the dense layer. - activation: str - Activation function used in the dense layer. - """ - - def __init__( - self, - atom_features, - edge_features, - out_features, - activation=None, - ): - super().__init__() - in_features = 2 * atom_features + edge_features - self.dense = Dense(in_features, out_features, activation=activation, bias=False) - - def forward( - self, - h, - m_rbf, - idx_s, - idx_t, - ): - """ - - Arguments - --------- - h - m_rbf: shape (nEdges, nFeatures) - in embedding block: m_rbf = rbf ; In interaction block: m_rbf = m_st - idx_s - idx_t - - Returns - ------- - m_st: torch.Tensor, shape=(nEdges, emb_size) - Edge embeddings. - """ - h_s = h[idx_s] # shape=(nEdges, emb_size) - h_t = h[idx_t] # shape=(nEdges, emb_size) - - m_st = torch.cat([h_s, h_t, m_rbf], dim=-1) # (nEdges, 2*emb_size+nFeatures) - m_st = self.dense(m_st) # (nEdges, emb_size) - return m_st diff --git a/ocpmodels/models/gemnet/layers/interaction_block.py b/ocpmodels/models/gemnet/layers/interaction_block.py deleted file mode 100644 index 033e25ee021cd1b914e679f03d87ed1d03ea7638..0000000000000000000000000000000000000000 --- a/ocpmodels/models/gemnet/layers/interaction_block.py +++ /dev/null @@ -1,341 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import math - -import torch - -from ocpmodels.modules.scaling.scale_factor import ScaleFactor - -from .atom_update_block import AtomUpdateBlock -from .base_layers import Dense, ResidualLayer -from .efficient import EfficientInteractionBilinear -from .embedding_block import EdgeEmbedding - - -class InteractionBlockTripletsOnly(torch.nn.Module): - """ - Interaction block for GemNet-T/dT. - - Parameters - ---------- - emb_size_atom: int - Embedding size of the atoms. - emb_size_edge: int - Embedding size of the edges. - emb_size_trip: int - (Down-projected) Embedding size in the triplet message passing block. - emb_size_rbf: int - Embedding size of the radial basis transformation. - emb_size_cbf: int - Embedding size of the circular basis transformation (one angle). - - emb_size_bil_trip: int - Embedding size of the edge embeddings in the triplet-based message passing block after the bilinear layer. - num_before_skip: int - Number of residual blocks before the first skip connection. - num_after_skip: int - Number of residual blocks after the first skip connection. - num_concat: int - Number of residual blocks after the concatenation. - num_atom: int - Number of residual blocks in the atom embedding blocks. - - activation: str - Name of the activation function to use in the dense layers except for the final dense layer. - """ - - def __init__( - self, - emb_size_atom, - emb_size_edge, - emb_size_trip, - emb_size_rbf, - emb_size_cbf, - emb_size_bil_trip, - num_before_skip, - num_after_skip, - num_concat, - num_atom, - activation=None, - name="Interaction", - ): - super().__init__() - self.name = name - - block_nr = name.split("_")[-1] - - ## -------------------------------------------- Message Passing ------------------------------------------- ## - # Dense transformation of skip connection - self.dense_ca = Dense( - emb_size_edge, - emb_size_edge, - activation=activation, - bias=False, - ) - - # Triplet Interaction - self.trip_interaction = TripletInteraction( - emb_size_edge=emb_size_edge, - emb_size_trip=emb_size_trip, - emb_size_bilinear=emb_size_bil_trip, - emb_size_rbf=emb_size_rbf, - emb_size_cbf=emb_size_cbf, - activation=activation, - name=f"TripInteraction_{block_nr}", - ) - - ## ---------------------------------------- Update Edge Embeddings ---------------------------------------- ## - # Residual layers before skip connection - self.layers_before_skip = torch.nn.ModuleList( - [ - ResidualLayer( - emb_size_edge, - activation=activation, - ) - for i in range(num_before_skip) - ] - ) - - # Residual layers after skip connection - self.layers_after_skip = torch.nn.ModuleList( - [ - ResidualLayer( - emb_size_edge, - activation=activation, - ) - for i in range(num_after_skip) - ] - ) - - ## ---------------------------------------- Update Atom Embeddings ---------------------------------------- ## - self.atom_update = AtomUpdateBlock( - emb_size_atom=emb_size_atom, - emb_size_edge=emb_size_edge, - emb_size_rbf=emb_size_rbf, - nHidden=num_atom, - activation=activation, - name=f"AtomUpdate_{block_nr}", - ) - - ## ------------------------------ Update Edge Embeddings with Atom Embeddings ----------------------------- ## - self.concat_layer = EdgeEmbedding( - emb_size_atom, - emb_size_edge, - emb_size_edge, - activation=activation, - ) - self.residual_m = torch.nn.ModuleList( - [ - ResidualLayer(emb_size_edge, activation=activation) - for _ in range(num_concat) - ] - ) - - self.inv_sqrt_2 = 1 / math.sqrt(2.0) - - def forward( - self, - h, - m, - rbf3, - cbf3, - id3_ragged_idx, - id_swap, - id3_ba, - id3_ca, - rbf_h, - idx_s, - idx_t, - ): - """ - Returns - ------- - h: torch.Tensor, shape=(nEdges, emb_size_atom) - Atom embeddings. - m: torch.Tensor, shape=(nEdges, emb_size_edge) - Edge embeddings (c->a). - """ - - # Initial transformation - x_ca_skip = self.dense_ca(m) # (nEdges, emb_size_edge) - - x3 = self.trip_interaction( - m, - rbf3, - cbf3, - id3_ragged_idx, - id_swap, - id3_ba, - id3_ca, - ) - - ## ----------------------------- Merge Embeddings after Triplet Interaction ------------------------------ ## - x = x_ca_skip + x3 # (nEdges, emb_size_edge) - x = x * self.inv_sqrt_2 - - ## ---------------------------------------- Update Edge Embeddings --------------------------------------- ## - # Transformations before skip connection - for i, layer in enumerate(self.layers_before_skip): - x = layer(x) # (nEdges, emb_size_edge) - - # Skip connection - m = m + x # (nEdges, emb_size_edge) - m = m * self.inv_sqrt_2 - - # Transformations after skip connection - for i, layer in enumerate(self.layers_after_skip): - m = layer(m) # (nEdges, emb_size_edge) - - ## ---------------------------------------- Update Atom Embeddings --------------------------------------- ## - h2 = self.atom_update(h, m, rbf_h, idx_t) - - # Skip connection - h = h + h2 # (nAtoms, emb_size_atom) - h = h * self.inv_sqrt_2 - - ## ----------------------------- Update Edge Embeddings with Atom Embeddings ----------------------------- ## - m2 = self.concat_layer(h, m, idx_s, idx_t) # (nEdges, emb_size_edge) - - for i, layer in enumerate(self.residual_m): - m2 = layer(m2) # (nEdges, emb_size_edge) - - # Skip connection - m = m + m2 # (nEdges, emb_size_edge) - m = m * self.inv_sqrt_2 - return h, m - - -class TripletInteraction(torch.nn.Module): - """ - Triplet-based message passing block. - - Parameters - ---------- - emb_size_edge: int - Embedding size of the edges. - emb_size_trip: int - (Down-projected) Embedding size of the edge embeddings after the hadamard product with rbf. - emb_size_bilinear: int - Embedding size of the edge embeddings after the bilinear layer. - emb_size_rbf: int - Embedding size of the radial basis transformation. - emb_size_cbf: int - Embedding size of the circular basis transformation (one angle). - - activation: str - Name of the activation function to use in the dense layers except for the final dense layer. - """ - - def __init__( - self, - emb_size_edge, - emb_size_trip, - emb_size_bilinear, - emb_size_rbf, - emb_size_cbf, - activation=None, - name="TripletInteraction", - **kwargs, - ): - super().__init__() - self.name = name - - # Dense transformation - self.dense_ba = Dense( - emb_size_edge, - emb_size_edge, - activation=activation, - bias=False, - ) - - # Up projections of basis representations, bilinear layer and scaling factors - self.mlp_rbf = Dense( - emb_size_rbf, - emb_size_edge, - activation=None, - bias=False, - ) - self.scale_rbf = ScaleFactor(name + "_had_rbf") - - self.mlp_cbf = EfficientInteractionBilinear( - emb_size_trip, emb_size_cbf, emb_size_bilinear - ) - - # combines scaling for bilinear layer and summation - self.scale_cbf_sum = ScaleFactor(name + "_sum_cbf") - - # Down and up projections - self.down_projection = Dense( - emb_size_edge, - emb_size_trip, - activation=activation, - bias=False, - ) - self.up_projection_ca = Dense( - emb_size_bilinear, - emb_size_edge, - activation=activation, - bias=False, - ) - self.up_projection_ac = Dense( - emb_size_bilinear, - emb_size_edge, - activation=activation, - bias=False, - ) - - self.inv_sqrt_2 = 1 / math.sqrt(2.0) - - def forward( - self, - m, - rbf3, - cbf3, - id3_ragged_idx, - id_swap, - id3_ba, - id3_ca, - ): - """ - Returns - ------- - m: torch.Tensor, shape=(nEdges, emb_size_edge) - Edge embeddings (c->a). - """ - - # Dense transformation - x_ba = self.dense_ba(m) # (nEdges, emb_size_edge) - - # Transform via radial bessel basis - rbf_emb = self.mlp_rbf(rbf3) # (nEdges, emb_size_edge) - x_ba2 = x_ba * rbf_emb - x_ba = self.scale_rbf(x_ba2, ref=x_ba) - - x_ba = self.down_projection(x_ba) # (nEdges, emb_size_trip) - - # Transform via circular spherical basis - x_ba = x_ba[id3_ba] - - # Efficient bilinear layer - x = self.mlp_cbf(cbf3, x_ba, id3_ca, id3_ragged_idx) - # (nEdges, emb_size_quad) - x = self.scale_cbf_sum(x, ref=x_ba) - - # => - # rbf(d_ba) - # cbf(d_ca, angle_cab) - - # Up project embeddings - x_ca = self.up_projection_ca(x) # (nEdges, emb_size_edge) - x_ac = self.up_projection_ac(x) # (nEdges, emb_size_edge) - - # Merge interaction of c->a and a->c - x_ac = x_ac[id_swap] # swap to add to edge a->c and not c->a - x3 = x_ca + x_ac - x3 = x3 * self.inv_sqrt_2 - return x3 diff --git a/ocpmodels/models/gemnet/layers/radial_basis.py b/ocpmodels/models/gemnet/layers/radial_basis.py deleted file mode 100644 index cad79f2093c1f5480f1de030f40760f04c779983..0000000000000000000000000000000000000000 --- a/ocpmodels/models/gemnet/layers/radial_basis.py +++ /dev/null @@ -1,200 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import math - -import numpy as np -import torch -from scipy.special import binom -from torch_geometric.nn.models.schnet import GaussianSmearing - - -class PolynomialEnvelope(torch.nn.Module): - """ - Polynomial envelope function that ensures a smooth cutoff. - - Parameters - ---------- - exponent: int - Exponent of the envelope function. - """ - - def __init__(self, exponent): - super().__init__() - assert exponent > 0 - self.p = exponent - self.a = -(self.p + 1) * (self.p + 2) / 2 - self.b = self.p * (self.p + 2) - self.c = -self.p * (self.p + 1) / 2 - - def forward(self, d_scaled): - env_val = ( - 1 - + self.a * d_scaled**self.p - + self.b * d_scaled ** (self.p + 1) - + self.c * d_scaled ** (self.p + 2) - ) - return torch.where(d_scaled < 1, env_val, torch.zeros_like(d_scaled)) - - -class ExponentialEnvelope(torch.nn.Module): - """ - Exponential envelope function that ensures a smooth cutoff, - as proposed in Unke, Chmiela, Gastegger, Schütt, Sauceda, Müller 2021. - SpookyNet: Learning Force Fields with Electronic Degrees of Freedom - and Nonlocal Effects - """ - - def __init__(self): - super().__init__() - - def forward(self, d_scaled): - env_val = torch.exp(-(d_scaled**2) / ((1 - d_scaled) * (1 + d_scaled))) - return torch.where(d_scaled < 1, env_val, torch.zeros_like(d_scaled)) - - -class SphericalBesselBasis(torch.nn.Module): - """ - 1D spherical Bessel basis - - Parameters - ---------- - num_radial: int - Controls maximum frequency. - cutoff: float - Cutoff distance in Angstrom. - """ - - def __init__( - self, - num_radial: int, - cutoff: float, - ): - super().__init__() - self.norm_const = math.sqrt(2 / (cutoff**3)) - # cutoff ** 3 to counteract dividing by d_scaled = d / cutoff - - # Initialize frequencies at canonical positions - self.frequencies = torch.nn.Parameter( - data=torch.tensor(np.pi * np.arange(1, num_radial + 1, dtype=np.float32)), - requires_grad=True, - ) - - def forward(self, d_scaled): - return ( - self.norm_const - / d_scaled[:, None] - * torch.sin(self.frequencies * d_scaled[:, None]) - ) # (num_edges, num_radial) - - -class BernsteinBasis(torch.nn.Module): - """ - Bernstein polynomial basis, - as proposed in Unke, Chmiela, Gastegger, Schütt, Sauceda, Müller 2021. - SpookyNet: Learning Force Fields with Electronic Degrees of Freedom - and Nonlocal Effects - - Parameters - ---------- - num_radial: int - Controls maximum frequency. - pregamma_initial: float - Initial value of exponential coefficient gamma. - Default: gamma = 0.5 * a_0**-1 = 0.94486, - inverse softplus -> pregamma = log e**gamma - 1 = 0.45264 - """ - - def __init__( - self, - num_radial: int, - pregamma_initial: float = 0.45264, - ): - super().__init__() - prefactor = binom(num_radial - 1, np.arange(num_radial)) - self.register_buffer( - "prefactor", - torch.tensor(prefactor, dtype=torch.float), - persistent=False, - ) - - self.pregamma = torch.nn.Parameter( - data=torch.tensor(pregamma_initial, dtype=torch.float), - requires_grad=True, - ) - self.softplus = torch.nn.Softplus() - - exp1 = torch.arange(num_radial) - self.register_buffer("exp1", exp1[None, :], persistent=False) - exp2 = num_radial - 1 - exp1 - self.register_buffer("exp2", exp2[None, :], persistent=False) - - def forward(self, d_scaled): - gamma = self.softplus(self.pregamma) # constrain to positive - exp_d = torch.exp(-gamma * d_scaled)[:, None] - return self.prefactor * (exp_d**self.exp1) * ((1 - exp_d) ** self.exp2) - - -class RadialBasis(torch.nn.Module): - """ - - Parameters - ---------- - num_radial: int - Controls maximum frequency. - cutoff: float - Cutoff distance in Angstrom. - rbf: dict = {"name": "gaussian"} - Basis function and its hyperparameters. - envelope: dict = {"name": "polynomial", "exponent": 5} - Envelope function and its hyperparameters. - """ - - def __init__( - self, - num_radial: int, - cutoff: float, - rbf: dict = {"name": "gaussian"}, - envelope: dict = {"name": "polynomial", "exponent": 5}, - ): - super().__init__() - self.inv_cutoff = 1 / cutoff - - env_name = envelope["name"].lower() - env_hparams = envelope.copy() - del env_hparams["name"] - - if env_name == "polynomial": - self.envelope = PolynomialEnvelope(**env_hparams) - elif env_name == "exponential": - self.envelope = ExponentialEnvelope(**env_hparams) - else: - raise ValueError(f"Unknown envelope function '{env_name}'.") - - rbf_name = rbf["name"].lower() - rbf_hparams = rbf.copy() - del rbf_hparams["name"] - - # RBFs get distances scaled to be in [0, 1] - if rbf_name == "gaussian": - self.rbf = GaussianSmearing( - start=0, stop=1, num_gaussians=num_radial, **rbf_hparams - ) - elif rbf_name == "spherical_bessel": - self.rbf = SphericalBesselBasis( - num_radial=num_radial, cutoff=cutoff, **rbf_hparams - ) - elif rbf_name == "bernstein": - self.rbf = BernsteinBasis(num_radial=num_radial, **rbf_hparams) - else: - raise ValueError(f"Unknown radial basis function '{rbf_name}'.") - - def forward(self, d): - d_scaled = d * self.inv_cutoff - - env = self.envelope(d_scaled) - return env[:, None] * self.rbf(d_scaled) # (nEdges, num_radial) diff --git a/ocpmodels/models/gemnet/layers/spherical_basis.py b/ocpmodels/models/gemnet/layers/spherical_basis.py deleted file mode 100644 index 6695a83adf5c2d9c4413b844a2e266aa55f2e726..0000000000000000000000000000000000000000 --- a/ocpmodels/models/gemnet/layers/spherical_basis.py +++ /dev/null @@ -1,89 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import sympy as sym -import torch -from torch_geometric.nn.models.schnet import GaussianSmearing - -from .basis_utils import real_sph_harm -from .radial_basis import RadialBasis - - -class CircularBasisLayer(torch.nn.Module): - """ - 2D Fourier Bessel Basis - - Parameters - ---------- - num_spherical: int - Controls maximum frequency. - radial_basis: RadialBasis - Radial basis functions - cbf: dict - Name and hyperparameters of the cosine basis function - efficient: bool - Whether to use the "efficient" summation order - """ - - def __init__( - self, - num_spherical: int, - radial_basis: RadialBasis, - cbf: str, - efficient: bool = False, - ): - super().__init__() - - self.radial_basis = radial_basis - self.efficient = efficient - - cbf_name = cbf["name"].lower() - cbf_hparams = cbf.copy() - del cbf_hparams["name"] - - if cbf_name == "gaussian": - self.cosφ_basis = GaussianSmearing( - start=-1, stop=1, num_gaussians=num_spherical, **cbf_hparams - ) - elif cbf_name == "spherical_harmonics": - Y_lm = real_sph_harm(num_spherical, use_theta=False, zero_m_only=True) - sph_funcs = [] # (num_spherical,) - - # convert to tensorflow functions - z = sym.symbols("z") - modules = {"sin": torch.sin, "cos": torch.cos, "sqrt": torch.sqrt} - m_order = 0 # only single angle - for l_degree in range(len(Y_lm)): # num_spherical - if ( - l_degree == 0 - ): # Y_00 is only a constant -> function returns value and not tensor - first_sph = sym.lambdify([z], Y_lm[l_degree][m_order], modules) - sph_funcs.append(lambda z: torch.zeros_like(z) + first_sph(z)) - else: - sph_funcs.append( - sym.lambdify([z], Y_lm[l_degree][m_order], modules) - ) - self.cosφ_basis = lambda cosφ: torch.stack( - [f(cosφ) for f in sph_funcs], dim=1 - ) - else: - raise ValueError(f"Unknown cosine basis function '{cbf_name}'.") - - def forward(self, D_ca, cosφ_cab, id3_ca): - rbf = self.radial_basis(D_ca) # (num_edges, num_radial) - cbf = self.cosφ_basis(cosφ_cab) # (num_triplets, num_spherical) - - if not self.efficient: - rbf = rbf[id3_ca] # (num_triplets, num_radial) - out = (rbf[:, None, :] * cbf[:, :, None]).view( - -1, rbf.shape[-1] * cbf.shape[-1] - ) - return (out,) - # (num_triplets, num_radial * num_spherical) - else: - return (rbf[None, :, :], cbf) - # (1, num_edges, num_radial), (num_edges, num_spherical) diff --git a/ocpmodels/models/gemnet/utils.py b/ocpmodels/models/gemnet/utils.py deleted file mode 100644 index 9ff5c6c1e5acf4651e2f1802eaf090fb432b3149..0000000000000000000000000000000000000000 --- a/ocpmodels/models/gemnet/utils.py +++ /dev/null @@ -1,275 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import json - -import torch -from torch_scatter import segment_csr - - -def read_json(path): - """""" - if not path.endswith(".json"): - raise UserWarning(f"Path {path} is not a json-path.") - - with open(path, "r") as f: - content = json.load(f) - return content - - -def update_json(path, data): - """""" - if not path.endswith(".json"): - raise UserWarning(f"Path {path} is not a json-path.") - - content = read_json(path) - content.update(data) - write_json(path, content) - - -def write_json(path, data): - """""" - if not path.endswith(".json"): - raise UserWarning(f"Path {path} is not a json-path.") - - with open(path, "w", encoding="utf-8") as f: - json.dump(data, f, ensure_ascii=False, indent=4) - - -def read_value_json(path, key): - """""" - content = read_json(path) - - if key in content.keys(): - return content[key] - else: - return None - - -def ragged_range(sizes): - """Multiple concatenated ranges. - - Examples - -------- - sizes = [1 4 2 3] - Return: [0 0 1 2 3 0 1 0 1 2] - """ - assert sizes.dim() == 1 - if sizes.sum() == 0: - return sizes.new_empty(0) - - # Remove 0 sizes - sizes_nonzero = sizes > 0 - if not torch.all(sizes_nonzero): - sizes = torch.masked_select(sizes, sizes_nonzero) - - # Initialize indexing array with ones as we need to setup incremental indexing - # within each group when cumulatively summed at the final stage. - id_steps = torch.ones(sizes.sum(), dtype=torch.long, device=sizes.device) - id_steps[0] = 0 - insert_index = sizes[:-1].cumsum(0) - insert_val = (1 - sizes)[:-1] - - # Assign index-offsetting values - id_steps[insert_index] = insert_val - - # Finally index into input array for the group repeated o/p - res = id_steps.cumsum(0) - return res - - -def repeat_blocks( - sizes, - repeats, - continuous_indexing=True, - start_idx=0, - block_inc=0, - repeat_inc=0, -): - """Repeat blocks of indices. - Adapted from https://stackoverflow.com/questions/51154989/numpy-vectorized-function-to-repeat-blocks-of-consecutive-elements - - continuous_indexing: Whether to keep increasing the index after each block - start_idx: Starting index - block_inc: Number to increment by after each block, - either global or per block. Shape: len(sizes) - 1 - repeat_inc: Number to increment by after each repetition, - either global or per block - - Examples - -------- - sizes = [1,3,2] ; repeats = [3,2,3] ; continuous_indexing = False - Return: [0 0 0 0 1 2 0 1 2 0 1 0 1 0 1] - sizes = [1,3,2] ; repeats = [3,2,3] ; continuous_indexing = True - Return: [0 0 0 1 2 3 1 2 3 4 5 4 5 4 5] - sizes = [1,3,2] ; repeats = [3,2,3] ; continuous_indexing = True ; - repeat_inc = 4 - Return: [0 4 8 1 2 3 5 6 7 4 5 8 9 12 13] - sizes = [1,3,2] ; repeats = [3,2,3] ; continuous_indexing = True ; - start_idx = 5 - Return: [5 5 5 6 7 8 6 7 8 9 10 9 10 9 10] - sizes = [1,3,2] ; repeats = [3,2,3] ; continuous_indexing = True ; - block_inc = 1 - Return: [0 0 0 2 3 4 2 3 4 6 7 6 7 6 7] - sizes = [0,3,2] ; repeats = [3,2,3] ; continuous_indexing = True - Return: [0 1 2 0 1 2 3 4 3 4 3 4] - sizes = [2,3,2] ; repeats = [2,0,2] ; continuous_indexing = True - Return: [0 1 0 1 5 6 5 6] - """ - assert sizes.dim() == 1 - assert all(sizes >= 0) - - # Remove 0 sizes - sizes_nonzero = sizes > 0 - if not torch.all(sizes_nonzero): - assert block_inc == 0 # Implementing this is not worth the effort - sizes = torch.masked_select(sizes, sizes_nonzero) - if isinstance(repeats, torch.Tensor): - repeats = torch.masked_select(repeats, sizes_nonzero) - if isinstance(repeat_inc, torch.Tensor): - repeat_inc = torch.masked_select(repeat_inc, sizes_nonzero) - - if isinstance(repeats, torch.Tensor): - assert all(repeats >= 0) - insert_dummy = repeats[0] == 0 - if insert_dummy: - one = sizes.new_ones(1) - zero = sizes.new_zeros(1) - sizes = torch.cat((one, sizes)) - repeats = torch.cat((one, repeats)) - if isinstance(block_inc, torch.Tensor): - block_inc = torch.cat((zero, block_inc)) - if isinstance(repeat_inc, torch.Tensor): - repeat_inc = torch.cat((zero, repeat_inc)) - else: - assert repeats >= 0 - insert_dummy = False - - # Get repeats for each group using group lengths/sizes - r1 = torch.repeat_interleave(torch.arange(len(sizes), device=sizes.device), repeats) - - # Get total size of output array, as needed to initialize output indexing array - N = (sizes * repeats).sum() - - # Initialize indexing array with ones as we need to setup incremental indexing - # within each group when cumulatively summed at the final stage. - # Two steps here: - # 1. Within each group, we have multiple sequences, so setup the offsetting - # at each sequence lengths by the seq. lengths preceding those. - id_ar = torch.ones(N, dtype=torch.long, device=sizes.device) - id_ar[0] = 0 - insert_index = sizes[r1[:-1]].cumsum(0) - insert_val = (1 - sizes)[r1[:-1]] - - if isinstance(repeats, torch.Tensor) and torch.any(repeats == 0): - diffs = r1[1:] - r1[:-1] - indptr = torch.cat((sizes.new_zeros(1), diffs.cumsum(0))) - if continuous_indexing: - # If a group was skipped (repeats=0) we need to add its size - insert_val += segment_csr(sizes[: r1[-1]], indptr, reduce="sum") - - # Add block increments - if isinstance(block_inc, torch.Tensor): - insert_val += segment_csr(block_inc[: r1[-1]], indptr, reduce="sum") - else: - insert_val += block_inc * (indptr[1:] - indptr[:-1]) - if insert_dummy: - insert_val[0] -= block_inc - else: - idx = r1[1:] != r1[:-1] - if continuous_indexing: - # 2. For each group, make sure the indexing starts from the next group's - # first element. So, simply assign 1s there. - insert_val[idx] = 1 - - # Add block increments - insert_val[idx] += block_inc - - # Add repeat_inc within each group - if isinstance(repeat_inc, torch.Tensor): - insert_val += repeat_inc[r1[:-1]] - if isinstance(repeats, torch.Tensor): - repeat_inc_inner = repeat_inc[repeats > 0][:-1] - else: - repeat_inc_inner = repeat_inc[:-1] - else: - insert_val += repeat_inc - repeat_inc_inner = repeat_inc - - # Subtract the increments between groups - if isinstance(repeats, torch.Tensor): - repeats_inner = repeats[repeats > 0][:-1] - else: - repeats_inner = repeats - insert_val[r1[1:] != r1[:-1]] -= repeat_inc_inner * repeats_inner - - # Assign index-offsetting values - id_ar[insert_index] = insert_val - - if insert_dummy: - id_ar = id_ar[1:] - if continuous_indexing: - id_ar[0] -= 1 - - # Set start index now, in case of insertion due to leading repeats=0 - id_ar[0] += start_idx - - # Finally index into input array for the group repeated o/p - res = id_ar.cumsum(0) - return res - - -def calculate_interatomic_vectors(R, id_s, id_t, offsets_st): - """ - Calculate the vectors connecting the given atom pairs, - considering offsets from periodic boundary conditions (PBC). - - Parameters - ---------- - R: Tensor, shape = (nAtoms, 3) - Atom positions. - id_s: Tensor, shape = (nEdges,) - Indices of the source atom of the edges. - id_t: Tensor, shape = (nEdges,) - Indices of the target atom of the edges. - offsets_st: Tensor, shape = (nEdges,) - PBC offsets of the edges. - Subtract this from the correct direction. - - Returns - ------- - (D_st, V_st): tuple - D_st: Tensor, shape = (nEdges,) - Distance from atom t to s. - V_st: Tensor, shape = (nEdges,) - Unit direction from atom t to s. - """ - Rs = R[id_s] - Rt = R[id_t] - # ReLU prevents negative numbers in sqrt - if offsets_st is None: - V_st = Rt - Rs # s -> t - else: - V_st = Rt - Rs + offsets_st # s -> t - D_st = torch.sqrt(torch.sum(V_st**2, dim=1)) - V_st = V_st / D_st[..., None] - return D_st, V_st - - -def inner_product_normalized(x, y): - """ - Calculate the inner product between the given normalized vectors, - giving a result between -1 and 1. - """ - return torch.sum(x * y, dim=-1).clamp(min=-1, max=1) - - -def mask_neighbors(neighbors, edge_mask): - neighbors_old_indptr = torch.cat([neighbors.new_zeros(1), neighbors]) - neighbors_old_indptr = torch.cumsum(neighbors_old_indptr, dim=0) - neighbors = segment_csr(edge_mask.long(), neighbors_old_indptr) - return neighbors diff --git a/ocpmodels/models/gemnet_gp/README.md b/ocpmodels/models/gemnet_gp/README.md deleted file mode 100644 index 1b47420c495b036e48a4dbea38f0c13a6350e7ed..0000000000000000000000000000000000000000 --- a/ocpmodels/models/gemnet_gp/README.md +++ /dev/null @@ -1,30 +0,0 @@ -# Towards Training Billion Parameter Graph Neural Networks for Atomic Simulations - -Anuroop Sriram, Abhishek Das, Brandon M. Wood, Siddharth Goyal, C. Lawrence Zitnick - -[[`arXiv:2203.09697`](https://arxiv.org/abs/2203.09697)] - - -To use graph parallel training, add `--gp-gpus N` to your command line, where N = number of GPUs to split the model over. This flag works for all tasks (`train`, `predict`, `validate` & `run-relaxations`). - -As an example, the Gemnet-XL model can be trained using: -```bash -python main.py --mode train --config-yml configs/s2ef/all/gp_gemnet/gp-gemnet-xl.yml \ - --distributed --num-nodes 32 --num-gpus 8 --gp-gpus 4 -``` -This trains the model on 256 GPUs (32 nodes x 8 GPUs each) with 4-way graph parallelism (i.e. the graph is distributed over 4 GPUs) and 64-way data parallelism (64 == 256 / 4). - -The Gemnet-XL model was trained without AMP as it led to unstable training. - -## Citing - -If you use Graph Parallelism in your work, please consider citing: - -```bibtex -@inproceedings{sriram_graphparallel_2022, - title={{Towards Training Billion Parameter Graph Neural Networks for Atomic Simulations}}, - author={Sriram, Anuroop and Das, Abhishek and Wood, Brandon M. and Goyal, Siddharth and Zitnick, C. Lawrence}, - booktitle={International Conference on Learning Representations (ICLR)}, - year={2022} -} -``` diff --git a/ocpmodels/models/gemnet_gp/gemnet.py b/ocpmodels/models/gemnet_gp/gemnet.py deleted file mode 100644 index 8a9658e9e5cec65c9d5bbbf5f108131ee5245dc3..0000000000000000000000000000000000000000 --- a/ocpmodels/models/gemnet_gp/gemnet.py +++ /dev/null @@ -1,649 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -from typing import Optional - -import numpy as np -import torch -from torch_cluster import radius_graph -from torch_scatter import scatter -from torch_sparse import SparseTensor - -from ocpmodels.common import distutils, gp_utils -from ocpmodels.common.registry import registry -from ocpmodels.common.utils import ( - compute_neighbors, - conditional_grad, - get_pbc_distances, - radius_graph_pbc, -) -from ocpmodels.models.base import BaseModel -from ocpmodels.modules.scaling.compat import load_scales_compat - -from .layers.atom_update_block import OutputBlock -from .layers.base_layers import Dense -from .layers.efficient import EfficientInteractionDownProjection -from .layers.embedding_block import AtomEmbedding, EdgeEmbedding -from .layers.interaction_block import InteractionBlockTripletsOnly -from .layers.radial_basis import RadialBasis -from .layers.spherical_basis import CircularBasisLayer -from .utils import ( - inner_product_normalized, - mask_neighbors, - ragged_range, - repeat_blocks, -) - - -@registry.register_model("gp_gemnet_t") -class GraphParallelGemNetT(BaseModel): - """ - GemNet-T, triplets-only variant of GemNet - - Parameters - ---------- - num_atoms (int): Unused argument - bond_feat_dim (int): Unused argument - num_targets: int - Number of prediction targets. - - num_spherical: int - Controls maximum frequency. - num_radial: int - Controls maximum frequency. - num_blocks: int - Number of building blocks to be stacked. - - emb_size_atom: int - Embedding size of the atoms. - emb_size_edge: int - Embedding size of the edges. - emb_size_trip: int - (Down-projected) Embedding size in the triplet message passing block. - emb_size_rbf: int - Embedding size of the radial basis transformation. - emb_size_cbf: int - Embedding size of the circular basis transformation (one angle). - emb_size_bil_trip: int - Embedding size of the edge embeddings in the triplet-based message passing block after the bilinear layer. - - num_before_skip: int - Number of residual blocks before the first skip connection. - num_after_skip: int - Number of residual blocks after the first skip connection. - num_concat: int - Number of residual blocks after the concatenation. - num_atom: int - Number of residual blocks in the atom embedding blocks. - - regress_forces: bool - Whether to predict forces. Default: True - direct_forces: bool - If True predict forces based on aggregation of interatomic directions. - If False predict forces based on negative gradient of energy potential. - - cutoff: float - Embedding cutoff for interactomic directions in Angstrom. - rbf: dict - Name and hyperparameters of the radial basis function. - envelope: dict - Name and hyperparameters of the envelope function. - cbf: dict - Name and hyperparameters of the cosine basis function. - extensive: bool - Whether the output should be extensive (proportional to the number of atoms) - output_init: str - Initialization method for the final dense layer. - activation: str - Name of the activation function. - scale_file: str - Path to the json file containing the scaling factors. - """ - - def __init__( - self, - num_atoms: Optional[int], - bond_feat_dim: int, - num_targets: int, - num_spherical: int, - num_radial: int, - num_blocks: int, - emb_size_atom: int, - emb_size_edge: int, - emb_size_trip: int, - emb_size_rbf: int, - emb_size_cbf: int, - emb_size_bil_trip: int, - num_before_skip: int, - num_after_skip: int, - num_concat: int, - num_atom: int, - regress_forces: bool = True, - direct_forces: bool = False, - cutoff: float = 6.0, - max_neighbors: int = 50, - rbf: dict = {"name": "gaussian"}, - envelope: dict = {"name": "polynomial", "exponent": 5}, - cbf: dict = {"name": "spherical_harmonics"}, - extensive: bool = True, - otf_graph: bool = False, - use_pbc: bool = True, - output_init: str = "HeOrthogonal", - activation: str = "swish", - scale_num_blocks: bool = False, - scatter_atoms: bool = True, - scale_file: Optional[str] = None, - ): - super().__init__() - self.num_targets = num_targets - assert num_blocks > 0 - self.num_blocks = num_blocks - self.extensive = extensive - self.scale_num_blocks = scale_num_blocks - self.scatter_atoms = scatter_atoms - - self.cutoff = cutoff - assert self.cutoff <= 6 or otf_graph - - self.max_neighbors = max_neighbors - assert self.max_neighbors == 50 or otf_graph - - self.regress_forces = regress_forces - self.otf_graph = otf_graph - self.use_pbc = use_pbc - - # GemNet variants - self.direct_forces = direct_forces - - ### ---------------------------------- Basis Functions ---------------------------------- ### - self.radial_basis = RadialBasis( - num_radial=num_radial, - cutoff=cutoff, - rbf=rbf, - envelope=envelope, - ) - - radial_basis_cbf3 = RadialBasis( - num_radial=num_radial, - cutoff=cutoff, - rbf=rbf, - envelope=envelope, - ) - self.cbf_basis3 = CircularBasisLayer( - num_spherical, - radial_basis=radial_basis_cbf3, - cbf=cbf, - efficient=True, - ) - ### ------------------------------------------------------------------------------------- ### - - ### ------------------------------- Share Down Projections ------------------------------ ### - # Share down projection across all interaction blocks - self.mlp_rbf3 = Dense( - num_radial, - emb_size_rbf, - activation=None, - bias=False, - ) - self.mlp_cbf3 = EfficientInteractionDownProjection( - num_spherical, num_radial, emb_size_cbf - ) - - # Share the dense Layer of the atom embedding block accross the interaction blocks - self.mlp_rbf_h = Dense( - num_radial, - emb_size_rbf, - activation=None, - bias=False, - ) - self.mlp_rbf_out = Dense( - num_radial, - emb_size_rbf, - activation=None, - bias=False, - ) - ### ------------------------------------------------------------------------------------- ### - - # Embedding block - self.atom_emb = AtomEmbedding(emb_size_atom) - self.edge_emb = EdgeEmbedding( - emb_size_atom, num_radial, emb_size_edge, activation=activation - ) - - out_blocks = [] - int_blocks = [] - - # Interaction Blocks - interaction_block = InteractionBlockTripletsOnly # GemNet-(d)T - for i in range(num_blocks): - int_blocks.append( - interaction_block( - emb_size_atom=emb_size_atom, - emb_size_edge=emb_size_edge, - emb_size_trip=emb_size_trip, - emb_size_rbf=emb_size_rbf, - emb_size_cbf=emb_size_cbf, - emb_size_bil_trip=emb_size_bil_trip, - num_before_skip=num_before_skip, - num_after_skip=num_after_skip, - num_concat=num_concat, - num_atom=num_atom, - activation=activation, - name=f"IntBlock_{i+1}", - ) - ) - - for i in range(num_blocks + 1): - out_blocks.append( - OutputBlock( - emb_size_atom=emb_size_atom, - emb_size_edge=emb_size_edge, - emb_size_rbf=emb_size_rbf, - nHidden=num_atom, - num_targets=num_targets, - activation=activation, - output_init=output_init, - direct_forces=direct_forces, - name=f"OutBlock_{i}", - ) - ) - - self.out_blocks = torch.nn.ModuleList(out_blocks) - self.int_blocks = torch.nn.ModuleList(int_blocks) - - load_scales_compat(self, scale_file) - - def get_triplets(self, edge_index, num_atoms): - """ - Get all b->a for each edge c->a. - It is possible that b=c, as long as the edges are distinct. - - Returns - ------- - id3_ba: torch.Tensor, shape (num_triplets,) - Indices of input edge b->a of each triplet b->a<-c - id3_ca: torch.Tensor, shape (num_triplets,) - Indices of output edge c->a of each triplet b->a<-c - id3_ragged_idx: torch.Tensor, shape (num_triplets,) - Indices enumerating the copies of id3_ca for creating a padded matrix - """ - idx_s, idx_t = edge_index # c->a (source=c, target=a) - - value = torch.arange(idx_s.size(0), device=idx_s.device, dtype=idx_s.dtype) - # Possibly contains multiple copies of the same edge (for periodic interactions) - adj = SparseTensor( - row=idx_t, - col=idx_s, - value=value, - sparse_sizes=(num_atoms, num_atoms), - ) - adj_edges = adj[idx_t] - - # Edge indices (b->a, c->a) for triplets. - id3_ba = adj_edges.storage.value() - id3_ca = adj_edges.storage.row() - - # Remove self-loop triplets - # Compare edge indices, not atom indices to correctly handle periodic interactions - mask = id3_ba != id3_ca - id3_ba = id3_ba[mask] - id3_ca = id3_ca[mask] - - # Get indices to reshape the neighbor indices b->a into a dense matrix. - # id3_ca has to be sorted for this to work. - num_triplets = torch.bincount(id3_ca, minlength=idx_s.size(0)) - id3_ragged_idx = ragged_range(num_triplets) - - return id3_ba, id3_ca, id3_ragged_idx - - def select_symmetric_edges(self, tensor, mask, reorder_idx, inverse_neg): - # Mask out counter-edges - tensor_directed = tensor[mask] - # Concatenate counter-edges after normal edges - sign = 1 - 2 * inverse_neg - tensor_cat = torch.cat([tensor_directed, sign * tensor_directed]) - # Reorder everything so the edges of every image are consecutive - tensor_ordered = tensor_cat[reorder_idx] - return tensor_ordered - - def reorder_symmetric_edges( - self, edge_index, cell_offsets, neighbors, edge_dist, edge_vector - ): - """ - Reorder edges to make finding counter-directional edges easier. - - Some edges are only present in one direction in the data, - since every atom has a maximum number of neighbors. Since we only use i->j - edges here, we lose some j->i edges and add others by - making it symmetric. - We could fix this by merging edge_index with its counter-edges, - including the cell_offsets, and then running torch.unique. - But this does not seem worth it. - """ - - # Generate mask - mask_sep_atoms = edge_index[0] < edge_index[1] - # Distinguish edges between the same (periodic) atom by ordering the cells - cell_earlier = ( - (cell_offsets[:, 0] < 0) - | ((cell_offsets[:, 0] == 0) & (cell_offsets[:, 1] < 0)) - | ( - (cell_offsets[:, 0] == 0) - & (cell_offsets[:, 1] == 0) - & (cell_offsets[:, 2] < 0) - ) - ) - mask_same_atoms = edge_index[0] == edge_index[1] - mask_same_atoms &= cell_earlier - mask = mask_sep_atoms | mask_same_atoms - - # Mask out counter-edges - edge_index_new = edge_index[mask[None, :].expand(2, -1)].view(2, -1) - - # Concatenate counter-edges after normal edges - edge_index_cat = torch.cat( - [ - edge_index_new, - torch.stack([edge_index_new[1], edge_index_new[0]], dim=0), - ], - dim=1, - ) - - # Count remaining edges per image - neighbors = neighbors.to(edge_index.device) - batch_edge = torch.repeat_interleave( - torch.arange(neighbors.size(0), device=edge_index.device), - neighbors, - ) - batch_edge = batch_edge[mask] - neighbors_new = 2 * torch.bincount(batch_edge, minlength=neighbors.size(0)) - - # Create indexing array - edge_reorder_idx = repeat_blocks( - neighbors_new // 2, - repeats=2, - continuous_indexing=True, - repeat_inc=edge_index_new.size(1), - ) - - # Reorder everything so the edges of every image are consecutive - edge_index_new = edge_index_cat[:, edge_reorder_idx] - cell_offsets_new = self.select_symmetric_edges( - cell_offsets, mask, edge_reorder_idx, True - ) - edge_dist_new = self.select_symmetric_edges( - edge_dist, mask, edge_reorder_idx, False - ) - edge_vector_new = self.select_symmetric_edges( - edge_vector, mask, edge_reorder_idx, True - ) - - return ( - edge_index_new, - cell_offsets_new, - neighbors_new, - edge_dist_new, - edge_vector_new, - ) - - def select_edges( - self, - data, - edge_index, - cell_offsets, - neighbors, - edge_dist, - edge_vector, - cutoff=None, - ): - if cutoff is not None: - edge_mask = edge_dist <= cutoff - - edge_index = edge_index[:, edge_mask] - cell_offsets = cell_offsets[edge_mask] - neighbors = mask_neighbors(neighbors, edge_mask) - edge_dist = edge_dist[edge_mask] - edge_vector = edge_vector[edge_mask] - - empty_image = neighbors == 0 - if torch.any(empty_image): - raise ValueError( - f"An image has no neighbors: id={data.id[empty_image]}, " - f"sid={data.sid[empty_image]}, fid={data.fid[empty_image]}" - ) - return edge_index, cell_offsets, neighbors, edge_dist, edge_vector - - def generate_interaction_graph(self, data): - num_atoms = data.atomic_numbers.size(0) - - ( - edge_index, - D_st, - distance_vec, - cell_offsets, - _, # cell offset distances - neighbors, - ) = self.generate_graph(data) - # These vectors actually point in the opposite direction. - # But we want to use col as idx_t for efficient aggregation. - V_st = -distance_vec / D_st[:, None] - - # Mask interaction edges if required - if self.otf_graph or np.isclose(self.cutoff, 6): - select_cutoff = None - else: - select_cutoff = self.cutoff - ( - edge_index, - cell_offsets, - neighbors, - D_st, - V_st, - ) = self.select_edges( - data=data, - edge_index=edge_index, - cell_offsets=cell_offsets, - neighbors=neighbors, - edge_dist=D_st, - edge_vector=V_st, - cutoff=select_cutoff, - ) - - ( - edge_index, - cell_offsets, - neighbors, - D_st, - V_st, - ) = self.reorder_symmetric_edges( - edge_index, cell_offsets, neighbors, D_st, V_st - ) - - # Indices for swapping c->a and a->c (for symmetric MP) - block_sizes = neighbors // 2 - id_swap = repeat_blocks( - block_sizes, - repeats=2, - continuous_indexing=False, - start_idx=block_sizes[0], - block_inc=block_sizes[:-1] + block_sizes[1:], - repeat_inc=-block_sizes, - ) - - id3_ba, id3_ca, id3_ragged_idx = self.get_triplets( - edge_index, num_atoms=num_atoms - ) - - return ( - edge_index, - neighbors, - D_st, - V_st, - id_swap, - id3_ba, - id3_ca, - id3_ragged_idx, - ) - - @conditional_grad(torch.enable_grad()) - def forward(self, data): - pos = data.pos - batch = data.batch - atomic_numbers = data.atomic_numbers.long() - - if self.regress_forces and not self.direct_forces: - pos.requires_grad_(True) - - ( - edge_index, - neighbors, - D_st, - V_st, - id_swap, - id3_ba, - id3_ca, - id3_ragged_idx, - ) = self.generate_interaction_graph(data) - idx_s, idx_t = edge_index - - # Graph Parallel: Precompute Kmax so all processes have the same value - Kmax = torch.max( - torch.max(id3_ragged_idx) + 1, - torch.tensor(0).to(id3_ragged_idx.device), - ) - - # Graph Parallel: Scatter triplets (consistent with edge splits) - edge_partition = gp_utils.scatter_to_model_parallel_region( - torch.arange(edge_index.size(1)) - ) - triplet_partition = torch.where( - torch.logical_and( - id3_ca >= edge_partition.min(), id3_ca <= edge_partition.max() - ) - )[0] - id3_ba = id3_ba[triplet_partition] - id3_ca = id3_ca[triplet_partition] - id3_ragged_idx = id3_ragged_idx[triplet_partition] - edge_offset = edge_partition.min() - - # Calculate triplet angles - cosφ_cab = inner_product_normalized(V_st[id3_ca], V_st[id3_ba]) - rad_cbf3, cbf3 = self.cbf_basis3(D_st, cosφ_cab, id3_ca) - - # TODO: Only do this for the partitioned edges - cbf3 = self.mlp_cbf3(rad_cbf3, cbf3, id3_ca, id3_ragged_idx, Kmax) - - # Graph Paralllel: Scatter edges - D_st = gp_utils.scatter_to_model_parallel_region(D_st, dim=0) - cbf3 = ( - gp_utils.scatter_to_model_parallel_region(cbf3[0], dim=0), - gp_utils.scatter_to_model_parallel_region(cbf3[1], dim=0), - ) - idx_s = gp_utils.scatter_to_model_parallel_region(idx_s, dim=0) - idx_t_full = idx_t - idx_t = gp_utils.scatter_to_model_parallel_region(idx_t, dim=0) - - rbf = self.radial_basis(D_st) - - # Graph Paralllel: Scatter Nodes - nAtoms = atomic_numbers.shape[0] - if self.scatter_atoms: - atomic_numbers = gp_utils.scatter_to_model_parallel_region( - atomic_numbers, dim=0 - ) - - # Embedding block - h = self.atom_emb(atomic_numbers) - # (nAtoms, emb_size_atom) - m = self.edge_emb(h, rbf, idx_s, idx_t) # (nEdges, emb_size_edge) - - rbf3 = self.mlp_rbf3(rbf) - - rbf_h = self.mlp_rbf_h(rbf) - rbf_out = self.mlp_rbf_out(rbf) - - E_t, F_st = self.out_blocks[0](nAtoms, m, rbf_out, idx_t) - # (nAtoms, num_targets), (nEdges, num_targets) - - for i in range(self.num_blocks): - # Interaction block - h, m = self.int_blocks[i]( - h=h, - m=m, - rbf3=rbf3, - cbf3=cbf3, - id3_ragged_idx=id3_ragged_idx, - id_swap=id_swap, - id3_ba=id3_ba, - id3_ca=id3_ca, - rbf_h=rbf_h, - idx_s=idx_s, - idx_t=idx_t, - edge_offset=edge_offset, - Kmax=Kmax, - nAtoms=nAtoms, - ) # (nAtoms, emb_size_atom), (nEdges, emb_size_edge) - - E, F = self.out_blocks[i + 1](nAtoms, m, rbf_out, idx_t) - # (nAtoms, num_targets), (nEdges, num_targets) - F_st += F - E_t += E - - if self.scale_num_blocks: - F_st = F_st / (self.num_blocks + 1) - E_t = E_t / (self.num_blocks + 1) - - # Graph Parallel: Gather F_st - F_st = gp_utils.gather_from_model_parallel_region(F_st, dim=0) - - nMolecules = torch.max(batch) + 1 - if self.extensive: - E_t = gp_utils.gather_from_model_parallel_region(E_t, dim=0) - E_t = scatter( - E_t, batch, dim=0, dim_size=nMolecules, reduce="add" - ) # (nMolecules, num_targets) - else: - E_t = scatter( - E_t, batch, dim=0, dim_size=nMolecules, reduce="mean" - ) # (nMolecules, num_targets) - - if self.regress_forces: - if self.direct_forces: - # map forces in edge directions - F_st_vec = F_st[:, :, None] * V_st[:, None, :] - # (nEdges, num_targets, 3) - F_t = scatter( - F_st_vec, - idx_t_full, - dim=0, - dim_size=data.atomic_numbers.size(0), - reduce="add", - ) # (nAtoms, num_targets, 3) - F_t = F_t.squeeze(1) # (nAtoms, 3) - else: - if self.num_targets > 1: - forces = [] - for i in range(self.num_targets): - # maybe this can be solved differently - forces += [ - -torch.autograd.grad( - E_t[:, i].sum(), pos, create_graph=True - )[0] - ] - F_t = torch.stack(forces, dim=1) - # (nAtoms, num_targets, 3) - else: - F_t = -torch.autograd.grad(E_t.sum(), pos, create_graph=True)[0] - # (nAtoms, 3) - - return E_t, F_t # (nMolecules, num_targets), (nAtoms, 3) - else: - return E_t - - @property - def num_params(self): - return sum(p.numel() for p in self.parameters()) diff --git a/ocpmodels/models/gemnet_gp/initializers.py b/ocpmodels/models/gemnet_gp/initializers.py deleted file mode 100644 index 39a07ce6a5daadb7dc9ce351199ae24b152b6a1d..0000000000000000000000000000000000000000 --- a/ocpmodels/models/gemnet_gp/initializers.py +++ /dev/null @@ -1,47 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import torch - - -def _standardize(kernel): - """ - Makes sure that N*Var(W) = 1 and E[W] = 0 - """ - eps = 1e-6 - - if len(kernel.shape) == 3: - axis = [0, 1] # last dimension is output dimension - else: - axis = 1 - - var, mean = torch.var_mean(kernel, dim=axis, unbiased=True, keepdim=True) - kernel = (kernel - mean) / (var + eps) ** 0.5 - return kernel - - -def he_orthogonal_init(tensor): - """ - Generate a weight matrix with variance according to He (Kaiming) initialization. - Based on a random (semi-)orthogonal matrix neural networks - are expected to learn better when features are decorrelated - (stated by eg. "Reducing overfitting in deep networks by decorrelating representations", - "Dropout: a simple way to prevent neural networks from overfitting", - "Exact solutions to the nonlinear dynamics of learning in deep linear neural networks") - """ - tensor = torch.nn.init.orthogonal_(tensor) - - if len(tensor.shape) == 3: - fan_in = tensor.shape[:-1].numel() - else: - fan_in = tensor.shape[1] - - with torch.no_grad(): - tensor.data = _standardize(tensor.data) - tensor.data *= (1 / fan_in) ** 0.5 - - return tensor diff --git a/ocpmodels/models/gemnet_gp/layers/atom_update_block.py b/ocpmodels/models/gemnet_gp/layers/atom_update_block.py deleted file mode 100644 index d956d4000796960ba2aa77c1a36a9415083b14d4..0000000000000000000000000000000000000000 --- a/ocpmodels/models/gemnet_gp/layers/atom_update_block.py +++ /dev/null @@ -1,238 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -from typing import Optional - -import torch -from torch_scatter import scatter -from torch_scatter.utils import broadcast - -from ocpmodels.common import gp_utils -from ocpmodels.modules.scaling import ScaleFactor - -from ..initializers import he_orthogonal_init -from .base_layers import Dense, ResidualLayer - - -def scatter_sum( - src: torch.Tensor, - index: torch.Tensor, - dim: int = -1, - out: Optional[torch.Tensor] = None, - dim_size: Optional[int] = None, -) -> torch.Tensor: - """ - Clone of torch_scatter.scatter_sum but without in-place operations - """ - index = broadcast(index, src, dim) - if out is None: - size = list(src.size()) - if dim_size is not None: - size[dim] = dim_size - elif index.numel() == 0: - size[dim] = 0 - else: - size[dim] = int(index.max()) + 1 - - out = torch.zeros(size, dtype=src.dtype, device=src.device) - return torch.scatter_add(out, dim, index, src) - else: - return out.scatter_add(dim, index, src) - - -class AtomUpdateBlock(torch.nn.Module): - """ - Aggregate the message embeddings of the atoms - - Parameters - ---------- - emb_size_atom: int - Embedding size of the atoms. - emb_size_atom: int - Embedding size of the edges. - nHidden: int - Number of residual blocks. - activation: callable/str - Name of the activation function to use in the dense layers. - """ - - def __init__( - self, - emb_size_atom: int, - emb_size_edge: int, - emb_size_rbf: int, - nHidden: int, - activation=None, - name: str = "atom_update", - ): - super().__init__() - self.name = name - - self.dense_rbf = Dense(emb_size_rbf, emb_size_edge, activation=None, bias=False) - self.scale_sum = ScaleFactor(name + "_sum") - - self.layers = self.get_mlp(emb_size_edge, emb_size_atom, nHidden, activation) - - def get_mlp(self, units_in, units, nHidden, activation): - dense1 = Dense(units_in, units, activation=activation, bias=False) - mlp = [dense1] - res = [ - ResidualLayer(units, nLayers=2, activation=activation) - for i in range(nHidden) - ] - mlp = mlp + res - return torch.nn.ModuleList(mlp) - - def forward(self, nAtoms, m, rbf, id_j): - """ - Returns - ------- - h: torch.Tensor, shape=(nAtoms, emb_size_atom) - Atom embedding. - """ - mlp_rbf = self.dense_rbf(rbf) # (nEdges, emb_size_edge) - x = m * mlp_rbf - - # Graph Parallel: Local node aggregation - x2 = scatter(x, id_j, dim=0, dim_size=nAtoms, reduce="sum") - - # Graph Parallel: Global node aggregation - x2 = gp_utils.reduce_from_model_parallel_region(x2) - x2 = gp_utils.scatter_to_model_parallel_region(x2, dim=0) - - # (nAtoms, emb_size_edge) - x = self.scale_sum(x2, ref=m) - - for layer in self.layers: - x = layer(x) # (nAtoms, emb_size_atom) - - return x - - -class OutputBlock(AtomUpdateBlock): - """ - Combines the atom update block and subsequent final dense layer. - - Parameters - ---------- - emb_size_atom: int - Embedding size of the atoms. - emb_size_atom: int - Embedding size of the edges. - nHidden: int - Number of residual blocks. - num_targets: int - Number of targets. - activation: str - Name of the activation function to use in the dense layers except for the final dense layer. - direct_forces: bool - If true directly predict forces without taking the gradient of the energy potential. - output_init: int - Kernel initializer of the final dense layer. - """ - - def __init__( - self, - emb_size_atom: int, - emb_size_edge: int, - emb_size_rbf: int, - nHidden: int, - num_targets: int, - activation=None, - direct_forces=True, - output_init="HeOrthogonal", - name: str = "output", - **kwargs, - ): - - super().__init__( - name=name, - emb_size_atom=emb_size_atom, - emb_size_edge=emb_size_edge, - emb_size_rbf=emb_size_rbf, - nHidden=nHidden, - activation=activation, - **kwargs, - ) - - assert isinstance(output_init, str) - self.output_init = output_init.lower() - self.direct_forces = direct_forces - - self.seq_energy = self.layers # inherited from parent class - self.out_energy = Dense(emb_size_atom, num_targets, bias=False, activation=None) - - if self.direct_forces: - self.scale_rbf_F = ScaleFactor(name + "_had") - self.seq_forces = self.get_mlp( - emb_size_edge, emb_size_edge, nHidden, activation - ) - self.out_forces = Dense( - emb_size_edge, num_targets, bias=False, activation=None - ) - self.dense_rbf_F = Dense( - emb_size_rbf, emb_size_edge, activation=None, bias=False - ) - - self.reset_parameters() - - def reset_parameters(self): - if self.output_init == "heorthogonal": - self.out_energy.reset_parameters(he_orthogonal_init) - if self.direct_forces: - self.out_forces.reset_parameters(he_orthogonal_init) - elif self.output_init == "zeros": - self.out_energy.reset_parameters(torch.nn.init.zeros_) - if self.direct_forces: - self.out_forces.reset_parameters(torch.nn.init.zeros_) - else: - raise UserWarning(f"Unknown output_init: {self.output_init}") - - def forward(self, nAtoms, m, rbf, id_j): - """ - Returns - ------- - (E, F): tuple - - E: torch.Tensor, shape=(nAtoms, num_targets) - - F: torch.Tensor, shape=(nEdges, num_targets) - Energy and force prediction - """ - - # -------------------------------------- Energy Prediction -------------------------------------- # - rbf_emb_E = self.dense_rbf(rbf) # (nEdges, emb_size_edge) - x = m * rbf_emb_E - - # Graph Parallel: Local Node aggregation - x_E = scatter(x, id_j, dim=0, dim_size=nAtoms, reduce="sum") - # Graph Parallel: Global Node aggregation - x_E = gp_utils.reduce_from_model_parallel_region(x_E) - x_E = gp_utils.scatter_to_model_parallel_region(x_E, dim=0) - - # (nAtoms, emb_size_edge) - x_E = self.scale_sum(x_E, ref=m) - - for layer in self.seq_energy: - x_E = layer(x_E) # (nAtoms, emb_size_atom) - - x_E = self.out_energy(x_E) # (nAtoms, num_targets) - - # --------------------------------------- Force Prediction -------------------------------------- # - if self.direct_forces: - x_F = m - for i, layer in enumerate(self.seq_forces): - x_F = layer(x_F) # (nEdges, emb_size_edge) - - rbf_emb_F = self.dense_rbf_F(rbf) # (nEdges, emb_size_edge) - x_F_rbf = x_F * rbf_emb_F - x_F = self.scale_rbf_F(x_F_rbf, ref=x_F) - - x_F = self.out_forces(x_F) # (nEdges, num_targets) - else: - x_F = 0 - # ----------------------------------------------------------------------------------------------- # - - return x_E, x_F diff --git a/ocpmodels/models/gemnet_gp/layers/base_layers.py b/ocpmodels/models/gemnet_gp/layers/base_layers.py deleted file mode 100644 index 948bb5607b6816e51c4f945affe824b47999eee1..0000000000000000000000000000000000000000 --- a/ocpmodels/models/gemnet_gp/layers/base_layers.py +++ /dev/null @@ -1,106 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import math - -import torch - -from ..initializers import he_orthogonal_init - - -class Dense(torch.nn.Module): - """ - Combines dense layer with scaling for swish activation. - - Parameters - ---------- - units: int - Output embedding size. - activation: str - Name of the activation function to use. - bias: bool - True if use bias. - """ - - def __init__(self, in_features, out_features, bias=False, activation=None): - super().__init__() - - self.linear = torch.nn.Linear(in_features, out_features, bias=bias) - self.reset_parameters() - - if isinstance(activation, str): - activation = activation.lower() - if activation in ["swish", "silu"]: - self._activation = ScaledSiLU() - elif activation == "siqu": - self._activation = SiQU() - elif activation is None: - self._activation = torch.nn.Identity() - else: - raise NotImplementedError( - "Activation function not implemented for GemNet (yet)." - ) - - def reset_parameters(self, initializer=he_orthogonal_init): - initializer(self.linear.weight) - if self.linear.bias is not None: - self.linear.bias.data.fill_(0) - - def forward(self, x): - x = self.linear(x) - x = self._activation(x) - return x - - -class ScaledSiLU(torch.nn.Module): - def __init__(self): - super().__init__() - self.scale_factor = 1 / 0.6 - self._activation = torch.nn.SiLU() - - def forward(self, x): - return self._activation(x) * self.scale_factor - - -class SiQU(torch.nn.Module): - def __init__(self): - super().__init__() - self._activation = torch.nn.SiLU() - - def forward(self, x): - return x * self._activation(x) - - -class ResidualLayer(torch.nn.Module): - """ - Residual block with output scaled by 1/sqrt(2). - - Parameters - ---------- - units: int - Output embedding size. - nLayers: int - Number of dense layers. - layer_kwargs: str - Keyword arguments for initializing the layers. - """ - - def __init__(self, units: int, nLayers: int = 2, layer=Dense, **layer_kwargs): - super().__init__() - self.dense_mlp = torch.nn.Sequential( - *[ - layer(in_features=units, out_features=units, bias=False, **layer_kwargs) - for _ in range(nLayers) - ] - ) - self.inv_sqrt_2 = 1 / math.sqrt(2) - - def forward(self, input): - x = self.dense_mlp(input) - x = input + x - x = x * self.inv_sqrt_2 - return x diff --git a/ocpmodels/models/gemnet_gp/layers/basis_utils.py b/ocpmodels/models/gemnet_gp/layers/basis_utils.py deleted file mode 100644 index 987bdd630ff7a89af37fe48996ee71665fc59d05..0000000000000000000000000000000000000000 --- a/ocpmodels/models/gemnet_gp/layers/basis_utils.py +++ /dev/null @@ -1,280 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import numpy as np -import sympy as sym -from scipy import special as sp -from scipy.optimize import brentq - - -def Jn(r, n): - """ - numerical spherical bessel functions of order n - """ - return sp.spherical_jn(n, r) - - -def Jn_zeros(n, k): - """ - Compute the first k zeros of the spherical bessel functions up to order n (excluded) - """ - zerosj = np.zeros((n, k), dtype="float32") - zerosj[0] = np.arange(1, k + 1) * np.pi - points = np.arange(1, k + n) * np.pi - racines = np.zeros(k + n - 1, dtype="float32") - for i in range(1, n): - for j in range(k + n - 1 - i): - foo = brentq(Jn, points[j], points[j + 1], (i,)) - racines[j] = foo - points = racines - zerosj[i][:k] = racines[:k] - - return zerosj - - -def spherical_bessel_formulas(n): - """ - Computes the sympy formulas for the spherical bessel functions up to order n (excluded) - """ - x = sym.symbols("x") - # j_i = (-x)^i * (1/x * d/dx)^î * sin(x)/x - j = [sym.sin(x) / x] # j_0 - a = sym.sin(x) / x - for i in range(1, n): - b = sym.diff(a, x) / x - j += [sym.simplify(b * (-x) ** i)] - a = sym.simplify(b) - return j - - -def bessel_basis(n, k): - """ - Compute the sympy formulas for the normalized and rescaled spherical bessel functions up to - order n (excluded) and maximum frequency k (excluded). - - Returns: - bess_basis: list - Bessel basis formulas taking in a single argument x. - Has length n where each element has length k. -> In total n*k many. - """ - zeros = Jn_zeros(n, k) - normalizer = [] - for order in range(n): - normalizer_tmp = [] - for i in range(k): - normalizer_tmp += [0.5 * Jn(zeros[order, i], order + 1) ** 2] - normalizer_tmp = ( - 1 / np.array(normalizer_tmp) ** 0.5 - ) # sqrt(2/(j_l+1)**2) , sqrt(1/c**3) not taken into account yet - normalizer += [normalizer_tmp] - - f = spherical_bessel_formulas(n) - x = sym.symbols("x") - bess_basis = [] - for order in range(n): - bess_basis_tmp = [] - for i in range(k): - bess_basis_tmp += [ - sym.simplify( - normalizer[order][i] * f[order].subs(x, zeros[order, i] * x) - ) - ] - bess_basis += [bess_basis_tmp] - return bess_basis - - -def sph_harm_prefactor(l_degree, m_order): - """Computes the constant pre-factor for the spherical harmonic of degree l and order m. - - Parameters - ---------- - l_degree: int - Degree of the spherical harmonic. l >= 0 - m_order: int - Order of the spherical harmonic. -l <= m <= l - - Returns - ------- - factor: float - - """ - # sqrt((2*l+1)/4*pi * (l-m)!/(l+m)! ) - return ( - (2 * l_degree + 1) - / (4 * np.pi) - * np.math.factorial(l_degree - abs(m_order)) - / np.math.factorial(l_degree + abs(m_order)) - ) ** 0.5 - - -def associated_legendre_polynomials(L_maxdegree, zero_m_only=True, pos_m_only=True): - """Computes string formulas of the associated legendre polynomials up to degree L (excluded). - - Parameters - ---------- - L_maxdegree: int - Degree up to which to calculate the associated legendre polynomials (degree L is excluded). - zero_m_only: bool - If True only calculate the polynomials for the polynomials where m=0. - pos_m_only: bool - If True only calculate the polynomials for the polynomials where m>=0. Overwritten by zero_m_only. - - Returns - ------- - polynomials: list - Contains the sympy functions of the polynomials (in total L many if zero_m_only is True else L^2 many). - """ - # calculations from http://web.cmb.usc.edu/people/alber/Software/tomominer/docs/cpp/group__legendre__polynomials.html - z = sym.symbols("z") - P_l_m = [ - [0] * (2 * l_degree + 1) for l_degree in range(L_maxdegree) - ] # for order l: -l <= m <= l - - P_l_m[0][0] = 1 - if L_maxdegree > 0: - if zero_m_only: - # m = 0 - P_l_m[1][0] = z - for l_degree in range(2, L_maxdegree): - P_l_m[l_degree][0] = sym.simplify( - ( - (2 * l_degree - 1) * z * P_l_m[l_degree - 1][0] - - (l_degree - 1) * P_l_m[l_degree - 2][0] - ) - / l_degree - ) - return P_l_m - else: - # for m >= 0 - for l_degree in range(1, L_maxdegree): - P_l_m[l_degree][l_degree] = sym.simplify( - (1 - 2 * l_degree) - * (1 - z**2) ** 0.5 - * P_l_m[l_degree - 1][l_degree - 1] - ) # P_00, P_11, P_22, P_33 - - for m_order in range(0, L_maxdegree - 1): - P_l_m[m_order + 1][m_order] = sym.simplify( - (2 * m_order + 1) * z * P_l_m[m_order][m_order] - ) # P_10, P_21, P_32, P_43 - - for l_degree in range(2, L_maxdegree): - for m_order in range(l_degree - 1): # P_20, P_30, P_31 - P_l_m[l_degree][m_order] = sym.simplify( - ( - (2 * l_degree - 1) * z * P_l_m[l_degree - 1][m_order] - - (l_degree + m_order - 1) * P_l_m[l_degree - 2][m_order] - ) - / (l_degree - m_order) - ) - - if not pos_m_only: - # for m < 0: P_l(-m) = (-1)^m * (l-m)!/(l+m)! * P_lm - for l_degree in range(1, L_maxdegree): - for m_order in range(1, l_degree + 1): # P_1(-1), P_2(-1) P_2(-2) - P_l_m[l_degree][-m_order] = sym.simplify( - (-1) ** m_order - * np.math.factorial(l_degree - m_order) - / np.math.factorial(l_degree + m_order) - * P_l_m[l_degree][m_order] - ) - - return P_l_m - - -def real_sph_harm(L_maxdegree, use_theta, use_phi=True, zero_m_only=True): - """ - Computes formula strings of the the real part of the spherical harmonics up to degree L (excluded). - Variables are either spherical coordinates phi and theta (or cartesian coordinates x,y,z) on the UNIT SPHERE. - - Parameters - ---------- - L_maxdegree: int - Degree up to which to calculate the spherical harmonics (degree L is excluded). - use_theta: bool - - True: Expects the input of the formula strings to contain theta. - - False: Expects the input of the formula strings to contain z. - use_phi: bool - - True: Expects the input of the formula strings to contain phi. - - False: Expects the input of the formula strings to contain x and y. - Does nothing if zero_m_only is True - zero_m_only: bool - If True only calculate the harmonics where m=0. - - Returns - ------- - Y_lm_real: list - Computes formula strings of the the real part of the spherical harmonics up - to degree L (where degree L is not excluded). - In total L^2 many sph harm exist up to degree L (excluded). However, if zero_m_only only is True then - the total count is reduced to be only L many. - """ - z = sym.symbols("z") - P_l_m = associated_legendre_polynomials(L_maxdegree, zero_m_only) - if zero_m_only: - # for all m != 0: Y_lm = 0 - Y_l_m = [[0] for l_degree in range(L_maxdegree)] - else: - Y_l_m = [ - [0] * (2 * l_degree + 1) for l_degree in range(L_maxdegree) - ] # for order l: -l <= m <= l - - # convert expressions to spherical coordiantes - if use_theta: - # replace z by cos(theta) - theta = sym.symbols("theta") - for l_degree in range(L_maxdegree): - for m_order in range(len(P_l_m[l_degree])): - if not isinstance(P_l_m[l_degree][m_order], int): - P_l_m[l_degree][m_order] = P_l_m[l_degree][m_order].subs( - z, sym.cos(theta) - ) - - ## calculate Y_lm - # Y_lm = N * P_lm(cos(theta)) * exp(i*m*phi) - # { sqrt(2) * (-1)^m * N * P_l|m| * sin(|m|*phi) if m < 0 - # Y_lm_real = { Y_lm if m = 0 - # { sqrt(2) * (-1)^m * N * P_lm * cos(m*phi) if m > 0 - - for l_degree in range(L_maxdegree): - Y_l_m[l_degree][0] = sym.simplify( - sph_harm_prefactor(l_degree, 0) * P_l_m[l_degree][0] - ) # Y_l0 - - if not zero_m_only: - phi = sym.symbols("phi") - for l_degree in range(1, L_maxdegree): - # m > 0 - for m_order in range(1, l_degree + 1): - Y_l_m[l_degree][m_order] = sym.simplify( - 2**0.5 - * (-1) ** m_order - * sph_harm_prefactor(l_degree, m_order) - * P_l_m[l_degree][m_order] - * sym.cos(m_order * phi) - ) - # m < 0 - for m_order in range(1, l_degree + 1): - Y_l_m[l_degree][-m_order] = sym.simplify( - 2**0.5 - * (-1) ** m_order - * sph_harm_prefactor(l_degree, -m_order) - * P_l_m[l_degree][m_order] - * sym.sin(m_order * phi) - ) - - # convert expressions to cartesian coordinates - if not use_phi: - # replace phi by atan2(y,x) - x = sym.symbols("x") - y = sym.symbols("y") - for l_degree in range(L_maxdegree): - for m_order in range(len(Y_l_m[l_degree])): - Y_l_m[l_degree][m_order] = sym.simplify( - Y_l_m[l_degree][m_order].subs(phi, sym.atan2(y, x)) - ) - return Y_l_m diff --git a/ocpmodels/models/gemnet_gp/layers/efficient.py b/ocpmodels/models/gemnet_gp/layers/efficient.py deleted file mode 100644 index 27197413919f62ec6dd7fa68ed97d8e93a686db3..0000000000000000000000000000000000000000 --- a/ocpmodels/models/gemnet_gp/layers/efficient.py +++ /dev/null @@ -1,156 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import torch - -from ..initializers import he_orthogonal_init - - -class EfficientInteractionDownProjection(torch.nn.Module): - """ - Down projection in the efficient reformulation. - - Parameters - ---------- - emb_size_interm: int - Intermediate embedding size (down-projection size). - kernel_initializer: callable - Initializer of the weight matrix. - """ - - def __init__( - self, - num_spherical: int, - num_radial: int, - emb_size_interm: int, - ): - super().__init__() - - self.num_spherical = num_spherical - self.num_radial = num_radial - self.emb_size_interm = emb_size_interm - - self.reset_parameters() - - def reset_parameters(self): - self.weight = torch.nn.Parameter( - torch.empty((self.num_spherical, self.num_radial, self.emb_size_interm)), - requires_grad=True, - ) - he_orthogonal_init(self.weight) - - def forward(self, rbf, sph, id_ca, id_ragged_idx, Kmax): - """ - - Arguments - --------- - rbf: torch.Tensor, shape=(1, nEdges, num_radial) - sph: torch.Tensor, shape=(nEdges, Kmax, num_spherical) - id_ca - id_ragged_idx - - Returns - ------- - rbf_W1: torch.Tensor, shape=(nEdges, emb_size_interm, num_spherical) - sph: torch.Tensor, shape=(nEdges, Kmax, num_spherical) - Kmax = maximum number of neighbors of the edges - """ - num_edges = rbf.shape[1] - - # MatMul: mul + sum over num_radial - rbf_W1 = torch.matmul(rbf, self.weight) - # (num_spherical, nEdges , emb_size_interm) - rbf_W1 = rbf_W1.permute(1, 2, 0) - # (nEdges, emb_size_interm, num_spherical) - - # Zero padded dense matrix - # maximum number of neighbors, catch empty id_ca with maximum - if sph.shape[0] == 0: - Kmax = 0 - - sph2 = sph.new_zeros(num_edges, Kmax, self.num_spherical) - sph2[id_ca, id_ragged_idx] = sph - - sph2 = torch.transpose(sph2, 1, 2) - # (nEdges, num_spherical/emb_size_interm, Kmax) - - return rbf_W1, sph2 - - -class EfficientInteractionBilinear(torch.nn.Module): - """ - Efficient reformulation of the bilinear layer and subsequent summation. - - Parameters - ---------- - units_out: int - Embedding output size of the bilinear layer. - kernel_initializer: callable - Initializer of the weight matrix. - """ - - def __init__( - self, - emb_size: int, - emb_size_interm: int, - units_out: int, - ): - super().__init__() - self.emb_size = emb_size - self.emb_size_interm = emb_size_interm - self.units_out = units_out - - self.reset_parameters() - - def reset_parameters(self): - self.weight = torch.nn.Parameter( - torch.empty( - (self.emb_size, self.emb_size_interm, self.units_out), - requires_grad=True, - ) - ) - he_orthogonal_init(self.weight) - - def forward(self, basis, m, id_reduce, id_ragged_idx, edge_offset, Kmax): - """ - - Arguments - --------- - basis - m: quadruplets: m = m_db , triplets: m = m_ba - id_reduce - id_ragged_idx - - Returns - ------- - m_ca: torch.Tensor, shape=(nEdges, units_out) - Edge embeddings. - """ - # num_spherical is actually num_spherical**2 for quadruplets - (rbf_W1, sph) = basis - # (nEdges, emb_size_interm, num_spherical), (nEdges, num_spherical, Kmax) - nEdges = rbf_W1.shape[0] - - # Create (zero-padded) dense matrix of the neighboring edge embeddings. - # maximum number of neighbors, catch empty id_reduce_ji with maximum - m2 = m.new_zeros(nEdges, Kmax, self.emb_size) - m2[id_reduce - edge_offset, id_ragged_idx] = m - # (num_quadruplets or num_triplets, emb_size) -> (nEdges, Kmax, emb_size) - - sum_k = torch.matmul(sph, m2) # (nEdges, num_spherical, emb_size) - - # MatMul: mul + sum over num_spherical - rbf_W1_sum_k = torch.matmul(rbf_W1, sum_k) - # (nEdges, emb_size_interm, emb_size) - - # Bilinear: Sum over emb_size_interm and emb_size - m_ca = torch.matmul(rbf_W1_sum_k.permute(2, 0, 1), self.weight) - # (emb_size, nEdges, units_out) - m_ca = torch.sum(m_ca, dim=0) - # (nEdges, units_out) - - return m_ca diff --git a/ocpmodels/models/gemnet_gp/layers/embedding_block.py b/ocpmodels/models/gemnet_gp/layers/embedding_block.py deleted file mode 100644 index e7de3976b58c466bd75d7cf841dfaa02d3e26b58..0000000000000000000000000000000000000000 --- a/ocpmodels/models/gemnet_gp/layers/embedding_block.py +++ /dev/null @@ -1,97 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import numpy as np -import torch - -from ocpmodels.common import gp_utils - -from .base_layers import Dense - - -class AtomEmbedding(torch.nn.Module): - """ - Initial atom embeddings based on the atom type - - Parameters - ---------- - emb_size: int - Atom embeddings size - """ - - def __init__(self, emb_size): - super().__init__() - self.emb_size = emb_size - - # Atom embeddings: We go up to Bi (83). - self.embeddings = torch.nn.Embedding(83, emb_size) - # init by uniform distribution - torch.nn.init.uniform_(self.embeddings.weight, a=-np.sqrt(3), b=np.sqrt(3)) - - def forward(self, Z): - """ - Returns - ------- - h: torch.Tensor, shape=(nAtoms, emb_size) - Atom embeddings. - """ - h = self.embeddings(Z - 1) # -1 because Z.min()=1 (==Hydrogen) - return h - - -class EdgeEmbedding(torch.nn.Module): - """ - Edge embedding based on the concatenation of atom embeddings and subsequent dense layer. - - Parameters - ---------- - emb_size: int - Embedding size after the dense layer. - activation: str - Activation function used in the dense layer. - """ - - def __init__( - self, - atom_features, - edge_features, - out_features, - activation=None, - ): - super().__init__() - in_features = 2 * atom_features + edge_features - self.dense = Dense(in_features, out_features, activation=activation, bias=False) - - def forward( - self, - h, - m_rbf, - idx_s, - idx_t, - ): - """ - - Arguments - --------- - h - m_rbf: shape (nEdges, nFeatures) - in embedding block: m_rbf = rbf ; In interaction block: m_rbf = m_st - idx_s - idx_t - - Returns - ------- - m_st: torch.Tensor, shape=(nEdges, emb_size) - Edge embeddings. - """ - h = gp_utils.gather_from_model_parallel_region(h, dim=0) - h_s = h[idx_s] # shape=(nEdges, emb_size) - h_t = h[idx_t] # shape=(nEdges, emb_size) - - m_st = torch.cat([h_s, h_t, m_rbf], dim=-1) # (nEdges, 2*emb_size+nFeatures) - m_st = self.dense(m_st) # (nEdges, emb_size) - return m_st diff --git a/ocpmodels/models/gemnet_gp/layers/interaction_block.py b/ocpmodels/models/gemnet_gp/layers/interaction_block.py deleted file mode 100644 index b77da760ff386d4333a0cb06d9b4afc977c024a6..0000000000000000000000000000000000000000 --- a/ocpmodels/models/gemnet_gp/layers/interaction_block.py +++ /dev/null @@ -1,360 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import math - -import torch - -from ocpmodels.common import gp_utils -from ocpmodels.modules.scaling import ScaleFactor - -from .atom_update_block import AtomUpdateBlock -from .base_layers import Dense, ResidualLayer -from .efficient import EfficientInteractionBilinear -from .embedding_block import EdgeEmbedding - - -class InteractionBlockTripletsOnly(torch.nn.Module): - """ - Interaction block for GemNet-T/dT. - - Parameters - ---------- - emb_size_atom: int - Embedding size of the atoms. - emb_size_edge: int - Embedding size of the edges. - emb_size_trip: int - (Down-projected) Embedding size in the triplet message passing block. - emb_size_rbf: int - Embedding size of the radial basis transformation. - emb_size_cbf: int - Embedding size of the circular basis transformation (one angle). - - emb_size_bil_trip: int - Embedding size of the edge embeddings in the triplet-based message passing block after the bilinear layer. - num_before_skip: int - Number of residual blocks before the first skip connection. - num_after_skip: int - Number of residual blocks after the first skip connection. - num_concat: int - Number of residual blocks after the concatenation. - num_atom: int - Number of residual blocks in the atom embedding blocks. - - activation: str - Name of the activation function to use in the dense layers except for the final dense layer. - """ - - def __init__( - self, - emb_size_atom, - emb_size_edge, - emb_size_trip, - emb_size_rbf, - emb_size_cbf, - emb_size_bil_trip, - num_before_skip, - num_after_skip, - num_concat, - num_atom, - activation=None, - name="Interaction", - ): - super().__init__() - self.name = name - - block_nr = name.split("_")[-1] - - ## -------------------------------------------- Message Passing ------------------------------------------- ## - # Dense transformation of skip connection - self.dense_ca = Dense( - emb_size_edge, - emb_size_edge, - activation=activation, - bias=False, - ) - - # Triplet Interaction - self.trip_interaction = TripletInteraction( - emb_size_edge=emb_size_edge, - emb_size_trip=emb_size_trip, - emb_size_bilinear=emb_size_bil_trip, - emb_size_rbf=emb_size_rbf, - emb_size_cbf=emb_size_cbf, - activation=activation, - name=f"TripInteraction_{block_nr}", - ) - - ## ---------------------------------------- Update Edge Embeddings ---------------------------------------- ## - # Residual layers before skip connection - self.layers_before_skip = torch.nn.ModuleList( - [ - ResidualLayer( - emb_size_edge, - activation=activation, - ) - for i in range(num_before_skip) - ] - ) - - # Residual layers after skip connection - self.layers_after_skip = torch.nn.ModuleList( - [ - ResidualLayer( - emb_size_edge, - activation=activation, - ) - for i in range(num_after_skip) - ] - ) - - ## ---------------------------------------- Update Atom Embeddings ---------------------------------------- ## - self.atom_update = AtomUpdateBlock( - emb_size_atom=emb_size_atom, - emb_size_edge=emb_size_edge, - emb_size_rbf=emb_size_rbf, - nHidden=num_atom, - activation=activation, - name=f"AtomUpdate_{block_nr}", - ) - - ## ------------------------------ Update Edge Embeddings with Atom Embeddings ----------------------------- ## - self.concat_layer = EdgeEmbedding( - emb_size_atom, - emb_size_edge, - emb_size_edge, - activation=activation, - ) - self.residual_m = torch.nn.ModuleList( - [ - ResidualLayer(emb_size_edge, activation=activation) - for _ in range(num_concat) - ] - ) - - self.inv_sqrt_2 = 1 / math.sqrt(2.0) - - def forward( - self, - h, - m, - rbf3, - cbf3, - id3_ragged_idx, - id_swap, - id3_ba, - id3_ca, - rbf_h, - idx_s, - idx_t, - edge_offset, - Kmax, - nAtoms, - ): - """ - Returns - ------- - h: torch.Tensor, shape=(nEdges, emb_size_atom) - Atom embeddings. - m: torch.Tensor, shape=(nEdges, emb_size_edge) - Edge embeddings (c->a). - Node: h - Edge: m, rbf3, id_swap, rbf_h, idx_s, idx_t, cbf3[0], cbf3[1] (dense) - Triplet: id3_ragged_idx, id3_ba, id3_ca - """ - # Initial transformation - x_ca_skip = self.dense_ca(m) # (nEdges, emb_size_edge) - - x3 = self.trip_interaction( - m, - rbf3, - cbf3, - id3_ragged_idx, - id_swap, - id3_ba, - id3_ca, - edge_offset, - Kmax, - ) - - ## ----------------------------- Merge Embeddings after Triplet Interaction ------------------------------ ## - x = x_ca_skip + x3 # (nEdges, emb_size_edge) - x = x * self.inv_sqrt_2 - - ## ---------------------------------------- Update Edge Embeddings --------------------------------------- ## - # Transformations before skip connection - for i, layer in enumerate(self.layers_before_skip): - x = layer(x) # (nEdges, emb_size_edge) - - # Skip connection - m = m + x # (nEdges, emb_size_edge) - m = m * self.inv_sqrt_2 - - # Transformations after skip connection - for i, layer in enumerate(self.layers_after_skip): - m = layer(m) # (nEdges, emb_size_edge) - - ## ---------------------------------------- Update Atom Embeddings --------------------------------------- ## - h2 = self.atom_update(nAtoms, m, rbf_h, idx_t) - - # Skip connection - h = h + h2 # (nAtoms, emb_size_atom) - h = h * self.inv_sqrt_2 - - ## ----------------------------- Update Edge Embeddings with Atom Embeddings ----------------------------- ## - m2 = self.concat_layer(h, m, idx_s, idx_t) # (nEdges, emb_size_edge) - - for i, layer in enumerate(self.residual_m): - m2 = layer(m2) # (nEdges, emb_size_edge) - - # Skip connection - m = m + m2 # (nEdges, emb_size_edge) - m = m * self.inv_sqrt_2 - return h, m - - -class TripletInteraction(torch.nn.Module): - """ - Triplet-based message passing block. - - Parameters - ---------- - emb_size_edge: int - Embedding size of the edges. - emb_size_trip: int - (Down-projected) Embedding size of the edge embeddings after the hadamard product with rbf. - emb_size_bilinear: int - Embedding size of the edge embeddings after the bilinear layer. - emb_size_rbf: int - Embedding size of the radial basis transformation. - emb_size_cbf: int - Embedding size of the circular basis transformation (one angle). - - activation: str - Name of the activation function to use in the dense layers except for the final dense layer. - """ - - def __init__( - self, - emb_size_edge, - emb_size_trip, - emb_size_bilinear, - emb_size_rbf, - emb_size_cbf, - activation=None, - name="TripletInteraction", - **kwargs, - ): - super().__init__() - self.name = name - - # Dense transformation - self.dense_ba = Dense( - emb_size_edge, - emb_size_edge, - activation=activation, - bias=False, - ) - - # Up projections of basis representations, bilinear layer and scaling factors - self.mlp_rbf = Dense( - emb_size_rbf, - emb_size_edge, - activation=None, - bias=False, - ) - self.scale_rbf = ScaleFactor(name + "_had_rbf") - - self.mlp_cbf = EfficientInteractionBilinear( - emb_size_trip, emb_size_cbf, emb_size_bilinear - ) - - # combines scaling for bilinear layer and summation - self.scale_cbf_sum = ScaleFactor(name + "_sum_cbf") - - # Down and up projections - self.down_projection = Dense( - emb_size_edge, - emb_size_trip, - activation=activation, - bias=False, - ) - self.up_projection_ca = Dense( - emb_size_bilinear, - emb_size_edge, - activation=activation, - bias=False, - ) - self.up_projection_ac = Dense( - emb_size_bilinear, - emb_size_edge, - activation=activation, - bias=False, - ) - - self.inv_sqrt_2 = 1 / math.sqrt(2.0) - - def forward( - self, - m, - rbf3, - cbf3, - id3_ragged_idx, - id_swap, - id3_ba, - id3_ca, - edge_offset, - Kmax, - ): - """ - Returns - ------- - m: torch.Tensor, shape=(nEdges, emb_size_edge) - Edge embeddings (c->a). - """ - - # Dense transformation - x_ba = self.dense_ba(m) # (nEdges, emb_size_edge) - - # Transform via radial bessel basis - rbf_emb = self.mlp_rbf(rbf3) # (nEdges, emb_size_edge) - x_ba2 = x_ba * rbf_emb - x_ba = self.scale_rbf(x_ba2, ref=x_ba) - - x_ba = self.down_projection(x_ba) # (nEdges, emb_size_trip) - - # Graph Parallel: Gather x_ba from all nodes - x_ba = gp_utils.gather_from_model_parallel_region(x_ba, dim=0) - - # Transform via circular spherical basis - x_ba = x_ba[id3_ba] - - # Efficient bilinear layer - x = self.mlp_cbf(cbf3, x_ba, id3_ca, id3_ragged_idx, edge_offset, Kmax) - # (nEdges, emb_size_quad) - x = self.scale_cbf_sum(x, ref=x_ba) - - # => - # rbf(d_ba) - # cbf(d_ca, angle_cab) - - # Up project embeddings - x_ca = self.up_projection_ca(x) # (nEdges, emb_size_edge) - x_ac = self.up_projection_ac(x) # (nEdges, emb_size_edge) - - # Graph Parallel: Gather x_ac from all nodes - x_ac = gp_utils.gather_from_model_parallel_region(x_ac, dim=0) - - # Merge interaction of c->a and a->c - x_ac = x_ac[id_swap] # swap to add to edge a->c and not c->a - x_ac = gp_utils.scatter_to_model_parallel_region(x_ac, dim=0) - - x3 = x_ca + x_ac - x3 = x3 * self.inv_sqrt_2 - - return x3 diff --git a/ocpmodels/models/gemnet_gp/layers/radial_basis.py b/ocpmodels/models/gemnet_gp/layers/radial_basis.py deleted file mode 100644 index cad79f2093c1f5480f1de030f40760f04c779983..0000000000000000000000000000000000000000 --- a/ocpmodels/models/gemnet_gp/layers/radial_basis.py +++ /dev/null @@ -1,200 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import math - -import numpy as np -import torch -from scipy.special import binom -from torch_geometric.nn.models.schnet import GaussianSmearing - - -class PolynomialEnvelope(torch.nn.Module): - """ - Polynomial envelope function that ensures a smooth cutoff. - - Parameters - ---------- - exponent: int - Exponent of the envelope function. - """ - - def __init__(self, exponent): - super().__init__() - assert exponent > 0 - self.p = exponent - self.a = -(self.p + 1) * (self.p + 2) / 2 - self.b = self.p * (self.p + 2) - self.c = -self.p * (self.p + 1) / 2 - - def forward(self, d_scaled): - env_val = ( - 1 - + self.a * d_scaled**self.p - + self.b * d_scaled ** (self.p + 1) - + self.c * d_scaled ** (self.p + 2) - ) - return torch.where(d_scaled < 1, env_val, torch.zeros_like(d_scaled)) - - -class ExponentialEnvelope(torch.nn.Module): - """ - Exponential envelope function that ensures a smooth cutoff, - as proposed in Unke, Chmiela, Gastegger, Schütt, Sauceda, Müller 2021. - SpookyNet: Learning Force Fields with Electronic Degrees of Freedom - and Nonlocal Effects - """ - - def __init__(self): - super().__init__() - - def forward(self, d_scaled): - env_val = torch.exp(-(d_scaled**2) / ((1 - d_scaled) * (1 + d_scaled))) - return torch.where(d_scaled < 1, env_val, torch.zeros_like(d_scaled)) - - -class SphericalBesselBasis(torch.nn.Module): - """ - 1D spherical Bessel basis - - Parameters - ---------- - num_radial: int - Controls maximum frequency. - cutoff: float - Cutoff distance in Angstrom. - """ - - def __init__( - self, - num_radial: int, - cutoff: float, - ): - super().__init__() - self.norm_const = math.sqrt(2 / (cutoff**3)) - # cutoff ** 3 to counteract dividing by d_scaled = d / cutoff - - # Initialize frequencies at canonical positions - self.frequencies = torch.nn.Parameter( - data=torch.tensor(np.pi * np.arange(1, num_radial + 1, dtype=np.float32)), - requires_grad=True, - ) - - def forward(self, d_scaled): - return ( - self.norm_const - / d_scaled[:, None] - * torch.sin(self.frequencies * d_scaled[:, None]) - ) # (num_edges, num_radial) - - -class BernsteinBasis(torch.nn.Module): - """ - Bernstein polynomial basis, - as proposed in Unke, Chmiela, Gastegger, Schütt, Sauceda, Müller 2021. - SpookyNet: Learning Force Fields with Electronic Degrees of Freedom - and Nonlocal Effects - - Parameters - ---------- - num_radial: int - Controls maximum frequency. - pregamma_initial: float - Initial value of exponential coefficient gamma. - Default: gamma = 0.5 * a_0**-1 = 0.94486, - inverse softplus -> pregamma = log e**gamma - 1 = 0.45264 - """ - - def __init__( - self, - num_radial: int, - pregamma_initial: float = 0.45264, - ): - super().__init__() - prefactor = binom(num_radial - 1, np.arange(num_radial)) - self.register_buffer( - "prefactor", - torch.tensor(prefactor, dtype=torch.float), - persistent=False, - ) - - self.pregamma = torch.nn.Parameter( - data=torch.tensor(pregamma_initial, dtype=torch.float), - requires_grad=True, - ) - self.softplus = torch.nn.Softplus() - - exp1 = torch.arange(num_radial) - self.register_buffer("exp1", exp1[None, :], persistent=False) - exp2 = num_radial - 1 - exp1 - self.register_buffer("exp2", exp2[None, :], persistent=False) - - def forward(self, d_scaled): - gamma = self.softplus(self.pregamma) # constrain to positive - exp_d = torch.exp(-gamma * d_scaled)[:, None] - return self.prefactor * (exp_d**self.exp1) * ((1 - exp_d) ** self.exp2) - - -class RadialBasis(torch.nn.Module): - """ - - Parameters - ---------- - num_radial: int - Controls maximum frequency. - cutoff: float - Cutoff distance in Angstrom. - rbf: dict = {"name": "gaussian"} - Basis function and its hyperparameters. - envelope: dict = {"name": "polynomial", "exponent": 5} - Envelope function and its hyperparameters. - """ - - def __init__( - self, - num_radial: int, - cutoff: float, - rbf: dict = {"name": "gaussian"}, - envelope: dict = {"name": "polynomial", "exponent": 5}, - ): - super().__init__() - self.inv_cutoff = 1 / cutoff - - env_name = envelope["name"].lower() - env_hparams = envelope.copy() - del env_hparams["name"] - - if env_name == "polynomial": - self.envelope = PolynomialEnvelope(**env_hparams) - elif env_name == "exponential": - self.envelope = ExponentialEnvelope(**env_hparams) - else: - raise ValueError(f"Unknown envelope function '{env_name}'.") - - rbf_name = rbf["name"].lower() - rbf_hparams = rbf.copy() - del rbf_hparams["name"] - - # RBFs get distances scaled to be in [0, 1] - if rbf_name == "gaussian": - self.rbf = GaussianSmearing( - start=0, stop=1, num_gaussians=num_radial, **rbf_hparams - ) - elif rbf_name == "spherical_bessel": - self.rbf = SphericalBesselBasis( - num_radial=num_radial, cutoff=cutoff, **rbf_hparams - ) - elif rbf_name == "bernstein": - self.rbf = BernsteinBasis(num_radial=num_radial, **rbf_hparams) - else: - raise ValueError(f"Unknown radial basis function '{rbf_name}'.") - - def forward(self, d): - d_scaled = d * self.inv_cutoff - - env = self.envelope(d_scaled) - return env[:, None] * self.rbf(d_scaled) # (nEdges, num_radial) diff --git a/ocpmodels/models/gemnet_gp/layers/spherical_basis.py b/ocpmodels/models/gemnet_gp/layers/spherical_basis.py deleted file mode 100644 index 6695a83adf5c2d9c4413b844a2e266aa55f2e726..0000000000000000000000000000000000000000 --- a/ocpmodels/models/gemnet_gp/layers/spherical_basis.py +++ /dev/null @@ -1,89 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import sympy as sym -import torch -from torch_geometric.nn.models.schnet import GaussianSmearing - -from .basis_utils import real_sph_harm -from .radial_basis import RadialBasis - - -class CircularBasisLayer(torch.nn.Module): - """ - 2D Fourier Bessel Basis - - Parameters - ---------- - num_spherical: int - Controls maximum frequency. - radial_basis: RadialBasis - Radial basis functions - cbf: dict - Name and hyperparameters of the cosine basis function - efficient: bool - Whether to use the "efficient" summation order - """ - - def __init__( - self, - num_spherical: int, - radial_basis: RadialBasis, - cbf: str, - efficient: bool = False, - ): - super().__init__() - - self.radial_basis = radial_basis - self.efficient = efficient - - cbf_name = cbf["name"].lower() - cbf_hparams = cbf.copy() - del cbf_hparams["name"] - - if cbf_name == "gaussian": - self.cosφ_basis = GaussianSmearing( - start=-1, stop=1, num_gaussians=num_spherical, **cbf_hparams - ) - elif cbf_name == "spherical_harmonics": - Y_lm = real_sph_harm(num_spherical, use_theta=False, zero_m_only=True) - sph_funcs = [] # (num_spherical,) - - # convert to tensorflow functions - z = sym.symbols("z") - modules = {"sin": torch.sin, "cos": torch.cos, "sqrt": torch.sqrt} - m_order = 0 # only single angle - for l_degree in range(len(Y_lm)): # num_spherical - if ( - l_degree == 0 - ): # Y_00 is only a constant -> function returns value and not tensor - first_sph = sym.lambdify([z], Y_lm[l_degree][m_order], modules) - sph_funcs.append(lambda z: torch.zeros_like(z) + first_sph(z)) - else: - sph_funcs.append( - sym.lambdify([z], Y_lm[l_degree][m_order], modules) - ) - self.cosφ_basis = lambda cosφ: torch.stack( - [f(cosφ) for f in sph_funcs], dim=1 - ) - else: - raise ValueError(f"Unknown cosine basis function '{cbf_name}'.") - - def forward(self, D_ca, cosφ_cab, id3_ca): - rbf = self.radial_basis(D_ca) # (num_edges, num_radial) - cbf = self.cosφ_basis(cosφ_cab) # (num_triplets, num_spherical) - - if not self.efficient: - rbf = rbf[id3_ca] # (num_triplets, num_radial) - out = (rbf[:, None, :] * cbf[:, :, None]).view( - -1, rbf.shape[-1] * cbf.shape[-1] - ) - return (out,) - # (num_triplets, num_radial * num_spherical) - else: - return (rbf[None, :, :], cbf) - # (1, num_edges, num_radial), (num_edges, num_spherical) diff --git a/ocpmodels/models/gemnet_gp/utils.py b/ocpmodels/models/gemnet_gp/utils.py deleted file mode 100644 index 9ff5c6c1e5acf4651e2f1802eaf090fb432b3149..0000000000000000000000000000000000000000 --- a/ocpmodels/models/gemnet_gp/utils.py +++ /dev/null @@ -1,275 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import json - -import torch -from torch_scatter import segment_csr - - -def read_json(path): - """""" - if not path.endswith(".json"): - raise UserWarning(f"Path {path} is not a json-path.") - - with open(path, "r") as f: - content = json.load(f) - return content - - -def update_json(path, data): - """""" - if not path.endswith(".json"): - raise UserWarning(f"Path {path} is not a json-path.") - - content = read_json(path) - content.update(data) - write_json(path, content) - - -def write_json(path, data): - """""" - if not path.endswith(".json"): - raise UserWarning(f"Path {path} is not a json-path.") - - with open(path, "w", encoding="utf-8") as f: - json.dump(data, f, ensure_ascii=False, indent=4) - - -def read_value_json(path, key): - """""" - content = read_json(path) - - if key in content.keys(): - return content[key] - else: - return None - - -def ragged_range(sizes): - """Multiple concatenated ranges. - - Examples - -------- - sizes = [1 4 2 3] - Return: [0 0 1 2 3 0 1 0 1 2] - """ - assert sizes.dim() == 1 - if sizes.sum() == 0: - return sizes.new_empty(0) - - # Remove 0 sizes - sizes_nonzero = sizes > 0 - if not torch.all(sizes_nonzero): - sizes = torch.masked_select(sizes, sizes_nonzero) - - # Initialize indexing array with ones as we need to setup incremental indexing - # within each group when cumulatively summed at the final stage. - id_steps = torch.ones(sizes.sum(), dtype=torch.long, device=sizes.device) - id_steps[0] = 0 - insert_index = sizes[:-1].cumsum(0) - insert_val = (1 - sizes)[:-1] - - # Assign index-offsetting values - id_steps[insert_index] = insert_val - - # Finally index into input array for the group repeated o/p - res = id_steps.cumsum(0) - return res - - -def repeat_blocks( - sizes, - repeats, - continuous_indexing=True, - start_idx=0, - block_inc=0, - repeat_inc=0, -): - """Repeat blocks of indices. - Adapted from https://stackoverflow.com/questions/51154989/numpy-vectorized-function-to-repeat-blocks-of-consecutive-elements - - continuous_indexing: Whether to keep increasing the index after each block - start_idx: Starting index - block_inc: Number to increment by after each block, - either global or per block. Shape: len(sizes) - 1 - repeat_inc: Number to increment by after each repetition, - either global or per block - - Examples - -------- - sizes = [1,3,2] ; repeats = [3,2,3] ; continuous_indexing = False - Return: [0 0 0 0 1 2 0 1 2 0 1 0 1 0 1] - sizes = [1,3,2] ; repeats = [3,2,3] ; continuous_indexing = True - Return: [0 0 0 1 2 3 1 2 3 4 5 4 5 4 5] - sizes = [1,3,2] ; repeats = [3,2,3] ; continuous_indexing = True ; - repeat_inc = 4 - Return: [0 4 8 1 2 3 5 6 7 4 5 8 9 12 13] - sizes = [1,3,2] ; repeats = [3,2,3] ; continuous_indexing = True ; - start_idx = 5 - Return: [5 5 5 6 7 8 6 7 8 9 10 9 10 9 10] - sizes = [1,3,2] ; repeats = [3,2,3] ; continuous_indexing = True ; - block_inc = 1 - Return: [0 0 0 2 3 4 2 3 4 6 7 6 7 6 7] - sizes = [0,3,2] ; repeats = [3,2,3] ; continuous_indexing = True - Return: [0 1 2 0 1 2 3 4 3 4 3 4] - sizes = [2,3,2] ; repeats = [2,0,2] ; continuous_indexing = True - Return: [0 1 0 1 5 6 5 6] - """ - assert sizes.dim() == 1 - assert all(sizes >= 0) - - # Remove 0 sizes - sizes_nonzero = sizes > 0 - if not torch.all(sizes_nonzero): - assert block_inc == 0 # Implementing this is not worth the effort - sizes = torch.masked_select(sizes, sizes_nonzero) - if isinstance(repeats, torch.Tensor): - repeats = torch.masked_select(repeats, sizes_nonzero) - if isinstance(repeat_inc, torch.Tensor): - repeat_inc = torch.masked_select(repeat_inc, sizes_nonzero) - - if isinstance(repeats, torch.Tensor): - assert all(repeats >= 0) - insert_dummy = repeats[0] == 0 - if insert_dummy: - one = sizes.new_ones(1) - zero = sizes.new_zeros(1) - sizes = torch.cat((one, sizes)) - repeats = torch.cat((one, repeats)) - if isinstance(block_inc, torch.Tensor): - block_inc = torch.cat((zero, block_inc)) - if isinstance(repeat_inc, torch.Tensor): - repeat_inc = torch.cat((zero, repeat_inc)) - else: - assert repeats >= 0 - insert_dummy = False - - # Get repeats for each group using group lengths/sizes - r1 = torch.repeat_interleave(torch.arange(len(sizes), device=sizes.device), repeats) - - # Get total size of output array, as needed to initialize output indexing array - N = (sizes * repeats).sum() - - # Initialize indexing array with ones as we need to setup incremental indexing - # within each group when cumulatively summed at the final stage. - # Two steps here: - # 1. Within each group, we have multiple sequences, so setup the offsetting - # at each sequence lengths by the seq. lengths preceding those. - id_ar = torch.ones(N, dtype=torch.long, device=sizes.device) - id_ar[0] = 0 - insert_index = sizes[r1[:-1]].cumsum(0) - insert_val = (1 - sizes)[r1[:-1]] - - if isinstance(repeats, torch.Tensor) and torch.any(repeats == 0): - diffs = r1[1:] - r1[:-1] - indptr = torch.cat((sizes.new_zeros(1), diffs.cumsum(0))) - if continuous_indexing: - # If a group was skipped (repeats=0) we need to add its size - insert_val += segment_csr(sizes[: r1[-1]], indptr, reduce="sum") - - # Add block increments - if isinstance(block_inc, torch.Tensor): - insert_val += segment_csr(block_inc[: r1[-1]], indptr, reduce="sum") - else: - insert_val += block_inc * (indptr[1:] - indptr[:-1]) - if insert_dummy: - insert_val[0] -= block_inc - else: - idx = r1[1:] != r1[:-1] - if continuous_indexing: - # 2. For each group, make sure the indexing starts from the next group's - # first element. So, simply assign 1s there. - insert_val[idx] = 1 - - # Add block increments - insert_val[idx] += block_inc - - # Add repeat_inc within each group - if isinstance(repeat_inc, torch.Tensor): - insert_val += repeat_inc[r1[:-1]] - if isinstance(repeats, torch.Tensor): - repeat_inc_inner = repeat_inc[repeats > 0][:-1] - else: - repeat_inc_inner = repeat_inc[:-1] - else: - insert_val += repeat_inc - repeat_inc_inner = repeat_inc - - # Subtract the increments between groups - if isinstance(repeats, torch.Tensor): - repeats_inner = repeats[repeats > 0][:-1] - else: - repeats_inner = repeats - insert_val[r1[1:] != r1[:-1]] -= repeat_inc_inner * repeats_inner - - # Assign index-offsetting values - id_ar[insert_index] = insert_val - - if insert_dummy: - id_ar = id_ar[1:] - if continuous_indexing: - id_ar[0] -= 1 - - # Set start index now, in case of insertion due to leading repeats=0 - id_ar[0] += start_idx - - # Finally index into input array for the group repeated o/p - res = id_ar.cumsum(0) - return res - - -def calculate_interatomic_vectors(R, id_s, id_t, offsets_st): - """ - Calculate the vectors connecting the given atom pairs, - considering offsets from periodic boundary conditions (PBC). - - Parameters - ---------- - R: Tensor, shape = (nAtoms, 3) - Atom positions. - id_s: Tensor, shape = (nEdges,) - Indices of the source atom of the edges. - id_t: Tensor, shape = (nEdges,) - Indices of the target atom of the edges. - offsets_st: Tensor, shape = (nEdges,) - PBC offsets of the edges. - Subtract this from the correct direction. - - Returns - ------- - (D_st, V_st): tuple - D_st: Tensor, shape = (nEdges,) - Distance from atom t to s. - V_st: Tensor, shape = (nEdges,) - Unit direction from atom t to s. - """ - Rs = R[id_s] - Rt = R[id_t] - # ReLU prevents negative numbers in sqrt - if offsets_st is None: - V_st = Rt - Rs # s -> t - else: - V_st = Rt - Rs + offsets_st # s -> t - D_st = torch.sqrt(torch.sum(V_st**2, dim=1)) - V_st = V_st / D_st[..., None] - return D_st, V_st - - -def inner_product_normalized(x, y): - """ - Calculate the inner product between the given normalized vectors, - giving a result between -1 and 1. - """ - return torch.sum(x * y, dim=-1).clamp(min=-1, max=1) - - -def mask_neighbors(neighbors, edge_mask): - neighbors_old_indptr = torch.cat([neighbors.new_zeros(1), neighbors]) - neighbors_old_indptr = torch.cumsum(neighbors_old_indptr, dim=0) - neighbors = segment_csr(edge_mask.long(), neighbors_old_indptr) - return neighbors diff --git a/ocpmodels/models/gemnet_oc/README.md b/ocpmodels/models/gemnet_oc/README.md deleted file mode 100644 index a5f47044bf45aa483a4c8a315cd80d6004d3c311..0000000000000000000000000000000000000000 --- a/ocpmodels/models/gemnet_oc/README.md +++ /dev/null @@ -1,29 +0,0 @@ -# GemNet-OC: Developing Graph Neural Networks for Large and Diverse Molecular Simulation Datasets - -Johannes Gasteiger, Muhammed Shuaibi, Anuroop Sriram, Stephan Günnemann, Zachary Ulissi, C. Lawrence Zitnick, Abhishek Das - -[[`arXiv:2204.02782`](https://arxiv.org/abs/2204.02782)] - -When running inference with a pretrained GemNet-OC model, make sure that the -`scale_file` path is correct in the config, otherwise predictions will be inaccurate. - -| Model | Val ID 30k Force MAE | Val ID 30k Energy MAE | Val ID 30k Force cos | Test metrics | Download | -| ----- | -------------------- | --------------------- | -------------------- | ------------ | -------- | -| gemnet_oc_2M | 0.0225 | 0.2299 | 0.6174 | [S2EF](https://evalai.s3.amazonaws.com/media/submission_files/submission_179229/062c037e-4f1f-49c2-9eeb-8e14681a70ee.json) \| [IS2RE](https://evalai.s3.amazonaws.com/media/submission_files/submission_179296/6688f44f-9d5a-4020-beca-8b804e0212fb.json) \| [IS2RS](https://evalai.s3.amazonaws.com/media/submission_files/submission_179257/0d02a349-0abe-44c0-a65c-29a9df75c886.json) | [config](https://github.com/Open-Catalyst-Project/ocp/blob/main/configs/s2ef/2M/gemnet/gemnet-oc.yml) \| [checkpoint](https://dl.fbaipublicfiles.com/opencatalystproject/models/2022_07/s2ef/gemnet_oc_base_s2ef_2M.pt) \| [scale file](https://github.com/Open-Catalyst-Project/ocp/blob/481f3a5a92dc787384ddae9fe3f50f5d932712fd/configs/s2ef/all/gemnet/scaling_factors/gemnet-oc.pt) | -| gemnet_oc_all | 0.0179 | 0.1668 | 0.6879 | [S2EF](https://evalai.s3.amazonaws.com/media/submission_files/submission_179008/6e731f20-17cf-417e-b0ad-97352be8cc37.json) \| [IS2RE]() \| [IS2RS](https://evalai.s3.amazonaws.com/media/submission_files/submission_160550/72a65a42-1fa9-44c5-8546-9eb691df8d2e.json) | [config](https://github.com/Open-Catalyst-Project/ocp/blob/main/configs/s2ef/all/gemnet/gemnet-oc.yml) \| [checkpoint](https://dl.fbaipublicfiles.com/opencatalystproject/models/2022_07/s2ef/gemnet_oc_base_s2ef_all.pt) \| [scale file](https://github.com/Open-Catalyst-Project/ocp/blob/481f3a5a92dc787384ddae9fe3f50f5d932712fd/configs/s2ef/all/gemnet/scaling_factors/gemnet-oc.pt) | -| gemnet_oc_large_all_md_energy | 0.0178 | 0.1504 | 0.6906 | [S2EF](https://evalai.s3.amazonaws.com/media/submission_files/submission_179143/40940149-6a4a-49a4-a2ce-38486215990f.json) | - | -| gemnet_oc_large_all_md_force | 0.0164 | 0.1665 | 0.7139 | [S2EF](https://evalai.s3.amazonaws.com/media/submission_files/submission_179042/ba160459-0de3-4583-a98b-12102138c61e.json) \| [IS2RS](https://evalai.s3.amazonaws.com/media/submission_files/submission_169243/10bc7c8d-5124-4338-aaf3-04a7d015c4a0.json) | [config](https://github.com/Open-Catalyst-Project/ocp/blob/main/configs/s2ef/all/gemnet/gemnet-oc-large.yml) \| [checkpoint](https://dl.fbaipublicfiles.com/opencatalystproject/models/2022_07/s2ef/gemnet_oc_large_s2ef_all_md.pt) \| [scale file](https://github.com/Open-Catalyst-Project/ocp/blob/481f3a5a92dc787384ddae9fe3f50f5d932712fd/configs/s2ef/all/gemnet/scaling_factors/gemnet-oc-large.pt) | -| gemnet_oc_large_all_md_energy + gemnet_oc_large_all_md_force | - | - | - | [IS2RE](https://evalai.s3.amazonaws.com/media/submission_files/submission_212962/6acc7cf7-e18b-4d6a-9082-b4a114110dbf.json) | - | - -## Citing - -If you use GemNet-OC in your work, please consider citing: - -```bibtex -@article{gasteiger_gemnet_oc_2022, - title = {{GemNet-OC: Developing Graph Neural Networks for Large and Diverse Molecular Simulation Datasets}}, - author = {Gasteiger, Johannes and Shuaibi, Muhammed and Sriram, Anuroop and G{\"u}nnemann, Stephan and Ulissi, Zachary and Zitnick, C Lawrence and Das, Abhishek}, - journal = {Transactions on Machine Learning Research (TMLR)}, - year = {2022}, -} -``` diff --git a/ocpmodels/models/gemnet_oc/gemnet_oc.py b/ocpmodels/models/gemnet_oc/gemnet_oc.py deleted file mode 100644 index 5e721421ab12c1b7b001f41af922da3b48161e95..0000000000000000000000000000000000000000 --- a/ocpmodels/models/gemnet_oc/gemnet_oc.py +++ /dev/null @@ -1,1338 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import logging -import os -from typing import Optional - -import numpy as np -import torch -from torch_geometric.nn import radius_graph -from torch_scatter import scatter, segment_coo - -from ocpmodels.common.registry import registry -from ocpmodels.common.utils import ( - compute_neighbors, - conditional_grad, - get_max_neighbors_mask, - get_pbc_distances, - radius_graph_pbc, - scatter_det, -) -from ocpmodels.models.base import BaseModel -from ocpmodels.modules.scaling.compat import load_scales_compat - -from .initializers import get_initializer -from .interaction_indices import ( - get_mixed_triplets, - get_quadruplets, - get_triplets, -) -from .layers.atom_update_block import OutputBlock -from .layers.base_layers import Dense, ResidualLayer -from .layers.efficient import BasisEmbedding -from .layers.embedding_block import AtomEmbedding, EdgeEmbedding -from .layers.force_scaler import ForceScaler -from .layers.interaction_block import InteractionBlock -from .layers.radial_basis import RadialBasis -from .layers.spherical_basis import CircularBasisLayer, SphericalBasisLayer -from .utils import ( - get_angle, - get_edge_id, - get_inner_idx, - inner_product_clamped, - mask_neighbors, - repeat_blocks, -) - - -@registry.register_model("gemnet_oc") -class GemNetOC(BaseModel): - """ - Arguments - --------- - num_atoms (int): Unused argument - bond_feat_dim (int): Unused argument - num_targets: int - Number of prediction targets. - - num_spherical: int - Controls maximum frequency. - num_radial: int - Controls maximum frequency. - num_blocks: int - Number of building blocks to be stacked. - - emb_size_atom: int - Embedding size of the atoms. - emb_size_edge: int - Embedding size of the edges. - emb_size_trip_in: int - (Down-projected) embedding size of the quadruplet edge embeddings - before the bilinear layer. - emb_size_trip_out: int - (Down-projected) embedding size of the quadruplet edge embeddings - after the bilinear layer. - emb_size_quad_in: int - (Down-projected) embedding size of the quadruplet edge embeddings - before the bilinear layer. - emb_size_quad_out: int - (Down-projected) embedding size of the quadruplet edge embeddings - after the bilinear layer. - emb_size_aint_in: int - Embedding size in the atom interaction before the bilinear layer. - emb_size_aint_out: int - Embedding size in the atom interaction after the bilinear layer. - emb_size_rbf: int - Embedding size of the radial basis transformation. - emb_size_cbf: int - Embedding size of the circular basis transformation (one angle). - emb_size_sbf: int - Embedding size of the spherical basis transformation (two angles). - - num_before_skip: int - Number of residual blocks before the first skip connection. - num_after_skip: int - Number of residual blocks after the first skip connection. - num_concat: int - Number of residual blocks after the concatenation. - num_atom: int - Number of residual blocks in the atom embedding blocks. - num_output_afteratom: int - Number of residual blocks in the output blocks - after adding the atom embedding. - num_atom_emb_layers: int - Number of residual blocks for transforming atom embeddings. - num_global_out_layers: int - Number of final residual blocks before the output. - - regress_forces: bool - Whether to predict forces. Default: True - direct_forces: bool - If True predict forces based on aggregation of interatomic directions. - If False predict forces based on negative gradient of energy potential. - use_pbc: bool - Whether to use periodic boundary conditions. - scale_backprop_forces: bool - Whether to scale up the energy and then scales down the forces - to prevent NaNs and infs in backpropagated forces. - - cutoff: float - Embedding cutoff for interatomic connections and embeddings in Angstrom. - cutoff_qint: float - Quadruplet interaction cutoff in Angstrom. - Optional. Uses cutoff per default. - cutoff_aeaint: float - Edge-to-atom and atom-to-edge interaction cutoff in Angstrom. - Optional. Uses cutoff per default. - cutoff_aint: float - Atom-to-atom interaction cutoff in Angstrom. - Optional. Uses maximum of all other cutoffs per default. - max_neighbors: int - Maximum number of neighbors for interatomic connections and embeddings. - max_neighbors_qint: int - Maximum number of quadruplet interactions per embedding. - Optional. Uses max_neighbors per default. - max_neighbors_aeaint: int - Maximum number of edge-to-atom and atom-to-edge interactions per embedding. - Optional. Uses max_neighbors per default. - max_neighbors_aint: int - Maximum number of atom-to-atom interactions per atom. - Optional. Uses maximum of all other neighbors per default. - - rbf: dict - Name and hyperparameters of the radial basis function. - rbf_spherical: dict - Name and hyperparameters of the radial basis function used as part of the - circular and spherical bases. - Optional. Uses rbf per default. - envelope: dict - Name and hyperparameters of the envelope function. - cbf: dict - Name and hyperparameters of the circular basis function. - sbf: dict - Name and hyperparameters of the spherical basis function. - extensive: bool - Whether the output should be extensive (proportional to the number of atoms) - forces_coupled: bool - If True, enforce that |F_st| = |F_ts|. No effect if direct_forces is False. - output_init: str - Initialization method for the final dense layer. - activation: str - Name of the activation function. - scale_file: str - Path to the pytorch file containing the scaling factors. - - quad_interaction: bool - Whether to use quadruplet interactions (with dihedral angles) - atom_edge_interaction: bool - Whether to use atom-to-edge interactions - edge_atom_interaction: bool - Whether to use edge-to-atom interactions - atom_interaction: bool - Whether to use atom-to-atom interactions - - scale_basis: bool - Whether to use a scaling layer in the raw basis function for better - numerical stability. - qint_tags: list - Which atom tags to use quadruplet interactions for. - 0=sub-surface bulk, 1=surface, 2=adsorbate atoms. - """ - - def __init__( - self, - num_atoms: Optional[int], - bond_feat_dim: int, - num_targets: int, - num_spherical: int, - num_radial: int, - num_blocks: int, - emb_size_atom: int, - emb_size_edge: int, - emb_size_trip_in: int, - emb_size_trip_out: int, - emb_size_quad_in: int, - emb_size_quad_out: int, - emb_size_aint_in: int, - emb_size_aint_out: int, - emb_size_rbf: int, - emb_size_cbf: int, - emb_size_sbf: int, - num_before_skip: int, - num_after_skip: int, - num_concat: int, - num_atom: int, - num_output_afteratom: int, - num_atom_emb_layers: int = 0, - num_global_out_layers: int = 2, - regress_forces: bool = True, - direct_forces: bool = False, - use_pbc: bool = True, - scale_backprop_forces: bool = False, - cutoff: float = 6.0, - cutoff_qint: Optional[float] = None, - cutoff_aeaint: Optional[float] = None, - cutoff_aint: Optional[float] = None, - max_neighbors: int = 50, - max_neighbors_qint: Optional[int] = None, - max_neighbors_aeaint: Optional[int] = None, - max_neighbors_aint: Optional[int] = None, - rbf: dict = {"name": "gaussian"}, - rbf_spherical: Optional[dict] = None, - envelope: dict = {"name": "polynomial", "exponent": 5}, - cbf: dict = {"name": "spherical_harmonics"}, - sbf: dict = {"name": "spherical_harmonics"}, - extensive: bool = True, - forces_coupled: bool = False, - output_init: str = "HeOrthogonal", - activation: str = "silu", - quad_interaction: bool = False, - atom_edge_interaction: bool = False, - edge_atom_interaction: bool = False, - atom_interaction: bool = False, - scale_basis: bool = False, - qint_tags: list = [0, 1, 2], - num_elements: int = 83, - otf_graph: bool = False, - scale_file: Optional[str] = None, - **kwargs, # backwards compatibility with deprecated arguments - ): - super().__init__() - if len(kwargs) > 0: - logging.warning(f"Unrecognized arguments: {list(kwargs.keys())}") - self.num_targets = num_targets - assert num_blocks > 0 - self.num_blocks = num_blocks - self.extensive = extensive - - self.atom_edge_interaction = atom_edge_interaction - self.edge_atom_interaction = edge_atom_interaction - self.atom_interaction = atom_interaction - self.quad_interaction = quad_interaction - self.qint_tags = torch.tensor(qint_tags) - self.otf_graph = otf_graph - if not rbf_spherical: - rbf_spherical = rbf - - self.set_cutoffs(cutoff, cutoff_qint, cutoff_aeaint, cutoff_aint) - self.set_max_neighbors( - max_neighbors, - max_neighbors_qint, - max_neighbors_aeaint, - max_neighbors_aint, - ) - self.use_pbc = use_pbc - - self.direct_forces = direct_forces - self.forces_coupled = forces_coupled - self.regress_forces = regress_forces - self.force_scaler = ForceScaler(enabled=scale_backprop_forces) - - self.init_basis_functions( - num_radial, - num_spherical, - rbf, - rbf_spherical, - envelope, - cbf, - sbf, - scale_basis, - ) - self.init_shared_basis_layers( - num_radial, num_spherical, emb_size_rbf, emb_size_cbf, emb_size_sbf - ) - - # Embedding blocks - self.atom_emb = AtomEmbedding(emb_size_atom, num_elements) - self.edge_emb = EdgeEmbedding( - emb_size_atom, num_radial, emb_size_edge, activation=activation - ) - - # Interaction Blocks - int_blocks = [] - for _ in range(num_blocks): - int_blocks.append( - InteractionBlock( - emb_size_atom=emb_size_atom, - emb_size_edge=emb_size_edge, - emb_size_trip_in=emb_size_trip_in, - emb_size_trip_out=emb_size_trip_out, - emb_size_quad_in=emb_size_quad_in, - emb_size_quad_out=emb_size_quad_out, - emb_size_a2a_in=emb_size_aint_in, - emb_size_a2a_out=emb_size_aint_out, - emb_size_rbf=emb_size_rbf, - emb_size_cbf=emb_size_cbf, - emb_size_sbf=emb_size_sbf, - num_before_skip=num_before_skip, - num_after_skip=num_after_skip, - num_concat=num_concat, - num_atom=num_atom, - num_atom_emb_layers=num_atom_emb_layers, - quad_interaction=quad_interaction, - atom_edge_interaction=atom_edge_interaction, - edge_atom_interaction=edge_atom_interaction, - atom_interaction=atom_interaction, - activation=activation, - ) - ) - self.int_blocks = torch.nn.ModuleList(int_blocks) - - out_blocks = [] - for _ in range(num_blocks + 1): - out_blocks.append( - OutputBlock( - emb_size_atom=emb_size_atom, - emb_size_edge=emb_size_edge, - emb_size_rbf=emb_size_rbf, - nHidden=num_atom, - nHidden_afteratom=num_output_afteratom, - activation=activation, - direct_forces=direct_forces, - ) - ) - self.out_blocks = torch.nn.ModuleList(out_blocks) - - out_mlp_E = [ - Dense( - emb_size_atom * (num_blocks + 1), - emb_size_atom, - activation=activation, - ) - ] - out_mlp_E += [ - ResidualLayer( - emb_size_atom, - activation=activation, - ) - for _ in range(num_global_out_layers) - ] - self.out_mlp_E = torch.nn.Sequential(*out_mlp_E) - self.out_energy = Dense(emb_size_atom, num_targets, bias=False, activation=None) - if direct_forces: - out_mlp_F = [ - Dense( - emb_size_edge * (num_blocks + 1), - emb_size_edge, - activation=activation, - ) - ] - out_mlp_F += [ - ResidualLayer( - emb_size_edge, - activation=activation, - ) - for _ in range(num_global_out_layers) - ] - self.out_mlp_F = torch.nn.Sequential(*out_mlp_F) - self.out_forces = Dense( - emb_size_edge, num_targets, bias=False, activation=None - ) - - out_initializer = get_initializer(output_init) - self.out_energy.reset_parameters(out_initializer) - if direct_forces: - self.out_forces.reset_parameters(out_initializer) - - load_scales_compat(self, scale_file) - - def set_cutoffs(self, cutoff, cutoff_qint, cutoff_aeaint, cutoff_aint): - self.cutoff = cutoff - - if ( - not (self.atom_edge_interaction or self.edge_atom_interaction) - or cutoff_aeaint is None - ): - self.cutoff_aeaint = self.cutoff - else: - self.cutoff_aeaint = cutoff_aeaint - if not self.quad_interaction or cutoff_qint is None: - self.cutoff_qint = self.cutoff - else: - self.cutoff_qint = cutoff_qint - if not self.atom_interaction or cutoff_aint is None: - self.cutoff_aint = max( - self.cutoff, - self.cutoff_aeaint, - self.cutoff_qint, - ) - else: - self.cutoff_aint = cutoff_aint - - assert self.cutoff <= self.cutoff_aint - assert self.cutoff_aeaint <= self.cutoff_aint - assert self.cutoff_qint <= self.cutoff_aint - - def set_max_neighbors( - self, - max_neighbors, - max_neighbors_qint, - max_neighbors_aeaint, - max_neighbors_aint, - ): - self.max_neighbors = max_neighbors - - if ( - not (self.atom_edge_interaction or self.edge_atom_interaction) - or max_neighbors_aeaint is None - ): - self.max_neighbors_aeaint = self.max_neighbors - else: - self.max_neighbors_aeaint = max_neighbors_aeaint - if not self.quad_interaction or max_neighbors_qint is None: - self.max_neighbors_qint = self.max_neighbors - else: - self.max_neighbors_qint = max_neighbors_qint - if not self.atom_interaction or max_neighbors_aint is None: - self.max_neighbors_aint = max( - self.max_neighbors, - self.max_neighbors_aeaint, - self.max_neighbors_qint, - ) - else: - self.max_neighbors_aint = max_neighbors_aint - - assert self.max_neighbors <= self.max_neighbors_aint - assert self.max_neighbors_aeaint <= self.max_neighbors_aint - assert self.max_neighbors_qint <= self.max_neighbors_aint - - def init_basis_functions( - self, - num_radial, - num_spherical, - rbf, - rbf_spherical, - envelope, - cbf, - sbf, - scale_basis, - ): - self.radial_basis = RadialBasis( - num_radial=num_radial, - cutoff=self.cutoff, - rbf=rbf, - envelope=envelope, - scale_basis=scale_basis, - ) - radial_basis_spherical = RadialBasis( - num_radial=num_radial, - cutoff=self.cutoff, - rbf=rbf_spherical, - envelope=envelope, - scale_basis=scale_basis, - ) - if self.quad_interaction: - radial_basis_spherical_qint = RadialBasis( - num_radial=num_radial, - cutoff=self.cutoff_qint, - rbf=rbf_spherical, - envelope=envelope, - scale_basis=scale_basis, - ) - self.cbf_basis_qint = CircularBasisLayer( - num_spherical, - radial_basis=radial_basis_spherical_qint, - cbf=cbf, - scale_basis=scale_basis, - ) - - self.sbf_basis_qint = SphericalBasisLayer( - num_spherical, - radial_basis=radial_basis_spherical, - sbf=sbf, - scale_basis=scale_basis, - ) - if self.atom_edge_interaction: - self.radial_basis_aeaint = RadialBasis( - num_radial=num_radial, - cutoff=self.cutoff_aeaint, - rbf=rbf, - envelope=envelope, - scale_basis=scale_basis, - ) - self.cbf_basis_aeint = CircularBasisLayer( - num_spherical, - radial_basis=radial_basis_spherical, - cbf=cbf, - scale_basis=scale_basis, - ) - if self.edge_atom_interaction: - self.radial_basis_aeaint = RadialBasis( - num_radial=num_radial, - cutoff=self.cutoff_aeaint, - rbf=rbf, - envelope=envelope, - scale_basis=scale_basis, - ) - radial_basis_spherical_aeaint = RadialBasis( - num_radial=num_radial, - cutoff=self.cutoff_aeaint, - rbf=rbf_spherical, - envelope=envelope, - scale_basis=scale_basis, - ) - self.cbf_basis_eaint = CircularBasisLayer( - num_spherical, - radial_basis=radial_basis_spherical_aeaint, - cbf=cbf, - scale_basis=scale_basis, - ) - if self.atom_interaction: - self.radial_basis_aint = RadialBasis( - num_radial=num_radial, - cutoff=self.cutoff_aint, - rbf=rbf, - envelope=envelope, - scale_basis=scale_basis, - ) - - self.cbf_basis_tint = CircularBasisLayer( - num_spherical, - radial_basis=radial_basis_spherical, - cbf=cbf, - scale_basis=scale_basis, - ) - - def init_shared_basis_layers( - self, - num_radial, - num_spherical, - emb_size_rbf, - emb_size_cbf, - emb_size_sbf, - ): - # Share basis down projections across all interaction blocks - if self.quad_interaction: - self.mlp_rbf_qint = Dense( - num_radial, - emb_size_rbf, - activation=None, - bias=False, - ) - self.mlp_cbf_qint = BasisEmbedding(num_radial, emb_size_cbf, num_spherical) - self.mlp_sbf_qint = BasisEmbedding( - num_radial, emb_size_sbf, num_spherical**2 - ) - - if self.atom_edge_interaction: - self.mlp_rbf_aeint = Dense( - num_radial, - emb_size_rbf, - activation=None, - bias=False, - ) - self.mlp_cbf_aeint = BasisEmbedding(num_radial, emb_size_cbf, num_spherical) - if self.edge_atom_interaction: - self.mlp_rbf_eaint = Dense( - num_radial, - emb_size_rbf, - activation=None, - bias=False, - ) - self.mlp_cbf_eaint = BasisEmbedding(num_radial, emb_size_cbf, num_spherical) - if self.atom_interaction: - self.mlp_rbf_aint = BasisEmbedding(num_radial, emb_size_rbf) - - self.mlp_rbf_tint = Dense( - num_radial, - emb_size_rbf, - activation=None, - bias=False, - ) - self.mlp_cbf_tint = BasisEmbedding(num_radial, emb_size_cbf, num_spherical) - - # Share the dense Layer of the atom embedding block accross the interaction blocks - self.mlp_rbf_h = Dense( - num_radial, - emb_size_rbf, - activation=None, - bias=False, - ) - self.mlp_rbf_out = Dense( - num_radial, - emb_size_rbf, - activation=None, - bias=False, - ) - - # Set shared parameters for better gradients - self.shared_parameters = [ - (self.mlp_rbf_tint.linear.weight, self.num_blocks), - (self.mlp_cbf_tint.weight, self.num_blocks), - (self.mlp_rbf_h.linear.weight, self.num_blocks), - (self.mlp_rbf_out.linear.weight, self.num_blocks + 1), - ] - if self.quad_interaction: - self.shared_parameters += [ - (self.mlp_rbf_qint.linear.weight, self.num_blocks), - (self.mlp_cbf_qint.weight, self.num_blocks), - (self.mlp_sbf_qint.weight, self.num_blocks), - ] - if self.atom_edge_interaction: - self.shared_parameters += [ - (self.mlp_rbf_aeint.linear.weight, self.num_blocks), - (self.mlp_cbf_aeint.weight, self.num_blocks), - ] - if self.edge_atom_interaction: - self.shared_parameters += [ - (self.mlp_rbf_eaint.linear.weight, self.num_blocks), - (self.mlp_cbf_eaint.weight, self.num_blocks), - ] - if self.atom_interaction: - self.shared_parameters += [ - (self.mlp_rbf_aint.weight, self.num_blocks), - ] - - def calculate_quad_angles( - self, - V_st, - V_qint_st, - quad_idx, - ): - """Calculate angles for quadruplet-based message passing. - - Arguments - --------- - V_st: Tensor, shape = (nAtoms, 3) - Normalized directions from s to t - V_qint_st: Tensor, shape = (nAtoms, 3) - Normalized directions from s to t for the quadruplet - interaction graph - quad_idx: dict of torch.Tensor - Indices relevant for quadruplet interactions. - - Returns - ------- - cosφ_cab: Tensor, shape = (num_triplets_inint,) - Cosine of angle between atoms c -> a <- b. - cosφ_abd: Tensor, shape = (num_triplets_qint,) - Cosine of angle between atoms a -> b -> d. - angle_cabd: Tensor, shape = (num_quadruplets,) - Dihedral angle between atoms c <- a-b -> d. - """ - # ---------------------------------- d -> b -> a ---------------------------------- # - V_ba = V_qint_st[quad_idx["triplet_in"]["out"]] - # (num_triplets_qint, 3) - V_db = V_st[quad_idx["triplet_in"]["in"]] - # (num_triplets_qint, 3) - cosφ_abd = inner_product_clamped(V_ba, V_db) - # (num_triplets_qint,) - - # Project for calculating dihedral angle - # Cross product is the same as projection, just 90° rotated - V_db_cross = torch.cross(V_db, V_ba, dim=-1) # a - b -| d - V_db_cross = V_db_cross[quad_idx["trip_in_to_quad"]] - # (num_quadruplets,) - - # --------------------------------- c -> a <- b ---------------------------------- # - V_ca = V_st[quad_idx["triplet_out"]["out"]] # (num_triplets_in, 3) - V_ba = V_qint_st[quad_idx["triplet_out"]["in"]] # (num_triplets_in, 3) - cosφ_cab = inner_product_clamped(V_ca, V_ba) # (n4Triplets,) - - # Project for calculating dihedral angle - # Cross product is the same as projection, just 90° rotated - V_ca_cross = torch.cross(V_ca, V_ba, dim=-1) # c |- a - b - V_ca_cross = V_ca_cross[quad_idx["trip_out_to_quad"]] - # (num_quadruplets,) - - # -------------------------------- c -> a - b <- d -------------------------------- # - half_angle_cabd = get_angle(V_ca_cross, V_db_cross) - # (num_quadruplets,) - angle_cabd = half_angle_cabd - # Ignore parity and just use the half angle. - - return cosφ_cab, cosφ_abd, angle_cabd - - def select_symmetric_edges(self, tensor, mask, reorder_idx, opposite_neg): - """Use a mask to remove values of removed edges and then - duplicate the values for the correct edge direction. - - Arguments - --------- - tensor: torch.Tensor - Values to symmetrize for the new tensor. - mask: torch.Tensor - Mask defining which edges go in the correct direction. - reorder_idx: torch.Tensor - Indices defining how to reorder the tensor values after - concatenating the edge values of both directions. - opposite_neg: bool - Whether the edge in the opposite direction should use the - negative tensor value. - - Returns - ------- - tensor_ordered: torch.Tensor - A tensor with symmetrized values. - """ - # Mask out counter-edges - tensor_directed = tensor[mask] - # Concatenate counter-edges after normal edges - sign = 1 - 2 * opposite_neg - tensor_cat = torch.cat([tensor_directed, sign * tensor_directed]) - # Reorder everything so the edges of every image are consecutive - tensor_ordered = tensor_cat[reorder_idx] - return tensor_ordered - - def symmetrize_edges( - self, - graph, - batch_idx, - ): - """ - Symmetrize edges to ensure existence of counter-directional edges. - - Some edges are only present in one direction in the data, - since every atom has a maximum number of neighbors. - We only use i->j edges here. So we lose some j->i edges - and add others by making it symmetric. - """ - num_atoms = batch_idx.shape[0] - new_graph = {} - - # Generate mask - mask_sep_atoms = graph["edge_index"][0] < graph["edge_index"][1] - # Distinguish edges between the same (periodic) atom by ordering the cells - cell_earlier = ( - (graph["cell_offset"][:, 0] < 0) - | ((graph["cell_offset"][:, 0] == 0) & (graph["cell_offset"][:, 1] < 0)) - | ( - (graph["cell_offset"][:, 0] == 0) - & (graph["cell_offset"][:, 1] == 0) - & (graph["cell_offset"][:, 2] < 0) - ) - ) - mask_same_atoms = graph["edge_index"][0] == graph["edge_index"][1] - mask_same_atoms &= cell_earlier - mask = mask_sep_atoms | mask_same_atoms - - # Mask out counter-edges - edge_index_directed = graph["edge_index"][mask[None, :].expand(2, -1)].view( - 2, -1 - ) - - # Concatenate counter-edges after normal edges - edge_index_cat = torch.cat( - [edge_index_directed, edge_index_directed.flip(0)], - dim=1, - ) - - # Count remaining edges per image - batch_edge = torch.repeat_interleave( - torch.arange( - graph["num_neighbors"].size(0), - device=graph["edge_index"].device, - ), - graph["num_neighbors"], - ) - batch_edge = batch_edge[mask] - # segment_coo assumes sorted batch_edge - # Factor 2 since this is only one half of the edges - ones = batch_edge.new_ones(1).expand_as(batch_edge) - new_graph["num_neighbors"] = 2 * segment_coo( - ones, batch_edge, dim_size=graph["num_neighbors"].size(0) - ) - - # Create indexing array - edge_reorder_idx = repeat_blocks( - torch.div(new_graph["num_neighbors"], 2, rounding_mode="floor"), - repeats=2, - continuous_indexing=True, - repeat_inc=edge_index_directed.size(1), - ) - - # Reorder everything so the edges of every image are consecutive - new_graph["edge_index"] = edge_index_cat[:, edge_reorder_idx] - new_graph["cell_offset"] = self.select_symmetric_edges( - graph["cell_offset"], mask, edge_reorder_idx, True - ) - new_graph["distance"] = self.select_symmetric_edges( - graph["distance"], mask, edge_reorder_idx, False - ) - new_graph["vector"] = self.select_symmetric_edges( - graph["vector"], mask, edge_reorder_idx, True - ) - - # Indices for swapping c->a and a->c (for symmetric MP) - # To obtain these efficiently and without any index assumptions, - # we get order the counter-edge IDs and then - # map this order back to the edge IDs. - # Double argsort gives the desired mapping - # from the ordered tensor to the original tensor. - edge_ids = get_edge_id( - new_graph["edge_index"], new_graph["cell_offset"], num_atoms - ) - order_edge_ids = torch.argsort(edge_ids) - inv_order_edge_ids = torch.argsort(order_edge_ids) - edge_ids_counter = get_edge_id( - new_graph["edge_index"].flip(0), - -new_graph["cell_offset"], - num_atoms, - ) - order_edge_ids_counter = torch.argsort(edge_ids_counter) - id_swap = order_edge_ids_counter[inv_order_edge_ids] - - return new_graph, id_swap - - def subselect_edges( - self, - data, - graph, - cutoff=None, - max_neighbors=None, - ): - """Subselect edges using a stricter cutoff and max_neighbors.""" - subgraph = graph.copy() - - if cutoff is not None: - edge_mask = subgraph["distance"] <= cutoff - - subgraph["edge_index"] = subgraph["edge_index"][:, edge_mask] - subgraph["cell_offset"] = subgraph["cell_offset"][edge_mask] - subgraph["num_neighbors"] = mask_neighbors( - subgraph["num_neighbors"], edge_mask - ) - subgraph["distance"] = subgraph["distance"][edge_mask] - subgraph["vector"] = subgraph["vector"][edge_mask] - - if max_neighbors is not None: - edge_mask, subgraph["num_neighbors"] = get_max_neighbors_mask( - natoms=data.natoms, - index=subgraph["edge_index"][1], - atom_distance=subgraph["distance"], - max_num_neighbors_threshold=max_neighbors, - ) - if not torch.all(edge_mask): - subgraph["edge_index"] = subgraph["edge_index"][:, edge_mask] - subgraph["cell_offset"] = subgraph["cell_offset"][edge_mask] - subgraph["distance"] = subgraph["distance"][edge_mask] - subgraph["vector"] = subgraph["vector"][edge_mask] - - empty_image = subgraph["num_neighbors"] == 0 - if torch.any(empty_image): - raise ValueError( - f"An image has no neighbors: id={data.id[empty_image]}, " - f"sid={data.sid[empty_image]}, fid={data.fid[empty_image]}" - ) - return subgraph - - def generate_graph_dict(self, data, cutoff, max_neighbors): - """Generate a radius/nearest neighbor graph.""" - otf_graph = cutoff > 6 or max_neighbors > 50 or self.otf_graph - - ( - edge_index, - edge_dist, - distance_vec, - cell_offsets, - _, # cell offset distances - num_neighbors, - ) = self.generate_graph( - data, - cutoff=cutoff, - max_neighbors=max_neighbors, - otf_graph=otf_graph, - ) - # These vectors actually point in the opposite direction. - # But we want to use col as idx_t for efficient aggregation. - edge_vector = -distance_vec / edge_dist[:, None] - cell_offsets = -cell_offsets # a - c + offset - - graph = { - "edge_index": edge_index, - "distance": edge_dist, - "vector": edge_vector, - "cell_offset": cell_offsets, - "num_neighbors": num_neighbors, - } - - # Mask interaction edges if required - if otf_graph or np.isclose(cutoff, 6): - select_cutoff = None - else: - select_cutoff = cutoff - if otf_graph or max_neighbors == 50: - select_neighbors = None - else: - select_neighbors = max_neighbors - graph = self.subselect_edges( - data=data, - graph=graph, - cutoff=select_cutoff, - max_neighbors=select_neighbors, - ) - - return graph - - def subselect_graph( - self, - data, - graph, - cutoff, - max_neighbors, - cutoff_orig, - max_neighbors_orig, - ): - """If the new cutoff and max_neighbors is different from the original, - subselect the edges of a given graph. - """ - # Check if embedding edges are different from interaction edges - if np.isclose(cutoff, cutoff_orig): - select_cutoff = None - else: - select_cutoff = cutoff - if max_neighbors == max_neighbors_orig: - select_neighbors = None - else: - select_neighbors = max_neighbors - - return self.subselect_edges( - data=data, - graph=graph, - cutoff=select_cutoff, - max_neighbors=select_neighbors, - ) - - def get_graphs_and_indices(self, data): - """ "Generate embedding and interaction graphs and indices.""" - num_atoms = data.atomic_numbers.size(0) - - # Atom interaction graph is always the largest - if ( - self.atom_edge_interaction - or self.edge_atom_interaction - or self.atom_interaction - ): - a2a_graph = self.generate_graph_dict( - data, self.cutoff_aint, self.max_neighbors_aint - ) - main_graph = self.subselect_graph( - data, - a2a_graph, - self.cutoff, - self.max_neighbors, - self.cutoff_aint, - self.max_neighbors_aint, - ) - a2ee2a_graph = self.subselect_graph( - data, - a2a_graph, - self.cutoff_aeaint, - self.max_neighbors_aeaint, - self.cutoff_aint, - self.max_neighbors_aint, - ) - else: - main_graph = self.generate_graph_dict(data, self.cutoff, self.max_neighbors) - a2a_graph = {} - a2ee2a_graph = {} - if self.quad_interaction: - if ( - self.atom_edge_interaction - or self.edge_atom_interaction - or self.atom_interaction - ): - qint_graph = self.subselect_graph( - data, - a2a_graph, - self.cutoff_qint, - self.max_neighbors_qint, - self.cutoff_aint, - self.max_neighbors_aint, - ) - else: - assert self.cutoff_qint <= self.cutoff - assert self.max_neighbors_qint <= self.max_neighbors - qint_graph = self.subselect_graph( - data, - main_graph, - self.cutoff_qint, - self.max_neighbors_qint, - self.cutoff, - self.max_neighbors, - ) - - # Only use quadruplets for certain tags - self.qint_tags = self.qint_tags.to(qint_graph["edge_index"].device) - tags_s = data.tags[qint_graph["edge_index"][0]] - tags_t = data.tags[qint_graph["edge_index"][1]] - qint_tag_mask_s = (tags_s[..., None] == self.qint_tags).any(dim=-1) - qint_tag_mask_t = (tags_t[..., None] == self.qint_tags).any(dim=-1) - qint_tag_mask = qint_tag_mask_s | qint_tag_mask_t - qint_graph["edge_index"] = qint_graph["edge_index"][:, qint_tag_mask] - qint_graph["cell_offset"] = qint_graph["cell_offset"][qint_tag_mask, :] - qint_graph["distance"] = qint_graph["distance"][qint_tag_mask] - qint_graph["vector"] = qint_graph["vector"][qint_tag_mask, :] - del qint_graph["num_neighbors"] - else: - qint_graph = {} - - # Symmetrize edges for swapping in symmetric message passing - main_graph, id_swap = self.symmetrize_edges(main_graph, data.batch) - - trip_idx_e2e = get_triplets(main_graph, num_atoms=num_atoms) - - # Additional indices for quadruplets - if self.quad_interaction: - quad_idx = get_quadruplets( - main_graph, - qint_graph, - num_atoms, - ) - else: - quad_idx = {} - - if self.atom_edge_interaction: - trip_idx_a2e = get_mixed_triplets( - a2ee2a_graph, - main_graph, - num_atoms=num_atoms, - return_agg_idx=True, - ) - else: - trip_idx_a2e = {} - if self.edge_atom_interaction: - trip_idx_e2a = get_mixed_triplets( - main_graph, - a2ee2a_graph, - num_atoms=num_atoms, - return_agg_idx=True, - ) - # a2ee2a_graph['edge_index'][1] has to be sorted for this - a2ee2a_graph["target_neighbor_idx"] = get_inner_idx( - a2ee2a_graph["edge_index"][1], dim_size=num_atoms - ) - else: - trip_idx_e2a = {} - if self.atom_interaction: - # a2a_graph['edge_index'][1] has to be sorted for this - a2a_graph["target_neighbor_idx"] = get_inner_idx( - a2a_graph["edge_index"][1], dim_size=num_atoms - ) - - return ( - main_graph, - a2a_graph, - a2ee2a_graph, - qint_graph, - id_swap, - trip_idx_e2e, - trip_idx_a2e, - trip_idx_e2a, - quad_idx, - ) - - def get_bases( - self, - main_graph, - a2a_graph, - a2ee2a_graph, - qint_graph, - trip_idx_e2e, - trip_idx_a2e, - trip_idx_e2a, - quad_idx, - num_atoms, - ): - """Calculate and transform basis functions.""" - basis_rad_main_raw = self.radial_basis(main_graph["distance"]) - - # Calculate triplet angles - cosφ_cab = inner_product_clamped( - main_graph["vector"][trip_idx_e2e["out"]], - main_graph["vector"][trip_idx_e2e["in"]], - ) - basis_rad_cir_e2e_raw, basis_cir_e2e_raw = self.cbf_basis_tint( - main_graph["distance"], cosφ_cab - ) - - if self.quad_interaction: - # Calculate quadruplet angles - cosφ_cab_q, cosφ_abd, angle_cabd = self.calculate_quad_angles( - main_graph["vector"], - qint_graph["vector"], - quad_idx, - ) - - basis_rad_cir_qint_raw, basis_cir_qint_raw = self.cbf_basis_qint( - qint_graph["distance"], cosφ_abd - ) - basis_rad_sph_qint_raw, basis_sph_qint_raw = self.sbf_basis_qint( - main_graph["distance"], - cosφ_cab_q[quad_idx["trip_out_to_quad"]], - angle_cabd, - ) - if self.atom_edge_interaction: - basis_rad_a2ee2a_raw = self.radial_basis_aeaint(a2ee2a_graph["distance"]) - cosφ_cab_a2e = inner_product_clamped( - main_graph["vector"][trip_idx_a2e["out"]], - a2ee2a_graph["vector"][trip_idx_a2e["in"]], - ) - basis_rad_cir_a2e_raw, basis_cir_a2e_raw = self.cbf_basis_aeint( - main_graph["distance"], cosφ_cab_a2e - ) - if self.edge_atom_interaction: - cosφ_cab_e2a = inner_product_clamped( - a2ee2a_graph["vector"][trip_idx_e2a["out"]], - main_graph["vector"][trip_idx_e2a["in"]], - ) - basis_rad_cir_e2a_raw, basis_cir_e2a_raw = self.cbf_basis_eaint( - a2ee2a_graph["distance"], cosφ_cab_e2a - ) - if self.atom_interaction: - basis_rad_a2a_raw = self.radial_basis_aint(a2a_graph["distance"]) - - # Shared Down Projections - bases_qint = {} - if self.quad_interaction: - bases_qint["rad"] = self.mlp_rbf_qint(basis_rad_main_raw) - bases_qint["cir"] = self.mlp_cbf_qint( - rad_basis=basis_rad_cir_qint_raw, - sph_basis=basis_cir_qint_raw, - idx_sph_outer=quad_idx["triplet_in"]["out"], - ) - bases_qint["sph"] = self.mlp_sbf_qint( - rad_basis=basis_rad_sph_qint_raw, - sph_basis=basis_sph_qint_raw, - idx_sph_outer=quad_idx["out"], - idx_sph_inner=quad_idx["out_agg"], - ) - - bases_a2e = {} - if self.atom_edge_interaction: - bases_a2e["rad"] = self.mlp_rbf_aeint(basis_rad_a2ee2a_raw) - bases_a2e["cir"] = self.mlp_cbf_aeint( - rad_basis=basis_rad_cir_a2e_raw, - sph_basis=basis_cir_a2e_raw, - idx_sph_outer=trip_idx_a2e["out"], - idx_sph_inner=trip_idx_a2e["out_agg"], - ) - bases_e2a = {} - if self.edge_atom_interaction: - bases_e2a["rad"] = self.mlp_rbf_eaint(basis_rad_main_raw) - bases_e2a["cir"] = self.mlp_cbf_eaint( - rad_basis=basis_rad_cir_e2a_raw, - sph_basis=basis_cir_e2a_raw, - idx_rad_outer=a2ee2a_graph["edge_index"][1], - idx_rad_inner=a2ee2a_graph["target_neighbor_idx"], - idx_sph_outer=trip_idx_e2a["out"], - idx_sph_inner=trip_idx_e2a["out_agg"], - num_atoms=num_atoms, - ) - if self.atom_interaction: - basis_a2a_rad = self.mlp_rbf_aint( - rad_basis=basis_rad_a2a_raw, - idx_rad_outer=a2a_graph["edge_index"][1], - idx_rad_inner=a2a_graph["target_neighbor_idx"], - num_atoms=num_atoms, - ) - else: - basis_a2a_rad = None - - bases_e2e = {} - bases_e2e["rad"] = self.mlp_rbf_tint(basis_rad_main_raw) - bases_e2e["cir"] = self.mlp_cbf_tint( - rad_basis=basis_rad_cir_e2e_raw, - sph_basis=basis_cir_e2e_raw, - idx_sph_outer=trip_idx_e2e["out"], - idx_sph_inner=trip_idx_e2e["out_agg"], - ) - - basis_atom_update = self.mlp_rbf_h(basis_rad_main_raw) - basis_output = self.mlp_rbf_out(basis_rad_main_raw) - - return ( - basis_rad_main_raw, - basis_atom_update, - basis_output, - bases_qint, - bases_e2e, - bases_a2e, - bases_e2a, - basis_a2a_rad, - ) - - @conditional_grad(torch.enable_grad()) - def forward(self, data): - pos = data.pos - batch = data.batch - atomic_numbers = data.atomic_numbers.long() - num_atoms = atomic_numbers.shape[0] - - if self.regress_forces and not self.direct_forces: - pos.requires_grad_(True) - - ( - main_graph, - a2a_graph, - a2ee2a_graph, - qint_graph, - id_swap, - trip_idx_e2e, - trip_idx_a2e, - trip_idx_e2a, - quad_idx, - ) = self.get_graphs_and_indices(data) - _, idx_t = main_graph["edge_index"] - - ( - basis_rad_raw, - basis_atom_update, - basis_output, - bases_qint, - bases_e2e, - bases_a2e, - bases_e2a, - basis_a2a_rad, - ) = self.get_bases( - main_graph=main_graph, - a2a_graph=a2a_graph, - a2ee2a_graph=a2ee2a_graph, - qint_graph=qint_graph, - trip_idx_e2e=trip_idx_e2e, - trip_idx_a2e=trip_idx_a2e, - trip_idx_e2a=trip_idx_e2a, - quad_idx=quad_idx, - num_atoms=num_atoms, - ) - - # Embedding block - h = self.atom_emb(atomic_numbers) - # (nAtoms, emb_size_atom) - m = self.edge_emb(h, basis_rad_raw, main_graph["edge_index"]) - # (nEdges, emb_size_edge) - - x_E, x_F = self.out_blocks[0](h, m, basis_output, idx_t) - # (nAtoms, emb_size_atom), (nEdges, emb_size_edge) - xs_E, xs_F = [x_E], [x_F] - - for i in range(self.num_blocks): - # Interaction block - h, m = self.int_blocks[i]( - h=h, - m=m, - bases_qint=bases_qint, - bases_e2e=bases_e2e, - bases_a2e=bases_a2e, - bases_e2a=bases_e2a, - basis_a2a_rad=basis_a2a_rad, - basis_atom_update=basis_atom_update, - edge_index_main=main_graph["edge_index"], - a2ee2a_graph=a2ee2a_graph, - a2a_graph=a2a_graph, - id_swap=id_swap, - trip_idx_e2e=trip_idx_e2e, - trip_idx_a2e=trip_idx_a2e, - trip_idx_e2a=trip_idx_e2a, - quad_idx=quad_idx, - ) # (nAtoms, emb_size_atom), (nEdges, emb_size_edge) - - x_E, x_F = self.out_blocks[i + 1](h, m, basis_output, idx_t) - # (nAtoms, emb_size_atom), (nEdges, emb_size_edge) - xs_E.append(x_E) - xs_F.append(x_F) - - # Global output block for final predictions - x_E = self.out_mlp_E(torch.cat(xs_E, dim=-1)) - if self.direct_forces: - x_F = self.out_mlp_F(torch.cat(xs_F, dim=-1)) - with torch.amp.autocast("cuda", False): - E_t = self.out_energy(x_E.float()) - if self.direct_forces: - F_st = self.out_forces(x_F.float()) - - nMolecules = torch.max(batch) + 1 - if self.extensive: - E_t = scatter_det( - E_t, batch, dim=0, dim_size=nMolecules, reduce="add" - ) # (nMolecules, num_targets) - else: - E_t = scatter_det( - E_t, batch, dim=0, dim_size=nMolecules, reduce="mean" - ) # (nMolecules, num_targets) - - if self.regress_forces: - if self.direct_forces: - if self.forces_coupled: # enforce F_st = F_ts - nEdges = idx_t.shape[0] - id_undir = repeat_blocks( - main_graph["num_neighbors"] // 2, - repeats=2, - continuous_indexing=True, - ) - F_st = scatter_det( - F_st, - id_undir, - dim=0, - dim_size=int(nEdges / 2), - reduce="mean", - ) # (nEdges/2, num_targets) - F_st = F_st[id_undir] # (nEdges, num_targets) - - # map forces in edge directions - F_st_vec = F_st[:, :, None] * main_graph["vector"][:, None, :] - # (nEdges, num_targets, 3) - F_t = scatter_det( - F_st_vec, - idx_t, - dim=0, - dim_size=num_atoms, - reduce="add", - ) # (nAtoms, num_targets, 3) - else: - F_t = self.force_scaler.calc_forces_and_update(E_t, pos) - - E_t = E_t.squeeze(1) # (num_molecules) - F_t = F_t.squeeze(1) # (num_atoms, 3) - return E_t, F_t - else: - E_t = E_t.squeeze(1) # (num_molecules) - return E_t - - @property - def num_params(self): - return sum(p.numel() for p in self.parameters()) diff --git a/ocpmodels/models/gemnet_oc/initializers.py b/ocpmodels/models/gemnet_oc/initializers.py deleted file mode 100644 index badbca1bf5e050a2747d56c3c4145e1ea1c61f22..0000000000000000000000000000000000000000 --- a/ocpmodels/models/gemnet_oc/initializers.py +++ /dev/null @@ -1,95 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -from functools import partial - -import torch - - -def _standardize(kernel): - """ - Makes sure that N*Var(W) = 1 and E[W] = 0 - """ - eps = 1e-6 - - if len(kernel.shape) == 3: - axis = [0, 1] # last dimension is output dimension - else: - axis = 1 - - var, mean = torch.var_mean(kernel, dim=axis, unbiased=True, keepdim=True) - kernel = (kernel - mean) / (var + eps) ** 0.5 - return kernel - - -def he_orthogonal_init(tensor): - """ - Generate a weight matrix with variance according to He (Kaiming) initialization. - Based on a random (semi-)orthogonal matrix neural networks - are expected to learn better when features are decorrelated - (stated by eg. "Reducing overfitting in deep networks by decorrelating representations", - "Dropout: a simple way to prevent neural networks from overfitting", - "Exact solutions to the nonlinear dynamics of learning in deep linear neural networks") - """ - tensor = torch.nn.init.orthogonal_(tensor) - - if len(tensor.shape) == 3: - fan_in = tensor.shape[:-1].numel() - else: - fan_in = tensor.shape[1] - - with torch.no_grad(): - tensor.data = _standardize(tensor.data) - tensor.data *= (1 / fan_in) ** 0.5 - - return tensor - - -def grid_init(tensor, start=-1, end=1): - """ - Generate a weight matrix so that each input value corresponds to one value on a regular grid between start and end. - """ - fan_in = tensor.shape[1] - - with torch.no_grad(): - data = torch.linspace( - start, end, fan_in, device=tensor.device, dtype=tensor.dtype - ).expand_as(tensor) - tensor.copy_(data) - - return tensor - - -def log_grid_init(tensor, start=-4, end=0): - """ - Generate a weight matrix so that each input value corresponds to one value on a regular logarithmic grid between 10^start and 10^end. - """ - fan_in = tensor.shape[1] - - with torch.no_grad(): - data = torch.logspace( - start, end, fan_in, device=tensor.device, dtype=tensor.dtype - ).expand_as(tensor) - tensor.copy_(data) - - return tensor - - -def get_initializer(name, **init_kwargs): - name = name.lower() - if name == "heorthogonal": - initializer = he_orthogonal_init - elif name == "zeros": - initializer = torch.nn.init.zeros_ - elif name == "grid": - initializer = grid_init - elif name == "loggrid": - initializer = log_grid_init - else: - raise UserWarning(f"Unknown initializer: {name}") - - initializer = partial(initializer, **init_kwargs) - return initializer diff --git a/ocpmodels/models/gemnet_oc/interaction_indices.py b/ocpmodels/models/gemnet_oc/interaction_indices.py deleted file mode 100644 index ba05e6093f9b2038702730fd2152ee81d0ceaba4..0000000000000000000000000000000000000000 --- a/ocpmodels/models/gemnet_oc/interaction_indices.py +++ /dev/null @@ -1,300 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import torch -from torch_scatter import segment_coo -from torch_sparse import SparseTensor - -from .utils import get_inner_idx, masked_select_sparsetensor_flat - - -def get_triplets(graph, num_atoms): - """ - Get all input edges b->a for each output edge c->a. - It is possible that b=c, as long as the edges are distinct - (i.e. atoms b and c stem from different unit cells). - - Arguments - --------- - graph: dict of torch.Tensor - Contains the graph's edge_index. - num_atoms: int - Total number of atoms. - - Returns - ------- - Dictionary containing the entries: - in: torch.Tensor, shape (num_triplets,) - Indices of input edge b->a of each triplet b->a<-c - out: torch.Tensor, shape (num_triplets,) - Indices of output edge c->a of each triplet b->a<-c - out_agg: torch.Tensor, shape (num_triplets,) - Indices enumerating the intermediate edges of each output edge. - Used for creating a padded matrix and aggregating via matmul. - """ - idx_s, idx_t = graph["edge_index"] # c->a (source=c, target=a) - num_edges = idx_s.size(0) - - value = torch.arange(num_edges, device=idx_s.device, dtype=idx_s.dtype) - # Possibly contains multiple copies of the same edge (for periodic interactions) - adj = SparseTensor( - row=idx_t, - col=idx_s, - value=value, - sparse_sizes=(num_atoms, num_atoms), - ) - adj_edges = adj[idx_t] - - # Edge indices (b->a, c->a) for triplets. - idx = {} - idx["in"] = adj_edges.storage.value() - idx["out"] = adj_edges.storage.row() - - # Remove self-loop triplets - # Compare edge indices, not atom indices to correctly handle periodic interactions - mask = idx["in"] != idx["out"] - idx["in"] = idx["in"][mask] - idx["out"] = idx["out"][mask] - - # idx['out'] has to be sorted for this - idx["out_agg"] = get_inner_idx(idx["out"], dim_size=num_edges) - - return idx - - -def get_mixed_triplets( - graph_in, - graph_out, - num_atoms, - to_outedge=False, - return_adj=False, - return_agg_idx=False, -): - """ - Get all output edges (ingoing or outgoing) for each incoming edge. - It is possible that in atom=out atom, as long as the edges are distinct - (i.e. they stem from different unit cells). In edges and out edges stem - from separate graphs (hence "mixed") with shared atoms. - - Arguments - --------- - graph_in: dict of torch.Tensor - Contains the input graph's edge_index and cell_offset. - graph_out: dict of torch.Tensor - Contains the output graph's edge_index and cell_offset. - Input and output graphs use the same atoms, but different edges. - num_atoms: int - Total number of atoms. - to_outedge: bool - Whether to map the output to the atom's outgoing edges a->c - instead of the ingoing edges c->a. - return_adj: bool - Whether to output the adjacency (incidence) matrix between output - edges and atoms adj_edges. - return_agg_idx: bool - Whether to output the indices enumerating the intermediate edges - of each output edge. - - Returns - ------- - Dictionary containing the entries: - in: torch.Tensor, shape (num_triplets,) - Indices of input edges - out: torch.Tensor, shape (num_triplets,) - Indices of output edges - adj_edges: SparseTensor, shape (num_edges, num_atoms) - Adjacency (incidence) matrix between output edges and atoms, - with values specifying the input edges. - Only returned if return_adj is True. - out_agg: torch.Tensor, shape (num_triplets,) - Indices enumerating the intermediate edges of each output edge. - Used for creating a padded matrix and aggregating via matmul. - Only returned if return_agg_idx is True. - """ - idx_out_s, idx_out_t = graph_out["edge_index"] - # c->a (source=c, target=a) - idx_in_s, idx_in_t = graph_in["edge_index"] - num_edges = idx_out_s.size(0) - - value_in = torch.arange( - idx_in_s.size(0), device=idx_in_s.device, dtype=idx_in_s.dtype - ) - # This exploits that SparseTensor can have multiple copies of the same edge! - adj_in = SparseTensor( - row=idx_in_t, - col=idx_in_s, - value=value_in, - sparse_sizes=(num_atoms, num_atoms), - ) - if to_outedge: - adj_edges = adj_in[idx_out_s] - else: - adj_edges = adj_in[idx_out_t] - - # Edge indices (b->a, c->a) for triplets. - idx_in = adj_edges.storage.value() - idx_out = adj_edges.storage.row() - - # Remove self-loop triplets c->a<-c or c<-a<-c - # Check atom as well as cell offset - if to_outedge: - idx_atom_in = idx_in_s[idx_in] - idx_atom_out = idx_out_t[idx_out] - cell_offsets_sum = ( - graph_out["cell_offset"][idx_out] + graph_in["cell_offset"][idx_in] - ) - else: - idx_atom_in = idx_in_s[idx_in] - idx_atom_out = idx_out_s[idx_out] - cell_offsets_sum = ( - graph_out["cell_offset"][idx_out] - graph_in["cell_offset"][idx_in] - ) - mask = (idx_atom_in != idx_atom_out) | torch.any(cell_offsets_sum != 0, dim=-1) - - idx = {} - if return_adj: - idx["adj_edges"] = masked_select_sparsetensor_flat(adj_edges, mask) - idx["in"] = idx["adj_edges"].storage.value().clone() - idx["out"] = idx["adj_edges"].storage.row() - else: - idx["in"] = idx_in[mask] - idx["out"] = idx_out[mask] - - if return_agg_idx: - # idx['out'] has to be sorted - idx["out_agg"] = get_inner_idx(idx["out"], dim_size=num_edges) - - return idx - - -def get_quadruplets( - main_graph, - qint_graph, - num_atoms, -): - """ - Get all d->b for each edge c->a and connection b->a - Careful about periodic images! - Separate interaction cutoff not supported. - - Arguments - --------- - main_graph: dict of torch.Tensor - Contains the main graph's edge_index and cell_offset. - The main graph defines which edges are embedded. - qint_graph: dict of torch.Tensor - Contains the quadruplet interaction graph's edge_index and - cell_offset. main_graph and qint_graph use the same atoms, - but different edges. - num_atoms: int - Total number of atoms. - - Returns - ------- - Dictionary containing the entries: - triplet_in['in']: torch.Tensor, shape (nTriplets,) - Indices of input edge d->b in triplet d->b->a. - triplet_in['out']: torch.Tensor, shape (nTriplets,) - Interaction indices of output edge b->a in triplet d->b->a. - triplet_out['in']: torch.Tensor, shape (nTriplets,) - Interaction indices of input edge b->a in triplet c->a<-b. - triplet_out['out']: torch.Tensor, shape (nTriplets,) - Indices of output edge c->a in triplet c->a<-b. - out: torch.Tensor, shape (nQuadruplets,) - Indices of output edge c->a in quadruplet - trip_in_to_quad: torch.Tensor, shape (nQuadruplets,) - Indices to map from input triplet d->b->a - to quadruplet d->b->a<-c. - trip_out_to_quad: torch.Tensor, shape (nQuadruplets,) - Indices to map from output triplet c->a<-b - to quadruplet d->b->a<-c. - out_agg: torch.Tensor, shape (num_triplets,) - Indices enumerating the intermediate edges of each output edge. - Used for creating a padded matrix and aggregating via matmul. - """ - idx_s, _ = main_graph["edge_index"] - idx_qint_s, _ = qint_graph["edge_index"] - # c->a (source=c, target=a) - num_edges = idx_s.size(0) - idx = {} - - idx["triplet_in"] = get_mixed_triplets( - main_graph, - qint_graph, - num_atoms, - to_outedge=True, - return_adj=True, - ) - # Input triplets d->b->a - - idx["triplet_out"] = get_mixed_triplets( - qint_graph, - main_graph, - num_atoms, - to_outedge=False, - ) - # Output triplets c->a<-b - - # ---------------- Quadruplets ----------------- - # Repeat indices by counting the number of input triplets per - # intermediate edge ba. segment_coo assumes sorted idx['triplet_in']['out'] - ones = idx["triplet_in"]["out"].new_ones(1).expand_as(idx["triplet_in"]["out"]) - num_trip_in_per_inter = segment_coo( - ones, idx["triplet_in"]["out"], dim_size=idx_qint_s.size(0) - ) - - num_trip_out_per_inter = num_trip_in_per_inter[idx["triplet_out"]["in"]] - idx["out"] = torch.repeat_interleave( - idx["triplet_out"]["out"], num_trip_out_per_inter - ) - idx_inter = torch.repeat_interleave( - idx["triplet_out"]["in"], num_trip_out_per_inter - ) - idx["trip_out_to_quad"] = torch.repeat_interleave( - torch.arange( - len(idx["triplet_out"]["out"]), - device=idx_s.device, - dtype=idx_s.dtype, - ), - num_trip_out_per_inter, - ) - - # Generate input indices by using the adjacency - # matrix idx['triplet_in']['adj_edges'] - idx["triplet_in"]["adj_edges"].set_value_( - torch.arange( - len(idx["triplet_in"]["in"]), - device=idx_s.device, - dtype=idx_s.dtype, - ), - layout="coo", - ) - adj_trip_in_per_trip_out = idx["triplet_in"]["adj_edges"][idx["triplet_out"]["in"]] - # Rows in adj_trip_in_per_trip_out are intermediate edges ba - idx["trip_in_to_quad"] = adj_trip_in_per_trip_out.storage.value() - idx_in = idx["triplet_in"]["in"][idx["trip_in_to_quad"]] - - # Remove quadruplets with c == d - # Triplets should already ensure that a != d and b != c - # Compare atom indices and cell offsets - idx_atom_c = idx_s[idx["out"]] - idx_atom_d = idx_s[idx_in] - - cell_offset_cd = ( - main_graph["cell_offset"][idx_in] - + qint_graph["cell_offset"][idx_inter] - - main_graph["cell_offset"][idx["out"]] - ) - mask_cd = (idx_atom_c != idx_atom_d) | torch.any(cell_offset_cd != 0, dim=-1) - - idx["out"] = idx["out"][mask_cd] - idx["trip_out_to_quad"] = idx["trip_out_to_quad"][mask_cd] - idx["trip_in_to_quad"] = idx["trip_in_to_quad"][mask_cd] - - # idx['out'] has to be sorted for this - idx["out_agg"] = get_inner_idx(idx["out"], dim_size=num_edges) - - return idx diff --git a/ocpmodels/models/gemnet_oc/layers/atom_update_block.py b/ocpmodels/models/gemnet_oc/layers/atom_update_block.py deleted file mode 100644 index c122c483873a8599b2bcf2732813266b27beb857..0000000000000000000000000000000000000000 --- a/ocpmodels/models/gemnet_oc/layers/atom_update_block.py +++ /dev/null @@ -1,192 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import math - -import torch -from torch_scatter import scatter - -from ocpmodels.common.utils import scatter_det -from ocpmodels.modules.scaling import ScaleFactor - -from ..initializers import get_initializer -from .base_layers import Dense, ResidualLayer - - -class AtomUpdateBlock(torch.nn.Module): - """ - Aggregate the message embeddings of the atoms - - Arguments - --------- - emb_size_atom: int - Embedding size of the atoms. - emb_size_edge: int - Embedding size of the edges. - emb_size_rbf: int - Embedding size of the radial basis. - nHidden: int - Number of residual blocks. - activation: callable/str - Name of the activation function to use in the dense layers. - """ - - def __init__( - self, - emb_size_atom: int, - emb_size_edge: int, - emb_size_rbf: int, - nHidden: int, - activation=None, - ): - super().__init__() - - self.dense_rbf = Dense(emb_size_rbf, emb_size_edge, activation=None, bias=False) - self.scale_sum = ScaleFactor() - - self.layers = self.get_mlp(emb_size_edge, emb_size_atom, nHidden, activation) - - def get_mlp(self, units_in, units, nHidden, activation): - if units_in != units: - dense1 = Dense(units_in, units, activation=activation, bias=False) - mlp = [dense1] - else: - mlp = [] - res = [ - ResidualLayer(units, nLayers=2, activation=activation) - for i in range(nHidden) - ] - mlp += res - return torch.nn.ModuleList(mlp) - - def forward(self, h, m, basis_rad, idx_atom): - """ - Returns - ------- - h: torch.Tensor, shape=(nAtoms, emb_size_atom) - Atom embedding. - """ - nAtoms = h.shape[0] - - bases_emb = self.dense_rbf(basis_rad) # (nEdges, emb_size_edge) - x = m * bases_emb - - x2 = scatter_det( - x, idx_atom, dim=0, dim_size=nAtoms, reduce="sum" - ) # (nAtoms, emb_size_edge) - x = self.scale_sum(x2, ref=m) - - for layer in self.layers: - x = layer(x) # (nAtoms, emb_size_atom) - - return x - - -class OutputBlock(AtomUpdateBlock): - """ - Combines the atom update block and subsequent final dense layer. - - Arguments - --------- - emb_size_atom: int - Embedding size of the atoms. - emb_size_edge: int - Embedding size of the edges. - emb_size_rbf: int - Embedding size of the radial basis. - nHidden: int - Number of residual blocks before adding the atom embedding. - nHidden_afteratom: int - Number of residual blocks after adding the atom embedding. - activation: str - Name of the activation function to use in the dense layers. - direct_forces: bool - If true directly predict forces, i.e. without taking the gradient - of the energy potential. - """ - - def __init__( - self, - emb_size_atom: int, - emb_size_edge: int, - emb_size_rbf: int, - nHidden: int, - nHidden_afteratom: int, - activation=None, - direct_forces=True, - ): - super().__init__( - emb_size_atom=emb_size_atom, - emb_size_edge=emb_size_edge, - emb_size_rbf=emb_size_rbf, - nHidden=nHidden, - activation=activation, - ) - - self.direct_forces = direct_forces - - self.seq_energy_pre = self.layers # inherited from parent class - if nHidden_afteratom >= 1: - self.seq_energy2 = self.get_mlp( - emb_size_atom, emb_size_atom, nHidden_afteratom, activation - ) - self.inv_sqrt_2 = 1 / math.sqrt(2.0) - else: - self.seq_energy2 = None - - if self.direct_forces: - self.scale_rbf_F = ScaleFactor() - self.seq_forces = self.get_mlp( - emb_size_edge, emb_size_edge, nHidden, activation - ) - self.dense_rbf_F = Dense( - emb_size_rbf, emb_size_edge, activation=None, bias=False - ) - - def forward(self, h, m, basis_rad, idx_atom): - """ - Returns - ------- - torch.Tensor, shape=(nAtoms, emb_size_atom) - Output atom embeddings. - torch.Tensor, shape=(nEdges, emb_size_edge) - Output edge embeddings. - """ - nAtoms = h.shape[0] - - # ------------------------ Atom embeddings ------------------------ # - basis_emb_E = self.dense_rbf(basis_rad) # (nEdges, emb_size_edge) - x = m * basis_emb_E - - x_E = scatter_det( - x, idx_atom, dim=0, dim_size=nAtoms, reduce="sum" - ) # (nAtoms, emb_size_edge) - x_E = self.scale_sum(x_E, ref=m) - - for layer in self.seq_energy_pre: - x_E = layer(x_E) # (nAtoms, emb_size_atom) - - if self.seq_energy2 is not None: - x_E = x_E + h - x_E = x_E * self.inv_sqrt_2 - for layer in self.seq_energy2: - x_E = layer(x_E) # (nAtoms, emb_size_atom) - - # ------------------------- Edge embeddings ------------------------ # - if self.direct_forces: - x_F = m - for i, layer in enumerate(self.seq_forces): - x_F = layer(x_F) # (nEdges, emb_size_edge) - - basis_emb_F = self.dense_rbf_F(basis_rad) - # (nEdges, emb_size_edge) - x_F_basis = x_F * basis_emb_F - x_F = self.scale_rbf_F(x_F_basis, ref=x_F) - else: - x_F = 0 - # ------------------------------------------------------------------ # - - return x_E, x_F diff --git a/ocpmodels/models/gemnet_oc/layers/base_layers.py b/ocpmodels/models/gemnet_oc/layers/base_layers.py deleted file mode 100644 index 4e6d78b5eb40b2f0bef17040c4801033ed0b951c..0000000000000000000000000000000000000000 --- a/ocpmodels/models/gemnet_oc/layers/base_layers.py +++ /dev/null @@ -1,98 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import math - -import torch - -from ..initializers import he_orthogonal_init - - -class Dense(torch.nn.Module): - """ - Combines dense layer with scaling for silu activation. - - Arguments - --------- - in_features: int - Input embedding size. - out_features: int - Output embedding size. - bias: bool - True if use bias. - activation: str - Name of the activation function to use. - """ - - def __init__(self, in_features, out_features, bias=False, activation=None): - super().__init__() - - self.linear = torch.nn.Linear(in_features, out_features, bias=bias) - self.reset_parameters() - - if isinstance(activation, str): - activation = activation.lower() - if activation in ["silu", "swish"]: - self._activation = ScaledSiLU() - elif activation is None: - self._activation = torch.nn.Identity() - else: - raise NotImplementedError( - "Activation function not implemented for GemNet (yet)." - ) - - def reset_parameters(self, initializer=he_orthogonal_init): - initializer(self.linear.weight) - if self.linear.bias is not None: - self.linear.bias.data.fill_(0) - - def forward(self, x): - x = self.linear(x) - x = self._activation(x) - return x - - -class ScaledSiLU(torch.nn.Module): - def __init__(self): - super().__init__() - self.scale_factor = 1 / 0.6 - self._activation = torch.nn.SiLU() - - def forward(self, x): - return self._activation(x) * self.scale_factor - - -class ResidualLayer(torch.nn.Module): - """ - Residual block with output scaled by 1/sqrt(2). - - Arguments - --------- - units: int - Input and output embedding size. - nLayers: int - Number of dense layers. - layer: torch.nn.Module - Class for the layers inside the residual block. - layer_kwargs: str - Keyword arguments for initializing the layers. - """ - - def __init__(self, units: int, nLayers: int = 2, layer=Dense, **layer_kwargs): - super().__init__() - self.dense_mlp = torch.nn.Sequential( - *[ - layer(in_features=units, out_features=units, bias=False, **layer_kwargs) - for _ in range(nLayers) - ] - ) - self.inv_sqrt_2 = 1 / math.sqrt(2) - - def forward(self, input): - x = self.dense_mlp(input) - x = input + x - x = x * self.inv_sqrt_2 - return x diff --git a/ocpmodels/models/gemnet_oc/layers/basis_utils.py b/ocpmodels/models/gemnet_oc/layers/basis_utils.py deleted file mode 100644 index 99279cb98256974385953fd3b34cb3c6c9063819..0000000000000000000000000000000000000000 --- a/ocpmodels/models/gemnet_oc/layers/basis_utils.py +++ /dev/null @@ -1,319 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import numpy as np -import sympy as sym -import torch -from scipy import special as sp -from scipy.optimize import brentq - - -def Jn(r, n): - """ - numerical spherical bessel functions of order n - """ - return sp.spherical_jn(n, r) - - -def Jn_zeros(n, k): - """ - Compute the first k zeros of the spherical bessel functions - up to order n (excluded) - """ - zerosj = np.zeros((n, k), dtype="float32") - zerosj[0] = np.arange(1, k + 1) * np.pi - points = np.arange(1, k + n) * np.pi - racines = np.zeros(k + n - 1, dtype="float32") - for i in range(1, n): - for j in range(k + n - 1 - i): - foo = brentq(Jn, points[j], points[j + 1], (i,)) - racines[j] = foo - points = racines - zerosj[i][:k] = racines[:k] - - return zerosj - - -def spherical_bessel_formulas(n): - """ - Computes the sympy formulas for the spherical bessel functions - up to order n (excluded) - """ - x = sym.symbols("x", real=True) - # j_i = (-x)^i * (1/x * d/dx)^î * sin(x)/x - j = [sym.sin(x) / x] # j_0 - a = sym.sin(x) / x - for i in range(1, n): - b = sym.diff(a, x) / x - j += [sym.simplify(b * (-x) ** i)] - a = sym.simplify(b) - return j - - -def bessel_basis(n, k): - """ - Compute the sympy formulas for the normalized and rescaled spherical bessel - functions up to order n (excluded) and maximum frequency k (excluded). - - Returns - ------- - bess_basis: list - Bessel basis formulas taking in a single argument x. - Has length n where each element has length k. -> In total n*k many. - """ - zeros = Jn_zeros(n, k) - normalizer = [] - for order in range(n): - normalizer_tmp = [] - for i in range(k): - normalizer_tmp += [0.5 * Jn(zeros[order, i], order + 1) ** 2] - normalizer_tmp = ( - 1 / np.array(normalizer_tmp) ** 0.5 - ) # sqrt(2/(j_l+1)**2) , sqrt(1/c**3) not taken into account yet - normalizer += [normalizer_tmp] - - f = spherical_bessel_formulas(n) - x = sym.symbols("x", real=True) - bess_basis = [] - for order in range(n): - bess_basis_tmp = [] - for i in range(k): - bess_basis_tmp += [ - sym.simplify( - normalizer[order][i] * f[order].subs(x, zeros[order, i] * x) - ) - ] - bess_basis += [bess_basis_tmp] - return bess_basis - - -def sph_harm_prefactor(l_degree, m_order): - """ - Computes the constant pre-factor for the spherical harmonic - of degree l and order m. - - Arguments - --------- - l_degree: int - Degree of the spherical harmonic. l >= 0 - m_order: int - Order of the spherical harmonic. -l <= m <= l - - Returns - ------- - factor: float - - """ - # sqrt((2*l+1)/4*pi * (l-m)!/(l+m)! ) - return ( - (2 * l_degree + 1) - / (4 * np.pi) - * np.math.factorial(l_degree - abs(m_order)) - / np.math.factorial(l_degree + abs(m_order)) - ) ** 0.5 - - -def associated_legendre_polynomials(L_maxdegree, zero_m_only=True, pos_m_only=True): - """ - Computes string formulas of the associated legendre polynomials - up to degree L (excluded). - - Arguments - --------- - L_maxdegree: int - Degree up to which to calculate the associated legendre polynomials - (degree L is excluded). - zero_m_only: bool - If True only calculate the polynomials for the polynomials where m=0. - pos_m_only: bool - If True only calculate the polynomials for the polynomials where m>=0. - Overwritten by zero_m_only. - - Returns - ------- - polynomials: list - Contains the sympy functions of the polynomials - (in total L many if zero_m_only is True else L^2 many). - """ - # calculations from http://web.cmb.usc.edu/people/alber/Software/tomominer/docs/cpp/group__legendre__polynomials.html - z = sym.symbols("z", real=True) - P_l_m = [ - [0] * (2 * l_degree + 1) for l_degree in range(L_maxdegree) - ] # for order l: -l <= m <= l - - P_l_m[0][0] = 1 - if L_maxdegree > 1: - if zero_m_only: - # m = 0 - P_l_m[1][0] = z - for l_degree in range(2, L_maxdegree): - P_l_m[l_degree][0] = sym.simplify( - ( - (2 * l_degree - 1) * z * P_l_m[l_degree - 1][0] - - (l_degree - 1) * P_l_m[l_degree - 2][0] - ) - / l_degree - ) - return P_l_m - else: - # for m >= 0 - for l_degree in range(1, L_maxdegree): - P_l_m[l_degree][l_degree] = sym.simplify( - (1 - 2 * l_degree) - * (1 - z**2) ** 0.5 - * P_l_m[l_degree - 1][l_degree - 1] - ) # P_00, P_11, P_22, P_33 - - for m_order in range(0, L_maxdegree - 1): - P_l_m[m_order + 1][m_order] = sym.simplify( - (2 * m_order + 1) * z * P_l_m[m_order][m_order] - ) # P_10, P_21, P_32, P_43 - - for l_degree in range(2, L_maxdegree): - for m_order in range(l_degree - 1): # P_20, P_30, P_31 - P_l_m[l_degree][m_order] = sym.simplify( - ( - (2 * l_degree - 1) * z * P_l_m[l_degree - 1][m_order] - - (l_degree + m_order - 1) * P_l_m[l_degree - 2][m_order] - ) - / (l_degree - m_order) - ) - - if not pos_m_only: - # for m < 0: P_l(-m) = (-1)^m * (l-m)!/(l+m)! * P_lm - for l_degree in range(1, L_maxdegree): - for m_order in range(1, l_degree + 1): # P_1(-1), P_2(-1) P_2(-2) - P_l_m[l_degree][-m_order] = sym.simplify( - (-1) ** m_order - * np.math.factorial(l_degree - m_order) - / np.math.factorial(l_degree + m_order) - * P_l_m[l_degree][m_order] - ) - - return P_l_m - - -def real_sph_harm(L_maxdegree, use_theta, use_phi=True, zero_m_only=True): - """ - Computes formula strings of the the real part of the spherical harmonics - up to degree L (excluded). Variables are either spherical coordinates phi - and theta (or cartesian coordinates x,y,z) on the UNIT SPHERE. - - Arguments - --------- - L_maxdegree: int - Degree up to which to calculate the spherical harmonics - (degree L is excluded). - use_theta: bool - - True: Expects the input of the formula strings to contain theta. - - False: Expects the input of the formula strings to contain z. - use_phi: bool - - True: Expects the input of the formula strings to contain phi. - - False: Expects the input of the formula strings to contain x and y. - Does nothing if zero_m_only is True - zero_m_only: bool - If True only calculate the harmonics where m=0. - - Returns - ------- - Y_lm_real: list - Computes formula strings of the the real part of the spherical - harmonics up to degree L (where degree L is not excluded). - In total L^2 many sph harm exist up to degree L (excluded). - However, if zero_m_only only is True then the total count - is reduced to L. - """ - z = sym.symbols("z", real=True) - P_l_m = associated_legendre_polynomials(L_maxdegree, zero_m_only) - if zero_m_only: - # for all m != 0: Y_lm = 0 - Y_l_m = [[0] for l_degree in range(L_maxdegree)] - else: - Y_l_m = [ - [0] * (2 * l_degree + 1) for l_degree in range(L_maxdegree) - ] # for order l: -l <= m <= l - - # convert expressions to spherical coordiantes - if use_theta: - # replace z by cos(theta) - theta = sym.symbols("theta", real=True) - for l_degree in range(L_maxdegree): - for m_order in range(len(P_l_m[l_degree])): - if not isinstance(P_l_m[l_degree][m_order], int): - P_l_m[l_degree][m_order] = P_l_m[l_degree][m_order].subs( - z, sym.cos(theta) - ) - - ## calculate Y_lm - # Y_lm = N * P_lm(cos(theta)) * exp(i*m*phi) - # { sqrt(2) * (-1)^m * N * P_l|m| * sin(|m|*phi) if m < 0 - # Y_lm_real = { Y_lm if m = 0 - # { sqrt(2) * (-1)^m * N * P_lm * cos(m*phi) if m > 0 - - for l_degree in range(L_maxdegree): - Y_l_m[l_degree][0] = sym.simplify( - sph_harm_prefactor(l_degree, 0) * P_l_m[l_degree][0] - ) # Y_l0 - - if not zero_m_only: - phi = sym.symbols("phi", real=True) - for l_degree in range(1, L_maxdegree): - # m > 0 - for m_order in range(1, l_degree + 1): - Y_l_m[l_degree][m_order] = sym.simplify( - 2**0.5 - * (-1) ** m_order - * sph_harm_prefactor(l_degree, m_order) - * P_l_m[l_degree][m_order] - * sym.cos(m_order * phi) - ) - # m < 0 - for m_order in range(1, l_degree + 1): - Y_l_m[l_degree][-m_order] = sym.simplify( - 2**0.5 - * (-1) ** m_order - * sph_harm_prefactor(l_degree, -m_order) - * P_l_m[l_degree][m_order] - * sym.sin(m_order * phi) - ) - - # convert expressions to cartesian coordinates - if not use_phi: - # replace phi by atan2(y,x) - x, y = sym.symbols("x y", real=True) - for l_degree in range(L_maxdegree): - for m_order in range(len(Y_l_m[l_degree])): - Y_l_m[l_degree][m_order] = sym.simplify( - Y_l_m[l_degree][m_order].subs(phi, sym.atan2(y, x)) - ) - return Y_l_m - - -def get_sph_harm_basis(L_maxdegree, zero_m_only=True): - """Get a function calculating the spherical harmonics basis from z and phi.""" - # retrieve equations - Y_lm = real_sph_harm( - L_maxdegree, use_theta=False, use_phi=True, zero_m_only=zero_m_only - ) - Y_lm_flat = [Y for Y_l in Y_lm for Y in Y_l] - - # convert to pytorch functions - z = sym.symbols("z", real=True) - variables = [z] - if not zero_m_only: - variables.append(sym.symbols("phi", real=True)) - - modules = {"sin": torch.sin, "cos": torch.cos, "sqrt": torch.sqrt} - sph_funcs = sym.lambdify(variables, Y_lm_flat, modules) - - # Return as a single function - # args are either [cosφ] or [cosφ, ϑ] - def basis_fn(*args): - basis = sph_funcs(*args) - basis[0] = args[0].new_tensor(basis[0]).expand_as(args[0]) - return torch.stack(basis, dim=1) - - return basis_fn diff --git a/ocpmodels/models/gemnet_oc/layers/efficient.py b/ocpmodels/models/gemnet_oc/layers/efficient.py deleted file mode 100644 index 9a78c8230d63d718fcf1c4b7a0348367349ff662..0000000000000000000000000000000000000000 --- a/ocpmodels/models/gemnet_oc/layers/efficient.py +++ /dev/null @@ -1,261 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -from typing import Optional - -import torch -from torch_scatter import scatter - -from ..initializers import he_orthogonal_init -from .base_layers import Dense - - -class BasisEmbedding(torch.nn.Module): - """ - Embed a basis (CBF, SBF), optionally using the efficient reformulation. - - Arguments - --------- - num_radial: int - Number of radial basis functions. - emb_size_interm: int - Intermediate embedding size of triplets/quadruplets. - num_spherical: int - Number of circular/spherical basis functions. - Only required if there is a circular/spherical basis. - """ - - def __init__( - self, - num_radial: int, - emb_size_interm: int, - num_spherical: Optional[int] = None, - ): - super().__init__() - self.num_radial = num_radial - self.num_spherical = num_spherical - if num_spherical is None: - self.weight = torch.nn.Parameter( - torch.empty(emb_size_interm, num_radial), - requires_grad=True, - ) - else: - self.weight = torch.nn.Parameter( - torch.empty(num_radial, num_spherical, emb_size_interm), - requires_grad=True, - ) - self.reset_parameters() - - def reset_parameters(self): - he_orthogonal_init(self.weight) - - def forward( - self, - rad_basis, - sph_basis=None, - idx_rad_outer=None, - idx_rad_inner=None, - idx_sph_outer=None, - idx_sph_inner=None, - num_atoms=None, - ): - """ - - Arguments - --------- - rad_basis: torch.Tensor, shape=(num_edges, num_radial or num_orders * num_radial) - Raw radial basis. - sph_basis: torch.Tensor, shape=(num_triplets or num_quadruplets, num_spherical) - Raw spherical or circular basis. - idx_rad_outer: torch.Tensor, shape=(num_edges) - Atom associated with each radial basis value. - Optional, used for efficient edge aggregation. - idx_rad_inner: torch.Tensor, shape=(num_edges) - Enumerates radial basis values per atom. - Optional, used for efficient edge aggregation. - idx_sph_outer: torch.Tensor, shape=(num_triplets or num_quadruplets) - Edge associated with each circular/spherical basis value. - Optional, used for efficient triplet/quadruplet aggregation. - idx_sph_inner: torch.Tensor, shape=(num_triplets or num_quadruplets) - Enumerates circular/spherical basis values per edge. - Optional, used for efficient triplet/quadruplet aggregation. - num_atoms: int - Total number of atoms. - Optional, used for efficient edge aggregation. - - Returns - ------- - rad_W1: torch.Tensor, shape=(num_edges, emb_size_interm, num_spherical) - sph: torch.Tensor, shape=(num_edges, Kmax, num_spherical) - Kmax = maximum number of neighbors of the edges - """ - num_edges = rad_basis.shape[0] - - if self.num_spherical is not None: - # MatMul: mul + sum over num_radial - rad_W1 = rad_basis @ self.weight.reshape(self.weight.shape[0], -1) - # (num_edges, emb_size_interm * num_spherical) - rad_W1 = rad_W1.reshape(num_edges, -1, sph_basis.shape[-1]) - # (num_edges, emb_size_interm, num_spherical) - else: - # MatMul: mul + sum over num_radial - rad_W1 = rad_basis @ self.weight.T - # (num_edges, emb_size_interm) - - if idx_rad_inner is not None: - # Zero padded dense matrix - # maximum number of neighbors - if idx_rad_outer.shape[0] == 0: - # catch empty idx_rad_outer - Kmax = 0 - else: - Kmax = torch.max(idx_rad_inner) + 1 - - rad_W1_padded = rad_W1.new_zeros([num_atoms, Kmax] + list(rad_W1.shape[1:])) - rad_W1_padded[idx_rad_outer, idx_rad_inner] = rad_W1 - # (num_atoms, Kmax, emb_size_interm, ...) - rad_W1_padded = torch.transpose(rad_W1_padded, 1, 2) - # (num_atoms, emb_size_interm, Kmax, ...) - rad_W1_padded = rad_W1_padded.reshape(num_atoms, rad_W1.shape[1], -1) - # (num_atoms, emb_size_interm, Kmax2 * ...) - rad_W1 = rad_W1_padded - - if idx_sph_inner is not None: - # Zero padded dense matrix - # maximum number of neighbors - if idx_sph_outer.shape[0] == 0: - # catch empty idx_sph_outer - Kmax = 0 - else: - Kmax = torch.max(idx_sph_inner) + 1 - - sph2 = sph_basis.new_zeros(num_edges, Kmax, sph_basis.shape[-1]) - sph2[idx_sph_outer, idx_sph_inner] = sph_basis - # (num_edges, Kmax, num_spherical) - sph2 = torch.transpose(sph2, 1, 2) - # (num_edges, num_spherical, Kmax) - - if sph_basis is None: - return rad_W1 - else: - if idx_sph_inner is None: - rad_W1 = rad_W1[idx_sph_outer] - # (num_triplets, emb_size_interm, num_spherical) - - sph_W1 = rad_W1 @ sph_basis[:, :, None] - # (num_triplets, emb_size_interm, num_spherical) - return sph_W1.squeeze(-1) - else: - return rad_W1, sph2 - - -class EfficientInteractionBilinear(torch.nn.Module): - """ - Efficient reformulation of the bilinear layer and subsequent summation. - - Arguments - --------- - emb_size_in: int - Embedding size of input triplets/quadruplets. - emb_size_interm: int - Intermediate embedding size of the basis transformation. - emb_size_out: int - Embedding size of output triplets/quadruplets. - """ - - def __init__( - self, - emb_size_in: int, - emb_size_interm: int, - emb_size_out: int, - ): - super().__init__() - self.emb_size_in = emb_size_in - self.emb_size_interm = emb_size_interm - self.emb_size_out = emb_size_out - - self.bilinear = Dense( - self.emb_size_in * self.emb_size_interm, - self.emb_size_out, - bias=False, - activation=None, - ) - - def forward( - self, - basis, - m, - idx_agg_outer, - idx_agg_inner, - idx_agg2_outer=None, - idx_agg2_inner=None, - agg2_out_size=None, - ): - """ - - Arguments - --------- - basis: Tuple (torch.Tensor, torch.Tensor), - shapes=((num_edges, emb_size_interm, num_spherical), - (num_edges, num_spherical, Kmax)) - First element: Radial basis multiplied with weight matrix - Second element: Circular/spherical basis - m: torch.Tensor, shape=(num_edges, emb_size_in) - Input edge embeddings - idx_agg_outer: torch.Tensor, shape=(num_triplets or num_quadruplets) - Output edge aggregating this intermediate triplet/quadruplet edge. - idx_agg_inner: torch.Tensor, shape=(num_triplets or num_quadruplets) - Enumerates intermediate edges per output edge. - idx_agg2_outer: torch.Tensor, shape=(num_edges) - Output atom aggregating this edge. - idx_agg2_inner: torch.Tensor, shape=(num_edges) - Enumerates edges per output atom. - agg2_out_size: int - Number of output embeddings when aggregating twice. Typically - the number of atoms. - - Returns - ------- - m_ca: torch.Tensor, shape=(num_edges, emb_size) - Aggregated edge/atom embeddings. - """ - # num_spherical is actually num_spherical**2 for quadruplets - (rad_W1, sph) = basis - # (num_edges, emb_size_interm, num_spherical), - # (num_edges, num_spherical, Kmax) - num_edges = sph.shape[0] - - # Create (zero-padded) dense matrix of the neighboring edge embeddings. - Kmax = torch.max(idx_agg_inner) + 1 - m_padded = m.new_zeros(num_edges, Kmax, self.emb_size_in) - m_padded[idx_agg_outer, idx_agg_inner] = m - # (num_quadruplets/num_triplets, emb_size_in) -> (num_edges, Kmax, emb_size_in) - - sph_m = torch.matmul(sph, m_padded) - # (num_edges, num_spherical, emb_size_in) - - if idx_agg2_outer is not None: - Kmax2 = torch.max(idx_agg2_inner) + 1 - sph_m_padded = sph_m.new_zeros( - agg2_out_size, Kmax2, sph_m.shape[1], sph_m.shape[2] - ) - sph_m_padded[idx_agg2_outer, idx_agg2_inner] = sph_m - # (num_atoms, Kmax2, num_spherical, emb_size_in) - sph_m_padded = sph_m_padded.reshape(agg2_out_size, -1, sph_m.shape[-1]) - # (num_atoms, Kmax2 * num_spherical, emb_size_in) - - rad_W1_sph_m = rad_W1 @ sph_m_padded - # (num_atoms, emb_size_interm, emb_size_in) - else: - # MatMul: mul + sum over num_spherical - rad_W1_sph_m = torch.matmul(rad_W1, sph_m) - # (num_edges, emb_size_interm, emb_size_in) - - # Bilinear: Sum over emb_size_interm and emb_size_in - m_ca = self.bilinear(rad_W1_sph_m.reshape(-1, rad_W1_sph_m.shape[1:].numel())) - # (num_edges/num_atoms, emb_size_out) - - return m_ca diff --git a/ocpmodels/models/gemnet_oc/layers/embedding_block.py b/ocpmodels/models/gemnet_oc/layers/embedding_block.py deleted file mode 100644 index 9dde5a9a7ed7c7326ba1f8f746b75817123fc807..0000000000000000000000000000000000000000 --- a/ocpmodels/models/gemnet_oc/layers/embedding_block.py +++ /dev/null @@ -1,95 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import numpy as np -import torch - -from .base_layers import Dense - - -class AtomEmbedding(torch.nn.Module): - """ - Initial atom embeddings based on the atom type - - Arguments - --------- - emb_size: int - Atom embeddings size - """ - - def __init__(self, emb_size, num_elements): - super().__init__() - self.emb_size = emb_size - - self.embeddings = torch.nn.Embedding(num_elements, emb_size) - # init by uniform distribution - torch.nn.init.uniform_(self.embeddings.weight, a=-np.sqrt(3), b=np.sqrt(3)) - - def forward(self, Z): - """ - Returns - ------- - h: torch.Tensor, shape=(nAtoms, emb_size) - Atom embeddings. - """ - h = self.embeddings(Z - 1) # -1 because Z.min()=1 (==Hydrogen) - return h - - -class EdgeEmbedding(torch.nn.Module): - """ - Edge embedding based on the concatenation of atom embeddings - and a subsequent dense layer. - - Arguments - --------- - atom_features: int - Embedding size of the atom embedding. - edge_features: int - Embedding size of the input edge embedding. - out_features: int - Embedding size after the dense layer. - activation: str - Activation function used in the dense layer. - """ - - def __init__( - self, - atom_features, - edge_features, - out_features, - activation=None, - ): - super().__init__() - in_features = 2 * atom_features + edge_features - self.dense = Dense(in_features, out_features, activation=activation, bias=False) - - def forward( - self, - h, - m, - edge_index, - ): - """ - Arguments - --------- - h: torch.Tensor, shape (num_atoms, atom_features) - Atom embeddings. - m: torch.Tensor, shape (num_edges, edge_features) - Radial basis in embedding block, - edge embedding in interaction block. - - Returns - ------- - m_st: torch.Tensor, shape=(nEdges, emb_size) - Edge embeddings. - """ - h_s = h[edge_index[0]] # shape=(nEdges, emb_size) - h_t = h[edge_index[1]] # shape=(nEdges, emb_size) - - m_st = torch.cat([h_s, h_t, m], dim=-1) # (nEdges, 2*emb_size+nFeatures) - m_st = self.dense(m_st) # (nEdges, emb_size) - return m_st diff --git a/ocpmodels/models/gemnet_oc/layers/force_scaler.py b/ocpmodels/models/gemnet_oc/layers/force_scaler.py deleted file mode 100644 index f444542ce5c1d4c5f6fea2a4b2daab6b5231d01e..0000000000000000000000000000000000000000 --- a/ocpmodels/models/gemnet_oc/layers/force_scaler.py +++ /dev/null @@ -1,94 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import logging - -import torch - - -class ForceScaler: - """ - Scales up the energy and then scales down the forces - to prevent NaNs and infs in calculations using AMP. - Inspired by torch.cuda.amp.GradScaler. - """ - - def __init__( - self, - init_scale=2.0**8, - growth_factor=2.0, - backoff_factor=0.5, - growth_interval=2000, - max_force_iters=50, - enabled=True, - ): - self.scale_factor = init_scale - self.growth_factor = growth_factor - self.backoff_factor = backoff_factor - self.growth_interval = growth_interval - self.max_force_iters = max_force_iters - self.enabled = enabled - self.finite_force_results = 0 - - def scale(self, energy): - return energy * self.scale_factor if self.enabled else energy - - def unscale(self, forces): - return forces / self.scale_factor if self.enabled else forces - - def calc_forces(self, energy, pos): - energy_scaled = self.scale(energy) - forces_scaled = -torch.autograd.grad( - energy_scaled, - pos, - grad_outputs=torch.ones_like(energy_scaled), - create_graph=True, - )[0] - # (nAtoms, 3) - forces = self.unscale(forces_scaled) - return forces - - def calc_forces_and_update(self, energy, pos): - if self.enabled: - found_nans_or_infs = True - force_iters = 0 - - # Re-calculate forces until everything is nice and finite. - while found_nans_or_infs: - forces = self.calc_forces(energy, pos) - - found_nans_or_infs = not torch.all(forces.isfinite()) - if found_nans_or_infs: - self.finite_force_results = 0 - - # Prevent infinite loop - force_iters += 1 - if force_iters == self.max_force_iters: - logging.warning( - "Too many non-finite force results in a batch. " - "Breaking scaling loop." - ) - break - else: - # Delete graph to save memory - del forces - else: - self.finite_force_results += 1 - self.update() - else: - forces = self.calc_forces(energy, pos) - return forces - - def update(self): - if self.finite_force_results == 0: - self.scale_factor *= self.backoff_factor - - if self.finite_force_results == self.growth_interval: - self.scale_factor *= self.growth_factor - self.finite_force_results = 0 - - logging.info(f"finite force step count: {self.finite_force_results}") - logging.info(f"scaling factor: {self.scale_factor}") diff --git a/ocpmodels/models/gemnet_oc/layers/interaction_block.py b/ocpmodels/models/gemnet_oc/layers/interaction_block.py deleted file mode 100644 index 6c61670b516c55ea27903ae35ebb5fa50c66e2b5..0000000000000000000000000000000000000000 --- a/ocpmodels/models/gemnet_oc/layers/interaction_block.py +++ /dev/null @@ -1,758 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import math - -import torch - -from ocpmodels.modules.scaling import ScaleFactor - -from .atom_update_block import AtomUpdateBlock -from .base_layers import Dense, ResidualLayer -from .efficient import EfficientInteractionBilinear -from .embedding_block import EdgeEmbedding - - -class InteractionBlock(torch.nn.Module): - """ - Interaction block for GemNet-Q/dQ. - - Arguments - --------- - emb_size_atom: int - Embedding size of the atoms. - emb_size_edge: int - Embedding size of the edges. - emb_size_trip_in: int - (Down-projected) embedding size of the quadruplet edge embeddings - before the bilinear layer. - emb_size_trip_out: int - (Down-projected) embedding size of the quadruplet edge embeddings - after the bilinear layer. - emb_size_quad_in: int - (Down-projected) embedding size of the quadruplet edge embeddings - before the bilinear layer. - emb_size_quad_out: int - (Down-projected) embedding size of the quadruplet edge embeddings - after the bilinear layer. - emb_size_a2a_in: int - Embedding size in the atom interaction before the bilinear layer. - emb_size_a2a_out: int - Embedding size in the atom interaction after the bilinear layer. - emb_size_rbf: int - Embedding size of the radial basis transformation. - emb_size_cbf: int - Embedding size of the circular basis transformation (one angle). - emb_size_sbf: int - Embedding size of the spherical basis transformation (two angles). - num_before_skip: int - Number of residual blocks before the first skip connection. - num_after_skip: int - Number of residual blocks after the first skip connection. - num_concat: int - Number of residual blocks after the concatenation. - num_atom: int - Number of residual blocks in the atom embedding blocks. - num_atom_emb_layers: int - Number of residual blocks for transforming atom embeddings. - quad_interaction: bool - Whether to use quadruplet interactions. - atom_edge_interaction: bool - Whether to use atom-to-edge interactions. - edge_atom_interaction: bool - Whether to use edge-to-atom interactions. - atom_interaction: bool - Whether to use atom-to-atom interactions. - activation: str - Name of the activation function to use in the dense layers. - """ - - def __init__( - self, - emb_size_atom, - emb_size_edge, - emb_size_trip_in, - emb_size_trip_out, - emb_size_quad_in, - emb_size_quad_out, - emb_size_a2a_in, - emb_size_a2a_out, - emb_size_rbf, - emb_size_cbf, - emb_size_sbf, - num_before_skip, - num_after_skip, - num_concat, - num_atom, - num_atom_emb_layers=0, - quad_interaction=False, - atom_edge_interaction=False, - edge_atom_interaction=False, - atom_interaction=False, - activation=None, - ): - super().__init__() - - ## ------------------------ Message Passing ----------------------- ## - # Dense transformation of skip connection - self.dense_ca = Dense( - emb_size_edge, - emb_size_edge, - activation=activation, - bias=False, - ) - - # Triplet Interaction - self.trip_interaction = TripletInteraction( - emb_size_in=emb_size_edge, - emb_size_out=emb_size_edge, - emb_size_trip_in=emb_size_trip_in, - emb_size_trip_out=emb_size_trip_out, - emb_size_rbf=emb_size_rbf, - emb_size_cbf=emb_size_cbf, - symmetric_mp=True, - swap_output=True, - activation=activation, - ) - - # Quadruplet Interaction - if quad_interaction: - self.quad_interaction = QuadrupletInteraction( - emb_size_edge=emb_size_edge, - emb_size_quad_in=emb_size_quad_in, - emb_size_quad_out=emb_size_quad_out, - emb_size_rbf=emb_size_rbf, - emb_size_cbf=emb_size_cbf, - emb_size_sbf=emb_size_sbf, - symmetric_mp=True, - activation=activation, - ) - else: - self.quad_interaction = None - - if atom_edge_interaction: - self.atom_edge_interaction = TripletInteraction( - emb_size_in=emb_size_atom, - emb_size_out=emb_size_edge, - emb_size_trip_in=emb_size_trip_in, - emb_size_trip_out=emb_size_trip_out, - emb_size_rbf=emb_size_rbf, - emb_size_cbf=emb_size_cbf, - symmetric_mp=True, - swap_output=True, - activation=activation, - ) - else: - self.atom_edge_interaction = None - if edge_atom_interaction: - self.edge_atom_interaction = TripletInteraction( - emb_size_in=emb_size_edge, - emb_size_out=emb_size_atom, - emb_size_trip_in=emb_size_trip_in, - emb_size_trip_out=emb_size_trip_out, - emb_size_rbf=emb_size_rbf, - emb_size_cbf=emb_size_cbf, - symmetric_mp=False, - swap_output=False, - activation=activation, - ) - else: - self.edge_atom_interaction = None - if atom_interaction: - self.atom_interaction = PairInteraction( - emb_size_atom=emb_size_atom, - emb_size_pair_in=emb_size_a2a_in, - emb_size_pair_out=emb_size_a2a_out, - emb_size_rbf=emb_size_rbf, - activation=activation, - ) - else: - self.atom_interaction = None - - ## -------------------- Update Edge Embeddings -------------------- ## - # Residual layers before skip connection - self.layers_before_skip = torch.nn.ModuleList( - [ - ResidualLayer( - emb_size_edge, - activation=activation, - ) - for i in range(num_before_skip) - ] - ) - - # Residual layers after skip connection - self.layers_after_skip = torch.nn.ModuleList( - [ - ResidualLayer( - emb_size_edge, - activation=activation, - ) - for i in range(num_after_skip) - ] - ) - - ## -------------------- Update Atom Embeddings -------------------- ## - self.atom_emb_layers = torch.nn.ModuleList( - [ - ResidualLayer( - emb_size_atom, - activation=activation, - ) - for _ in range(num_atom_emb_layers) - ] - ) - - self.atom_update = AtomUpdateBlock( - emb_size_atom=emb_size_atom, - emb_size_edge=emb_size_edge, - emb_size_rbf=emb_size_rbf, - nHidden=num_atom, - activation=activation, - ) - - ## ---------- Update Edge Embeddings with Atom Embeddings --------- ## - self.concat_layer = EdgeEmbedding( - emb_size_atom, - emb_size_edge, - emb_size_edge, - activation=activation, - ) - self.residual_m = torch.nn.ModuleList( - [ - ResidualLayer(emb_size_edge, activation=activation) - for _ in range(num_concat) - ] - ) - - self.inv_sqrt_2 = 1 / math.sqrt(2.0) - num_eint = 2.0 + quad_interaction + atom_edge_interaction - self.inv_sqrt_num_eint = 1 / math.sqrt(num_eint) - num_aint = 1.0 + edge_atom_interaction + atom_interaction - self.inv_sqrt_num_aint = 1 / math.sqrt(num_aint) - - def forward( - self, - h, - m, - bases_qint, - bases_e2e, - bases_a2e, - bases_e2a, - basis_a2a_rad, - basis_atom_update, - edge_index_main, - a2ee2a_graph, - a2a_graph, - id_swap, - trip_idx_e2e, - trip_idx_a2e, - trip_idx_e2a, - quad_idx, - ): - """ - Returns - ------- - h: torch.Tensor, shape=(nEdges, emb_size_atom) - Atom embeddings. - m: torch.Tensor, shape=(nEdges, emb_size_edge) - Edge embeddings (c->a). - """ - num_atoms = h.shape[0] - - # Initial transformation - x_ca_skip = self.dense_ca(m) # (nEdges, emb_size_edge) - - x_e2e = self.trip_interaction( - m, - bases_e2e, - trip_idx_e2e, - id_swap, - ) - if self.quad_interaction is not None: - x_qint = self.quad_interaction( - m, - bases_qint, - quad_idx, - id_swap, - ) - if self.atom_edge_interaction is not None: - x_a2e = self.atom_edge_interaction( - h, - bases_a2e, - trip_idx_a2e, - id_swap, - expand_idx=a2ee2a_graph["edge_index"][0], - ) - if self.edge_atom_interaction is not None: - h_e2a = self.edge_atom_interaction( - m, - bases_e2a, - trip_idx_e2a, - id_swap, - idx_agg2=a2ee2a_graph["edge_index"][1], - idx_agg2_inner=a2ee2a_graph["target_neighbor_idx"], - agg2_out_size=num_atoms, - ) - if self.atom_interaction is not None: - h_a2a = self.atom_interaction( - h, - basis_a2a_rad, - a2a_graph["edge_index"], - a2a_graph["target_neighbor_idx"], - ) - - ## -------------- Merge Embeddings after interactions ------------- ## - x = x_ca_skip + x_e2e # (nEdges, emb_size_edge) - if self.quad_interaction is not None: - x += x_qint # (nEdges, emb_size_edge) - if self.atom_edge_interaction is not None: - x += x_a2e # (nEdges, emb_size_edge) - x = x * self.inv_sqrt_num_eint - - # Merge atom embeddings after interactions - if self.edge_atom_interaction is not None: - h = h + h_e2a # (nEdges, emb_size_edge) - if self.atom_interaction is not None: - h = h + h_a2a # (nEdges, emb_size_edge) - h = h * self.inv_sqrt_num_aint - - ## -------------------- Update Edge Embeddings -------------------- ## - # Transformations before skip connection - for i, layer in enumerate(self.layers_before_skip): - x = layer(x) # (nEdges, emb_size_edge) - - # Skip connection - m = m + x # (nEdges, emb_size_edge) - m = m * self.inv_sqrt_2 - - # Transformations after skip connection - for i, layer in enumerate(self.layers_after_skip): - m = layer(m) # (nEdges, emb_size_edge) - - ## -------------------- Update Atom Embeddings -------------------- ## - for layer in self.atom_emb_layers: - h = layer(h) # (nAtoms, emb_size_atom) - - h2 = self.atom_update(h, m, basis_atom_update, edge_index_main[1]) - - # Skip connection - h = h + h2 # (nAtoms, emb_size_atom) - h = h * self.inv_sqrt_2 - - ## ---------- Update Edge Embeddings with Atom Embeddings --------- ## - m2 = self.concat_layer(h, m, edge_index_main) - # (nEdges, emb_size_edge) - - for i, layer in enumerate(self.residual_m): - m2 = layer(m2) # (nEdges, emb_size_edge) - - # Skip connection - m = m + m2 # (nEdges, emb_size_edge) - m = m * self.inv_sqrt_2 - return h, m - - -class QuadrupletInteraction(torch.nn.Module): - """ - Quadruplet-based message passing block. - - Arguments - --------- - emb_size_edge: int - Embedding size of the edges. - emb_size_quad_in: int - (Down-projected) embedding size of the quadruplet edge embeddings - before the bilinear layer. - emb_size_quad_out: int - (Down-projected) embedding size of the quadruplet edge embeddings - after the bilinear layer. - emb_size_rbf: int - Embedding size of the radial basis transformation. - emb_size_cbf: int - Embedding size of the circular basis transformation (one angle). - emb_size_sbf: int - Embedding size of the spherical basis transformation (two angles). - symmetric_mp: bool - Whether to use symmetric message passing and - update the edges in both directions. - activation: str - Name of the activation function to use in the dense layers. - """ - - def __init__( - self, - emb_size_edge, - emb_size_quad_in, - emb_size_quad_out, - emb_size_rbf, - emb_size_cbf, - emb_size_sbf, - symmetric_mp=True, - activation=None, - ): - super().__init__() - self.symmetric_mp = symmetric_mp - - # Dense transformation - self.dense_db = Dense( - emb_size_edge, - emb_size_edge, - activation=activation, - bias=False, - ) - - # Up projections of basis representations, - # bilinear layer and scaling factors - self.mlp_rbf = Dense( - emb_size_rbf, - emb_size_edge, - activation=None, - bias=False, - ) - self.scale_rbf = ScaleFactor() - - self.mlp_cbf = Dense( - emb_size_cbf, - emb_size_quad_in, - activation=None, - bias=False, - ) - self.scale_cbf = ScaleFactor() - - self.mlp_sbf = EfficientInteractionBilinear( - emb_size_quad_in, emb_size_sbf, emb_size_quad_out - ) - self.scale_sbf_sum = ScaleFactor() - # combines scaling for bilinear layer and summation - - # Down and up projections - self.down_projection = Dense( - emb_size_edge, - emb_size_quad_in, - activation=activation, - bias=False, - ) - self.up_projection_ca = Dense( - emb_size_quad_out, - emb_size_edge, - activation=activation, - bias=False, - ) - if self.symmetric_mp: - self.up_projection_ac = Dense( - emb_size_quad_out, - emb_size_edge, - activation=activation, - bias=False, - ) - - self.inv_sqrt_2 = 1 / math.sqrt(2.0) - - def forward( - self, - m, - bases, - idx, - id_swap, - ): - """ - Returns - ------- - m: torch.Tensor, shape=(nEdges, emb_size_edge) - Edge embeddings (c->a). - """ - - x_db = self.dense_db(m) # (nEdges, emb_size_edge) - - # Transform via radial basis - x_db2 = x_db * self.mlp_rbf(bases["rad"]) # (nEdges, emb_size_edge) - x_db = self.scale_rbf(x_db2, ref=x_db) - - # Down project embeddings - x_db = self.down_projection(x_db) # (nEdges, emb_size_quad_in) - - # Transform via circular basis - x_db = x_db[idx["triplet_in"]["in"]] - # (num_triplets_int, emb_size_quad_in) - - x_db2 = x_db * self.mlp_cbf(bases["cir"]) - # (num_triplets_int, emb_size_quad_in) - x_db = self.scale_cbf(x_db2, ref=x_db) - - # Transform via spherical basis - x_db = x_db[idx["trip_in_to_quad"]] - # (num_quadruplets, emb_size_quad_in) - x = self.mlp_sbf(bases["sph"], x_db, idx["out"], idx["out_agg"]) - # (nEdges, emb_size_quad_out) - x = self.scale_sbf_sum(x, ref=x_db) - - # => - # rbf(d_db) - # cbf(d_ba, angle_abd) - # sbf(d_ca, angle_cab, angle_cabd) - - if self.symmetric_mp: - # Upproject embeddings - x_ca = self.up_projection_ca(x) # (nEdges, emb_size_edge) - x_ac = self.up_projection_ac(x) # (nEdges, emb_size_edge) - - # Merge interaction of c->a and a->c - x_ac = x_ac[id_swap] # swap to add to edge a->c and not c->a - x_res = x_ca + x_ac - x_res = x_res * self.inv_sqrt_2 - return x_res - else: - x_res = self.up_projection_ca(x) - return x_res - - -class TripletInteraction(torch.nn.Module): - """ - Triplet-based message passing block. - - Arguments - --------- - emb_size_in: int - Embedding size of the input embeddings. - emb_size_out: int - Embedding size of the output embeddings. - emb_size_trip_in: int - (Down-projected) embedding size of the quadruplet edge embeddings - before the bilinear layer. - emb_size_trip_out: int - (Down-projected) embedding size of the quadruplet edge embeddings - after the bilinear layer. - emb_size_rbf: int - Embedding size of the radial basis transformation. - emb_size_cbf: int - Embedding size of the circular basis transformation (one angle). - symmetric_mp: bool - Whether to use symmetric message passing and - update the edges in both directions. - swap_output: bool - Whether to swap the output embedding directions. - Only relevant if symmetric_mp is False. - activation: str - Name of the activation function to use in the dense layers. - """ - - def __init__( - self, - emb_size_in, - emb_size_out, - emb_size_trip_in, - emb_size_trip_out, - emb_size_rbf, - emb_size_cbf, - symmetric_mp=True, - swap_output=True, - activation=None, - ): - super().__init__() - self.symmetric_mp = symmetric_mp - self.swap_output = swap_output - - # Dense transformation - self.dense_ba = Dense( - emb_size_in, - emb_size_in, - activation=activation, - bias=False, - ) - - # Up projections of basis representations, bilinear layer and scaling factors - self.mlp_rbf = Dense( - emb_size_rbf, - emb_size_in, - activation=None, - bias=False, - ) - self.scale_rbf = ScaleFactor() - - self.mlp_cbf = EfficientInteractionBilinear( - emb_size_trip_in, emb_size_cbf, emb_size_trip_out - ) - self.scale_cbf_sum = ScaleFactor() - # combines scaling for bilinear layer and summation - - # Down and up projections - self.down_projection = Dense( - emb_size_in, - emb_size_trip_in, - activation=activation, - bias=False, - ) - self.up_projection_ca = Dense( - emb_size_trip_out, - emb_size_out, - activation=activation, - bias=False, - ) - if self.symmetric_mp: - self.up_projection_ac = Dense( - emb_size_trip_out, - emb_size_out, - activation=activation, - bias=False, - ) - - self.inv_sqrt_2 = 1 / math.sqrt(2.0) - - def forward( - self, - m, - bases, - idx, - id_swap, - expand_idx=None, - idx_agg2=None, - idx_agg2_inner=None, - agg2_out_size=None, - ): - """ - Returns - ------- - m: torch.Tensor, shape=(nEdges, emb_size_edge) - Edge embeddings. - """ - - # Dense transformation - x_ba = self.dense_ba(m) # (nEdges, emb_size_edge) - - if expand_idx is not None: - x_ba = x_ba[expand_idx] - - # Transform via radial basis - rad_emb = self.mlp_rbf(bases["rad"]) # (nEdges, emb_size_edge) - x_ba2 = x_ba * rad_emb - x_ba = self.scale_rbf(x_ba2, ref=x_ba) - - x_ba = self.down_projection(x_ba) # (nEdges, emb_size_trip_in) - - # Transform via circular spherical basis - x_ba = x_ba[idx["in"]] - - # Efficient bilinear layer - x = self.mlp_cbf( - basis=bases["cir"], - m=x_ba, - idx_agg_outer=idx["out"], - idx_agg_inner=idx["out_agg"], - idx_agg2_outer=idx_agg2, - idx_agg2_inner=idx_agg2_inner, - agg2_out_size=agg2_out_size, - ) - # (num_atoms, emb_size_trip_out) - x = self.scale_cbf_sum(x, ref=x_ba) - - # => - # rbf(d_ba) - # cbf(d_ca, angle_cab) - - if self.symmetric_mp: - # Up project embeddings - x_ca = self.up_projection_ca(x) # (nEdges, emb_size_edge) - x_ac = self.up_projection_ac(x) # (nEdges, emb_size_edge) - - # Merge interaction of c->a and a->c - x_ac = x_ac[id_swap] # swap to add to edge a->c and not c->a - x_res = x_ca + x_ac - x_res = x_res * self.inv_sqrt_2 - return x_res - else: - if self.swap_output: - x = x[id_swap] - x_res = self.up_projection_ca(x) # (nEdges, emb_size_edge) - return x_res - - -class PairInteraction(torch.nn.Module): - """ - Pair-based message passing block. - - Arguments - --------- - emb_size_atom: int - Embedding size of the atoms. - emb_size_pair_in: int - Embedding size of the atom pairs before the bilinear layer. - emb_size_pair_out: int - Embedding size of the atom pairs after the bilinear layer. - emb_size_rbf: int - Embedding size of the radial basis transformation. - activation: str - Name of the activation function to use in the dense layers. - """ - - def __init__( - self, - emb_size_atom, - emb_size_pair_in, - emb_size_pair_out, - emb_size_rbf, - activation=None, - ): - super().__init__() - - # Bilinear layer and scaling factor - self.bilinear = Dense( - emb_size_rbf * emb_size_pair_in, - emb_size_pair_out, - activation=None, - bias=False, - ) - self.scale_rbf_sum = ScaleFactor() - - # Down and up projections - self.down_projection = Dense( - emb_size_atom, - emb_size_pair_in, - activation=activation, - bias=False, - ) - self.up_projection = Dense( - emb_size_pair_out, - emb_size_atom, - activation=activation, - bias=False, - ) - - self.inv_sqrt_2 = 1 / math.sqrt(2.0) - - def forward( - self, - h, - rad_basis, - edge_index, - target_neighbor_idx, - ): - """ - Returns - ------- - h: torch.Tensor, shape=(num_atoms, emb_size_atom) - Atom embeddings. - """ - num_atoms = h.shape[0] - - x_b = self.down_projection(h) # (num_atoms, emb_size_edge) - x_ba = x_b[edge_index[0]] # (num_edges, emb_size_edge) - - Kmax = torch.max(target_neighbor_idx) + 1 - x2 = x_ba.new_zeros(num_atoms, Kmax, x_ba.shape[-1]) - x2[edge_index[1], target_neighbor_idx] = x_ba - # (num_atoms, Kmax, emb_size_edge) - - x_ba2 = rad_basis @ x2 - # (num_atoms, emb_size_interm, emb_size_edge) - h_out = self.bilinear(x_ba2.reshape(num_atoms, -1)) - - h_out = self.scale_rbf_sum(h_out, ref=x_ba) - # (num_atoms, emb_size_edge) - - h_out = self.up_projection(h_out) # (num_atoms, emb_size_atom) - - return h_out diff --git a/ocpmodels/models/gemnet_oc/layers/radial_basis.py b/ocpmodels/models/gemnet_oc/layers/radial_basis.py deleted file mode 100644 index d2ec06fe737b296a81ba03de64f06ac41e05073f..0000000000000000000000000000000000000000 --- a/ocpmodels/models/gemnet_oc/layers/radial_basis.py +++ /dev/null @@ -1,231 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import math - -import numpy as np -import sympy as sym -import torch -from scipy.special import binom - -from ocpmodels.modules.scaling import ScaleFactor - -from .basis_utils import bessel_basis - - -class PolynomialEnvelope(torch.nn.Module): - """ - Polynomial envelope function that ensures a smooth cutoff. - - Arguments - --------- - exponent: int - Exponent of the envelope function. - """ - - def __init__(self, exponent): - super().__init__() - assert exponent > 0 - self.p = exponent - self.a = -(self.p + 1) * (self.p + 2) / 2 - self.b = self.p * (self.p + 2) - self.c = -self.p * (self.p + 1) / 2 - - def forward(self, d_scaled): - env_val = ( - 1 - + self.a * d_scaled**self.p - + self.b * d_scaled ** (self.p + 1) - + self.c * d_scaled ** (self.p + 2) - ) - return torch.where(d_scaled < 1, env_val, torch.zeros_like(d_scaled)) - - -class ExponentialEnvelope(torch.nn.Module): - """ - Exponential envelope function that ensures a smooth cutoff, - as proposed in Unke, Chmiela, Gastegger, Schütt, Sauceda, Müller 2021. - SpookyNet: Learning Force Fields with Electronic Degrees of Freedom - and Nonlocal Effects - """ - - def __init__(self): - super().__init__() - - def forward(self, d_scaled): - env_val = torch.exp(-(d_scaled**2) / ((1 - d_scaled) * (1 + d_scaled))) - return torch.where(d_scaled < 1, env_val, torch.zeros_like(d_scaled)) - - -class GaussianBasis(torch.nn.Module): - def __init__(self, start=0.0, stop=5.0, num_gaussians=50, trainable=False): - super().__init__() - offset = torch.linspace(start, stop, num_gaussians) - if trainable: - self.offset = torch.nn.Parameter(offset, requires_grad=True) - else: - self.register_buffer("offset", offset) - self.coeff = -0.5 / ((stop - start) / (num_gaussians - 1)) ** 2 - - def forward(self, dist): - dist = dist[:, None] - self.offset[None, :] - return torch.exp(self.coeff * torch.pow(dist, 2)) - - -class SphericalBesselBasis(torch.nn.Module): - """ - First-order spherical Bessel basis - - Arguments - --------- - num_radial: int - Number of basis functions. Controls the maximum frequency. - cutoff: float - Cutoff distance in Angstrom. - """ - - def __init__( - self, - num_radial: int, - cutoff: float, - ): - super().__init__() - self.norm_const = math.sqrt(2 / (cutoff**3)) - # cutoff ** 3 to counteract dividing by d_scaled = d / cutoff - - # Initialize frequencies at canonical positions - self.frequencies = torch.nn.Parameter( - data=torch.tensor(np.pi * np.arange(1, num_radial + 1, dtype=np.float32)), - requires_grad=True, - ) - - def forward(self, d_scaled): - return ( - self.norm_const - / d_scaled[:, None] - * torch.sin(self.frequencies * d_scaled[:, None]) - ) # (num_edges, num_radial) - - -class BernsteinBasis(torch.nn.Module): - """ - Bernstein polynomial basis, - as proposed in Unke, Chmiela, Gastegger, Schütt, Sauceda, Müller 2021. - SpookyNet: Learning Force Fields with Electronic Degrees of Freedom - and Nonlocal Effects - - Arguments - --------- - num_radial: int - Number of basis functions. Controls the maximum frequency. - pregamma_initial: float - Initial value of exponential coefficient gamma. - Default: gamma = 0.5 * a_0**-1 = 0.94486, - inverse softplus -> pregamma = log e**gamma - 1 = 0.45264 - """ - - def __init__( - self, - num_radial: int, - pregamma_initial: float = 0.45264, - ): - super().__init__() - prefactor = binom(num_radial - 1, np.arange(num_radial)) - self.register_buffer( - "prefactor", - torch.tensor(prefactor, dtype=torch.float), - persistent=False, - ) - - self.pregamma = torch.nn.Parameter( - data=torch.tensor(pregamma_initial, dtype=torch.float), - requires_grad=True, - ) - self.softplus = torch.nn.Softplus() - - exp1 = torch.arange(num_radial) - self.register_buffer("exp1", exp1[None, :], persistent=False) - exp2 = num_radial - 1 - exp1 - self.register_buffer("exp2", exp2[None, :], persistent=False) - - def forward(self, d_scaled): - gamma = self.softplus(self.pregamma) # constrain to positive - exp_d = torch.exp(-gamma * d_scaled)[:, None] - return self.prefactor * (exp_d**self.exp1) * ((1 - exp_d) ** self.exp2) - - -class RadialBasis(torch.nn.Module): - """ - - Arguments - --------- - num_radial: int - Number of basis functions. Controls the maximum frequency. - cutoff: float - Cutoff distance in Angstrom. - rbf: dict = {"name": "gaussian"} - Basis function and its hyperparameters. - envelope: dict = {"name": "polynomial", "exponent": 5} - Envelope function and its hyperparameters. - scale_basis: bool - Whether to scale the basis values for better numerical stability. - """ - - def __init__( - self, - num_radial: int, - cutoff: float, - rbf: dict = {"name": "gaussian"}, - envelope: dict = {"name": "polynomial", "exponent": 5}, - scale_basis: bool = False, - ): - super().__init__() - self.inv_cutoff = 1 / cutoff - - self.scale_basis = scale_basis - if self.scale_basis: - self.scale_rbf = ScaleFactor() - - env_name = envelope["name"].lower() - env_hparams = envelope.copy() - del env_hparams["name"] - - if env_name == "polynomial": - self.envelope = PolynomialEnvelope(**env_hparams) - elif env_name == "exponential": - self.envelope = ExponentialEnvelope(**env_hparams) - else: - raise ValueError(f"Unknown envelope function '{env_name}'.") - - rbf_name = rbf["name"].lower() - rbf_hparams = rbf.copy() - del rbf_hparams["name"] - - # RBFs get distances scaled to be in [0, 1] - if rbf_name == "gaussian": - self.rbf = GaussianBasis( - start=0, stop=1, num_gaussians=num_radial, **rbf_hparams - ) - elif rbf_name == "spherical_bessel": - self.rbf = SphericalBesselBasis( - num_radial=num_radial, cutoff=cutoff, **rbf_hparams - ) - elif rbf_name == "bernstein": - self.rbf = BernsteinBasis(num_radial=num_radial, **rbf_hparams) - else: - raise ValueError(f"Unknown radial basis function '{rbf_name}'.") - - def forward(self, d): - d_scaled = d * self.inv_cutoff - - env = self.envelope(d_scaled) - res = env[:, None] * self.rbf(d_scaled) - - if self.scale_basis: - res = self.scale_rbf(res) - - return res - # (num_edges, num_radial) or (num_edges, num_orders * num_radial) diff --git a/ocpmodels/models/gemnet_oc/layers/spherical_basis.py b/ocpmodels/models/gemnet_oc/layers/spherical_basis.py deleted file mode 100644 index faa738e51fb8495e00afc33dbdb4f9c60f3f4815..0000000000000000000000000000000000000000 --- a/ocpmodels/models/gemnet_oc/layers/spherical_basis.py +++ /dev/null @@ -1,137 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import torch - -from ocpmodels.modules.scaling import ScaleFactor - -from .basis_utils import get_sph_harm_basis -from .radial_basis import GaussianBasis, RadialBasis - - -class CircularBasisLayer(torch.nn.Module): - """ - 2D Fourier Bessel Basis - - Arguments - --------- - num_spherical: int - Number of basis functions. Controls the maximum frequency. - radial_basis: RadialBasis - Radial basis function. - cbf: dict - Name and hyperparameters of the circular basis function. - scale_basis: bool - Whether to scale the basis values for better numerical stability. - """ - - def __init__( - self, - num_spherical: int, - radial_basis: RadialBasis, - cbf: dict, - scale_basis: bool = False, - ): - super().__init__() - - self.radial_basis = radial_basis - - self.scale_basis = scale_basis - if self.scale_basis: - self.scale_cbf = ScaleFactor() - - cbf_name = cbf["name"].lower() - cbf_hparams = cbf.copy() - del cbf_hparams["name"] - - if cbf_name == "gaussian": - self.cosφ_basis = GaussianBasis( - start=-1, stop=1, num_gaussians=num_spherical, **cbf_hparams - ) - elif cbf_name == "spherical_harmonics": - self.cosφ_basis = get_sph_harm_basis(num_spherical, zero_m_only=True) - else: - raise ValueError(f"Unknown cosine basis function '{cbf_name}'.") - - def forward(self, D_ca, cosφ_cab): - rad_basis = self.radial_basis(D_ca) # (num_edges, num_radial) - cir_basis = self.cosφ_basis(cosφ_cab) # (num_triplets, num_spherical) - - if self.scale_basis: - cir_basis = self.scale_cbf(cir_basis) - - return rad_basis, cir_basis - # (num_edges, num_radial), (num_triplets, num_spherical) - - -class SphericalBasisLayer(torch.nn.Module): - """ - 3D Fourier Bessel Basis - - Arguments - --------- - num_spherical: int - Number of basis functions. Controls the maximum frequency. - radial_basis: RadialBasis - Radial basis functions. - sbf: dict - Name and hyperparameters of the spherical basis function. - scale_basis: bool - Whether to scale the basis values for better numerical stability. - """ - - def __init__( - self, - num_spherical: int, - radial_basis: RadialBasis, - sbf: dict, - scale_basis: bool = False, - ): - super().__init__() - - self.num_spherical = num_spherical - self.radial_basis = radial_basis - - self.scale_basis = scale_basis - if self.scale_basis: - self.scale_sbf = ScaleFactor() - - sbf_name = sbf["name"].lower() - sbf_hparams = sbf.copy() - del sbf_hparams["name"] - - if sbf_name == "spherical_harmonics": - self.spherical_basis = get_sph_harm_basis(num_spherical, zero_m_only=False) - - elif sbf_name == "legendre_outer": - circular_basis = get_sph_harm_basis(num_spherical, zero_m_only=True) - self.spherical_basis = lambda cosφ, ϑ: ( - circular_basis(cosφ)[:, :, None] - * circular_basis(torch.cos(ϑ))[:, None, :] - ).reshape(cosφ.shape[0], -1) - - elif sbf_name == "gaussian_outer": - self.circular_basis = GaussianBasis( - start=-1, stop=1, num_gaussians=num_spherical, **sbf_hparams - ) - self.spherical_basis = lambda cosφ, ϑ: ( - self.circular_basis(cosφ)[:, :, None] - * self.circular_basis(torch.cos(ϑ))[:, None, :] - ).reshape(cosφ.shape[0], -1) - - else: - raise ValueError(f"Unknown spherical basis function '{sbf_name}'.") - - def forward(self, D_ca, cosφ_cab, θ_cabd): - rad_basis = self.radial_basis(D_ca) - sph_basis = self.spherical_basis(cosφ_cab, θ_cabd) - # (num_quadruplets, num_spherical**2) - - if self.scale_basis: - sph_basis = self.scale_sbf(sph_basis) - - return rad_basis, sph_basis - # (num_edges, num_radial), (num_quadruplets, num_spherical**2) diff --git a/ocpmodels/models/gemnet_oc/utils.py b/ocpmodels/models/gemnet_oc/utils.py deleted file mode 100644 index 37addbcee8dbc8e061cbb1863196da5ba7b61841..0000000000000000000000000000000000000000 --- a/ocpmodels/models/gemnet_oc/utils.py +++ /dev/null @@ -1,413 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import numpy as np -import torch -from torch_scatter import segment_coo, segment_csr -from torch_sparse import SparseTensor - - -def ragged_range(sizes): - """Multiple concatenated ranges. - - Examples - -------- - sizes = [1 4 2 3] - Return: [0 0 1 2 3 0 1 0 1 2] - """ - assert sizes.dim() == 1 - if sizes.sum() == 0: - return sizes.new_empty(0) - - # Remove 0 sizes - sizes_nonzero = sizes > 0 - if not torch.all(sizes_nonzero): - sizes = torch.masked_select(sizes, sizes_nonzero) - - # Initialize indexing array with ones as we need to setup incremental indexing - # within each group when cumulatively summed at the final stage. - id_steps = torch.ones(sizes.sum(), dtype=torch.long, device=sizes.device) - id_steps[0] = 0 - insert_index = sizes[:-1].cumsum(0) - insert_val = (1 - sizes)[:-1] - - # Assign index-offsetting values - id_steps[insert_index] = insert_val - - # Finally index into input array for the group repeated o/p - res = id_steps.cumsum(0) - return res - - -def repeat_blocks( - sizes, - repeats, - continuous_indexing=True, - start_idx=0, - block_inc=0, - repeat_inc=0, -): - """Repeat blocks of indices. - Adapted from https://stackoverflow.com/questions/51154989/numpy-vectorized-function-to-repeat-blocks-of-consecutive-elements - - continuous_indexing: Whether to keep increasing the index after each block - start_idx: Starting index - block_inc: Number to increment by after each block, - either global or per block. Shape: len(sizes) - 1 - repeat_inc: Number to increment by after each repetition, - either global or per block - - Examples - -------- - sizes = [1,3,2] ; repeats = [3,2,3] ; continuous_indexing = False - Return: [0 0 0 0 1 2 0 1 2 0 1 0 1 0 1] - sizes = [1,3,2] ; repeats = [3,2,3] ; continuous_indexing = True - Return: [0 0 0 1 2 3 1 2 3 4 5 4 5 4 5] - sizes = [1,3,2] ; repeats = [3,2,3] ; continuous_indexing = True ; - repeat_inc = 4 - Return: [0 4 8 1 2 3 5 6 7 4 5 8 9 12 13] - sizes = [1,3,2] ; repeats = [3,2,3] ; continuous_indexing = True ; - start_idx = 5 - Return: [5 5 5 6 7 8 6 7 8 9 10 9 10 9 10] - sizes = [1,3,2] ; repeats = [3,2,3] ; continuous_indexing = True ; - block_inc = 1 - Return: [0 0 0 2 3 4 2 3 4 6 7 6 7 6 7] - sizes = [0,3,2] ; repeats = [3,2,3] ; continuous_indexing = True - Return: [0 1 2 0 1 2 3 4 3 4 3 4] - sizes = [2,3,2] ; repeats = [2,0,2] ; continuous_indexing = True - Return: [0 1 0 1 5 6 5 6] - """ - assert sizes.dim() == 1 - assert all(sizes >= 0) - - # Remove 0 sizes - sizes_nonzero = sizes > 0 - if not torch.all(sizes_nonzero): - assert block_inc == 0 # Implementing this is not worth the effort - sizes = torch.masked_select(sizes, sizes_nonzero) - if isinstance(repeats, torch.Tensor): - repeats = torch.masked_select(repeats, sizes_nonzero) - if isinstance(repeat_inc, torch.Tensor): - repeat_inc = torch.masked_select(repeat_inc, sizes_nonzero) - - if isinstance(repeats, torch.Tensor): - assert all(repeats >= 0) - insert_dummy = repeats[0] == 0 - if insert_dummy: - one = sizes.new_ones(1) - zero = sizes.new_zeros(1) - sizes = torch.cat((one, sizes)) - repeats = torch.cat((one, repeats)) - if isinstance(block_inc, torch.Tensor): - block_inc = torch.cat((zero, block_inc)) - if isinstance(repeat_inc, torch.Tensor): - repeat_inc = torch.cat((zero, repeat_inc)) - else: - assert repeats >= 0 - insert_dummy = False - - # Get repeats for each group using group lengths/sizes - r1 = torch.repeat_interleave(torch.arange(len(sizes), device=sizes.device), repeats) - - # Get total size of output array, as needed to initialize output indexing array - N = (sizes * repeats).sum() - - # Initialize indexing array with ones as we need to setup incremental indexing - # within each group when cumulatively summed at the final stage. - # Two steps here: - # 1. Within each group, we have multiple sequences, so setup the offsetting - # at each sequence lengths by the seq. lengths preceding those. - id_ar = torch.ones(N, dtype=torch.long, device=sizes.device) - id_ar[0] = 0 - insert_index = sizes[r1[:-1]].cumsum(0) - insert_val = (1 - sizes)[r1[:-1]] - - if isinstance(repeats, torch.Tensor) and torch.any(repeats == 0): - diffs = r1[1:] - r1[:-1] - indptr = torch.cat((sizes.new_zeros(1), diffs.cumsum(0))) - if continuous_indexing: - # If a group was skipped (repeats=0) we need to add its size - insert_val += segment_csr(sizes[: r1[-1]], indptr, reduce="sum") - - # Add block increments - if isinstance(block_inc, torch.Tensor): - insert_val += segment_csr(block_inc[: r1[-1]], indptr, reduce="sum") - else: - insert_val += block_inc * (indptr[1:] - indptr[:-1]) - if insert_dummy: - insert_val[0] -= block_inc - else: - idx = r1[1:] != r1[:-1] - if continuous_indexing: - # 2. For each group, make sure the indexing starts from the next group's - # first element. So, simply assign 1s there. - insert_val[idx] = 1 - - # Add block increments - insert_val[idx] += block_inc - - # Add repeat_inc within each group - if isinstance(repeat_inc, torch.Tensor): - insert_val += repeat_inc[r1[:-1]] - if isinstance(repeats, torch.Tensor): - repeat_inc_inner = repeat_inc[repeats > 0][:-1] - else: - repeat_inc_inner = repeat_inc[:-1] - else: - insert_val += repeat_inc - repeat_inc_inner = repeat_inc - - # Subtract the increments between groups - if isinstance(repeats, torch.Tensor): - repeats_inner = repeats[repeats > 0][:-1] - else: - repeats_inner = repeats - insert_val[r1[1:] != r1[:-1]] -= repeat_inc_inner * repeats_inner - - # Assign index-offsetting values - id_ar[insert_index] = insert_val - - if insert_dummy: - id_ar = id_ar[1:] - if continuous_indexing: - id_ar[0] -= 1 - - # Set start index now, in case of insertion due to leading repeats=0 - id_ar[0] += start_idx - - # Finally index into input array for the group repeated o/p - res = id_ar.cumsum(0) - return res - - -def masked_select_sparsetensor_flat(src, mask): - row, col, value = src.coo() - row = row[mask] - col = col[mask] - value = value[mask] - return SparseTensor(row=row, col=col, value=value, sparse_sizes=src.sparse_sizes()) - - -def calculate_interatomic_vectors(R, id_s, id_t, offsets_st): - """ - Calculate the vectors connecting the given atom pairs, - considering offsets from periodic boundary conditions (PBC). - - Arguments - --------- - R: Tensor, shape = (nAtoms, 3) - Atom positions. - id_s: Tensor, shape = (nEdges,) - Indices of the source atom of the edges. - id_t: Tensor, shape = (nEdges,) - Indices of the target atom of the edges. - offsets_st: Tensor, shape = (nEdges,) - PBC offsets of the edges. - Subtract this from the correct direction. - - Returns - ------- - (D_st, V_st): tuple - D_st: Tensor, shape = (nEdges,) - Distance from atom t to s. - V_st: Tensor, shape = (nEdges,) - Unit direction from atom t to s. - """ - Rs = R[id_s] - Rt = R[id_t] - # ReLU prevents negative numbers in sqrt - if offsets_st is None: - V_st = Rt - Rs # s -> t - else: - V_st = Rt - Rs + offsets_st # s -> t - D_st = torch.sqrt(torch.sum(V_st**2, dim=1)) - V_st = V_st / D_st[..., None] - return D_st, V_st - - -def inner_product_clamped(x, y): - """ - Calculate the inner product between the given normalized vectors, - giving a result between -1 and 1. - """ - return torch.sum(x * y, dim=-1).clamp(min=-1, max=1) - - -def get_angle(R_ac, R_ab): - """Calculate angles between atoms c -> a <- b. - - Arguments - --------- - R_ac: Tensor, shape = (N, 3) - Vector from atom a to c. - R_ab: Tensor, shape = (N, 3) - Vector from atom a to b. - - Returns - ------- - angle_cab: Tensor, shape = (N,) - Angle between atoms c <- a -> b. - """ - # cos(alpha) = (u * v) / (|u|*|v|) - x = torch.sum(R_ac * R_ab, dim=-1) # shape = (N,) - # sin(alpha) = |u x v| / (|u|*|v|) - y = torch.cross(R_ac, R_ab, dim=-1).norm(dim=-1) # shape = (N,) - y = y.clamp(min=1e-9) # Avoid NaN gradient for y = (0,0,0) - - angle = torch.atan2(y, x) - return angle - - -def vector_rejection(R_ab, P_n): - """ - Project the vector R_ab onto a plane with normal vector P_n. - - Arguments - --------- - R_ab: Tensor, shape = (N, 3) - Vector from atom a to b. - P_n: Tensor, shape = (N, 3) - Normal vector of a plane onto which to project R_ab. - - Returns - ------- - R_ab_proj: Tensor, shape = (N, 3) - Projected vector (orthogonal to P_n). - """ - a_x_b = torch.sum(R_ab * P_n, dim=-1) - b_x_b = torch.sum(P_n * P_n, dim=-1) - return R_ab - (a_x_b / b_x_b)[:, None] * P_n - - -def get_projected_angle(R_ab, P_n, eps=1e-4): - """ - Project the vector R_ab onto a plane with normal vector P_n, - then calculate the angle w.r.t. the (x [cross] P_n), - or (y [cross] P_n) if the former would be ill-defined/numerically unstable. - - Arguments - --------- - R_ab: Tensor, shape = (N, 3) - Vector from atom a to b. - P_n: Tensor, shape = (N, 3) - Normal vector of a plane onto which to project R_ab. - eps: float - Norm of projection below which to use the y-axis instead of x. - - Returns - ------- - angle_ab: Tensor, shape = (N) - Angle on plane w.r.t. x- or y-axis. - """ - R_ab_proj = torch.cross(R_ab, P_n, dim=-1) - - # Obtain axis defining the angle=0 - x = P_n.new_tensor([[1, 0, 0]]).expand_as(P_n) - zero_angle = torch.cross(x, P_n, dim=-1) - - use_y = torch.norm(zero_angle, dim=-1) < eps - P_n_y = P_n[use_y] - y = P_n_y.new_tensor([[0, 1, 0]]).expand_as(P_n_y) - y_cross = torch.cross(y, P_n_y, dim=-1) - zero_angle[use_y] = y_cross - - angle = get_angle(zero_angle, R_ab_proj) - - # Flip sign of angle if necessary to obtain clock-wise angles - cross = torch.cross(zero_angle, R_ab_proj, dim=-1) - flip_sign = torch.sum(cross * P_n, dim=-1) < 0 - angle[flip_sign] = -angle[flip_sign] - - return angle - - -def mask_neighbors(neighbors, edge_mask): - neighbors_old_indptr = torch.cat([neighbors.new_zeros(1), neighbors]) - neighbors_old_indptr = torch.cumsum(neighbors_old_indptr, dim=0) - neighbors = segment_csr(edge_mask.long(), neighbors_old_indptr) - return neighbors - - -def get_neighbor_order(num_atoms, index, atom_distance): - """ - Give a mask that filters out edges so that each atom has at most - `max_num_neighbors_threshold` neighbors. - """ - device = index.device - - # Get sorted index and inverse sorting - # Necessary for index_sort_map - index_sorted, index_order = torch.sort(index) - index_order_inverse = torch.argsort(index_order) - - # Get number of neighbors - ones = index_sorted.new_ones(1).expand_as(index_sorted) - num_neighbors = segment_coo(ones, index_sorted, dim_size=num_atoms) - max_num_neighbors = num_neighbors.max() - - # Create a tensor of size [num_atoms, max_num_neighbors] to sort the distances of the neighbors. - # Fill with infinity so we can easily remove unused distances later. - distance_sort = torch.full([num_atoms * max_num_neighbors], np.inf, device=device) - - # Create an index map to map distances from atom_distance to distance_sort - index_neighbor_offset = torch.cumsum(num_neighbors, dim=0) - num_neighbors - index_neighbor_offset_expand = torch.repeat_interleave( - index_neighbor_offset, num_neighbors - ) - index_sort_map = ( - index_sorted * max_num_neighbors - + torch.arange(len(index_sorted), device=device) - - index_neighbor_offset_expand - ) - distance_sort.index_copy_(0, index_sort_map, atom_distance) - distance_sort = distance_sort.view(num_atoms, max_num_neighbors) - - # Sort neighboring atoms based on distance - distance_sort, index_sort = torch.sort(distance_sort, dim=1) - - # Offset index_sort so that it indexes into index_sorted - index_sort = index_sort + index_neighbor_offset.view(-1, 1).expand( - -1, max_num_neighbors - ) - # Remove "unused pairs" with infinite distances - mask_finite = torch.isfinite(distance_sort) - index_sort = torch.masked_select(index_sort, mask_finite) - - # Create indices specifying the order in index_sort - order_peratom = torch.arange(max_num_neighbors, device=device)[None, :].expand_as( - mask_finite - ) - order_peratom = torch.masked_select(order_peratom, mask_finite) - - # Re-index to obtain order value of each neighbor in index_sorted - order = torch.zeros(len(index), device=device, dtype=torch.long) - order[index_sort] = order_peratom - - return order[index_order_inverse] - - -def get_inner_idx(idx, dim_size): - """ - Assign an inner index to each element (neighbor) with the same index. - For example, with idx=[0 0 0 1 1 1 1 2 2] this returns [0 1 2 0 1 2 3 0 1]. - These indices allow reshape neighbor indices into a dense matrix. - idx has to be sorted for this to work. - """ - ones = idx.new_ones(1).expand_as(idx) - num_neighbors = segment_coo(ones, idx, dim_size=dim_size) - inner_idx = ragged_range(num_neighbors) - return inner_idx - - -def get_edge_id(edge_idx, cell_offsets, num_atoms): - cell_basis = cell_offsets.max() - cell_offsets.min() + 1 - cell_id = ( - (cell_offsets * cell_offsets.new_tensor([[1, cell_basis, cell_basis**2]])) - .sum(-1) - .long() - ) - edge_id = edge_idx[0] + edge_idx[1] * num_atoms + cell_id * num_atoms**2 - return edge_id diff --git a/ocpmodels/models/painn/README.md b/ocpmodels/models/painn/README.md deleted file mode 100644 index 8e127dc1958f46749af04e961b5bffad84c50129..0000000000000000000000000000000000000000 --- a/ocpmodels/models/painn/README.md +++ /dev/null @@ -1,39 +0,0 @@ -# Polarizable Atom Interaction Neural Network (PaiNN) - -Kristof T. Schütt, Oliver T. Unke, Michael Gastegger - -[[`arXiv:2102.03150`](https://arxiv.org/abs/2102.03150)] - -This is our independent reimplementation of the original PaiNN architecture -with the difference that forces are predicted directly from vectorial features -via a gated equivariant block instead of gradients of the energy output. -This breaks energy conservation but is essential for good performance on OC20. - -All PaiNN models were trained without AMP, as using AMP led to unstable training. - -## IS2RE - -Trained only using IS2RE data, no auxiliary losses and/or S2EF data. - -| Model | Val ID Energy MAE | Test metrics | Download | -| ----- | ----------------- | ------------ | -------- | -| painn_h1024_bs4x8 | 0.5728 | [IS2RE](https://evalai.s3.amazonaws.com/media/submission_files/submission_200972/45d289fc-8de9-45cc-aed4-6cd1753cb56d.json) | [config](https://github.com/Open-Catalyst-Project/ocp/tree/main/configs/is2re/all/painn/painn_h1024_bs8x4.yml) \| [checkpoint](https://dl.fbaipublicfiles.com/opencatalystproject/models/2022_05/is2re/painn_h1024_bs4x8_is2re_all.pt) | - -## S2EF - -| Model | Val ID 30k Force MAE | Val ID 30k Energy MAE | Val ID 30k Force cos | Test metrics | Download | -| ----- | -------------------- | --------------------- | -------------------- | ------------ | -------- | -| painn_h512 | 0.02945 | 0.2459 | 0.5143 | [S2EF](https://evalai.s3.amazonaws.com/media/submission_files/submission_200711/2f487981-051d-445e-a7cd-6eb00ebe0735.json) \| [IS2RE](https://evalai.s3.amazonaws.com/media/submission_files/submission_200710/7fe29c4c-c203-434d-a6d4-9ea992d3bb5c.json) \| [IS2RS](https://evalai.s3.amazonaws.com/media/submission_files/submission_200700/8fd419e6-bab3-49be-a936-ae31979b4866.json) | [config](https://github.com/Open-Catalyst-Project/ocp/tree/main/configs/s2ef/all/painn/painn_h512.yml) \| [checkpoint](https://dl.fbaipublicfiles.com/opencatalystproject/models/2022_05/s2ef/painn_h512_s2ef_all.pt) | - -## Citing - -If you use PaiNN in your work, please consider citing the original paper: - -```bibtex -@inproceedings{schutt_painn_2021, - title = {Equivariant message passing for the prediction of tensorial properties and molecular spectra}, - author = {Sch{\"u}tt, Kristof and Unke, Oliver and Gastegger, Michael}, - booktitle = {Proceedings of the International Conference on Machine Learning (ICML)}, - year = {2021}, -} -``` diff --git a/ocpmodels/models/painn/painn.py b/ocpmodels/models/painn/painn.py deleted file mode 100644 index ee234758f4457d9429d5ee7bab2a70f0815584f8..0000000000000000000000000000000000000000 --- a/ocpmodels/models/painn/painn.py +++ /dev/null @@ -1,628 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. - ---- - -MIT License - -Copyright (c) 2021 www.compscience.org - -Permission is hereby granted, free of charge, to any person obtaining a copy -of this software and associated documentation files (the "Software"), to deal -in the Software without restriction, including without limitation the rights -to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -copies of the Software, and to permit persons to whom the Software is -furnished to do so, subject to the following conditions: - -The above copyright notice and this permission notice shall be included in all -copies or substantial portions of the Software. - -THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -SOFTWARE. -""" - -import logging -import math -import os -from typing import Optional, Tuple - -import torch -from torch import nn -from torch_geometric.nn import MessagePassing, radius_graph -from torch_scatter import scatter, segment_coo - -from ocpmodels.common.registry import registry -from ocpmodels.common.utils import ( - compute_neighbors, - conditional_grad, - get_pbc_distances, - radius_graph_pbc, -) -from ocpmodels.models.base import BaseModel -from ocpmodels.models.gemnet.layers.base_layers import ScaledSiLU -from ocpmodels.models.gemnet.layers.embedding_block import AtomEmbedding -from ocpmodels.models.gemnet.layers.radial_basis import RadialBasis -from ocpmodels.modules.scaling import ScaleFactor -from ocpmodels.modules.scaling.compat import load_scales_compat - -from .utils import get_edge_id, repeat_blocks - - -@registry.register_model("painn") -class PaiNN(BaseModel): - r"""PaiNN model based on the description in Schütt et al. (2021): - Equivariant message passing for the prediction of tensorial properties - and molecular spectra, https://arxiv.org/abs/2102.03150. - """ - - def __init__( - self, - num_atoms, - bond_feat_dim, - num_targets, - hidden_channels=512, - num_layers=6, - num_rbf=128, - cutoff=12.0, - max_neighbors=50, - rbf: dict = {"name": "gaussian"}, - envelope: dict = {"name": "polynomial", "exponent": 5}, - regress_forces=True, - direct_forces=True, - use_pbc=True, - otf_graph=True, - num_elements=83, - scale_file: Optional[str] = None, - ): - super(PaiNN, self).__init__() - - self.hidden_channels = hidden_channels - self.num_layers = num_layers - self.num_rbf = num_rbf - self.cutoff = cutoff - self.max_neighbors = max_neighbors - self.regress_forces = regress_forces - self.direct_forces = direct_forces - self.otf_graph = otf_graph - self.use_pbc = use_pbc - - # Borrowed from GemNet. - self.symmetric_edge_symmetrization = False - - #### Learnable parameters ############################################# - - self.atom_emb = AtomEmbedding(hidden_channels, num_elements) - - self.radial_basis = RadialBasis( - num_radial=num_rbf, - cutoff=self.cutoff, - rbf=rbf, - envelope=envelope, - ) - - self.message_layers = nn.ModuleList() - self.update_layers = nn.ModuleList() - - for i in range(num_layers): - self.message_layers.append( - PaiNNMessage(hidden_channels, num_rbf).jittable() - ) - self.update_layers.append(PaiNNUpdate(hidden_channels)) - setattr(self, "upd_out_scalar_scale_%d" % i, ScaleFactor()) - - self.out_energy = nn.Sequential( - nn.Linear(hidden_channels, hidden_channels // 2), - ScaledSiLU(), - nn.Linear(hidden_channels // 2, 1), - ) - - if self.regress_forces is True and self.direct_forces is True: - self.out_forces = PaiNNOutput(hidden_channels) - - self.inv_sqrt_2 = 1 / math.sqrt(2.0) - - self.reset_parameters() - - load_scales_compat(self, scale_file) - - def reset_parameters(self): - nn.init.xavier_uniform_(self.out_energy[0].weight) - self.out_energy[0].bias.data.fill_(0) - nn.init.xavier_uniform_(self.out_energy[2].weight) - self.out_energy[2].bias.data.fill_(0) - - # Borrowed from GemNet. - def select_symmetric_edges(self, tensor, mask, reorder_idx, inverse_neg): - # Mask out counter-edges - tensor_directed = tensor[mask] - # Concatenate counter-edges after normal edges - sign = 1 - 2 * inverse_neg - tensor_cat = torch.cat([tensor_directed, sign * tensor_directed]) - # Reorder everything so the edges of every image are consecutive - tensor_ordered = tensor_cat[reorder_idx] - return tensor_ordered - - # Borrowed from GemNet. - def symmetrize_edges( - self, - edge_index, - cell_offsets, - neighbors, - batch_idx, - reorder_tensors, - reorder_tensors_invneg, - ): - """ - Symmetrize edges to ensure existence of counter-directional edges. - - Some edges are only present in one direction in the data, - since every atom has a maximum number of neighbors. - If `symmetric_edge_symmetrization` is False, - we only use i->j edges here. So we lose some j->i edges - and add others by making it symmetric. - If `symmetric_edge_symmetrization` is True, - we always use both directions. - """ - num_atoms = batch_idx.shape[0] - - if self.symmetric_edge_symmetrization: - edge_index_bothdir = torch.cat( - [edge_index, edge_index.flip(0)], - dim=1, - ) - cell_offsets_bothdir = torch.cat( - [cell_offsets, -cell_offsets], - dim=0, - ) - - # Filter for unique edges - edge_ids = get_edge_id(edge_index_bothdir, cell_offsets_bothdir, num_atoms) - unique_ids, unique_inv = torch.unique(edge_ids, return_inverse=True) - perm = torch.arange( - unique_inv.size(0), - dtype=unique_inv.dtype, - device=unique_inv.device, - ) - unique_idx = scatter( - perm, - unique_inv, - dim=0, - dim_size=unique_ids.shape[0], - reduce="min", - ) - edge_index_new = edge_index_bothdir[:, unique_idx] - - # Order by target index - edge_index_order = torch.argsort(edge_index_new[1]) - edge_index_new = edge_index_new[:, edge_index_order] - unique_idx = unique_idx[edge_index_order] - - # Subindex remaining tensors - cell_offsets_new = cell_offsets_bothdir[unique_idx] - reorder_tensors = [ - self.symmetrize_tensor(tensor, unique_idx, False) - for tensor in reorder_tensors - ] - reorder_tensors_invneg = [ - self.symmetrize_tensor(tensor, unique_idx, True) - for tensor in reorder_tensors_invneg - ] - - # Count edges per image - # segment_coo assumes sorted edge_index_new[1] and batch_idx - ones = edge_index_new.new_ones(1).expand_as(edge_index_new[1]) - neighbors_per_atom = segment_coo( - ones, edge_index_new[1], dim_size=num_atoms - ) - neighbors_per_image = segment_coo( - neighbors_per_atom, batch_idx, dim_size=neighbors.shape[0] - ) - else: - # Generate mask - mask_sep_atoms = edge_index[0] < edge_index[1] - # Distinguish edges between the same (periodic) atom by ordering the cells - cell_earlier = ( - (cell_offsets[:, 0] < 0) - | ((cell_offsets[:, 0] == 0) & (cell_offsets[:, 1] < 0)) - | ( - (cell_offsets[:, 0] == 0) - & (cell_offsets[:, 1] == 0) - & (cell_offsets[:, 2] < 0) - ) - ) - mask_same_atoms = edge_index[0] == edge_index[1] - mask_same_atoms &= cell_earlier - mask = mask_sep_atoms | mask_same_atoms - - # Mask out counter-edges - edge_index_new = edge_index[mask[None, :].expand(2, -1)].view(2, -1) - - # Concatenate counter-edges after normal edges - edge_index_cat = torch.cat( - [edge_index_new, edge_index_new.flip(0)], - dim=1, - ) - - # Count remaining edges per image - batch_edge = torch.repeat_interleave( - torch.arange(neighbors.size(0), device=edge_index.device), - neighbors, - ) - batch_edge = batch_edge[mask] - # segment_coo assumes sorted batch_edge - # Factor 2 since this is only one half of the edges - ones = batch_edge.new_ones(1).expand_as(batch_edge) - neighbors_per_image = 2 * segment_coo( - ones, batch_edge, dim_size=neighbors.size(0) - ) - - # Create indexing array - edge_reorder_idx = repeat_blocks( - torch.div(neighbors_per_image, 2, rounding_mode="floor"), - repeats=2, - continuous_indexing=True, - repeat_inc=edge_index_new.size(1), - ) - - # Reorder everything so the edges of every image are consecutive - edge_index_new = edge_index_cat[:, edge_reorder_idx] - cell_offsets_new = self.select_symmetric_edges( - cell_offsets, mask, edge_reorder_idx, True - ) - reorder_tensors = [ - self.select_symmetric_edges(tensor, mask, edge_reorder_idx, False) - for tensor in reorder_tensors - ] - reorder_tensors_invneg = [ - self.select_symmetric_edges(tensor, mask, edge_reorder_idx, True) - for tensor in reorder_tensors_invneg - ] - - # Indices for swapping c->a and a->c (for symmetric MP) - # To obtain these efficiently and without any index assumptions, - # we get order the counter-edge IDs and then - # map this order back to the edge IDs. - # Double argsort gives the desired mapping - # from the ordered tensor to the original tensor. - edge_ids = get_edge_id(edge_index_new, cell_offsets_new, num_atoms) - order_edge_ids = torch.argsort(edge_ids) - inv_order_edge_ids = torch.argsort(order_edge_ids) - edge_ids_counter = get_edge_id( - edge_index_new.flip(0), -cell_offsets_new, num_atoms - ) - order_edge_ids_counter = torch.argsort(edge_ids_counter) - id_swap = order_edge_ids_counter[inv_order_edge_ids] - - return ( - edge_index_new, - cell_offsets_new, - neighbors_per_image, - reorder_tensors, - reorder_tensors_invneg, - id_swap, - ) - - def generate_graph_values(self, data): - ( - edge_index, - edge_dist, - distance_vec, - cell_offsets, - _, # cell offset distances - neighbors, - ) = self.generate_graph(data) - - # Unit vectors pointing from edge_index[1] to edge_index[0], - # i.e., edge_index[0] - edge_index[1] divided by the norm. - # make sure that the distances are not close to zero before dividing - mask_zero = torch.isclose(edge_dist, torch.tensor(0.0), atol=1e-6) - edge_dist[mask_zero] = 1.0e-6 - edge_vector = distance_vec / edge_dist[:, None] - - empty_image = neighbors == 0 - if torch.any(empty_image): - raise ValueError( - f"An image has no neighbors: id={data.id[empty_image]}, " - f"sid={data.sid[empty_image]}, fid={data.fid[empty_image]}" - ) - - # Symmetrize edges for swapping in symmetric message passing - ( - edge_index, - cell_offsets, - neighbors, - [edge_dist], - [edge_vector], - id_swap, - ) = self.symmetrize_edges( - edge_index, - cell_offsets, - neighbors, - data.batch, - [edge_dist], - [edge_vector], - ) - - return ( - edge_index, - neighbors, - edge_dist, - edge_vector, - id_swap, - ) - - @conditional_grad(torch.enable_grad()) - def forward(self, data): - pos = data.pos - batch = data.batch - z = data.atomic_numbers.long() - - if self.regress_forces and not self.direct_forces: - pos = pos.requires_grad_(True) - - ( - edge_index, - neighbors, - edge_dist, - edge_vector, - id_swap, - ) = self.generate_graph_values(data) - - assert z.dim() == 1 and z.dtype == torch.long - - edge_rbf = self.radial_basis(edge_dist) # rbf * envelope - - x = self.atom_emb(z) - vec = torch.zeros(x.size(0), 3, x.size(1), device=x.device) - - #### Interaction blocks ############################################### - - for i in range(self.num_layers): - dx, dvec = self.message_layers[i](x, vec, edge_index, edge_rbf, edge_vector) - - x = x + dx - vec = vec + dvec - x = x * self.inv_sqrt_2 - - dx, dvec = self.update_layers[i](x, vec) - - x = x + dx - vec = vec + dvec - x = getattr(self, "upd_out_scalar_scale_%d" % i)(x) - - #### Output block ##################################################### - - per_atom_energy = self.out_energy(x).squeeze(1) - energy = scatter(per_atom_energy, batch, dim=0) - - if self.regress_forces: - if self.direct_forces: - forces = self.out_forces(x, vec) - return energy, forces - else: - forces = ( - -1 - * torch.autograd.grad( - x, - pos, - grad_outputs=torch.ones_like(x), - create_graph=True, - )[0] - ) - return energy, forces - else: - return energy - - @property - def num_params(self): - return sum(p.numel() for p in self.parameters()) - - def __repr__(self): - return ( - f"{self.__class__.__name__}(" - f"hidden_channels={self.hidden_channels}, " - f"num_layers={self.num_layers}, " - f"num_rbf={self.num_rbf}, " - f"max_neighbors={self.max_neighbors}, " - f"cutoff={self.cutoff})" - ) - - -class PaiNNMessage(MessagePassing): - def __init__( - self, - hidden_channels, - num_rbf, - ): - super(PaiNNMessage, self).__init__(aggr="add", node_dim=0) - - self.hidden_channels = hidden_channels - - self.x_proj = nn.Sequential( - nn.Linear(hidden_channels, hidden_channels), - ScaledSiLU(), - nn.Linear(hidden_channels, hidden_channels * 3), - ) - self.rbf_proj = nn.Linear(num_rbf, hidden_channels * 3) - - self.inv_sqrt_3 = 1 / math.sqrt(3.0) - self.inv_sqrt_h = 1 / math.sqrt(hidden_channels) - self.x_layernorm = nn.LayerNorm(hidden_channels) - - self.reset_parameters() - - def reset_parameters(self): - nn.init.xavier_uniform_(self.x_proj[0].weight) - self.x_proj[0].bias.data.fill_(0) - nn.init.xavier_uniform_(self.x_proj[2].weight) - self.x_proj[2].bias.data.fill_(0) - nn.init.xavier_uniform_(self.rbf_proj.weight) - self.rbf_proj.bias.data.fill_(0) - self.x_layernorm.reset_parameters() - - def forward(self, x, vec, edge_index, edge_rbf, edge_vector): - xh = self.x_proj(self.x_layernorm(x)) - - # TODO(@abhshkdz): Nans out with AMP here during backprop. Debug / fix. - rbfh = self.rbf_proj(edge_rbf) - - # propagate_type: (xh: Tensor, vec: Tensor, rbfh_ij: Tensor, r_ij: Tensor) - dx, dvec = self.propagate( - edge_index, - xh=xh, - vec=vec, - rbfh_ij=rbfh, - r_ij=edge_vector, - size=None, - ) - - return dx, dvec - - def message(self, xh_j, vec_j, rbfh_ij, r_ij): - x, xh2, xh3 = torch.split(xh_j * rbfh_ij, self.hidden_channels, dim=-1) - xh2 = xh2 * self.inv_sqrt_3 - - vec = vec_j * xh2.unsqueeze(1) + xh3.unsqueeze(1) * r_ij.unsqueeze(2) - vec = vec * self.inv_sqrt_h - - return x, vec - - def aggregate( - self, - features: Tuple[torch.Tensor, torch.Tensor], - index: torch.Tensor, - ptr: Optional[torch.Tensor], - dim_size: Optional[int], - ) -> Tuple[torch.Tensor, torch.Tensor]: - x, vec = features - x = scatter(x, index, dim=self.node_dim, dim_size=dim_size) - vec = scatter(vec, index, dim=self.node_dim, dim_size=dim_size) - return x, vec - - def update( - self, inputs: Tuple[torch.Tensor, torch.Tensor] - ) -> Tuple[torch.Tensor, torch.Tensor]: - return inputs - - -class PaiNNUpdate(nn.Module): - def __init__(self, hidden_channels): - super().__init__() - self.hidden_channels = hidden_channels - - self.vec_proj = nn.Linear(hidden_channels, hidden_channels * 2, bias=False) - self.xvec_proj = nn.Sequential( - nn.Linear(hidden_channels * 2, hidden_channels), - ScaledSiLU(), - nn.Linear(hidden_channels, hidden_channels * 3), - ) - - self.inv_sqrt_2 = 1 / math.sqrt(2.0) - self.inv_sqrt_h = 1 / math.sqrt(hidden_channels) - - self.reset_parameters() - - def reset_parameters(self): - nn.init.xavier_uniform_(self.vec_proj.weight) - nn.init.xavier_uniform_(self.xvec_proj[0].weight) - self.xvec_proj[0].bias.data.fill_(0) - nn.init.xavier_uniform_(self.xvec_proj[2].weight) - self.xvec_proj[2].bias.data.fill_(0) - - def forward(self, x, vec): - vec1, vec2 = torch.split(self.vec_proj(vec), self.hidden_channels, dim=-1) - vec_dot = (vec1 * vec2).sum(dim=1) * self.inv_sqrt_h - - # NOTE: Can't use torch.norm because the gradient is NaN for input = 0. - # Add an epsilon offset to make sure sqrt is always positive. - x_vec_h = self.xvec_proj( - torch.cat([x, torch.sqrt(torch.sum(vec2**2, dim=-2) + 1e-8)], dim=-1) - ) - xvec1, xvec2, xvec3 = torch.split(x_vec_h, self.hidden_channels, dim=-1) - - dx = xvec1 + xvec2 * vec_dot - dx = dx * self.inv_sqrt_2 - - dvec = xvec3.unsqueeze(1) * vec1 - - return dx, dvec - - -class PaiNNOutput(nn.Module): - def __init__(self, hidden_channels): - super().__init__() - self.hidden_channels = hidden_channels - - self.output_network = nn.ModuleList( - [ - GatedEquivariantBlock( - hidden_channels, - hidden_channels // 2, - ), - GatedEquivariantBlock(hidden_channels // 2, 1), - ] - ) - - self.reset_parameters() - - def reset_parameters(self): - for layer in self.output_network: - layer.reset_parameters() - - def forward(self, x, vec): - for layer in self.output_network: - x, vec = layer(x, vec) - return vec.squeeze() - - -# Borrowed from TorchMD-Net -class GatedEquivariantBlock(nn.Module): - """Gated Equivariant Block as defined in Schütt et al. (2021): - Equivariant message passing for the prediction of tensorial properties and molecular spectra - """ - - def __init__( - self, - hidden_channels, - out_channels, - ): - super(GatedEquivariantBlock, self).__init__() - self.out_channels = out_channels - - self.vec1_proj = nn.Linear(hidden_channels, hidden_channels, bias=False) - self.vec2_proj = nn.Linear(hidden_channels, out_channels, bias=False) - - self.update_net = nn.Sequential( - nn.Linear(hidden_channels * 2, hidden_channels), - ScaledSiLU(), - nn.Linear(hidden_channels, out_channels * 2), - ) - - self.act = ScaledSiLU() - - def reset_parameters(self): - nn.init.xavier_uniform_(self.vec1_proj.weight) - nn.init.xavier_uniform_(self.vec2_proj.weight) - nn.init.xavier_uniform_(self.update_net[0].weight) - self.update_net[0].bias.data.fill_(0) - nn.init.xavier_uniform_(self.update_net[2].weight) - self.update_net[2].bias.data.fill_(0) - - def forward(self, x, v): - vec1 = torch.norm(self.vec1_proj(v), dim=-2) - vec2 = self.vec2_proj(v) - - x = torch.cat([x, vec1], dim=-1) - x, v = torch.split(self.update_net(x), self.out_channels, dim=-1) - v = v.unsqueeze(1) * vec2 - - x = self.act(x) - return x, v diff --git a/ocpmodels/models/painn/utils.py b/ocpmodels/models/painn/utils.py deleted file mode 100644 index 78eb7b22dd38a9b2f455e14f1964e9fe0adc8e75..0000000000000000000000000000000000000000 --- a/ocpmodels/models/painn/utils.py +++ /dev/null @@ -1,161 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import torch -from torch_scatter import segment_csr - - -def repeat_blocks( - sizes, - repeats, - continuous_indexing=True, - start_idx=0, - block_inc=0, - repeat_inc=0, -): - """Repeat blocks of indices. - Adapted from https://stackoverflow.com/questions/51154989/numpy-vectorized-function-to-repeat-blocks-of-consecutive-elements - - continuous_indexing: Whether to keep increasing the index after each block - start_idx: Starting index - block_inc: Number to increment by after each block, - either global or per block. Shape: len(sizes) - 1 - repeat_inc: Number to increment by after each repetition, - either global or per block - - Examples - -------- - sizes = [1,3,2] ; repeats = [3,2,3] ; continuous_indexing = False - Return: [0 0 0 0 1 2 0 1 2 0 1 0 1 0 1] - sizes = [1,3,2] ; repeats = [3,2,3] ; continuous_indexing = True - Return: [0 0 0 1 2 3 1 2 3 4 5 4 5 4 5] - sizes = [1,3,2] ; repeats = [3,2,3] ; continuous_indexing = True ; - repeat_inc = 4 - Return: [0 4 8 1 2 3 5 6 7 4 5 8 9 12 13] - sizes = [1,3,2] ; repeats = [3,2,3] ; continuous_indexing = True ; - start_idx = 5 - Return: [5 5 5 6 7 8 6 7 8 9 10 9 10 9 10] - sizes = [1,3,2] ; repeats = [3,2,3] ; continuous_indexing = True ; - block_inc = 1 - Return: [0 0 0 2 3 4 2 3 4 6 7 6 7 6 7] - sizes = [0,3,2] ; repeats = [3,2,3] ; continuous_indexing = True - Return: [0 1 2 0 1 2 3 4 3 4 3 4] - sizes = [2,3,2] ; repeats = [2,0,2] ; continuous_indexing = True - Return: [0 1 0 1 5 6 5 6] - """ - assert sizes.dim() == 1 - assert all(sizes >= 0) - - # Remove 0 sizes - sizes_nonzero = sizes > 0 - if not torch.all(sizes_nonzero): - assert block_inc == 0 # Implementing this is not worth the effort - sizes = torch.masked_select(sizes, sizes_nonzero) - if isinstance(repeats, torch.Tensor): - repeats = torch.masked_select(repeats, sizes_nonzero) - if isinstance(repeat_inc, torch.Tensor): - repeat_inc = torch.masked_select(repeat_inc, sizes_nonzero) - - if isinstance(repeats, torch.Tensor): - assert all(repeats >= 0) - insert_dummy = repeats[0] == 0 - if insert_dummy: - one = sizes.new_ones(1) - zero = sizes.new_zeros(1) - sizes = torch.cat((one, sizes)) - repeats = torch.cat((one, repeats)) - if isinstance(block_inc, torch.Tensor): - block_inc = torch.cat((zero, block_inc)) - if isinstance(repeat_inc, torch.Tensor): - repeat_inc = torch.cat((zero, repeat_inc)) - else: - assert repeats >= 0 - insert_dummy = False - - # Get repeats for each group using group lengths/sizes - r1 = torch.repeat_interleave(torch.arange(len(sizes), device=sizes.device), repeats) - - # Get total size of output array, as needed to initialize output indexing array - N = (sizes * repeats).sum() - - # Initialize indexing array with ones as we need to setup incremental indexing - # within each group when cumulatively summed at the final stage. - # Two steps here: - # 1. Within each group, we have multiple sequences, so setup the offsetting - # at each sequence lengths by the seq. lengths preceding those. - id_ar = torch.ones(N, dtype=torch.long, device=sizes.device) - id_ar[0] = 0 - insert_index = sizes[r1[:-1]].cumsum(0) - insert_val = (1 - sizes)[r1[:-1]] - - if isinstance(repeats, torch.Tensor) and torch.any(repeats == 0): - diffs = r1[1:] - r1[:-1] - indptr = torch.cat((sizes.new_zeros(1), diffs.cumsum(0))) - if continuous_indexing: - # If a group was skipped (repeats=0) we need to add its size - insert_val += segment_csr(sizes[: r1[-1]], indptr, reduce="sum") - - # Add block increments - if isinstance(block_inc, torch.Tensor): - insert_val += segment_csr(block_inc[: r1[-1]], indptr, reduce="sum") - else: - insert_val += block_inc * (indptr[1:] - indptr[:-1]) - if insert_dummy: - insert_val[0] -= block_inc - else: - idx = r1[1:] != r1[:-1] - if continuous_indexing: - # 2. For each group, make sure the indexing starts from the next group's - # first element. So, simply assign 1s there. - insert_val[idx] = 1 - - # Add block increments - insert_val[idx] += block_inc - - # Add repeat_inc within each group - if isinstance(repeat_inc, torch.Tensor): - insert_val += repeat_inc[r1[:-1]] - if isinstance(repeats, torch.Tensor): - repeat_inc_inner = repeat_inc[repeats > 0][:-1] - else: - repeat_inc_inner = repeat_inc[:-1] - else: - insert_val += repeat_inc - repeat_inc_inner = repeat_inc - - # Subtract the increments between groups - if isinstance(repeats, torch.Tensor): - repeats_inner = repeats[repeats > 0][:-1] - else: - repeats_inner = repeats - insert_val[r1[1:] != r1[:-1]] -= repeat_inc_inner * repeats_inner - - # Assign index-offsetting values - id_ar[insert_index] = insert_val - - if insert_dummy: - id_ar = id_ar[1:] - if continuous_indexing: - id_ar[0] -= 1 - - # Set start index now, in case of insertion due to leading repeats=0 - id_ar[0] += start_idx - - # Finally index into input array for the group repeated o/p - res = id_ar.cumsum(0) - return res - - -def get_edge_id(edge_idx, cell_offsets, num_atoms): - cell_basis = cell_offsets.max() - cell_offsets.min() + 1 - cell_id = ( - (cell_offsets * cell_offsets.new_tensor([[1, cell_basis, cell_basis**2]])) - .sum(-1) - .long() - ) - edge_id = edge_idx[0] + edge_idx[1] * num_atoms + cell_id * num_atoms**2 - return edge_id diff --git a/ocpmodels/models/schnet.py b/ocpmodels/models/schnet.py deleted file mode 100644 index 81e012caee8ab9c67ad15aeebf17f1e6d1aae63e..0000000000000000000000000000000000000000 --- a/ocpmodels/models/schnet.py +++ /dev/null @@ -1,142 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import torch -from torch_geometric.nn import SchNet -from torch_scatter import scatter - -from ocpmodels.common.registry import registry -from ocpmodels.common.utils import ( - conditional_grad, - get_pbc_distances, - radius_graph_pbc, -) -from ocpmodels.models.base import BaseModel - - -@registry.register_model("schnet") -class SchNetWrap(SchNet, BaseModel): - r"""Wrapper around the continuous-filter convolutional neural network SchNet from the - `"SchNet: A Continuous-filter Convolutional Neural Network for Modeling - Quantum Interactions" `_. Each layer uses interaction - block of the form: - - .. math:: - \mathbf{x}^{\prime}_i = \sum_{j \in \mathcal{N}(i)} \mathbf{x}_j \odot - h_{\mathbf{\Theta}} ( \exp(-\gamma(\mathbf{e}_{j,i} - \mathbf{\mu}))), - - Args: - num_atoms (int): Unused argument - bond_feat_dim (int): Unused argument - num_targets (int): Number of targets to predict. - use_pbc (bool, optional): If set to :obj:`True`, account for periodic boundary conditions. - (default: :obj:`True`) - regress_forces (bool, optional): If set to :obj:`True`, predict forces by differentiating - energy with respect to positions. - (default: :obj:`True`) - otf_graph (bool, optional): If set to :obj:`True`, compute graph edges on the fly. - (default: :obj:`False`) - hidden_channels (int, optional): Number of hidden channels. - (default: :obj:`128`) - num_filters (int, optional): Number of filters to use. - (default: :obj:`128`) - num_interactions (int, optional): Number of interaction blocks - (default: :obj:`6`) - num_gaussians (int, optional): The number of gaussians :math:`\mu`. - (default: :obj:`50`) - cutoff (float, optional): Cutoff distance for interatomic interactions. - (default: :obj:`10.0`) - readout (string, optional): Whether to apply :obj:`"add"` or - :obj:`"mean"` global aggregation. (default: :obj:`"add"`) - """ - - def __init__( - self, - num_atoms, # not used - bond_feat_dim, # not used - num_targets, - use_pbc=True, - regress_forces=True, - otf_graph=False, - hidden_channels=128, - num_filters=128, - num_interactions=6, - num_gaussians=50, - cutoff=10.0, - readout="add", - ): - self.num_targets = num_targets - self.regress_forces = regress_forces - self.use_pbc = use_pbc - self.cutoff = cutoff - self.otf_graph = otf_graph - self.max_neighbors = 50 - self.reduce = readout - super(SchNetWrap, self).__init__( - hidden_channels=hidden_channels, - num_filters=num_filters, - num_interactions=num_interactions, - num_gaussians=num_gaussians, - cutoff=cutoff, - readout=readout, - ) - - @conditional_grad(torch.enable_grad()) - def _forward(self, data): - z = data.atomic_numbers.long() - pos = data.pos - batch = data.batch - - ( - edge_index, - edge_weight, - distance_vec, - cell_offsets, - _, # cell offset distances - neighbors, - ) = self.generate_graph(data) - - if self.use_pbc: - assert z.dim() == 1 and z.dtype == torch.long - - edge_attr = self.distance_expansion(edge_weight) - - h = self.embedding(z) - for interaction in self.interactions: - h = h + interaction(h, edge_index, edge_weight, edge_attr) - - h = self.lin1(h) - h = self.act(h) - h = self.lin2(h) - - batch = torch.zeros_like(z) if batch is None else batch - energy = scatter(h, batch, dim=0, reduce=self.reduce) - else: - energy = super(SchNetWrap, self).forward(z, pos, batch) - return energy - - def forward(self, data): - if self.regress_forces: - data.pos.requires_grad_(True) - energy = self._forward(data) - - if self.regress_forces: - forces = -1 * ( - torch.autograd.grad( - energy, - data.pos, - grad_outputs=torch.ones_like(energy), - create_graph=True, - )[0] - ) - return energy, forces - else: - return energy - - @property - def num_params(self): - return sum(p.numel() for p in self.parameters()) diff --git a/ocpmodels/models/scn/Jd.pt b/ocpmodels/models/scn/Jd.pt deleted file mode 100644 index 9bf1f88c6fc9186c7df212787ad6ee53c96907b8..0000000000000000000000000000000000000000 --- a/ocpmodels/models/scn/Jd.pt +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:b4059c45be246dcb6c49c545670b65c56550eb0c2e7a9c92b4b50a92d370dbe2 -size 21697 diff --git a/ocpmodels/models/scn/README.md b/ocpmodels/models/scn/README.md deleted file mode 100644 index 0fbfe9e804c2ead8ab7e199c3737e57815dcb653..0000000000000000000000000000000000000000 --- a/ocpmodels/models/scn/README.md +++ /dev/null @@ -1,22 +0,0 @@ -# Spherical Channels for Modeling Atomic Interactions - -C. Lawrence Zitnick, Abhishek Das, Adeesh Kolluru, Janice Lan, Muhammed Shuaibi, Anuroop Sriram, Zachary Ulissi, Brandon Wood - -[[`arXiv:2206.14331`](https://arxiv.org/abs/2206.14331)] - -To run the Spherical Channel Network (SCN), install [e3nn](https://github.com/e3nn/e3nn/) with `pip install e3nn==0.2.6`. - -SCN was developed with e3nn v0.2.6, and might run slower with later versions [[1](https://github.com/Open-Catalyst-Project/ocp/issues/397), [2](https://github.com/Open-Catalyst-Project/ocp/pull/402)]. - -## Citing - -If you use SCN in your work, please consider citing: - -```bibtex -@inproceedings{zitnick_scn_2022, - title = {{Spherical Channels for Modeling Atomic Interactions}}, - author = {Zitnick, C. Lawrence and Das, Abhishek and Kolluru, Adeesh and Lan, Janice and Shuaibi, Muhammed and Sriram, Anuroop and Ulissi, Zachary and Wood, Brandon}, - booktitle = {Advances in Neural Information Processing Systems (NeurIPS)}, - year = {2022}, -} -``` diff --git a/ocpmodels/models/scn/sampling.py b/ocpmodels/models/scn/sampling.py deleted file mode 100644 index ba7600e5a1a1b64ff51d203dedd9bdfaf025b9a9..0000000000000000000000000000000000000000 --- a/ocpmodels/models/scn/sampling.py +++ /dev/null @@ -1,47 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import math - -import torch - -### Methods for sample points on a sphere - - -def CalcSpherePoints(num_points, device): - goldenRatio = (1 + 5**0.5) / 2 - i = torch.arange(num_points, device=device).view(-1, 1) - theta = 2 * math.pi * i / goldenRatio - phi = torch.arccos(1 - 2 * (i + 0.5) / num_points) - points = torch.cat( - [ - torch.cos(theta) * torch.sin(phi), - torch.sin(theta) * torch.sin(phi), - torch.cos(phi), - ], - dim=1, - ) - - # weight the points by their density - pt_cross = points.view(1, -1, 3) - points.view(-1, 1, 3) - pt_cross = torch.sum(pt_cross**2, dim=2) - pt_cross = torch.exp(-pt_cross / (0.5 * 0.3)) - scalar = 1.0 / torch.sum(pt_cross, dim=1) - scalar = num_points * scalar / torch.sum(scalar) - return points * (scalar.view(-1, 1)) - - -def CalcSpherePointsRandom(num_points, device): - pts = 2.0 * (torch.rand(num_points, 3, device=device) - 0.5) - radius = torch.sum(pts**2, dim=1) - while torch.max(radius) > 1.0: - replace_pts = 2.0 * (torch.rand(num_points, 3, device=device) - 0.5) - replace_mask = radius.gt(0.99) - pts.masked_scatter_(replace_mask.view(-1, 1).repeat(1, 3), replace_pts) - radius = torch.sum(pts**2, dim=1) - - return pts / radius.view(-1, 1) diff --git a/ocpmodels/models/scn/scn.py b/ocpmodels/models/scn/scn.py deleted file mode 100644 index 21c43c7c03a2a590b6f7f5fe65f5dbf11e6dc5cf..0000000000000000000000000000000000000000 --- a/ocpmodels/models/scn/scn.py +++ /dev/null @@ -1,783 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import logging -import sys -import time - -import numpy as np -import torch -import torch.nn as nn -from torch_geometric.nn import radius_graph - -from ocpmodels.common.registry import registry -from ocpmodels.common.utils import ( - conditional_grad, - get_pbc_distances, - radius_graph_pbc, -) -from ocpmodels.models.base import BaseModel -from ocpmodels.models.scn.sampling import CalcSpherePoints -from ocpmodels.models.scn.smearing import ( - GaussianSmearing, - LinearSigmoidSmearing, - SigmoidSmearing, - SiLUSmearing, -) -from ocpmodels.models.scn.spherical_harmonics import SphericalHarmonicsHelper - -try: - import e3nn - from e3nn import o3 -except ImportError: - pass - - -@registry.register_model("scn") -class SphericalChannelNetwork(BaseModel): - """Spherical Channel Network - Paper: Spherical Channels for Modeling Atomic Interactions - - Args: - use_pbc (bool): Use periodic boundary conditions - regress_forces (bool): Compute forces - otf_graph (bool): Compute graph On The Fly (OTF) - max_num_neighbors (int): Maximum number of neighbors per atom - cutoff (float): Maximum distance between nieghboring atoms in Angstroms - max_num_elements (int): Maximum atomic number - - num_interactions (int): Number of layers in the GNN - lmax (int): Maximum degree of the spherical harmonics (1 to 10) - mmax (int): Maximum order of the spherical harmonics (0 or 1) - num_resolutions (int): Number of resolutions used to compute messages, further away atoms has lower resolution (1 or 2) - sphere_channels (int): Number of spherical channels - sphere_channels_reduce (int): Number of spherical channels used during message passing (downsample or upsample) - hidden_channels (int): Number of hidden units in message passing - num_taps (int): Number of taps or rotations used during message passing (1 or otherwise set automatically based on mmax) - - use_grid (bool): Use non-linear pointwise convolution during aggregation - num_bands (int): Number of bands used during message aggregation for the 1x1 pointwise convolution (1 or 2) - - num_sphere_samples (int): Number of samples used to approximate the integration of the sphere in the output blocks - num_basis_functions (int): Number of basis functions used for distance and atomic number blocks - distance_function ("gaussian", "sigmoid", "linearsigmoid", "silu"): Basis function used for distances - basis_width_scalar (float): Width of distance basis function - distance_resolution (float): Distance between distance basis functions in Angstroms - - show_timing_info (bool): Show timing and memory info - """ - - def __init__( - self, - num_atoms, # not used - bond_feat_dim, # not used - num_targets, # not used - use_pbc=True, - regress_forces=True, - otf_graph=False, - max_num_neighbors=20, - cutoff=8.0, - max_num_elements=90, - num_interactions=8, - lmax=6, - mmax=1, - num_resolutions=2, - sphere_channels=128, - sphere_channels_reduce=128, - hidden_channels=256, - num_taps=-1, - use_grid=True, - num_bands=1, - num_sphere_samples=128, - num_basis_functions=128, - distance_function="gaussian", - basis_width_scalar=1.0, - distance_resolution=0.02, - show_timing_info=False, - direct_forces=True, - ): - super().__init__() - - if "e3nn" not in sys.modules: - logging.error("You need to install e3nn v0.2.6 to use the SCN model") - raise ImportError - - assert e3nn.__version__ == "0.2.6" - - self.regress_forces = regress_forces - self.use_pbc = use_pbc - self.cutoff = cutoff - self.otf_graph = otf_graph - self.show_timing_info = show_timing_info - self.max_num_elements = max_num_elements - self.hidden_channels = hidden_channels - self.num_interactions = num_interactions - self.num_atoms = 0 - self.num_sphere_samples = num_sphere_samples - self.sphere_channels = sphere_channels - self.sphere_channels_reduce = sphere_channels_reduce - self.max_num_neighbors = self.max_neighbors = max_num_neighbors - self.num_basis_functions = num_basis_functions - self.distance_resolution = distance_resolution - self.grad_forces = False - self.lmax = lmax - self.mmax = mmax - self.basis_width_scalar = basis_width_scalar - self.sphere_basis = (self.lmax + 1) ** 2 - self.use_grid = use_grid - self.distance_function = distance_function - - # variables used for display purposes - self.counter = 0 - - self.act = nn.SiLU() - - # Weights for message initialization - self.sphere_embedding = nn.Embedding( - self.max_num_elements, self.sphere_channels - ) - - assert self.distance_function in [ - "gaussian", - "sigmoid", - "linearsigmoid", - "silu", - ] - - self.num_gaussians = int(cutoff / self.distance_resolution) - if self.distance_function == "gaussian": - self.distance_expansion = GaussianSmearing( - 0.0, - cutoff, - self.num_gaussians, - basis_width_scalar, - ) - if self.distance_function == "sigmoid": - self.distance_expansion = SigmoidSmearing( - 0.0, - cutoff, - self.num_gaussians, - basis_width_scalar, - ) - if self.distance_function == "linearsigmoid": - self.distance_expansion = LinearSigmoidSmearing( - 0.0, - cutoff, - self.num_gaussians, - basis_width_scalar, - ) - - if self.distance_function == "silu": - self.distance_expansion = SiLUSmearing( - 0.0, - cutoff, - self.num_gaussians, - basis_width_scalar, - ) - - if num_resolutions == 1: - self.num_resolutions = 1 - self.hidden_channels_list = torch.tensor([self.hidden_channels]) - self.lmax_list = torch.tensor([self.lmax, -1]) # always end with -1 - self.cutoff_list = torch.tensor([self.max_num_neighbors - 0.01]) - if num_resolutions == 2: - self.num_resolutions = 2 - self.hidden_channels_list = torch.tensor( - [self.hidden_channels, self.hidden_channels // 4] - ) - self.lmax_list = torch.tensor([self.lmax, max(4, self.lmax - 2)]) - self.cutoff_list = torch.tensor([12 - 0.01, self.max_num_neighbors - 0.01]) - - self.sphharm_list = [] - for i in range(self.num_resolutions): - self.sphharm_list.append( - SphericalHarmonicsHelper( - self.lmax_list[i], - self.mmax, - num_taps, - num_bands, - ) - ) - - self.edge_blocks = nn.ModuleList() - for _ in range(self.num_interactions): - block = EdgeBlock( - self.num_resolutions, - self.sphere_channels_reduce, - self.hidden_channels_list, - self.cutoff_list, - self.sphharm_list, - self.sphere_channels, - self.distance_expansion, - self.max_num_elements, - self.num_basis_functions, - self.num_gaussians, - self.use_grid, - self.act, - ) - self.edge_blocks.append(block) - - # Energy estimation - self.energy_fc1 = nn.Linear(self.sphere_channels, self.sphere_channels) - self.energy_fc2 = nn.Linear(self.sphere_channels, self.sphere_channels_reduce) - self.energy_fc3 = nn.Linear(self.sphere_channels_reduce, 1) - - # Force estimation - if self.regress_forces: - self.force_fc1 = nn.Linear(self.sphere_channels, self.sphere_channels) - self.force_fc2 = nn.Linear( - self.sphere_channels, self.sphere_channels_reduce - ) - self.force_fc3 = nn.Linear(self.sphere_channels_reduce, 1) - - @conditional_grad(torch.enable_grad()) - def forward(self, data): - self.device = data.pos.device - self.num_atoms = len(data.batch) - self.batch_size = len(data.natoms) - # torch.autograd.set_detect_anomaly(True) - - start_time = time.time() - - outputs = self._forward_helper( - data, - ) - - if self.show_timing_info is True: - torch.cuda.synchronize() - print( - "{} Time: {}\tMemory: {}\t{}".format( - self.counter, - time.time() - start_time, - len(data.pos), - torch.cuda.max_memory_allocated() / 1000000, - ) - ) - - self.counter = self.counter + 1 - - return outputs - - # restructure forward helper for conditional grad - def _forward_helper(self, data): - atomic_numbers = data.atomic_numbers.long() - num_atoms = len(atomic_numbers) - pos = data.pos - - ( - edge_index, - edge_distance, - edge_distance_vec, - cell_offsets, - _, # cell offset distances - neighbors, - ) = self.generate_graph(data) - - ############################################################### - # Initialize data structures - ############################################################### - - # Calculate which message block each edge should use. Based on edge distance rank. - edge_rank = self._rank_edge_distances( - edge_distance, edge_index, self.max_num_neighbors - ) - - # Reorder edges so that they are grouped by distance rank (lowest to highest) - last_cutoff = -0.1 - message_block_idx = torch.zeros(len(edge_distance), device=pos.device) - edge_distance_reorder = torch.tensor([], device=self.device) - edge_index_reorder = torch.tensor([], device=self.device) - edge_distance_vec_reorder = torch.tensor([], device=self.device) - cutoff_index = torch.tensor([0], device=self.device) - for i in range(self.num_resolutions): - mask = torch.logical_and( - edge_rank.gt(last_cutoff), edge_rank.le(self.cutoff_list[i]) - ) - last_cutoff = self.cutoff_list[i] - message_block_idx.masked_fill_(mask, i) - edge_distance_reorder = torch.cat( - [ - edge_distance_reorder, - torch.masked_select(edge_distance, mask), - ], - dim=0, - ) - edge_index_reorder = torch.cat( - [ - edge_index_reorder, - torch.masked_select(edge_index, mask.view(1, -1).repeat(2, 1)).view( - 2, -1 - ), - ], - dim=1, - ) - edge_distance_vec_mask = torch.masked_select( - edge_distance_vec, mask.view(-1, 1).repeat(1, 3) - ).view(-1, 3) - edge_distance_vec_reorder = torch.cat( - [edge_distance_vec_reorder, edge_distance_vec_mask], dim=0 - ) - cutoff_index = torch.cat( - [ - cutoff_index, - torch.tensor([len(edge_distance_reorder)], device=self.device), - ], - dim=0, - ) - - edge_index = edge_index_reorder.long() - edge_distance = edge_distance_reorder - edge_distance_vec = edge_distance_vec_reorder - - # Compute 3x3 rotation matrix per edge - edge_rot_mat = self._init_edge_rot_mat(data, edge_index, edge_distance_vec) - - # Initialize the WignerD matrices and other values for spherical harmonic calculations - for i in range(self.num_resolutions): - self.sphharm_list[i].InitWignerDMatrix( - edge_rot_mat[cutoff_index[i] : cutoff_index[i + 1]], - ) - - ############################################################### - # Initialize node embeddings - ############################################################### - - # Init per node representations using an atomic number based embedding - x = torch.zeros( - num_atoms, - self.sphere_basis, - self.sphere_channels, - device=pos.device, - ) - x[:, 0, :] = self.sphere_embedding(atomic_numbers) - - ############################################################### - # Update spherical node embeddings - ############################################################### - for i, interaction in enumerate(self.edge_blocks): - if i > 0: - x = x + interaction( - x, atomic_numbers, edge_distance, edge_index, cutoff_index - ) - else: - x = interaction( - x, atomic_numbers, edge_distance, edge_index, cutoff_index - ) - - ############################################################### - # Estimate energy and forces using the node embeddings - ############################################################### - - # Create a roughly evenly distributed point sampling of the sphere - sphere_points = CalcSpherePoints(self.num_sphere_samples, x.device).detach() - sphharm_weights = o3.spherical_harmonics( - torch.arange(0, self.lmax + 1).tolist(), sphere_points, False - ).detach() - - # Energy estimation - node_energy = torch.einsum("abc, pb->apc", x, sphharm_weights).contiguous() - node_energy = node_energy.view(-1, self.sphere_channels) - node_energy = self.act(self.energy_fc1(node_energy)) - node_energy = self.act(self.energy_fc2(node_energy)) - node_energy = self.energy_fc3(node_energy) - node_energy = node_energy.view(-1, self.num_sphere_samples, 1) - node_energy = torch.sum(node_energy, dim=1) / self.num_sphere_samples - energy = torch.zeros(len(data.natoms), device=pos.device) - energy.index_add_(0, data.batch, node_energy.view(-1)) - - # Force estimation - if self.regress_forces: - forces = torch.einsum("abc, pb->apc", x, sphharm_weights).contiguous() - forces = forces.view(-1, self.sphere_channels) - forces = self.act(self.force_fc1(forces)) - forces = self.act(self.force_fc2(forces)) - forces = self.force_fc3(forces) - forces = forces.view(-1, self.num_sphere_samples, 1) - forces = forces * sphere_points.view(1, self.num_sphere_samples, 3) - forces = torch.sum(forces, dim=1) / self.num_sphere_samples - - if not self.regress_forces: - return energy - else: - return energy, forces - - def _init_edge_rot_mat(self, data, edge_index, edge_distance_vec): - edge_vec_0 = edge_distance_vec - edge_vec_0_distance = torch.sqrt(torch.sum(edge_vec_0**2, dim=1)) - - if torch.min(edge_vec_0_distance) < 0.0001: - print( - "Error edge_vec_0_distance: {}".format(torch.min(edge_vec_0_distance)) - ) - (minval, minidx) = torch.min(edge_vec_0_distance, 0) - print( - "Error edge_vec_0_distance: {} {} {} {} {}".format( - minidx, - edge_index[0, minidx], - edge_index[1, minidx], - data.pos[edge_index[0, minidx]], - data.pos[edge_index[1, minidx]], - ) - ) - - norm_x = edge_vec_0 / (edge_vec_0_distance.view(-1, 1)) - - edge_vec_2 = torch.rand_like(edge_vec_0) - 0.5 - edge_vec_2 = edge_vec_2 / ( - torch.sqrt(torch.sum(edge_vec_2**2, dim=1)).view(-1, 1) - ) - # Create two rotated copys of the random vectors in case the random vector is aligned with norm_x - # With two 90 degree rotated vectors, at least one should not be aligned with norm_x - edge_vec_2b = edge_vec_2.clone() - edge_vec_2b[:, 0] = -edge_vec_2[:, 1] - edge_vec_2b[:, 1] = edge_vec_2[:, 0] - edge_vec_2c = edge_vec_2.clone() - edge_vec_2c[:, 1] = -edge_vec_2[:, 2] - edge_vec_2c[:, 2] = edge_vec_2[:, 1] - vec_dot_b = torch.abs(torch.sum(edge_vec_2b * norm_x, dim=1)).view(-1, 1) - vec_dot_c = torch.abs(torch.sum(edge_vec_2c * norm_x, dim=1)).view(-1, 1) - - vec_dot = torch.abs(torch.sum(edge_vec_2 * norm_x, dim=1)).view(-1, 1) - edge_vec_2 = torch.where(torch.gt(vec_dot, vec_dot_b), edge_vec_2b, edge_vec_2) - vec_dot = torch.abs(torch.sum(edge_vec_2 * norm_x, dim=1)).view(-1, 1) - edge_vec_2 = torch.where(torch.gt(vec_dot, vec_dot_c), edge_vec_2c, edge_vec_2) - - vec_dot = torch.abs(torch.sum(edge_vec_2 * norm_x, dim=1)) - # Check the vectors aren't aligned - assert torch.max(vec_dot) < 0.99 - - norm_z = torch.cross(norm_x, edge_vec_2, dim=1) - norm_z = norm_z / (torch.sqrt(torch.sum(norm_z**2, dim=1, keepdim=True))) - norm_z = norm_z / (torch.sqrt(torch.sum(norm_z**2, dim=1)).view(-1, 1)) - norm_y = torch.cross(norm_x, norm_z, dim=1) - norm_y = norm_y / (torch.sqrt(torch.sum(norm_y**2, dim=1, keepdim=True))) - - norm_x = norm_x.view(-1, 3, 1) - norm_y = -norm_y.view(-1, 3, 1) - norm_z = norm_z.view(-1, 3, 1) - - edge_rot_mat_inv = torch.cat([norm_z, norm_x, norm_y], dim=2) - edge_rot_mat = torch.transpose(edge_rot_mat_inv, 1, 2) - - return edge_rot_mat.detach() - - def _rank_edge_distances(self, edge_distance, edge_index, max_num_neighbors): - device = edge_distance.device - # Create an index map to map distances from atom_distance to distance_sort - # index_sort_map assumes index to be sorted - output, num_neighbors = torch.unique(edge_index[1], return_counts=True) - index_neighbor_offset = torch.cumsum(num_neighbors, dim=0) - num_neighbors - index_neighbor_offset_expand = torch.repeat_interleave( - index_neighbor_offset, num_neighbors - ) - - index_sort_map = ( - edge_index[1] * max_num_neighbors - + torch.arange(len(edge_distance), device=device) - - index_neighbor_offset_expand - ) - - num_atoms = torch.max(edge_index) + 1 - distance_sort = torch.full( - [num_atoms * max_num_neighbors], np.inf, device=device - ) - distance_sort.index_copy_(0, index_sort_map, edge_distance) - distance_sort = distance_sort.view(num_atoms, max_num_neighbors) - no_op, index_sort = torch.sort(distance_sort, dim=1) - - index_map = ( - torch.arange(max_num_neighbors, device=device) - .view(1, -1) - .repeat(num_atoms, 1) - .view(-1) - ) - index_sort = index_sort + ( - torch.arange(num_atoms, device=device) * max_num_neighbors - ).view(-1, 1).repeat(1, max_num_neighbors) - edge_rank = torch.zeros_like(index_map) - edge_rank.index_copy_(0, index_sort.view(-1), index_map) - edge_rank = edge_rank.view(num_atoms, max_num_neighbors) - - index_sort_mask = distance_sort.lt(1000.0) - edge_rank = torch.masked_select(edge_rank, index_sort_mask) - - return edge_rank - - @property - def num_params(self): - return sum(p.numel() for p in self.parameters()) - - -class EdgeBlock(torch.nn.Module): - def __init__( - self, - num_resolutions, - sphere_channels_reduce, - hidden_channels_list, - cutoff_list, - sphharm_list, - sphere_channels, - distance_expansion, - max_num_elements, - num_basis_functions, - num_gaussians, - use_grid, - act, - ): - super(EdgeBlock, self).__init__() - self.num_resolutions = num_resolutions - self.act = act - self.hidden_channels_list = hidden_channels_list - self.sphere_channels = sphere_channels - self.sphere_channels_reduce = sphere_channels_reduce - self.distance_expansion = distance_expansion - self.cutoff_list = cutoff_list - self.sphharm_list = sphharm_list - self.max_num_elements = max_num_elements - self.num_basis_functions = num_basis_functions - self.use_grid = use_grid - self.num_gaussians = num_gaussians - - # Edge features - self.dist_block = DistanceBlock( - self.num_gaussians, - self.num_basis_functions, - self.distance_expansion, - self.max_num_elements, - self.act, - ) - - # Create a message block for each cutoff - self.message_blocks = nn.ModuleList() - for i in range(self.num_resolutions): - block = MessageBlock( - self.sphere_channels_reduce, - int(self.hidden_channels_list[i]), - self.num_basis_functions, - self.sphharm_list[i], - self.act, - ) - self.message_blocks.append(block) - - # Downsampling number of sphere channels - # Make sure bias is false unless equivariance is lost - if self.sphere_channels != self.sphere_channels_reduce: - self.downsample = nn.Linear( - self.sphere_channels, - self.sphere_channels_reduce, - bias=False, - ) - self.upsample = nn.Linear( - self.sphere_channels_reduce, - self.sphere_channels, - bias=False, - ) - - # Use non-linear message aggregation? - if self.use_grid: - # Network for each node to combine edge messages - self.fc1_sphere = nn.Linear( - self.sphharm_list[0].num_bands * 2 * self.sphere_channels_reduce, - self.sphharm_list[0].num_bands * 2 * self.sphere_channels_reduce, - ) - - self.fc2_sphere = nn.Linear( - self.sphharm_list[0].num_bands * 2 * self.sphere_channels_reduce, - 2 * self.sphere_channels_reduce, - ) - - self.fc3_sphere = nn.Linear( - 2 * self.sphere_channels_reduce, self.sphere_channels_reduce - ) - - def forward( - self, - x, - atomic_numbers, - edge_distance, - edge_index, - cutoff_index, - ): - - ############################################################### - # Update spherical node embeddings - ############################################################### - - x_edge = self.dist_block( - edge_distance, - atomic_numbers[edge_index[0]], - atomic_numbers[edge_index[1]], - ) - x_new = torch.zeros( - len(x), - self.sphharm_list[0].sphere_basis, - self.sphere_channels_reduce, - dtype=x.dtype, - device=x.device, - ) - - if self.sphere_channels != self.sphere_channels_reduce: - x_down = self.downsample(x.view(-1, self.sphere_channels)) - else: - x_down = x - x_down = x_down.view( - -1, self.sphharm_list[0].sphere_basis, self.sphere_channels_reduce - ) - - for i, interaction in enumerate(self.message_blocks): - start_idx = cutoff_index[i] - end_idx = cutoff_index[i + 1] - - x_message = interaction( - x_down[:, 0 : self.sphharm_list[i].sphere_basis, :], - x_edge[start_idx:end_idx], - edge_index[:, start_idx:end_idx], - ) - - # Sum all incoming edges to the target nodes - x_new[:, 0 : self.sphharm_list[i].sphere_basis, :].index_add_( - 0, edge_index[1, start_idx:end_idx], x_message.to(x_new.dtype) - ) - - if self.use_grid: - # Feed in the spherical functions from the previous time step - x_grid = self.sphharm_list[0].ToGrid(x_down, self.sphere_channels_reduce) - x_grid = torch.cat( - [ - x_grid, - self.sphharm_list[0].ToGrid(x_new, self.sphere_channels_reduce), - ], - dim=1, - ) - - x_grid = self.act(self.fc1_sphere(x_grid)) - x_grid = self.act(self.fc2_sphere(x_grid)) - x_grid = self.fc3_sphere(x_grid) - x_new = self.sphharm_list[0].FromGrid(x_grid, self.sphere_channels_reduce) - - if self.sphere_channels != self.sphere_channels_reduce: - x_new = x_new.view(-1, self.sphere_channels_reduce) - x_new = self.upsample(x_new) - x_new = x_new.view(-1, self.sphharm_list[0].sphere_basis, self.sphere_channels) - - return x_new - - -class MessageBlock(torch.nn.Module): - def __init__( - self, - sphere_channels_reduce, - hidden_channels, - num_basis_functions, - sphharm, - act, - ): - super(MessageBlock, self).__init__() - self.act = act - self.hidden_channels = hidden_channels - self.sphere_channels_reduce = sphere_channels_reduce - self.sphharm = sphharm - - self.fc1_dist = nn.Linear(num_basis_functions, self.hidden_channels) - - # Network for each edge to compute edge messages - self.fc1_edge_proj = nn.Linear( - 2 * self.sphharm.sphere_basis_reduce * self.sphere_channels_reduce, - self.hidden_channels, - ) - - self.fc1_edge = nn.Linear(self.hidden_channels, self.hidden_channels) - - self.fc2_edge = nn.Linear( - self.hidden_channels, - self.sphharm.sphere_basis_reduce * self.sphere_channels_reduce, - ) - - def forward( - self, - x, - x_edge, - edge_index, - ): - - ############################################################### - # Compute messages - ############################################################### - x_edge = self.act(self.fc1_dist(x_edge)) - - x_source = x[edge_index[0, :]] - x_target = x[edge_index[1, :]] - - # Rotate the spherical harmonic basis functions to align with the edge - x_msg_source = self.sphharm.Rotate(x_source) - x_msg_target = self.sphharm.Rotate(x_target) - - # Compute messages - x_message = torch.cat([x_msg_source, x_msg_target], dim=1) - x_message = self.act(self.fc1_edge_proj(x_message)) - x_message = ( - x_message.view(-1, self.sphharm.num_y_rotations, self.hidden_channels) - ) * x_edge.view(-1, 1, self.hidden_channels) - x_message = x_message.view(-1, self.hidden_channels) - - x_message = self.act(self.fc1_edge(x_message)) - x_message = self.act(self.fc2_edge(x_message)) - - # Combine the rotated versions of the messages - x_message = x_message.view(-1, self.sphere_channels_reduce) - x_message = self.sphharm.CombineYRotations(x_message) - - # Rotate the spherical harmonic basis functions back to global coordinate frame - x_message = self.sphharm.RotateInv(x_message) - - return x_message - - -class DistanceBlock(torch.nn.Module): - def __init__( - self, - in_channels, - num_basis_functions, - distance_expansion, - max_num_elements, - act, - ): - super(DistanceBlock, self).__init__() - self.in_channels = in_channels - self.distance_expansion = distance_expansion - self.act = act - self.num_basis_functions = num_basis_functions - self.max_num_elements = max_num_elements - self.num_edge_channels = self.num_basis_functions - - self.fc1_dist = nn.Linear(self.in_channels, self.num_basis_functions) - - self.source_embedding = nn.Embedding( - self.max_num_elements, self.num_basis_functions - ) - self.target_embedding = nn.Embedding( - self.max_num_elements, self.num_basis_functions - ) - nn.init.uniform_(self.source_embedding.weight.data, -0.001, 0.001) - nn.init.uniform_(self.target_embedding.weight.data, -0.001, 0.001) - - self.fc1_edge_attr = nn.Linear( - self.num_edge_channels, - self.num_edge_channels, - ) - - def forward(self, edge_distance, source_element, target_element): - x_dist = self.distance_expansion(edge_distance) - x_dist = self.fc1_dist(x_dist) - - source_embedding = self.source_embedding(source_element) - target_embedding = self.target_embedding(target_element) - - x_edge = self.act(source_embedding + target_embedding + x_dist) - x_edge = self.act(self.fc1_edge_attr(x_edge)) - - return x_edge diff --git a/ocpmodels/models/scn/smearing.py b/ocpmodels/models/scn/smearing.py deleted file mode 100644 index b0197eba50346e5c70d7735520632d71660cd19a..0000000000000000000000000000000000000000 --- a/ocpmodels/models/scn/smearing.py +++ /dev/null @@ -1,64 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import torch -import torch.nn as nn - - -# Different encodings for the atom distance embeddings -class GaussianSmearing(torch.nn.Module): - def __init__(self, start=-5.0, stop=5.0, num_gaussians=50, basis_width_scalar=1.0): - super(GaussianSmearing, self).__init__() - self.num_output = num_gaussians - offset = torch.linspace(start, stop, num_gaussians) - self.coeff = -0.5 / (basis_width_scalar * (offset[1] - offset[0])).item() ** 2 - self.register_buffer("offset", offset) - - def forward(self, dist): - dist = dist.view(-1, 1) - self.offset.view(1, -1) - return torch.exp(self.coeff * torch.pow(dist, 2)) - - -class SigmoidSmearing(torch.nn.Module): - def __init__(self, start=-5.0, stop=5.0, num_sigmoid=50, basis_width_scalar=1.0): - super(SigmoidSmearing, self).__init__() - self.num_output = num_sigmoid - offset = torch.linspace(start, stop, num_sigmoid) - self.coeff = (basis_width_scalar / (offset[1] - offset[0])).item() - self.register_buffer("offset", offset) - - def forward(self, dist): - exp_dist = self.coeff * (dist.view(-1, 1) - self.offset.view(1, -1)) - return torch.sigmoid(exp_dist) - - -class LinearSigmoidSmearing(torch.nn.Module): - def __init__(self, start=-5.0, stop=5.0, num_sigmoid=50, basis_width_scalar=1.0): - super(LinearSigmoidSmearing, self).__init__() - self.num_output = num_sigmoid - offset = torch.linspace(start, stop, num_sigmoid) - self.coeff = (basis_width_scalar / (offset[1] - offset[0])).item() - self.register_buffer("offset", offset) - - def forward(self, dist): - exp_dist = self.coeff * (dist.view(-1, 1) - self.offset.view(1, -1)) - x_dist = torch.sigmoid(exp_dist) + 0.001 * exp_dist - return x_dist - - -class SiLUSmearing(torch.nn.Module): - def __init__(self, start=-5.0, stop=5.0, num_output=50, basis_width_scalar=1.0): - super(SiLUSmearing, self).__init__() - self.num_output = num_output - self.fc1 = nn.Linear(2, num_output) - self.act = nn.SiLU() - - def forward(self, dist): - x_dist = dist.view(-1, 1) - x_dist = torch.cat([x_dist, torch.ones_like(x_dist)], dim=1) - x_dist = self.act(self.fc1(x_dist)) - return x_dist diff --git a/ocpmodels/models/scn/spherical_harmonics.py b/ocpmodels/models/scn/spherical_harmonics.py deleted file mode 100644 index e9525af5ee6c0c801a977762dd23e00e963361db..0000000000000000000000000000000000000000 --- a/ocpmodels/models/scn/spherical_harmonics.py +++ /dev/null @@ -1,370 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import logging -import math -import os - -import torch - -try: - from e3nn import o3 - from e3nn.o3 import FromS2Grid, ToS2Grid -except ImportError: - pass - -# Borrowed from e3nn @ 0.4.0: -# https://github.com/e3nn/e3nn/blob/0.4.0/e3nn/o3/_wigner.py#L10 -# _Jd is a list of tensors of shape (2l+1, 2l+1) -_Jd = torch.load(os.path.join(os.path.dirname(__file__), "Jd.pt"), weights_only=True) - - -class SphericalHarmonicsHelper: - """ - Helper functions for spherical harmonics calculations and representations - - Args: - lmax (int): Maximum degree of the spherical harmonics - mmax (int): Maximum order of the spherical harmonics - num_taps (int): Number of taps or rotations (1 or otherwise set automatically based on mmax) - num_bands (int): Number of bands used during message aggregation for the 1x1 pointwise convolution (1 or 2) - """ - - def __init__( - self, - lmax, - mmax, - num_taps, - num_bands, - ): - import sys - - if "e3nn" not in sys.modules: - logging.error( - "You need to install the e3nn library to use Spherical Harmonics" - ) - raise ImportError - - super().__init__() - self.lmax = lmax - self.mmax = mmax - self.num_taps = num_taps - self.num_bands = num_bands - - # Make sure lmax is large enough to support the num_bands - assert self.lmax - (self.num_bands - 1) >= 0 - - self.sphere_basis = (self.lmax + 1) ** 2 - self.sphere_basis = int(self.sphere_basis) - - self.sphere_basis_reduce = self.lmax + 1 - for i in range(1, self.mmax + 1): - self.sphere_basis_reduce = self.sphere_basis_reduce + 2 * ( - self.lmax + 1 - i - ) - self.sphere_basis_reduce = int(self.sphere_basis_reduce) - - def InitWignerDMatrix(self, edge_rot_mat): - self.device = edge_rot_mat.device - - # Initialize matrix to combine the y-axis rotations during message passing - self.mapping_y_rot, self.y_rotations = self.InitYRotMapping() - self.num_y_rotations = len(self.y_rotations) - - # Conversion from basis to grid respresentations - self.grid_res = (self.lmax + 1) * 2 - self.to_grid_shb = torch.tensor([], device=self.device) - self.to_grid_sha = torch.tensor([], device=self.device) - - for b in range(self.num_bands): - l = self.lmax - b # noqa: E741 - togrid = ToS2Grid( - l, - (self.grid_res, self.grid_res + 1), - normalization="integral", - device=self.device, - ) - shb = togrid.shb - sha = togrid.sha - - padding = torch.zeros( - shb.size()[0], - shb.size()[1], - self.sphere_basis - shb.size()[2], - device=self.device, - ) - shb = torch.cat([shb, padding], dim=2) - self.to_grid_shb = torch.cat([self.to_grid_shb, shb], dim=0) - if b == 0: - self.to_grid_sha = sha - else: - self.to_grid_sha = torch.block_diag(self.to_grid_sha, sha) - - self.to_grid_sha = self.to_grid_sha.view(self.num_bands, self.grid_res + 1, -1) - self.to_grid_sha = torch.transpose(self.to_grid_sha, 0, 1).contiguous() - self.to_grid_sha = self.to_grid_sha.view( - (self.grid_res + 1) * self.num_bands, -1 - ) - - self.to_grid_shb = self.to_grid_shb.detach() - self.to_grid_sha = self.to_grid_sha.detach() - - self.from_grid = FromS2Grid( - (self.grid_res, self.grid_res + 1), - self.lmax, - normalization="integral", - device=self.device, - ) - for p in self.from_grid.parameters(): - p.detach() - - # Compute subsets of Wigner matrices to use for messages - wigner = torch.tensor([], device=self.device) - wigner_inv = torch.tensor([], device=self.device) - - for y_rot in self.y_rotations: - - # Compute rotation about y-axis - y_rot_mat = self.RotationMatrix(0, y_rot, 0) - y_rot_mat = y_rot_mat.repeat(len(edge_rot_mat), 1, 1) - # Add additional rotation about y-axis - rot_mat = torch.bmm(y_rot_mat, edge_rot_mat) - - # Compute Wigner matrices corresponding to the 3x3 rotation matrices - wignerD = self.RotationToWignerDMatrix(rot_mat, 0, self.lmax) - - basis_in = torch.tensor([], device=self.device) - basis_out = torch.tensor([], device=self.device) - start_l = 0 - end_l = self.lmax + 1 - for l in range(start_l, end_l): # noqa: E741 - offset = l**2 - basis_in = torch.cat( - [ - basis_in, - torch.arange(2 * l + 1, device=self.device) + offset, - ], - dim=0, - ) - m_max = min(l, self.mmax) - basis_out = torch.cat( - [ - basis_out, - torch.arange(-m_max, m_max + 1, device=self.device) - + offset - + l, - ], - dim=0, - ) - - # Only keep the rows/columns of the matrices used given lmax and mmax - wignerD_reduce = wignerD[:, basis_out.long(), :] - wignerD_reduce = wignerD_reduce[:, :, basis_in.long()] - - if y_rot == 0.0: - wigner_inv = torch.transpose(wignerD_reduce, 1, 2).contiguous().detach() - - wigner = torch.cat([wigner, wignerD_reduce.unsqueeze(1)], dim=1) - - wigner = wigner.view(-1, self.sphere_basis_reduce, self.sphere_basis) - - self.wigner = wigner.detach() - self.wigner_inv = wigner_inv.detach() - - # If num_taps is greater than 1, calculate how to combine the different samples. - # Note the e3nn code flips the y-axis with the z-axis in the SCN paper description. - def InitYRotMapping(self): - - if self.mmax == 0: - y_rotations = torch.tensor([0.0], device=self.device) - num_y_rotations = 1 - mapping_y_rot = torch.eye(self.sphere_basis_reduce, device=self.device) - - if self.mmax == 1: - - if self.num_taps == 1: - y_rotations = torch.tensor([0.0], device=self.device) - num_y_rotations = len(y_rotations) - mapping_y_rot = torch.eye( - len(y_rotations) * self.sphere_basis_reduce, - self.sphere_basis_reduce, - device=self.device, - ) - else: - y_rotations = torch.tensor( - [0.0, 0.5 * math.pi, math.pi, 1.5 * math.pi], - device=self.device, - ) - num_y_rotations = len(y_rotations) - mapping_y_rot = torch.zeros( - len(y_rotations) * self.sphere_basis_reduce, - self.sphere_basis_reduce, - device=self.device, - ) - - # m = 0 - for l in range(0, self.lmax + 1): # noqa: E741 - offset = (l - 1) * 3 + 2 - if l == 0: # noqa: E741 - offset = 0 - for y in range(num_y_rotations): - mapping_y_rot[offset + y * self.sphere_basis_reduce, offset] = ( - 1.0 / num_y_rotations - ) - - # m = -1 - for l in range(1, self.lmax + 1): # noqa: E741 - offset = (l - 1) * 3 + 1 - for y in range(num_y_rotations): - mapping_y_rot[offset + y * self.sphere_basis_reduce, offset] = ( - math.cos(y_rotations[y]) / num_y_rotations - ) - mapping_y_rot[ - (offset + 2) + y * self.sphere_basis_reduce, offset - ] = (math.sin(y_rotations[y]) / num_y_rotations) - - # m = 1 - for l in range(1, self.lmax + 1): # noqa: E741 - offset = (l - 1) * 3 + 3 - for y in range(num_y_rotations): - mapping_y_rot[offset + y * self.sphere_basis_reduce, offset] = ( - math.cos(y_rotations[y]) / num_y_rotations - ) - mapping_y_rot[ - offset - 2 + y * self.sphere_basis_reduce, offset - ] = (-math.sin(y_rotations[y]) / num_y_rotations) - - return mapping_y_rot.detach(), y_rotations - - # Simplified version of function from e3nn - def ToGrid(self, x, channels): - x = x.view(-1, self.sphere_basis, channels) - x_grid = torch.einsum("mbi,zic->zbmc", self.to_grid_shb, x) - x_grid = torch.einsum("am,zbmc->zbac", self.to_grid_sha, x_grid).contiguous() - x_grid = x_grid.view(-1, self.num_bands * channels) - return x_grid - - # Simplified version of function from e3nn - def FromGrid(self, x_grid, channels): - x_grid = x_grid.view(-1, self.grid_res, (self.grid_res + 1), channels) - x = torch.einsum("am,zbac->zbmc", self.from_grid.sha, x_grid) - x = torch.einsum("mbi,zbmc->zic", self.from_grid.shb, x).contiguous() - x = x.view(-1, channels) - return x - - def CombineYRotations(self, x): - num_channels = x.size()[-1] - x = x.view(-1, self.num_y_rotations * self.sphere_basis_reduce, num_channels) - x = torch.einsum("abc, bd->adc", x, self.mapping_y_rot).contiguous() - return x - - def Rotate(self, x): - num_channels = x.size()[2] - x = x.view(-1, 1, self.sphere_basis, num_channels).repeat( - 1, self.num_y_rotations, 1, 1 - ) - x = x.view(-1, self.sphere_basis, num_channels) - # print('{} {}'.format(self.wigner.size(), x.size())) - x_rot = torch.bmm(self.wigner, x) - x_rot = x_rot.view(-1, self.sphere_basis_reduce * num_channels) - return x_rot - - def FlipGrid(self, grid, num_channels): - # lat long - long_res = self.grid_res - grid = grid.view(-1, self.grid_res, self.grid_res, num_channels) - grid = torch.roll(grid, int(long_res // 2), 2) - flip_grid = torch.flip(grid, [1]) - return flip_grid.view(-1, num_channels) - - def RotateInv(self, x): - x_rot = torch.bmm(self.wigner_inv, x) - return x_rot - - def RotateWigner(self, x, wigner): - x_rot = torch.bmm(wigner, x) - return x_rot - - def RotationMatrix(self, rot_x, rot_y, rot_z): - m1, m2, m3 = ( - torch.eye(3, device=self.device), - torch.eye(3, device=self.device), - torch.eye(3, device=self.device), - ) - if rot_x: - degree = rot_x - sin, cos = math.sin(degree), math.cos(degree) - m1 = torch.tensor( - [[1, 0, 0], [0, cos, sin], [0, -sin, cos]], device=self.device - ) - if rot_y: - degree = rot_y - sin, cos = math.sin(degree), math.cos(degree) - m2 = torch.tensor( - [[cos, 0, -sin], [0, 1, 0], [sin, 0, cos]], device=self.device - ) - if rot_z: - degree = rot_z - sin, cos = math.sin(degree), math.cos(degree) - m3 = torch.tensor( - [[cos, sin, 0], [-sin, cos, 0], [0, 0, 1]], device=self.device - ) - - matrix = torch.mm(torch.mm(m1, m2), m3) - matrix = matrix.view(1, 3, 3) - - return matrix - - def RotationToWignerDMatrix(self, edge_rot_mat, start_lmax, end_lmax): - x = edge_rot_mat @ edge_rot_mat.new_tensor([0.0, 1.0, 0.0]) - alpha, beta = o3.xyz_to_angles(x) - R = ( - o3.angles_to_matrix(alpha, beta, torch.zeros_like(alpha)).transpose(-1, -2) - @ edge_rot_mat - ) - gamma = torch.atan2(R[..., 0, 2], R[..., 0, 0]) - - size = (end_lmax + 1) ** 2 - (start_lmax) ** 2 - wigner = torch.zeros(len(alpha), size, size, device=self.device) - start = 0 - for lmax in range(start_lmax, end_lmax + 1): - block = wigner_D(lmax, alpha, beta, gamma) - end = start + block.size()[1] - wigner[:, start:end, start:end] = block - start = end - - return wigner.detach() - - -# Borrowed from e3nn @ 0.4.0: -# https://github.com/e3nn/e3nn/blob/0.4.0/e3nn/o3/_wigner.py#L37 -# -# In 0.5.0, e3nn shifted to torch.matrix_exp which is significantly slower: -# https://github.com/e3nn/e3nn/blob/0.5.0/e3nn/o3/_wigner.py#L92 -def wigner_D(l, alpha, beta, gamma): - if not l < len(_Jd): - raise NotImplementedError( - f"wigner D maximum l implemented is {len(_Jd) - 1}, send us an email to ask for more" - ) - - alpha, beta, gamma = torch.broadcast_tensors(alpha, beta, gamma) - J = _Jd[l].to(dtype=alpha.dtype, device=alpha.device) - Xa = _z_rot_mat(alpha, l) - Xb = _z_rot_mat(beta, l) - Xc = _z_rot_mat(gamma, l) - return Xa @ J @ Xb @ J @ Xc - - -def _z_rot_mat(angle, l): - shape, device, dtype = angle.shape, angle.device, angle.dtype - M = angle.new_zeros((*shape, 2 * l + 1, 2 * l + 1)) - inds = torch.arange(0, 2 * l + 1, 1, device=device) - reversed_inds = torch.arange(2 * l, -1, -1, device=device) - frequencies = torch.arange(l, -l - 1, -1, dtype=dtype, device=device) - M[..., inds, reversed_inds] = torch.sin(frequencies * angle[..., None]) - M[..., inds, inds] = torch.cos(frequencies * angle[..., None]) - return M diff --git a/ocpmodels/models/spinconv.py b/ocpmodels/models/spinconv.py deleted file mode 100644 index 0bbfc23dc2f7fb75767041c05e65dd3c685754d3..0000000000000000000000000000000000000000 --- a/ocpmodels/models/spinconv.py +++ /dev/null @@ -1,1186 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import math -import time -from math import pi as PI - -import numpy as np -import torch -import torch.nn as nn -import torch.nn.functional as F -from torch.nn import Embedding, Linear, ModuleList, Sequential -from torch_geometric.nn import MessagePassing, SchNet, radius_graph -from torch_scatter import scatter - -from ocpmodels.common.registry import registry -from ocpmodels.common.transforms import RandomRotate -from ocpmodels.common.utils import ( - compute_neighbors, - conditional_grad, - get_pbc_distances, - radius_graph_pbc, -) -from ocpmodels.models.base import BaseModel - -try: - from e3nn import o3 - from e3nn.io import SphericalTensor - from e3nn.o3 import FromS2Grid, SphericalHarmonics, ToS2Grid -except Exception: - pass - - -@registry.register_model("spinconv") -class spinconv(BaseModel): - def __init__( - self, - num_atoms, # not used - bond_feat_dim, # not used - num_targets, - use_pbc=True, - regress_forces=True, - otf_graph=False, - hidden_channels=32, - mid_hidden_channels=200, - num_interactions=1, - num_basis_functions=200, - basis_width_scalar=1.0, - max_num_neighbors=20, - sphere_size_lat=15, - sphere_size_long=9, - cutoff=10.0, - distance_block_scalar_max=2.0, - max_num_elements=90, - embedding_size=32, - show_timing_info=False, - sphere_message="fullconv", # message block sphere representation - output_message="fullconv", # output block sphere representation - lmax=False, - force_estimator="random", - model_ref_number=0, - readout="add", - num_rand_rotations=5, - scale_distances=True, - ): - super(spinconv, self).__init__() - - self.num_targets = num_targets - self.num_random_rotations = num_rand_rotations - self.regress_forces = regress_forces - self.use_pbc = use_pbc - self.cutoff = cutoff - self.otf_graph = otf_graph - self.show_timing_info = show_timing_info - self.max_num_elements = max_num_elements - self.mid_hidden_channels = mid_hidden_channels - self.sphere_size_lat = sphere_size_lat - self.sphere_size_long = sphere_size_long - self.num_atoms = 0 - self.hidden_channels = hidden_channels - self.embedding_size = embedding_size - self.max_num_neighbors = self.max_neighbors = max_num_neighbors - self.sphere_message = sphere_message - self.output_message = output_message - self.force_estimator = force_estimator - self.num_basis_functions = num_basis_functions - self.distance_block_scalar_max = distance_block_scalar_max - self.grad_forces = False - self.num_embedding_basis = 8 - self.lmax = lmax - self.scale_distances = scale_distances - self.basis_width_scalar = basis_width_scalar - - if self.sphere_message in ["spharm", "rotspharmroll", "rotspharmwd"]: - assert self.lmax, "lmax must be defined for spherical harmonics" - if self.output_message in ["spharm", "rotspharmroll", "rotspharmwd"]: - assert self.lmax, "lmax must be defined for spherical harmonics" - - # variables used for display purposes - self.counter = 0 - self.start_time = time.time() - self.total_time = 0 - self.model_ref_number = model_ref_number - - if self.force_estimator == "grad": - self.grad_forces = True - - # self.act = ShiftedSoftplus() - self.act = Swish() - - self.distance_expansion_forces = GaussianSmearing( - 0.0, - cutoff, - num_basis_functions, - basis_width_scalar, - ) - - # Weights for message initialization - self.embeddingblock2 = EmbeddingBlock( - self.mid_hidden_channels, - self.hidden_channels, - self.mid_hidden_channels, - self.embedding_size, - self.num_embedding_basis, - self.max_num_elements, - self.act, - ) - self.distfc1 = nn.Linear(self.mid_hidden_channels, self.mid_hidden_channels) - self.distfc2 = nn.Linear(self.mid_hidden_channels, self.mid_hidden_channels) - - self.dist_block = DistanceBlock( - self.num_basis_functions, - self.mid_hidden_channels, - self.max_num_elements, - self.distance_block_scalar_max, - self.distance_expansion_forces, - self.scale_distances, - ) - - self.message_blocks = ModuleList() - for _ in range(num_interactions): - block = MessageBlock( - hidden_channels, - hidden_channels, - mid_hidden_channels, - embedding_size, - self.sphere_size_lat, - self.sphere_size_long, - self.max_num_elements, - self.sphere_message, - self.act, - self.lmax, - ) - self.message_blocks.append(block) - - self.energyembeddingblock = EmbeddingBlock( - hidden_channels, - 1, - mid_hidden_channels, - embedding_size, - 8, - self.max_num_elements, - self.act, - ) - - if force_estimator == "random": - self.force_output_block = ForceOutputBlock( - hidden_channels, - 2, - mid_hidden_channels, - embedding_size, - self.sphere_size_lat, - self.sphere_size_long, - self.max_num_elements, - self.output_message, - self.act, - self.lmax, - ) - - @conditional_grad(torch.enable_grad()) - def forward(self, data): - self.device = data.pos.device - self.num_atoms = len(data.batch) - self.batch_size = len(data.natoms) - - pos = data.pos - if self.regress_forces: - pos = pos.requires_grad_(True) - - ( - edge_index, - edge_distance, - edge_distance_vec, - cell_offsets, - _, # cell offset distances - neighbors, - ) = self.generate_graph(data) - - edge_index, edge_distance, edge_distance_vec = self._filter_edges( - edge_index, - edge_distance, - edge_distance_vec, - self.max_num_neighbors, - ) - - outputs = self._forward_helper( - data, edge_index, edge_distance, edge_distance_vec - ) - if self.show_timing_info is True: - torch.cuda.synchronize() - print( - "Memory: {}\t{}\t{}".format( - len(edge_index[0]), - torch.cuda.memory_allocated() / (1000 * len(edge_index[0])), - torch.cuda.max_memory_allocated() / 1000000, - ) - ) - - return outputs - - # restructure forward helper for conditional grad - def _forward_helper(self, data, edge_index, edge_distance, edge_distance_vec): - ############################################################### - # Initialize messages - ############################################################### - - source_element = data.atomic_numbers[edge_index[0, :]].long() - target_element = data.atomic_numbers[edge_index[1, :]].long() - - x_dist = self.dist_block(edge_distance, source_element, target_element) - - x = x_dist - x = self.distfc1(x) - x = self.act(x) - x = self.distfc2(x) - x = self.act(x) - x = self.embeddingblock2(x, source_element, target_element) - - ############################################################### - # Update messages using block interactions - ############################################################### - - edge_rot_mat = self._init_edge_rot_mat(data, edge_index, edge_distance_vec) - ( - proj_edges_index, - proj_edges_delta, - proj_edges_src_index, - ) = self._project2D_edges_init(edge_rot_mat, edge_index, edge_distance_vec) - - for block_index, interaction in enumerate(self.message_blocks): - x_out = interaction( - x, - x_dist, - source_element, - target_element, - proj_edges_index, - proj_edges_delta, - proj_edges_src_index, - ) - - if block_index > 0: - x = x + x_out - else: - x = x_out - - ############################################################### - # Decoder - # Compute the forces and energies from the messages - ############################################################### - assert self.force_estimator in ["random", "grad"] - - energy = scatter(x, edge_index[1], dim=0, dim_size=data.num_nodes) / ( - self.max_num_neighbors / 2.0 + 1.0 - ) - atomic_numbers = data.atomic_numbers.long() - energy = self.energyembeddingblock(energy, atomic_numbers, atomic_numbers) - energy = scatter(energy, data.batch, dim=0) - - if self.regress_forces: - if self.force_estimator == "grad": - forces = -1 * ( - torch.autograd.grad( - energy, - data.pos, - grad_outputs=torch.ones_like(energy), - create_graph=True, - )[0] - ) - if self.force_estimator == "random": - forces = self._compute_forces_random_rotations( - x, - self.num_random_rotations, - data.atomic_numbers.long(), - edge_index, - edge_distance_vec, - data.batch, - ) - - if not self.regress_forces: - return energy - else: - return energy, forces - - def _compute_forces_random_rotations( - self, - x, - num_random_rotations, - target_element, - edge_index, - edge_distance_vec, - batch, - ): - # Compute the forces and energy by randomly rotating the system and taking the average - - device = x.device - - rot_mat_x = torch.zeros(3, 3, device=device) - rot_mat_x[0][0] = 1.0 - rot_mat_x[1][1] = 1.0 - rot_mat_x[2][2] = 1.0 - - rot_mat_y = torch.zeros(3, 3, device=device) - rot_mat_y[0][1] = 1.0 - rot_mat_y[1][0] = -1.0 - rot_mat_y[2][2] = 1.0 - - rot_mat_z = torch.zeros(3, 3, device=device) - rot_mat_z[0][2] = 1.0 - rot_mat_z[1][1] = 1.0 - rot_mat_z[2][0] = -1.0 - - rot_mat_x = rot_mat_x.view(-1, 3, 3).repeat(self.num_atoms, 1, 1) - rot_mat_y = rot_mat_y.view(-1, 3, 3).repeat(self.num_atoms, 1, 1) - rot_mat_z = rot_mat_z.view(-1, 3, 3).repeat(self.num_atoms, 1, 1) - - # compute the random rotations - random_rot_mat = self._random_rot_mat( - self.num_atoms * num_random_rotations, device - ) - random_rot_mat = random_rot_mat.view(num_random_rotations, self.num_atoms, 3, 3) - - # the first matrix is the identity with the rest being random - # atom_rot_mat = torch.cat([torch.eye(3, device=device).view(1, 1, 3, 3).repeat(1, self.num_atoms, 1, 1), random_rot_mat], dim=0) - # or they are all random - atom_rot_mat = random_rot_mat - - forces = torch.zeros(self.num_atoms, 3, device=device) - - for rot_index in range(num_random_rotations): - - rot_mat_x_perturb = torch.bmm(rot_mat_x, atom_rot_mat[rot_index]) - rot_mat_y_perturb = torch.bmm(rot_mat_y, atom_rot_mat[rot_index]) - rot_mat_z_perturb = torch.bmm(rot_mat_z, atom_rot_mat[rot_index]) - - # project neighbors using the random rotations - ( - proj_nodes_index_x, - proj_nodes_delta_x, - proj_nodes_src_index_x, - ) = self._project2D_nodes_init( - rot_mat_x_perturb, edge_index, edge_distance_vec - ) - ( - proj_nodes_index_y, - proj_nodes_delta_y, - proj_nodes_src_index_y, - ) = self._project2D_nodes_init( - rot_mat_y_perturb, edge_index, edge_distance_vec - ) - ( - proj_nodes_index_z, - proj_nodes_delta_z, - proj_nodes_src_index_z, - ) = self._project2D_nodes_init( - rot_mat_z_perturb, edge_index, edge_distance_vec - ) - - # estimate the force in each perpendicular direction - force_x = self.force_output_block( - x, - self.num_atoms, - target_element, - proj_nodes_index_x, - proj_nodes_delta_x, - proj_nodes_src_index_x, - ) - force_y = self.force_output_block( - x, - self.num_atoms, - target_element, - proj_nodes_index_y, - proj_nodes_delta_y, - proj_nodes_src_index_y, - ) - force_z = self.force_output_block( - x, - self.num_atoms, - target_element, - proj_nodes_index_z, - proj_nodes_delta_z, - proj_nodes_src_index_z, - ) - forces_perturb = torch.cat( - [force_x[:, 0:1], force_y[:, 0:1], force_z[:, 0:1]], dim=1 - ) - - # rotate the predicted forces back into the global reference frame - rot_mat_inv = torch.transpose(rot_mat_x_perturb, 1, 2) - forces_perturb = torch.bmm(rot_mat_inv, forces_perturb.view(-1, 3, 1)).view( - -1, 3 - ) - - forces = forces + forces_perturb - - forces = forces / (num_random_rotations) - - return forces - - def _filter_edges( - self, edge_index, edge_distance, edge_distance_vec, max_num_neighbors - ): - # Remove edges that aren't within the closest max_num_neighbors from either the target or source atom. - # This ensures all edges occur in pairs, i.e., if X -> Y exists then Y -> X is included. - # However, if both X -> Y and Y -> X don't both exist in the original list, this isn't guaranteed. - # Since some edges may have exactly the same distance, this function is not deterministic - device = edge_index.device - length = len(edge_distance) - - # Assuming the edges are consecutive based on the target index - target_node_index, neigh_count = torch.unique_consecutive( - edge_index[1], return_counts=True - ) - max_neighbors = torch.max(neigh_count) - - # handle special case where an atom doesn't have any neighbors - target_neigh_count = torch.zeros(self.num_atoms, device=device).long() - target_neigh_count.index_copy_(0, target_node_index.long(), neigh_count) - - # Create a list of edges for each atom - index_offset = torch.cumsum(target_neigh_count, dim=0) - target_neigh_count - neigh_index = torch.arange(length, device=device) - neigh_index = neigh_index - index_offset[edge_index[1]] - - edge_map_index = (edge_index[1] * max_neighbors + neigh_index).long() - target_lookup = ( - torch.zeros(self.num_atoms * max_neighbors, device=device) - 1 - ).long() - target_lookup.index_copy_( - 0, edge_map_index, torch.arange(length, device=device).long() - ) - - # Get the length of each edge - distance_lookup = ( - torch.zeros(self.num_atoms * max_neighbors, device=device) + 1000000.0 - ) - distance_lookup.index_copy_(0, edge_map_index, edge_distance) - distance_lookup = distance_lookup.view(self.num_atoms, max_neighbors) - - # Sort the distances - distance_sorted_no_op, indices = torch.sort(distance_lookup, dim=1) - - # Create a hash that maps edges that go from X -> Y and Y -> X in the same bin - edge_index_min, no_op = torch.min(edge_index, dim=0) - edge_index_max, no_op = torch.max(edge_index, dim=0) - edge_index_hash = edge_index_min * self.num_atoms + edge_index_max - edge_count_start = torch.zeros(self.num_atoms * self.num_atoms, device=device) - edge_count_start.index_add_( - 0, edge_index_hash, torch.ones(len(edge_index_hash), device=device) - ) - - # Find index into the original edge_index - indices = indices + ( - torch.arange(len(indices), device=device) * max_neighbors - ).view(-1, 1).repeat(1, max_neighbors) - indices = indices.view(-1) - target_lookup_sorted = ( - torch.zeros(self.num_atoms * max_neighbors, device=device) - 1 - ).long() - target_lookup_sorted = target_lookup[indices] - target_lookup_sorted = target_lookup_sorted.view(self.num_atoms, max_neighbors) - - # Select the closest max_num_neighbors for each edge and remove the unused entries - target_lookup_below_thres = ( - target_lookup_sorted[:, 0:max_num_neighbors].contiguous().view(-1) - ) - target_lookup_below_thres = target_lookup_below_thres.view(-1) - mask_unused = target_lookup_below_thres.ge(0) - target_lookup_below_thres = torch.masked_select( - target_lookup_below_thres, mask_unused - ) - - # Find edges that are used at least once and create a mask to keep - edge_count = torch.zeros(self.num_atoms * self.num_atoms, device=device) - edge_count.index_add_( - 0, - edge_index_hash[target_lookup_below_thres], - torch.ones(len(target_lookup_below_thres), device=device), - ) - edge_count_mask = edge_count.ne(0) - edge_keep = edge_count_mask[edge_index_hash] - - # Finally remove all edges that are too long in distance as indicated by the mask - edge_index_mask = edge_keep.view(1, -1).repeat(2, 1) - edge_index = torch.masked_select(edge_index, edge_index_mask).view(2, -1) - edge_distance = torch.masked_select(edge_distance, edge_keep) - edge_distance_vec_mask = edge_keep.view(-1, 1).repeat(1, 3) - edge_distance_vec = torch.masked_select( - edge_distance_vec, edge_distance_vec_mask - ).view(-1, 3) - - return edge_index, edge_distance, edge_distance_vec - - def _random_rot_mat(self, num_matrices, device): - ang_a = 2.0 * math.pi * torch.rand(num_matrices, device=device) - ang_b = 2.0 * math.pi * torch.rand(num_matrices, device=device) - ang_c = 2.0 * math.pi * torch.rand(num_matrices, device=device) - - cos_a = torch.cos(ang_a) - cos_b = torch.cos(ang_b) - cos_c = torch.cos(ang_c) - sin_a = torch.sin(ang_a) - sin_b = torch.sin(ang_b) - sin_c = torch.sin(ang_c) - - rot_a = torch.eye(3, device=device).view(1, 3, 3).repeat(num_matrices, 1, 1) - rot_b = torch.eye(3, device=device).view(1, 3, 3).repeat(num_matrices, 1, 1) - rot_c = torch.eye(3, device=device).view(1, 3, 3).repeat(num_matrices, 1, 1) - - rot_a[:, 1, 1] = cos_a - rot_a[:, 1, 2] = sin_a - rot_a[:, 2, 1] = -sin_a - rot_a[:, 2, 2] = cos_a - - rot_b[:, 0, 0] = cos_b - rot_b[:, 0, 2] = -sin_b - rot_b[:, 2, 0] = sin_b - rot_b[:, 2, 2] = cos_b - - rot_c[:, 0, 0] = cos_c - rot_c[:, 0, 1] = sin_c - rot_c[:, 1, 0] = -sin_c - rot_c[:, 1, 1] = cos_c - - return torch.bmm(torch.bmm(rot_a, rot_b), rot_c) - - def _init_edge_rot_mat(self, data, edge_index, edge_distance_vec): - device = data.pos.device - num_atoms = len(data.batch) - - edge_vec_0 = edge_distance_vec - edge_vec_0_distance = torch.sqrt(torch.sum(edge_vec_0**2, dim=1)) - - if torch.min(edge_vec_0_distance) < 0.0001: - print( - "Error edge_vec_0_distance: {}".format(torch.min(edge_vec_0_distance)) - ) - (minval, minidx) = torch.min(edge_vec_0_distance, 0) - print( - "Error edge_vec_0_distance: {} {} {} {} {}".format( - minidx, - edge_index[0, minidx], - edge_index[1, minidx], - data.pos[edge_index[0, minidx]], - data.pos[edge_index[1, minidx]], - ) - ) - - avg_vector = torch.zeros(num_atoms, 3, device=device) - weight = 0.5 * (torch.cos(edge_vec_0_distance * PI / self.cutoff) + 1.0) - avg_vector.index_add_( - 0, edge_index[1, :], edge_vec_0 * weight.view(-1, 1).expand(-1, 3) - ) - - edge_vec_2 = avg_vector[edge_index[1, :]] + 0.0001 - edge_vec_2_distance = torch.sqrt(torch.sum(edge_vec_2**2, dim=1)) - - if torch.min(edge_vec_2_distance) < 0.000001: - print( - "Error edge_vec_2_distance: {}".format(torch.min(edge_vec_2_distance)) - ) - - norm_x = edge_vec_0 / (edge_vec_0_distance.view(-1, 1)) - norm_0_2 = edge_vec_2 / (edge_vec_2_distance.view(-1, 1)) - norm_z = torch.cross(norm_x, norm_0_2, dim=1) - norm_z = norm_z / ( - torch.sqrt(torch.sum(norm_z**2, dim=1, keepdim=True)) + 0.0000001 - ) - norm_y = torch.cross(norm_x, norm_z, dim=1) - norm_y = norm_y / ( - torch.sqrt(torch.sum(norm_y**2, dim=1, keepdim=True)) + 0.0000001 - ) - - norm_x = norm_x.view(-1, 3, 1) - norm_y = norm_y.view(-1, 3, 1) - norm_z = norm_z.view(-1, 3, 1) - - edge_rot_mat_inv = torch.cat([norm_x, norm_y, norm_z], dim=2) - edge_rot_mat = torch.transpose(edge_rot_mat_inv, 1, 2) - - return edge_rot_mat - - def _project2D_edges_init(self, rot_mat, edge_index, edge_distance_vec): - torch.set_printoptions(sci_mode=False) - length = len(edge_distance_vec) - device = edge_distance_vec.device - - # Assuming the edges are consecutive based on the target index - target_node_index, neigh_count = torch.unique_consecutive( - edge_index[1], return_counts=True - ) - max_neighbors = torch.max(neigh_count) - target_neigh_count = torch.zeros(self.num_atoms, device=device).long() - target_neigh_count.index_copy_(0, target_node_index.long(), neigh_count) - - index_offset = torch.cumsum(target_neigh_count, dim=0) - target_neigh_count - neigh_index = torch.arange(length, device=device) - neigh_index = neigh_index - index_offset[edge_index[1]] - - edge_map_index = edge_index[1] * max_neighbors + neigh_index - target_lookup = ( - torch.zeros(self.num_atoms * max_neighbors, device=device) - 1 - ).long() - target_lookup.index_copy_( - 0, - edge_map_index.long(), - torch.arange(length, device=device).long(), - ) - target_lookup = target_lookup.view(self.num_atoms, max_neighbors) - - # target_lookup - For each target node, a list of edge indices - # target_neigh_count - number of neighbors for each target node - source_edge = target_lookup[edge_index[0]] - target_edge = ( - torch.arange(length, device=device) - .long() - .view(-1, 1) - .repeat(1, max_neighbors) - ) - - source_edge = source_edge.view(-1) - target_edge = target_edge.view(-1) - - mask_unused = source_edge.ge(0) - source_edge = torch.masked_select(source_edge, mask_unused) - target_edge = torch.masked_select(target_edge, mask_unused) - - return self._project2D_init( - source_edge, target_edge, rot_mat, edge_distance_vec - ) - - def _project2D_nodes_init(self, rot_mat, edge_index, edge_distance_vec): - torch.set_printoptions(sci_mode=False) - length = len(edge_distance_vec) - device = edge_distance_vec.device - - target_node = edge_index[1] - source_edge = torch.arange(length, device=device) - - return self._project2D_init( - source_edge, target_node, rot_mat, edge_distance_vec - ) - - def _project2D_init(self, source_edge, target_edge, rot_mat, edge_distance_vec): - edge_distance_norm = F.normalize(edge_distance_vec) - source_edge_offset = edge_distance_norm[source_edge] - - source_edge_offset_rot = torch.bmm( - rot_mat[target_edge], source_edge_offset.view(-1, 3, 1) - ) - - source_edge_X = torch.atan2( - source_edge_offset_rot[:, 1], source_edge_offset_rot[:, 2] - ).view(-1) - - # source_edge_X ranges from -pi to pi - source_edge_X = (source_edge_X + math.pi) / (2.0 * math.pi) - - # source_edge_Y ranges from -1 to 1 - source_edge_Y = source_edge_offset_rot[:, 0].view(-1) - source_edge_Y = torch.clamp(source_edge_Y, min=-1.0, max=1.0) - source_edge_Y = (source_edge_Y.asin() + (math.pi / 2.0)) / ( - math.pi - ) # bin by angle - # source_edge_Y = (source_edge_Y + 1.0) / 2.0 # bin by sin - source_edge_Y = 0.99 * (source_edge_Y) + 0.005 - - source_edge_X = source_edge_X * self.sphere_size_long - source_edge_Y = source_edge_Y * ( - self.sphere_size_lat - 1.0 - ) # not circular so pad by one - - source_edge_X_0 = torch.floor(source_edge_X).long() - source_edge_X_del = source_edge_X - source_edge_X_0 - source_edge_X_0 = source_edge_X_0 % self.sphere_size_long - source_edge_X_1 = (source_edge_X_0 + 1) % self.sphere_size_long - - source_edge_Y_0 = torch.floor(source_edge_Y).long() - source_edge_Y_del = source_edge_Y - source_edge_Y_0 - source_edge_Y_0 = source_edge_Y_0 % self.sphere_size_lat - source_edge_Y_1 = (source_edge_Y_0 + 1) % self.sphere_size_lat - - # Compute the values needed to bilinearly splat the values onto the spheres - index_0_0 = ( - target_edge * self.sphere_size_lat * self.sphere_size_long - + source_edge_Y_0 * self.sphere_size_long - + source_edge_X_0 - ) - index_0_1 = ( - target_edge * self.sphere_size_lat * self.sphere_size_long - + source_edge_Y_0 * self.sphere_size_long - + source_edge_X_1 - ) - index_1_0 = ( - target_edge * self.sphere_size_lat * self.sphere_size_long - + source_edge_Y_1 * self.sphere_size_long - + source_edge_X_0 - ) - index_1_1 = ( - target_edge * self.sphere_size_lat * self.sphere_size_long - + source_edge_Y_1 * self.sphere_size_long - + source_edge_X_1 - ) - - delta_0_0 = (1.0 - source_edge_X_del) * (1.0 - source_edge_Y_del) - delta_0_1 = (source_edge_X_del) * (1.0 - source_edge_Y_del) - delta_1_0 = (1.0 - source_edge_X_del) * (source_edge_Y_del) - delta_1_1 = (source_edge_X_del) * (source_edge_Y_del) - - index_0_0 = index_0_0.view(1, -1) - index_0_1 = index_0_1.view(1, -1) - index_1_0 = index_1_0.view(1, -1) - index_1_1 = index_1_1.view(1, -1) - - # NaNs otherwise - if self.grad_forces: - with torch.no_grad(): - delta_0_0 = delta_0_0.view(1, -1) - delta_0_1 = delta_0_1.view(1, -1) - delta_1_0 = delta_1_0.view(1, -1) - delta_1_1 = delta_1_1.view(1, -1) - else: - delta_0_0 = delta_0_0.view(1, -1) - delta_0_1 = delta_0_1.view(1, -1) - delta_1_0 = delta_1_0.view(1, -1) - delta_1_1 = delta_1_1.view(1, -1) - - return ( - torch.cat([index_0_0, index_0_1, index_1_0, index_1_1]), - torch.cat([delta_0_0, delta_0_1, delta_1_0, delta_1_1]), - source_edge, - ) - - @property - def num_params(self): - return sum(p.numel() for p in self.parameters()) - - -class MessageBlock(torch.nn.Module): - def __init__( - self, - in_hidden_channels, - out_hidden_channels, - mid_hidden_channels, - embedding_size, - sphere_size_lat, - sphere_size_long, - max_num_elements, - sphere_message, - act, - lmax, - ): - super(MessageBlock, self).__init__() - self.in_hidden_channels = in_hidden_channels - self.out_hidden_channels = out_hidden_channels - self.act = act - self.lmax = lmax - self.embedding_size = embedding_size - self.mid_hidden_channels = mid_hidden_channels - self.sphere_size_lat = sphere_size_lat - self.sphere_size_long = sphere_size_long - self.sphere_message = sphere_message - self.max_num_elements = max_num_elements - self.num_embedding_basis = 8 - - self.spinconvblock = SpinConvBlock( - self.in_hidden_channels, - self.mid_hidden_channels, - self.sphere_size_lat, - self.sphere_size_long, - self.sphere_message, - self.act, - self.lmax, - ) - - self.embeddingblock1 = EmbeddingBlock( - self.mid_hidden_channels, - self.mid_hidden_channels, - self.mid_hidden_channels, - self.embedding_size, - self.num_embedding_basis, - self.max_num_elements, - self.act, - ) - self.embeddingblock2 = EmbeddingBlock( - self.mid_hidden_channels, - self.out_hidden_channels, - self.mid_hidden_channels, - self.embedding_size, - self.num_embedding_basis, - self.max_num_elements, - self.act, - ) - - self.distfc1 = nn.Linear(self.mid_hidden_channels, self.mid_hidden_channels) - self.distfc2 = nn.Linear(self.mid_hidden_channels, self.mid_hidden_channels) - - def forward( - self, - x, - x_dist, - source_element, - target_element, - proj_index, - proj_delta, - proj_src_index, - ): - out_size = len(x) - - x = self.spinconvblock(x, out_size, proj_index, proj_delta, proj_src_index) - - x = self.embeddingblock1(x, source_element, target_element) - - x_dist = self.distfc1(x_dist) - x_dist = self.act(x_dist) - x_dist = self.distfc2(x_dist) - x = x + x_dist - - x = self.act(x) - x = self.embeddingblock2(x, source_element, target_element) - - return x - - -class ForceOutputBlock(torch.nn.Module): - def __init__( - self, - in_hidden_channels, - out_hidden_channels, - mid_hidden_channels, - embedding_size, - sphere_size_lat, - sphere_size_long, - max_num_elements, - sphere_message, - act, - lmax, - ): - super(ForceOutputBlock, self).__init__() - self.in_hidden_channels = in_hidden_channels - self.out_hidden_channels = out_hidden_channels - self.act = act - self.lmax = lmax - self.embedding_size = embedding_size - self.mid_hidden_channels = mid_hidden_channels - self.sphere_size_lat = sphere_size_lat - self.sphere_size_long = sphere_size_long - self.sphere_message = sphere_message - self.max_num_elements = max_num_elements - self.num_embedding_basis = 8 - - self.spinconvblock = SpinConvBlock( - self.in_hidden_channels, - self.mid_hidden_channels, - self.sphere_size_lat, - self.sphere_size_long, - self.sphere_message, - self.act, - self.lmax, - ) - - self.block1 = EmbeddingBlock( - self.mid_hidden_channels, - self.mid_hidden_channels, - self.mid_hidden_channels, - self.embedding_size, - self.num_embedding_basis, - self.max_num_elements, - self.act, - ) - self.block2 = EmbeddingBlock( - self.mid_hidden_channels, - self.out_hidden_channels, - self.mid_hidden_channels, - self.embedding_size, - self.num_embedding_basis, - self.max_num_elements, - self.act, - ) - - def forward( - self, - x, - out_size, - target_element, - proj_index, - proj_delta, - proj_src_index, - ): - x = self.spinconvblock(x, out_size, proj_index, proj_delta, proj_src_index) - - x = self.block1(x, target_element, target_element) - x = self.act(x) - x = self.block2(x, target_element, target_element) - - return x - - -class SpinConvBlock(torch.nn.Module): - def __init__( - self, - in_hidden_channels, - mid_hidden_channels, - sphere_size_lat, - sphere_size_long, - sphere_message, - act, - lmax, - ): - super(SpinConvBlock, self).__init__() - self.in_hidden_channels = in_hidden_channels - self.mid_hidden_channels = mid_hidden_channels - self.sphere_size_lat = sphere_size_lat - self.sphere_size_long = sphere_size_long - self.sphere_message = sphere_message - self.act = act - self.lmax = lmax - self.num_groups = self.in_hidden_channels // 8 - - self.ProjectLatLongSphere = ProjectLatLongSphere( - sphere_size_lat, sphere_size_long - ) - assert self.sphere_message in [ - "fullconv", - "rotspharmwd", - ] - if self.sphere_message in ["rotspharmwd"]: - self.sph_froms2grid = FromS2Grid( - (self.sphere_size_lat, self.sphere_size_long), self.lmax - ) - self.mlp = nn.Linear( - self.in_hidden_channels * (self.lmax + 1) ** 2, - self.mid_hidden_channels, - ) - self.sphlength = (self.lmax + 1) ** 2 - rotx = torch.zeros(self.sphere_size_long) + ( - 2 * math.pi / self.sphere_size_long - ) - roty = torch.zeros(self.sphere_size_long) - rotz = torch.zeros(self.sphere_size_long) - - self.wigner = [] - for xrot, yrot, zrot in zip(rotx, roty, rotz): - _blocks = [] - for l_degree in range(self.lmax + 1): - _blocks.append(o3.wigner_D(l_degree, xrot, yrot, zrot)) - self.wigner.append(torch.block_diag(*_blocks)) - - if self.sphere_message == "fullconv": - padding = self.sphere_size_long // 2 - self.conv1 = nn.Conv1d( - self.in_hidden_channels * self.sphere_size_lat, - self.mid_hidden_channels, - self.sphere_size_long, - groups=self.in_hidden_channels // 8, - padding=padding, - padding_mode="circular", - ) - self.pool = nn.AvgPool1d(sphere_size_long) - - self.GroupNorm = nn.GroupNorm(self.num_groups, self.mid_hidden_channels) - - def forward(self, x, out_size, proj_index, proj_delta, proj_src_index): - x = self.ProjectLatLongSphere( - x, out_size, proj_index, proj_delta, proj_src_index - ) - if self.sphere_message == "rotspharmwd": - sph_harm_calc = torch.zeros( - ((x.shape[0], self.mid_hidden_channels)), - device=x.device, - ) - - sph_harm = self.sph_froms2grid(x) - sph_harm = sph_harm.view(-1, self.sphlength, 1) - for wD_diag in self.wigner: - wD_diag = wD_diag.to(x.device) - sph_harm_calc += self.act(self.mlp(sph_harm.reshape(x.shape[0], -1))) - wd = wD_diag.view(1, self.sphlength, self.sphlength).expand( - len(x) * self.in_hidden_channels, -1, -1 - ) - sph_harm = torch.bmm(wd, sph_harm) - x = sph_harm_calc - - if self.sphere_message in ["fullconv"]: - x = x.view( - -1, - self.in_hidden_channels * self.sphere_size_lat, - self.sphere_size_long, - ) - x = self.conv1(x) - x = self.act(x) - # Pool in the longitudal direction - x = self.pool(x[:, :, 0 : self.sphere_size_long]) - x = x.view(out_size, -1) - - x = self.GroupNorm(x) - - return x - - -class EmbeddingBlock(torch.nn.Module): - def __init__( - self, - in_hidden_channels, - out_hidden_channels, - mid_hidden_channels, - embedding_size, - num_embedding_basis, - max_num_elements, - act, - ): - super(EmbeddingBlock, self).__init__() - self.in_hidden_channels = in_hidden_channels - self.out_hidden_channels = out_hidden_channels - self.act = act - self.embedding_size = embedding_size - self.mid_hidden_channels = mid_hidden_channels - self.num_embedding_basis = num_embedding_basis - self.max_num_elements = max_num_elements - - self.fc1 = nn.Linear(self.in_hidden_channels, self.mid_hidden_channels) - self.fc2 = nn.Linear( - self.mid_hidden_channels, - self.num_embedding_basis * self.mid_hidden_channels, - ) - self.fc3 = nn.Linear(self.mid_hidden_channels, self.out_hidden_channels) - - self.source_embedding = nn.Embedding(max_num_elements, self.embedding_size) - self.target_embedding = nn.Embedding(max_num_elements, self.embedding_size) - nn.init.uniform_(self.source_embedding.weight.data, -0.0001, 0.0001) - nn.init.uniform_(self.target_embedding.weight.data, -0.0001, 0.0001) - - self.embed_fc1 = nn.Linear(2 * self.embedding_size, self.num_embedding_basis) - - self.softmax = nn.Softmax(dim=1) - - def forward(self, x, source_element, target_element): - source_embedding = self.source_embedding(source_element) - target_embedding = self.target_embedding(target_element) - embedding = torch.cat([source_embedding, target_embedding], dim=1) - embedding = self.embed_fc1(embedding) - embedding = self.softmax(embedding) - - x = self.fc1(x) - x = self.act(x) - x = self.fc2(x) - x = self.act(x) - x = (x.view(-1, self.num_embedding_basis, self.mid_hidden_channels)) * ( - embedding.view(-1, self.num_embedding_basis, 1) - ) - x = torch.sum(x, dim=1) - x = self.fc3(x) - - return x - - -class DistanceBlock(torch.nn.Module): - def __init__( - self, - in_channels, - out_channels, - max_num_elements, - scalar_max, - distance_expansion, - scale_distances, - ): - super(DistanceBlock, self).__init__() - self.in_channels = in_channels - self.out_channels = out_channels - self.max_num_elements = max_num_elements - self.distance_expansion = distance_expansion - self.scalar_max = scalar_max - self.scale_distances = scale_distances - - if self.scale_distances: - self.dist_scalar = nn.Embedding( - self.max_num_elements * self.max_num_elements, 1 - ) - self.dist_offset = nn.Embedding( - self.max_num_elements * self.max_num_elements, 1 - ) - nn.init.uniform_(self.dist_scalar.weight.data, -0.0001, 0.0001) - nn.init.uniform_(self.dist_offset.weight.data, -0.0001, 0.0001) - - self.fc1 = nn.Linear(self.in_channels, self.out_channels) - - def forward(self, edge_distance, source_element, target_element): - if self.scale_distances: - embedding_index = source_element * self.max_num_elements + target_element - - # Restrict the scalar to range from 1 / self.scalar_max to self.scalar_max - scalar_max = math.log(self.scalar_max) - scalar = ( - 2.0 * torch.sigmoid(self.dist_scalar(embedding_index).view(-1)) - 1.0 - ) - scalar = torch.exp(scalar_max * scalar) - offset = self.dist_offset(embedding_index).view(-1) - x = self.distance_expansion(scalar * edge_distance + offset) - else: - x = self.distance_expansion(edge_distance) - - x = self.fc1(x) - - return x - - -class ProjectLatLongSphere(torch.nn.Module): - def __init__(self, sphere_size_lat, sphere_size_long): - super(ProjectLatLongSphere, self).__init__() - self.sphere_size_lat = sphere_size_lat - self.sphere_size_long = sphere_size_long - - def forward(self, x, length, index, delta, source_edge_index): - device = x.device - hidden_channels = len(x[0]) - - x_proj = torch.zeros( - length * self.sphere_size_lat * self.sphere_size_long, - hidden_channels, - device=device, - ) - splat_values = x[source_edge_index] - - # Perform bilinear splatting - x_proj.index_add_(0, index[0], splat_values * (delta[0].view(-1, 1))) - x_proj.index_add_(0, index[1], splat_values * (delta[1].view(-1, 1))) - x_proj.index_add_(0, index[2], splat_values * (delta[2].view(-1, 1))) - x_proj.index_add_(0, index[3], splat_values * (delta[3].view(-1, 1))) - - x_proj = x_proj.view( - length, - self.sphere_size_lat * self.sphere_size_long, - hidden_channels, - ) - x_proj = torch.transpose(x_proj, 1, 2).contiguous() - x_proj = x_proj.view( - length, - hidden_channels, - self.sphere_size_lat, - self.sphere_size_long, - ) - - return x_proj - - -class Swish(torch.nn.Module): - def __init__(self): - super(Swish, self).__init__() - - def forward(self, x): - return x * torch.sigmoid(x) - - -class GaussianSmearing(torch.nn.Module): - def __init__(self, start=-5.0, stop=5.0, num_gaussians=50, basis_width_scalar=1.0): - super(GaussianSmearing, self).__init__() - offset = torch.linspace(start, stop, num_gaussians) - self.coeff = -0.5 / (basis_width_scalar * (offset[1] - offset[0])).item() ** 2 - self.register_buffer("offset", offset) - - def forward(self, dist): - dist = dist.view(-1, 1) - self.offset.view(1, -1) - return torch.exp(self.coeff * torch.pow(dist, 2)) diff --git a/ocpmodels/models/utils/__init__.py b/ocpmodels/models/utils/__init__.py deleted file mode 100644 index 6264236915a7269a4d920ee8213004374dd86a9a..0000000000000000000000000000000000000000 --- a/ocpmodels/models/utils/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -# Copyright (c) Facebook, Inc. and its affiliates. -# -# This source code is licensed under the MIT license found in the -# LICENSE file in the root directory of this source tree. diff --git a/ocpmodels/models/utils/activations.py b/ocpmodels/models/utils/activations.py deleted file mode 100644 index 256830acbce5128acf2c3a7a6c8e31862540eee1..0000000000000000000000000000000000000000 --- a/ocpmodels/models/utils/activations.py +++ /dev/null @@ -1,43 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import torch -import torch.nn.functional as F - - -class Act(torch.nn.Module): - def __init__(self, act, slope=0.05): - super(Act, self).__init__() - self.act = act - self.slope = slope - self.shift = torch.log(torch.tensor(2.0)).item() - - def forward(self, input): - if self.act == "relu": - return F.relu(input) - elif self.act == "leaky_relu": - return F.leaky_relu(input) - elif self.act == "sp": - return F.softplus(input, beta=1) - elif self.act == "leaky_sp": - return F.softplus(input, beta=1) - self.slope * F.relu(-input) - elif self.act == "elu": - return F.elu(input, alpha=1) - elif self.act == "leaky_elu": - return F.elu(input, alpha=1) - self.slope * F.relu(-input) - elif self.act == "ssp": - return F.softplus(input, beta=1) - self.shift - elif self.act == "leaky_ssp": - return F.softplus(input, beta=1) - self.slope * F.relu(-input) - self.shift - elif self.act == "tanh": - return torch.tanh(input) - elif self.act == "leaky_tanh": - return torch.tanh(input) + self.slope * input - elif self.act == "swish": - return torch.sigmoid(input) * input - else: - raise RuntimeError(f"Undefined activation called {self.act}") diff --git a/ocpmodels/models/utils/basis.py b/ocpmodels/models/utils/basis.py deleted file mode 100644 index 27d24ab73b65b1550a9f866697ac3dc46a81b40f..0000000000000000000000000000000000000000 --- a/ocpmodels/models/utils/basis.py +++ /dev/null @@ -1,281 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import math -from math import pi as PI -from typing import List - -import numpy as np -import torch -import torch.nn as nn -from scipy.special import sph_harm -from torch.nn.init import _calculate_correct_fan - -from .activations import Act - - -class Sine(nn.Module): - def __init__(self, w0: float = 30.0): - - super(Sine, self).__init__() - self.w0 = w0 - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return torch.sin(self.w0 * x) - - -class SIREN(nn.Module): - def __init__( - self, - layers: List[int], - in_features: int, - out_features: int, - w0: float = 30.0, - initializer: str = "siren", - c: float = 6, - ): - - super(SIREN, self).__init__() - self.layers = [nn.Linear(in_features, layers[0]), Sine(w0=w0)] - - for index in range(len(layers) - 1): - self.layers.extend( - [nn.Linear(layers[index], layers[index + 1]), Sine(w0=1)] - ) - - self.layers.append(nn.Linear(layers[-1], out_features)) - self.network = nn.Sequential(*self.layers) - - if initializer is not None and initializer == "siren": - for m in self.network: - if isinstance(m, nn.Linear): - num_input = float(m.weight.size(-1)) - with torch.no_grad(): - m.weight.uniform_( - -math.sqrt(6.0 / num_input), - math.sqrt(6.0 / num_input), - ) - - def forward(self, X): - return self.network(X) - - -class SINESmearing(nn.Module): - def __init__(self, in_features, num_freqs=40, use_cosine=False): - - super(SINESmearing, self).__init__() - - self.num_freqs = num_freqs - self.out_dim = in_features * self.num_freqs - self.use_cosine = use_cosine - - freq = torch.arange(num_freqs).float() - freq = torch.pow(torch.ones_like(freq) * 1.1, freq) - self.freq_filter = nn.Parameter( - freq.view(-1, 1).repeat(1, in_features).view(1, -1), - requires_grad=False, - ) - - def forward(self, x): - x = x.repeat(1, self.num_freqs) - x = x * self.freq_filter - - if self.use_cosine: - return torch.cos(x) - else: - return torch.sin(x) - - -class GaussianSmearing(nn.Module): - def __init__(self, in_features, start=0, end=1, num_freqs=50): - super(GaussianSmearing, self).__init__() - self.num_freqs = num_freqs - offset = torch.linspace(start, end, num_freqs) - self.coeff = -0.5 / (offset[1] - offset[0]).item() ** 2 - self.offset = nn.Parameter( - offset.view(-1, 1).repeat(1, in_features).view(1, -1), - requires_grad=False, - ) - - def forward(self, x): - x = x.repeat(1, self.num_freqs) - x = x - self.offset - return torch.exp(self.coeff * torch.pow(x, 2)) - - -class FourierSmearing(nn.Module): - def __init__(self, in_features, num_freqs=40, use_cosine=False): - - super(FourierSmearing, self).__init__() - - self.num_freqs = num_freqs - self.out_dim = in_features * self.num_freqs - self.use_cosine = use_cosine - - freq = torch.arange(num_freqs).to(torch.float32) - self.freq_filter = nn.Parameter( - freq.view(-1, 1).repeat(1, in_features).view(1, -1), - requires_grad=False, - ) - - def forward(self, x): - x = x.repeat(1, self.num_freqs) - x = x * self.freq_filter - - if self.use_cosine: - return torch.cos(x) - else: - return torch.sin(x) - - -class Basis(nn.Module): - def __init__( - self, - in_features, - num_freqs=50, - basis_type="powersine", - act="ssp", - sph=None, - ): - super(Basis, self).__init__() - - self.num_freqs = num_freqs - self.basis_type = basis_type - - if basis_type == "powersine": - self.smearing = SINESmearing(in_features, num_freqs) - self.out_dim = in_features * num_freqs - elif basis_type == "powercosine": - self.smearing = SINESmearing(in_features, num_freqs, use_cosine=True) - self.out_dim = in_features * num_freqs - elif basis_type == "fouriersine": - self.smearing = FourierSmearing(in_features, num_freqs) - self.out_dim = in_features * num_freqs - elif basis_type == "gauss": - self.smearing = GaussianSmearing( - in_features, start=0, end=1, num_freqs=num_freqs - ) - self.out_dim = in_features * num_freqs - elif basis_type == "linact": - self.smearing = torch.nn.Sequential( - torch.nn.Linear(in_features, num_freqs * in_features), Act(act) - ) - self.out_dim = in_features * num_freqs - elif basis_type == "raw" or basis_type == "rawcat": - self.out_dim = in_features - elif "sph" in basis_type: - # by default, we use sine function to encode distance - # sph must be given here - assert sph is not None - # assumes the first three columns are normalizaed xyz - # the rest of the columns are distances - if "cat" in basis_type: - # concatenate - self.smearing_sine = SINESmearing(in_features - 3, num_freqs) - self.out_dim = sph.out_dim + (in_features - 3) * num_freqs - elif "mul" in basis_type: - self.smearing_sine = SINESmearing(in_features - 3, num_freqs) - self.lin = torch.nn.Linear(self.smearing_sine.out_dim, in_features - 3) - self.out_dim = (in_features - 3) * sph.out_dim - elif "m40" in basis_type: - dim = 40 - self.smearing_sine = SINESmearing(in_features - 3, num_freqs) - self.lin = torch.nn.Linear( - self.smearing_sine.out_dim, dim - ) # make the output dimensionality comparable. - self.out_dim = dim * sph.out_dim - elif "nosine" in basis_type: - # does not use sine smearing for encoding distance - self.out_dim = (in_features - 3) * sph.out_dim - else: - raise ValueError("cat or mul not specified for spherical harnomics.") - else: - raise RuntimeError("Undefined basis type.") - - def forward(self, x, edge_attr_sph=None): - if "sph" in self.basis_type: - if "nosine" not in self.basis_type: - x_sine = self.smearing_sine( - x[:, 3:] - ) # the first three features correspond to edge_vec_normalized, so we ignore - if "cat" in self.basis_type: - # just concatenate spherical edge feature and sined node features - return torch.cat([edge_attr_sph, x_sine], dim=1) - elif "mul" in self.basis_type or "m40" in self.basis_type: - # multiply sined node features into spherical edge feature (inspired by theory in spherical harmonics) - r = self.lin(x_sine) - outer = torch.einsum("ik,ij->ikj", edge_attr_sph, r) - return torch.flatten(outer, start_dim=1) - else: - raise RuntimeError(f"Unknown basis type called {self.basis_type}") - else: - outer = torch.einsum("ik,ij->ikj", edge_attr_sph, x[:, 3:]) - return torch.flatten(outer, start_dim=1) - - elif "raw" in self.basis_type: - # do nothing, just return node features - pass - else: - x = self.smearing(x) - return x - - -class SphericalSmearing(nn.Module): - def __init__(self, max_n=10, option="all"): - super(SphericalSmearing, self).__init__() - - self.max_n = max_n - - m = [] - n = [] - for i in range(max_n): - for j in range(0, i + 1): - n.append(i) - m.append(j) - - m = np.array(m) - n = np.array(n) - - if option == "all": - self.m = m - self.n = n - elif option == "sine": - self.m = m[n % 2 == 1] - self.n = n[n % 2 == 1] - elif option == "cosine": - self.m = m[n % 2 == 0] - self.n = n[n % 2 == 0] - - self.out_dim = int(np.sum(self.m == 0) + 2 * np.sum(self.m != 0)) - - def forward(self, xyz): - # assuming input is already normalized - assert xyz.size(1) == 3 - - xyz = xyz / xyz.norm(dim=-1).view(-1, 1) - - phi = torch.acos(xyz[:, 2]) - theta = torch.atan2(-xyz[:, 1], -xyz[:, 0]) + math.pi - - phi = phi.cpu().numpy() - theta = theta.cpu().numpy() - - m_tile = np.tile(self.m, (len(xyz), 1)) - n_tile = np.tile(self.n, (len(xyz), 1)) - theta_tile = np.tile(theta.reshape(len(xyz), 1), (1, len(self.m))) - phi_tile = np.tile(phi.reshape(len(xyz), 1), (1, len(self.m))) - - harm = sph_harm(m_tile, n_tile, theta_tile, phi_tile) - - harm_mzero = harm[:, self.m == 0] - harm_mnonzero = harm[:, self.m != 0] - - harm_real = np.concatenate( - [harm_mzero.real, harm_mnonzero.real, harm_mnonzero.imag], axis=1 - ) - - return torch.from_numpy(harm_real).to(torch.float32).to(xyz.device) diff --git a/ocpmodels/modules/__init__.py b/ocpmodels/modules/__init__.py deleted file mode 100644 index c17674b3f0bcb88e2dd7ad3faeb589434b1356a5..0000000000000000000000000000000000000000 --- a/ocpmodels/modules/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" diff --git a/ocpmodels/modules/evaluator.py b/ocpmodels/modules/evaluator.py deleted file mode 100644 index 365e28a7bda79807917a67a9a4c1c0c2eca4b61d..0000000000000000000000000000000000000000 --- a/ocpmodels/modules/evaluator.py +++ /dev/null @@ -1,292 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import numpy as np -import torch - - -""" -An evaluation module for use with the OCP dataset and suite of tasks. It should -be possible to import this independently of the rest of the codebase, e.g: - -``` -from ocpmodels.modules import Evaluator - -evaluator = Evaluator(task="is2re") -perf = evaluator.eval(prediction, target) -``` - -task: "s2ef", "is2rs", "is2re". - -We specify a default set of metrics for each task, but should be easy to extend -to add more metrics. `evaluator.eval` takes as input two dictionaries, one for -predictions and another for targets to check against. It returns a dictionary -with the relevant metrics computed. -""" - - -class Evaluator: - task_metrics = { - "s2ef": [ - "forcesx_mae", - "forcesy_mae", - "forcesz_mae", - "forces_mae", - "forces_cos", - "forces_magnitude", - "energy_mae", - "energy_force_within_threshold", - ], - "is2rs": [ - "average_distance_within_threshold", - "positions_mae", - "positions_mse", - ], - "is2re": ["energy_mae", "energy_mse", "energy_within_threshold"], - } - - task_attributes = { - "s2ef": ["energy", "forces", "natoms"], - "is2rs": ["positions", "cell", "pbc", "natoms"], - "is2re": ["energy"], - } - - task_primary_metric = { - "s2ef": "energy_force_within_threshold", - "is2rs": "average_distance_within_threshold", - "is2re": "energy_mae", - } - - def __init__(self, task=None): - assert task in ["s2ef", "is2rs", "is2re"] - self.task = task - self.metric_fn = self.task_metrics[task] - - def eval(self, prediction, target, prev_metrics={}): - for attr in self.task_attributes[self.task]: - assert attr in prediction - assert attr in target - assert prediction[attr].shape == target[attr].shape - - metrics = prev_metrics - - for fn in self.task_metrics[self.task]: - res = eval(fn)(prediction, target) - metrics = self.update(fn, res, metrics) - - return metrics - - def update(self, key, stat, metrics): - if key not in metrics: - metrics[key] = { - "metric": None, - "total": 0, - "numel": 0, - } - - if isinstance(stat, dict): - # If dictionary, we expect it to have `metric`, `total`, `numel`. - metrics[key]["total"] += stat["total"] - metrics[key]["numel"] += stat["numel"] - metrics[key]["metric"] = metrics[key]["total"] / metrics[key]["numel"] - elif isinstance(stat, float) or isinstance(stat, int): - # If float or int, just add to the total and increment numel by 1. - metrics[key]["total"] += stat - metrics[key]["numel"] += 1 - metrics[key]["metric"] = metrics[key]["total"] / metrics[key]["numel"] - elif torch.is_tensor(stat): - raise NotImplementedError - - return metrics - - -def energy_mae(prediction, target): - return absolute_error(prediction["energy"], target["energy"]) - - -def energy_mse(prediction, target): - return squared_error(prediction["energy"], target["energy"]) - - -def forcesx_mae(prediction, target): - return absolute_error(prediction["forces"][:, 0], target["forces"][:, 0]) - - -def forcesx_mse(prediction, target): - return squared_error(prediction["forces"][:, 0], target["forces"][:, 0]) - - -def forcesy_mae(prediction, target): - return absolute_error(prediction["forces"][:, 1], target["forces"][:, 1]) - - -def forcesy_mse(prediction, target): - return squared_error(prediction["forces"][:, 1], target["forces"][:, 1]) - - -def forcesz_mae(prediction, target): - return absolute_error(prediction["forces"][:, 2], target["forces"][:, 2]) - - -def forcesz_mse(prediction, target): - return squared_error(prediction["forces"][:, 2], target["forces"][:, 2]) - - -def forces_mae(prediction, target): - return absolute_error(prediction["forces"], target["forces"]) - - -def forces_mse(prediction, target): - return squared_error(prediction["forces"], target["forces"]) - - -def forces_cos(prediction, target): - return cosine_similarity(prediction["forces"], target["forces"]) - - -def forces_magnitude(prediction, target): - return magnitude_error(prediction["forces"], target["forces"], p=2) - - -def positions_mae(prediction, target): - return absolute_error(prediction["positions"], target["positions"]) - - -def positions_mse(prediction, target): - return squared_error(prediction["positions"], target["positions"]) - - -def energy_force_within_threshold(prediction, target): - # Note that this natoms should be the count of free atoms we evaluate over. - assert target["natoms"].sum() == prediction["forces"].size(0) - assert target["natoms"].size(0) == prediction["energy"].size(0) - - # compute absolute error on per-atom forces and energy per system. - # then count the no. of systems where max force error is < 0.03 and max - # energy error is < 0.02. - f_thresh = 0.03 - e_thresh = 0.02 - - success, total = 0.0, target["natoms"].size(0) - - error_forces = torch.abs(target["forces"] - prediction["forces"]) - error_energy = torch.abs(target["energy"] - prediction["energy"]) - - start_idx = 0 - for i, n in enumerate(target["natoms"]): - if ( - error_energy[i] < e_thresh - and error_forces[start_idx : start_idx + n].max() < f_thresh - ): - success += 1 - start_idx += n - - return { - "metric": success / total, - "total": success, - "numel": total, - } - - -def energy_within_threshold(prediction, target): - # compute absolute error on energy per system. - # then count the no. of systems where max energy error is < 0.02. - e_thresh = 0.02 - error_energy = torch.abs(target["energy"] - prediction["energy"]) - - success = (error_energy < e_thresh).sum().item() - total = target["energy"].size(0) - - return { - "metric": success / total, - "total": success, - "numel": total, - } - - -def average_distance_within_threshold(prediction, target): - pred_pos = torch.split(prediction["positions"], prediction["natoms"].tolist()) - target_pos = torch.split(target["positions"], target["natoms"].tolist()) - - mean_distance = [] - for idx, ml_pos in enumerate(pred_pos): - mean_distance.append( - np.mean( - np.linalg.norm( - min_diff( - ml_pos.detach().cpu().numpy(), - target_pos[idx].detach().cpu().numpy(), - target["cell"][idx].detach().cpu().numpy(), - target["pbc"].tolist(), - ), - axis=1, - ) - ) - ) - - success = 0 - intv = np.arange(0.01, 0.5, 0.001) - for i in intv: - success += sum(np.array(mean_distance) < i) - - total = len(mean_distance) * len(intv) - - return {"metric": success / total, "total": success, "numel": total} - - -def min_diff(pred_pos, dft_pos, cell, pbc): - pos_diff = pred_pos - dft_pos - fractional = np.linalg.solve(cell.T, pos_diff.T).T - - for i, periodic in enumerate(pbc): - # Yes, we need to do it twice - if periodic: - fractional[:, i] %= 1.0 - fractional[:, i] %= 1.0 - - fractional[fractional > 0.5] -= 1 - - return np.matmul(fractional, cell) - - -def cosine_similarity(prediction, target): - error = torch.cosine_similarity(prediction, target) - return { - "metric": torch.mean(error).item(), - "total": torch.sum(error).item(), - "numel": error.numel(), - } - - -def absolute_error(prediction, target): - error = torch.abs(target - prediction) - return { - "metric": torch.mean(error).item(), - "total": torch.sum(error).item(), - "numel": prediction.numel(), - } - - -def squared_error(prediction, target): - error = (target - prediction) ** 2 - return { - "metric": torch.mean(error).item(), - "total": torch.sum(error).item(), - "numel": prediction.numel(), - } - - -def magnitude_error(prediction, target, p=2): - assert prediction.shape[1] > 1 - error = torch.abs( - torch.norm(prediction, p=p, dim=-1) - torch.norm(target, p=p, dim=-1) - ) - return { - "metric": torch.mean(error).item(), - "total": torch.sum(error).item(), - "numel": error.numel(), - } diff --git a/ocpmodels/modules/exponential_moving_average.py b/ocpmodels/modules/exponential_moving_average.py deleted file mode 100644 index a38d831f0aa5cbf81884257520a5aaad8c8e824b..0000000000000000000000000000000000000000 --- a/ocpmodels/modules/exponential_moving_average.py +++ /dev/null @@ -1,192 +0,0 @@ -""" -Copied (and improved) from: -https://github.com/fadel/pytorch_ema/blob/master/torch_ema/ema.py (MIT license) -""" - -from __future__ import division, unicode_literals - -import copy -import weakref -from typing import Iterable, Optional - -import torch - - -# Partially based on: -# https://github.com/tensorflow/tensorflow/blob/r1.13/tensorflow/python/training/moving_averages.py -class ExponentialMovingAverage: - """ - Maintains (exponential) moving average of a set of parameters. - - Args: - parameters: Iterable of `torch.nn.Parameter` (typically from - `model.parameters()`). - decay: The exponential decay. - use_num_updates: Whether to use number of updates when computing - averages. - """ - - def __init__( - self, - parameters: Iterable[torch.nn.Parameter], - decay: float, - use_num_updates: bool = False, - ): - if decay < 0.0 or decay > 1.0: - raise ValueError("Decay must be between 0 and 1") - self.decay = decay - self.num_updates = 0 if use_num_updates else None - parameters = list(parameters) - self.shadow_params = [p.clone().detach() for p in parameters if p.requires_grad] - self.collected_params = [] - # By maintaining only a weakref to each parameter, - # we maintain the old GC behaviour of ExponentialMovingAverage: - # if the model goes out of scope but the ExponentialMovingAverage - # is kept, no references to the model or its parameters will be - # maintained, and the model will be cleaned up. - self._params_refs = [weakref.ref(p) for p in parameters if p.requires_grad] - - def _get_parameters( - self, parameters: Optional[Iterable[torch.nn.Parameter]] - ) -> Iterable[torch.nn.Parameter]: - if parameters is None: - parameters = [p() for p in self._params_refs] - if any(p is None for p in parameters): - raise ValueError( - "(One of) the parameters with which this " - "ExponentialMovingAverage " - "was initialized no longer exists (was garbage collected);" - " please either provide `parameters` explicitly or keep " - "the model to which they belong from being garbage " - "collected." - ) - return parameters - else: - return [p for p in parameters if p.requires_grad] - - def update(self, parameters: Optional[Iterable[torch.nn.Parameter]] = None) -> None: - """ - Update currently maintained parameters. - - Call this every time the parameters are updated, such as the result of - the `optimizer.step()` call. - - Args: - parameters: Iterable of `torch.nn.Parameter`; usually the same set of - parameters used to initialize this object. If `None`, the - parameters with which this `ExponentialMovingAverage` was - initialized will be used. - """ - parameters = self._get_parameters(parameters) - decay = self.decay - if self.num_updates is not None: - self.num_updates += 1 - decay = min(decay, (1 + self.num_updates) / (10 + self.num_updates)) - one_minus_decay = 1.0 - decay - with torch.no_grad(): - for s_param, param in zip(self.shadow_params, parameters): - tmp = param - s_param - s_param.add_(tmp, alpha=one_minus_decay) - - def copy_to( - self, parameters: Optional[Iterable[torch.nn.Parameter]] = None - ) -> None: - """ - Copy current parameters into given collection of parameters. - - Args: - parameters: Iterable of `torch.nn.Parameter`; the parameters to be - updated with the stored moving averages. If `None`, the - parameters with which this `ExponentialMovingAverage` was - initialized will be used. - """ - parameters = self._get_parameters(parameters) - for s_param, param in zip(self.shadow_params, parameters): - param.data.copy_(s_param.data) - - def store(self, parameters: Optional[Iterable[torch.nn.Parameter]] = None) -> None: - """ - Save the current parameters for restoring later. - - Args: - parameters: Iterable of `torch.nn.Parameter`; the parameters to be - temporarily stored. If `None`, the parameters of with which this - `ExponentialMovingAverage` was initialized will be used. - """ - parameters = self._get_parameters(parameters) - self.collected_params = [param.clone() for param in parameters] - - def restore( - self, parameters: Optional[Iterable[torch.nn.Parameter]] = None - ) -> None: - """ - Restore the parameters stored with the `store` method. - Useful to validate the model with EMA parameters without affecting the - original optimization process. Store the parameters before the - `copy_to` method. After validation (or model saving), use this to - restore the former parameters. - - Args: - parameters: Iterable of `torch.nn.Parameter`; the parameters to be - updated with the stored parameters. If `None`, the - parameters with which this `ExponentialMovingAverage` was - initialized will be used. - """ - parameters = self._get_parameters(parameters) - for c_param, param in zip(self.collected_params, parameters): - param.data.copy_(c_param.data) - - def state_dict(self) -> dict: - r"""Returns the state of the ExponentialMovingAverage as a dict.""" - # Following PyTorch conventions, references to tensors are returned: - # "returns a reference to the state and not its copy!" - - # https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict - return { - "decay": self.decay, - "num_updates": self.num_updates, - "shadow_params": self.shadow_params, - "collected_params": self.collected_params, - } - - def load_state_dict(self, state_dict: dict) -> None: - r"""Loads the ExponentialMovingAverage state. - - Args: - state_dict (dict): EMA state. Should be an object returned - from a call to :meth:`state_dict`. - """ - # deepcopy, to be consistent with module API - state_dict = copy.deepcopy(state_dict) - - self.decay = state_dict["decay"] - if self.decay < 0.0 or self.decay > 1.0: - raise ValueError("Decay must be between 0 and 1") - - self.num_updates = state_dict["num_updates"] - assert self.num_updates is None or isinstance( - self.num_updates, int - ), "Invalid num_updates" - - assert isinstance( - state_dict["shadow_params"], list - ), "shadow_params must be a list" - self.shadow_params = [ - p.to(self.shadow_params[i].device) - for i, p in enumerate(state_dict["shadow_params"]) - ] - assert all( - isinstance(p, torch.Tensor) for p in self.shadow_params - ), "shadow_params must all be Tensors" - - assert isinstance( - state_dict["collected_params"], list - ), "collected_params must be a list" - # collected_params is empty at initialization, - # so use shadow_params for device instead - self.collected_params = [ - p.to(self.shadow_params[i].device) - for i, p in enumerate(state_dict["collected_params"]) - ] - assert all( - isinstance(p, torch.Tensor) for p in self.collected_params - ), "collected_params must all be Tensors" diff --git a/ocpmodels/modules/loss.py b/ocpmodels/modules/loss.py deleted file mode 100644 index 3873225cdc8a830fbed8b54958d08df65257776b..0000000000000000000000000000000000000000 --- a/ocpmodels/modules/loss.py +++ /dev/null @@ -1,79 +0,0 @@ -import logging - -import torch -from torch import nn - -from ocpmodels.common import distutils - - -class L2MAELoss(nn.Module): - def __init__(self, reduction="mean"): - super().__init__() - self.reduction = reduction - assert reduction in ["mean", "sum"] - - def forward(self, input: torch.Tensor, target: torch.Tensor): - dists = torch.norm(input - target, p=2, dim=-1) - if self.reduction == "mean": - return torch.mean(dists) - elif self.reduction == "sum": - return torch.sum(dists) - - -class AtomwiseL2Loss(nn.Module): - def __init__(self, reduction="mean"): - super().__init__() - self.reduction = reduction - assert reduction in ["mean", "sum"] - - def forward( - self, - input: torch.Tensor, - target: torch.Tensor, - natoms: torch.Tensor, - ): - assert natoms.shape[0] == input.shape[0] == target.shape[0] - assert len(natoms.shape) == 1 # (nAtoms, ) - - dists = torch.norm(input - target, p=2, dim=-1) - loss = natoms * dists - - if self.reduction == "mean": - return torch.mean(loss) - elif self.reduction == "sum": - return torch.sum(loss) - - -class DDPLoss(nn.Module): - def __init__(self, loss_fn, reduction="mean"): - super().__init__() - self.loss_fn = loss_fn - self.loss_fn.reduction = "sum" - self.reduction = reduction - assert reduction in ["mean", "sum"] - - def forward( - self, - input: torch.Tensor, - target: torch.Tensor, - natoms: torch.Tensor = None, - batch_size: int = None, - ): - # zero out nans, if any - found_nans_or_infs = not torch.all(input.isfinite()) - if found_nans_or_infs is True: - logging.warning("Found nans while computing loss") - input = torch.nan_to_num(input, nan=0.0) - - if natoms is None: - loss = self.loss_fn(input, target) - else: # atom-wise loss - loss = self.loss_fn(input, target, natoms) - if self.reduction == "mean": - num_samples = batch_size if batch_size is not None else input.shape[0] - num_samples = distutils.all_reduce(num_samples, device=input.device) - # Multiply by world size since gradients are averaged - # across DDP replicas - return loss * distutils.get_world_size() / num_samples - else: - return loss diff --git a/ocpmodels/modules/normalizer.py b/ocpmodels/modules/normalizer.py deleted file mode 100644 index 302f0d6cef04567be6e89a5e250ed67f75f4ca8b..0000000000000000000000000000000000000000 --- a/ocpmodels/modules/normalizer.py +++ /dev/null @@ -1,46 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import torch - - -class Normalizer(object): - """Normalize a Tensor and restore it later.""" - - def __init__(self, tensor=None, mean=None, std=None, device=None): - """tensor is taken as a sample to calculate the mean and std""" - if tensor is None and mean is None: - return - - if device is None: - device = "cpu" - - if tensor is not None: - self.mean = torch.mean(tensor, dim=0).to(device) - self.std = torch.std(tensor, dim=0).to(device) - return - - if mean is not None and std is not None: - self.mean = torch.tensor(mean).to(device) - self.std = torch.tensor(std).to(device) - - def to(self, device): - self.mean = self.mean.to(device) - self.std = self.std.to(device) - - def norm(self, tensor): - return (tensor - self.mean) / self.std - - def denorm(self, normed_tensor): - return normed_tensor * self.std + self.mean - - def state_dict(self): - return {"mean": self.mean, "std": self.std} - - def load_state_dict(self, state_dict): - self.mean = state_dict["mean"].to(self.mean.device) - self.std = state_dict["std"].to(self.mean.device) diff --git a/ocpmodels/modules/scaling/__init__.py b/ocpmodels/modules/scaling/__init__.py deleted file mode 100644 index 807416b066c7ccf88f169f74686690b60362bcf8..0000000000000000000000000000000000000000 --- a/ocpmodels/modules/scaling/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from .scale_factor import ScaleFactor - -__all__ = ["ScaleFactor"] diff --git a/ocpmodels/modules/scaling/compat.py b/ocpmodels/modules/scaling/compat.py deleted file mode 100644 index 56ef12e3693810227d403165a3765d8c5c58fe79..0000000000000000000000000000000000000000 --- a/ocpmodels/modules/scaling/compat.py +++ /dev/null @@ -1,72 +0,0 @@ -import json -import logging -from pathlib import Path -from typing import Dict, Optional, Union - -import torch -import torch.nn as nn - -from .scale_factor import ScaleFactor - -ScaleDict = Union[Dict[str, float], Dict[str, torch.Tensor]] - - -def _load_scale_dict(scale_file: Optional[Union[str, ScaleDict]]): - """ - Loads scale factors from either: - - a JSON file mapping scale factor names to scale values - - a python dictionary pickled object (loaded using `torch.load`) mapping scale factor names to scale values - - a dictionary mapping scale factor names to scale values - """ - if not scale_file: - return None - - if isinstance(scale_file, dict): - if not scale_file: - logging.warning("Empty scale dictionary provided to model.") - return scale_file - - path = Path(scale_file) - if not path.exists(): - raise ValueError(f"Scale file {path} does not exist.") - - scale_dict: Optional[ScaleDict] = None - if path.suffix == ".pt": - scale_dict = torch.load(path) - elif path.suffix == ".json": - with open(path, "r") as f: - scale_dict = json.load(f) - - if isinstance(scale_dict, dict): - # old json scale factors have a comment field that has the model name - scale_dict.pop("comment", None) - else: - raise ValueError(f"Unsupported scale file extension: {path.suffix}") - - if not scale_dict: - return None - - return scale_dict - - -def load_scales_compat(module: nn.Module, scale_file: Optional[Union[str, ScaleDict]]): - scale_dict = _load_scale_dict(scale_file) - if not scale_dict: - return - - scale_factors = { - module.name or name: (module, name) - for name, module in module.named_modules() - if isinstance(module, ScaleFactor) - } - logging.debug( - f"Found the following scale factors: {[(k, name) for k, (_, name) in scale_factors.items()]}" - ) - for name, scale in scale_dict.items(): - if name not in scale_factors: - logging.warning(f"Scale factor {name} not found in model") - continue - - scale_module, module_name = scale_factors[name] - logging.debug(f"Loading scale factor {scale} for ({name} => {module_name})") - scale_module.set_(scale) diff --git a/ocpmodels/modules/scaling/fit.py b/ocpmodels/modules/scaling/fit.py deleted file mode 100644 index b1038e9d3bbf1293f1726ade79ea0aba0556e5c0..0000000000000000000000000000000000000000 --- a/ocpmodels/modules/scaling/fit.py +++ /dev/null @@ -1,233 +0,0 @@ -import logging -import math -import readline -import sys -from itertools import islice -from pathlib import Path -from typing import TYPE_CHECKING, Dict, Literal - -import torch -import torch.nn as nn -from torch.nn.parallel.distributed import DistributedDataParallel - -from ocpmodels.common.data_parallel import OCPDataParallel -from ocpmodels.common.flags import flags -from ocpmodels.common.utils import ( - build_config, - new_trainer_context, - setup_logging, -) -from ocpmodels.modules.scaling import ScaleFactor -from ocpmodels.modules.scaling.compat import load_scales_compat - -if TYPE_CHECKING: - from ocpmodels.trainers.base_trainer import BaseTrainer - - -def _prefilled_input(prompt: str, prefill: str = ""): - readline.set_startup_hook(lambda: readline.insert_text(prefill)) - try: - return input(prompt) - finally: - readline.set_startup_hook() - - -def _train_batch(trainer: "BaseTrainer", batch): - with torch.no_grad(): - with torch.amp.autocast("cuda", enabled=trainer.scaler is not None): - out = trainer._forward(batch) - loss = trainer._compute_loss(out, batch) - del out, loss - - -def main(*, num_batches: int = 16): - # region args/config setup - setup_logging() - - parser = flags.get_parser() - args, override_args = parser.parse_known_args() - _config = build_config(args, override_args) - _config["logger"] = "tensorboard" - # endregion - - assert not args.distributed, "This doesn't work with DDP" - with new_trainer_context(args=args, config=_config) as ctx: - config = ctx.config - trainer = ctx.trainer - - ckpt_file = config.get("checkpoint", None) - assert ( - ckpt_file is not None - ), "Checkpoint file not specified. Please specify --checkpoint " - ckpt_file = Path(ckpt_file) - - logging.info(f"Input checkpoint path: {ckpt_file}, {ckpt_file.exists()=}") - - model: nn.Module = trainer.model - val_loader = trainer.val_loader - assert val_loader is not None, "Val dataset is required for making predictions" - - if ckpt_file.exists(): - trainer.load_checkpoint(str(ckpt_file)) - - # region reoad scale file contents if necessary - # unwrap module from DP/DDP - unwrapped_model = model - while isinstance(unwrapped_model, (DistributedDataParallel, OCPDataParallel)): - unwrapped_model = unwrapped_model.module - assert isinstance(unwrapped_model, nn.Module), "Model is not a nn.Module" - load_scales_compat(unwrapped_model, config.get("scale_file", None)) - # endregion - - model.eval() - - # recursively go through the submodules and get the ScaleFactor modules - scale_factors: Dict[str, ScaleFactor] = { - name: module - for name, module in model.named_modules() - if isinstance(module, ScaleFactor) - } - - mode: Literal["all", "unfitted"] = "all" - - # region detect fitted/unfitted factors - fitted_scale_factors = [ - f"{name}: {module.scale_factor.item():.3f}" - for name, module in scale_factors.items() - if module.fitted - ] - unfitted_scale_factors = [ - name for name, module in scale_factors.items() if not module.fitted - ] - fitted_scale_factors_str = ", ".join(fitted_scale_factors) - logging.info(f"Fitted scale factors: [{fitted_scale_factors_str}]") - unfitted_scale_factors_str = ", ".join(unfitted_scale_factors) - logging.info(f"Unfitted scale factors: [{unfitted_scale_factors_str}]") - - if fitted_scale_factors: - flag = input( - "Do you want to continue and fit all scale factors (1), " - "only fit the variables not fitted yet (2), or exit (3)? " - ) - if str(flag) == "1": - mode = "all" - logging.info("Fitting all scale factors.") - elif str(flag) == "2": - mode = "unfitted" - logging.info("Only fitting unfitted variables.") - else: - print(flag) - logging.info("Exiting script") - sys.exit() - # endregion - - # region get the output path - out_path = Path( - _prefilled_input( - "Enter output path for fitted scale factors: ", - prefill=str(ckpt_file), - ) - ) - if out_path.exists(): - logging.warning(f"Already found existing file: {out_path}") - flag = input( - "Do you want to continue and overwrite existing file (1), " - "or exit (2)? " - ) - if str(flag) == "1": - logging.info("Overwriting existing file.") - else: - logging.info("Exiting script") - sys.exit() - - logging.info( - f"Output path for fitted scale factors: {out_path}, {out_path.exists()=}" - ) - # endregion - - # region reset the scale factors if mode == "all" - if mode == "all": - logging.info("Fitting all scale factors.") - for name, scale_factor in scale_factors.items(): - if scale_factor.fitted: - logging.info( - f"{name} is already fitted in the checkpoint, resetting it. {scale_factor.scale_factor}" - ) - scale_factor.reset_() - # endregion - - # region we do a single pass through the network to get the correct execution order of the scale factors - scale_factor_indices: Dict[str, int] = {} - max_idx = 0 - - # initialize all scale factors - for name, module in scale_factors.items(): - - def index_fn(name=name): - nonlocal max_idx - assert name is not None - if name not in scale_factor_indices: - scale_factor_indices[name] = max_idx - logging.debug(f"Scale factor for {name} = {max_idx}") - max_idx += 1 - - module.initialize_(index_fn=index_fn) - - # single pass through network - _train_batch(trainer, next(iter(val_loader))) - - # sort the scale factors by their computation order - sorted_factors = sorted( - scale_factors.items(), - key=lambda x: scale_factor_indices.get(x[0], math.inf), - ) - - logging.info("Sorted scale factors by computation order:") - for name, _ in sorted_factors: - logging.info(f"{name}: {scale_factor_indices[name]}") - - # endregion - - # loop over the scale factors in the computation order - # and fit them one by one - logging.info("Start fitting") - - for name, module in sorted_factors: - if mode == "unfitted" and module.fitted: - logging.info(f"Skipping {name} (already fitted)") - continue - - logging.info(f"Fitting {name}") - with module.fit_context_(): - for batch in islice(val_loader, num_batches): - _train_batch(trainer, batch) - stats, ratio, value = module.fit_() - - logging.info( - f"Variable: {name}, " - f"Var_in: {stats['variance_in']:.3f}, " - f"Var_out: {stats['variance_out']:.3f}, " - f"Ratio: {ratio:.3f} => Scaling factor: {value:.3f}" - ) - - # make sure all scale factors are fitted - for name, module in sorted_factors: - assert module.fitted, f"{name} is not fitted" - - # region save the scale factors to the checkpoint file - trainer.config["cmd"]["checkpoint_dir"] = out_path.parent - trainer.is_debug = False - out_file = trainer.save( - metrics=None, - checkpoint_file=out_path.name, - training_state=False, - ) - assert out_file is not None, "Failed to save checkpoint" - out_file = Path(out_file) - assert out_file.exists(), f"Failed to save checkpoint to {out_file}" - # endregion - logging.info(f"Saved results to: {out_file}") - - -if __name__ == "__main__": - main() diff --git a/ocpmodels/modules/scaling/scale_factor.py b/ocpmodels/modules/scaling/scale_factor.py deleted file mode 100644 index 1587040a007bb60397c6b8a81bfc39e05b4cf7b3..0000000000000000000000000000000000000000 --- a/ocpmodels/modules/scaling/scale_factor.py +++ /dev/null @@ -1,166 +0,0 @@ -import itertools -import logging -import math -from contextlib import contextmanager -from typing import Callable, Optional, TypedDict, Union - -import torch -import torch.nn as nn - - -class _Stats(TypedDict): - variance_in: float - variance_out: float - n_samples: int - - -IndexFn = Callable[[], None] - - -def _check_consistency(old: torch.Tensor, new: torch.Tensor, key: str): - if not torch.allclose(old, new): - raise ValueError( - f"Scale factor parameter {key} is inconsistent with the loaded state dict.\n" - f"Old: {old}\n" - f"Actual: {new}" - ) - - -class ScaleFactor(nn.Module): - scale_factor: torch.Tensor - - name: Optional[str] = None - index_fn: Optional[IndexFn] = None - stats: Optional[_Stats] = None - - def __init__( - self, - name: Optional[str] = None, - enforce_consistency: bool = True, - ): - super().__init__() - - self.name = name - self.index_fn = None - self.stats = None - - self.scale_factor = nn.parameter.Parameter( - torch.tensor(0.0), requires_grad=False - ) - if enforce_consistency: - self._register_load_state_dict_pre_hook(self._enforce_consistency) - - def _enforce_consistency( - self, - state_dict, - prefix, - _local_metadata, - _strict, - _missing_keys, - _unexpected_keys, - _error_msgs, - ): - if not self.fitted: - return - - persistent_buffers = { - k: v - for k, v in self._buffers.items() - if k not in self._non_persistent_buffers_set - } - local_name_params = itertools.chain( - self._parameters.items(), persistent_buffers.items() - ) - local_state = {k: v for k, v in local_name_params if v is not None} - - for name, param in local_state.items(): - key = prefix + name - if key not in state_dict: - continue - - input_param = state_dict[key] - _check_consistency(old=param, new=input_param, key=key) - - @property - def fitted(self): - return bool((self.scale_factor != 0.0).item()) - - @torch.jit.unused - def reset_(self): - self.scale_factor.zero_() - - @torch.jit.unused - def set_(self, scale: Union[float, torch.Tensor]): - if self.fitted: - _check_consistency( - old=self.scale_factor, - new=torch.tensor(scale) if isinstance(scale, float) else scale, - key="scale_factor", - ) - self.scale_factor.fill_(scale) - - @torch.jit.unused - def initialize_(self, *, index_fn: Optional[IndexFn] = None): - self.index_fn = index_fn - - @contextmanager - @torch.jit.unused - def fit_context_(self): - self.stats = _Stats(variance_in=0.0, variance_out=0.0, n_samples=0) - yield - del self.stats - self.stats = None - - @torch.jit.unused - def fit_(self): - assert self.stats, "Stats not set" - for k, v in self.stats.items(): - assert v > 0, f"{k} is {v}" - - self.stats["variance_in"] = self.stats["variance_in"] / self.stats["n_samples"] - self.stats["variance_out"] = ( - self.stats["variance_out"] / self.stats["n_samples"] - ) - - ratio = self.stats["variance_out"] / self.stats["variance_in"] - value = math.sqrt(1 / ratio) - - self.set_(value) - - stats = dict(**self.stats) - return stats, ratio, value - - @torch.no_grad() - @torch.jit.unused - def _observe(self, x: torch.Tensor, ref: Optional[torch.Tensor] = None): - if self.stats is None: - logging.debug("Observer not initialized but self.observe() called") - return - - n_samples = x.shape[0] - self.stats["variance_out"] += torch.mean(torch.var(x, dim=0)).item() * n_samples - - if ref is None: - self.stats["variance_in"] += n_samples - else: - self.stats["variance_in"] += ( - torch.mean(torch.var(ref, dim=0)).item() * n_samples - ) - self.stats["n_samples"] += n_samples - - def forward( - self, - x: torch.Tensor, - *, - ref: Optional[torch.Tensor] = None, - ): - if self.index_fn is not None: - self.index_fn() - - if self.fitted: - x = x * self.scale_factor - - if not torch.jit.is_scripting(): - self._observe(x, ref=ref) - - return x diff --git a/ocpmodels/modules/scaling/util.py b/ocpmodels/modules/scaling/util.py deleted file mode 100644 index 15c58b5d42f9a492ea2faaf70a79d6a8dea9b129..0000000000000000000000000000000000000000 --- a/ocpmodels/modules/scaling/util.py +++ /dev/null @@ -1,23 +0,0 @@ -import logging - -import torch.nn as nn - -from .scale_factor import ScaleFactor - - -def ensure_fitted(module: nn.Module, warn: bool = False): - for name, child in module.named_modules(): - if not isinstance(child, ScaleFactor) or child.fitted: - continue - if child.name is not None: - name = f"{child.name} ({name})" - msg = ( - f"Scale factor {name} is not fitted. " - "Please make sure that you either (1) load a checkpoint with fitted scale factors, " - "(2) explicitly load scale factors using the `model.scale_file` attribute, or " - "(3) fit the scale factors using the `fit.py` script." - ) - if warn: - logging.warning(msg) - else: - raise ValueError(msg) diff --git a/ocpmodels/modules/scheduler.py b/ocpmodels/modules/scheduler.py deleted file mode 100644 index 223ec24447b423fbb80f51bd3a040a58b2cc4aa9..0000000000000000000000000000000000000000 --- a/ocpmodels/modules/scheduler.py +++ /dev/null @@ -1,64 +0,0 @@ -import inspect - -import torch.optim.lr_scheduler as lr_scheduler - -from ocpmodels.common.utils import warmup_lr_lambda - - -class LRScheduler: - """ - Learning rate scheduler class for torch.optim learning rate schedulers - - Notes: - If no learning rate scheduler is specified in the config the default - scheduler is warmup_lr_lambda (ocpmodels.common.utils) not no scheduler, - this is for backward-compatibility reasons. To run without a lr scheduler - specify scheduler: "Null" in the optim section of the config. - - Args: - config (dict): Optim dict from the input config - optimizer (obj): torch optim object - """ - - def __init__(self, optimizer, config): - self.optimizer = optimizer - self.config = config.copy() - if "scheduler" in self.config: - self.scheduler_type = self.config["scheduler"] - else: - self.scheduler_type = "LambdaLR" - scheduler_lambda_fn = lambda x: warmup_lr_lambda(x, self.config) - self.config["lr_lambda"] = scheduler_lambda_fn - - if self.scheduler_type != "Null": - self.scheduler = getattr(lr_scheduler, self.scheduler_type) - scheduler_args = self.filter_kwargs(config) - self.scheduler = self.scheduler(optimizer, **scheduler_args) - - def step(self, metrics=None, epoch=None): - if self.scheduler_type == "Null": - return - if self.scheduler_type == "ReduceLROnPlateau": - if metrics is None: - raise Exception("Validation set required for ReduceLROnPlateau.") - self.scheduler.step(metrics) - else: - self.scheduler.step() - - def filter_kwargs(self, config): - # adapted from https://stackoverflow.com/questions/26515595/ - sig = inspect.signature(self.scheduler) - filter_keys = [ - param.name - for param in sig.parameters.values() - if param.kind == param.POSITIONAL_OR_KEYWORD - ] - filter_keys.remove("optimizer") - scheduler_args = { - arg: self.config[arg] for arg in self.config if arg in filter_keys - } - return scheduler_args - - def get_lr(self): - for group in self.optimizer.param_groups: - return group["lr"] diff --git a/ocpmodels/preprocessing/__init__.py b/ocpmodels/preprocessing/__init__.py deleted file mode 100644 index 1139688707a2037b501e9b417fc7df2ab544b3dd..0000000000000000000000000000000000000000 --- a/ocpmodels/preprocessing/__init__.py +++ /dev/null @@ -1,8 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -from .atoms_to_graphs import AtomsToGraphs diff --git a/ocpmodels/preprocessing/atoms_to_graphs.py b/ocpmodels/preprocessing/atoms_to_graphs.py deleted file mode 100644 index 7434a6f7ac76472b1d9a6fe5bee8df0b5dcaec23..0000000000000000000000000000000000000000 --- a/ocpmodels/preprocessing/atoms_to_graphs.py +++ /dev/null @@ -1,245 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import ase.db.sqlite -import ase.io.trajectory -import numpy as np -import torch -from torch_geometric.data import Data - -from ocpmodels.common.utils import collate - -try: - from pymatgen.io.ase import AseAtomsAdaptor -except Exception: - pass - - -try: - shell = get_ipython().__class__.__name__ - if shell == "ZMQInteractiveShell": - from tqdm.notebook import tqdm - else: - from tqdm import tqdm -except NameError: - from tqdm import tqdm - - -class AtomsToGraphs: - """A class to help convert periodic atomic structures to graphs. - - The AtomsToGraphs class takes in periodic atomic structures in form of ASE atoms objects and converts - them into graph representations for use in PyTorch. The primary purpose of this class is to determine the - nearest neighbors within some radius around each individual atom, taking into account PBC, and set the - pair index and distance between atom pairs appropriately. Lastly, atomic properties and the graph information - are put into a PyTorch geometric data object for use with PyTorch. - - Args: - max_neigh (int): Maximum number of neighbors to consider. - radius (int or float): Cutoff radius in Angstroms to search for neighbors. - r_energy (bool): Return the energy with other properties. Default is False, so the energy will not be returned. - r_forces (bool): Return the forces with other properties. Default is False, so the forces will not be returned. - r_distances (bool): Return the distances with other properties. - Default is False, so the distances will not be returned. - r_edges (bool): Return interatomic edges with other properties. Default is True, so edges will be returned. - r_fixed (bool): Return a binary vector with flags for fixed (1) vs free (0) atoms. - Default is True, so the fixed indices will be returned. - r_pbc (bool): Return the periodic boundary conditions with other properties. - Default is False, so the periodic boundary conditions will not be returned. - - Attributes: - max_neigh (int): Maximum number of neighbors to consider. - radius (int or float): Cutoff radius in Angstoms to search for neighbors. - r_energy (bool): Return the energy with other properties. Default is False, so the energy will not be returned. - r_forces (bool): Return the forces with other properties. Default is False, so the forces will not be returned. - r_distances (bool): Return the distances with other properties. - Default is False, so the distances will not be returned. - r_edges (bool): Return interatomic edges with other properties. Default is True, so edges will be returned. - r_fixed (bool): Return a binary vector with flags for fixed (1) vs free (0) atoms. - Default is True, so the fixed indices will be returned. - r_pbc (bool): Return the periodic boundary conditions with other properties. - Default is False, so the periodic boundary conditions will not be returned. - - """ - - def __init__( - self, - max_neigh=200, - radius=6, - r_energy=False, - r_forces=False, - r_distances=False, - r_edges=True, - r_fixed=True, - r_pbc=False, - ): - self.max_neigh = max_neigh - self.radius = radius - self.r_energy = r_energy - self.r_forces = r_forces - self.r_distances = r_distances - self.r_fixed = r_fixed - self.r_edges = r_edges - self.r_pbc = r_pbc - - def _get_neighbors_pymatgen(self, atoms): - """Preforms nearest neighbor search and returns edge index, distances, - and cell offsets""" - struct = AseAtomsAdaptor.get_structure(atoms) - _c_index, _n_index, _offsets, n_distance = struct.get_neighbor_list( - r=self.radius, numerical_tol=0, exclude_self=True - ) - - _nonmax_idx = [] - for i in range(len(atoms)): - idx_i = (_c_index == i).nonzero()[0] - # sort neighbors by distance, remove edges larger than max_neighbors - idx_sorted = np.argsort(n_distance[idx_i])[: self.max_neigh] - _nonmax_idx.append(idx_i[idx_sorted]) - _nonmax_idx = np.concatenate(_nonmax_idx) - - _c_index = _c_index[_nonmax_idx] - _n_index = _n_index[_nonmax_idx] - n_distance = n_distance[_nonmax_idx] - _offsets = _offsets[_nonmax_idx] - - return _c_index, _n_index, n_distance, _offsets - - def _reshape_features(self, c_index, n_index, n_distance, offsets): - """Stack center and neighbor index and reshapes distances, - takes in np.arrays and returns torch tensors""" - edge_index = torch.LongTensor(np.vstack((n_index, c_index))) - edge_distances = torch.FloatTensor(n_distance) - cell_offsets = torch.LongTensor(offsets) - - # remove distances smaller than a tolerance ~ 0. The small tolerance is - # needed to correct for pymatgen's neighbor_list returning self atoms - # in a few edge cases. - nonzero = torch.where(edge_distances >= 1e-8)[0] - edge_index = edge_index[:, nonzero] - edge_distances = edge_distances[nonzero] - cell_offsets = cell_offsets[nonzero] - - return edge_index, edge_distances, cell_offsets - - def convert( - self, - atoms, - ): - """Convert a single atomic stucture to a graph. - - Args: - atoms (ase.atoms.Atoms): An ASE atoms object. - - Returns: - data (torch_geometric.data.Data): A torch geometic data object with positions, atomic_numbers, tags, - and optionally, energy, forces, distances, edges, and periodic boundary conditions. - Optional properties can included by setting r_property=True when constructing the class. - """ - - # set the atomic numbers, positions, and cell - atomic_numbers = torch.Tensor(atoms.get_atomic_numbers()) - positions = torch.Tensor(atoms.get_positions()) - cell = torch.Tensor(atoms.get_cell()).view(1, 3, 3) - natoms = positions.shape[0] - # initialized to torch.zeros(natoms) if tags missing. - # https://wiki.fysik.dtu.dk/ase/_modules/ase/atoms.html#Atoms.get_tags - tags = torch.Tensor(atoms.get_tags()) - - # put the minimum data in torch geometric data object - data = Data( - cell=cell, - pos=positions, - atomic_numbers=atomic_numbers, - natoms=natoms, - tags=tags, - ) - - # optionally include other properties - if self.r_edges: - # run internal functions to get padded indices and distances - split_idx_dist = self._get_neighbors_pymatgen(atoms) - edge_index, edge_distances, cell_offsets = self._reshape_features( - *split_idx_dist - ) - - data.edge_index = edge_index - data.cell_offsets = cell_offsets - if self.r_energy: - energy = atoms.get_potential_energy(apply_constraint=False) - data.y = energy - if self.r_forces: - forces = torch.Tensor(atoms.get_forces(apply_constraint=False)) - data.force = forces - if self.r_distances and self.r_edges: - data.distances = edge_distances - if self.r_fixed: - fixed_idx = torch.zeros(natoms) - if hasattr(atoms, "constraints"): - from ase.constraints import FixAtoms - - for constraint in atoms.constraints: - if isinstance(constraint, FixAtoms): - fixed_idx[constraint.index] = 1 - data.fixed = fixed_idx - if self.r_pbc: - data.pbc = torch.tensor(atoms.pbc) - - return data - - def convert_all( - self, - atoms_collection, - processed_file_path=None, - collate_and_save=False, - disable_tqdm=False, - ): - """Convert all atoms objects in a list or in an ase.db to graphs. - - Args: - atoms_collection (list of ase.atoms.Atoms or ase.db.sqlite.SQLite3Database): - Either a list of ASE atoms objects or an ASE database. - processed_file_path (str): - A string of the path to where the processed file will be written. Default is None. - collate_and_save (bool): A boolean to collate and save or not. Default is False, so will not write a file. - - Returns: - data_list (list of torch_geometric.data.Data): - A list of torch geometric data objects containing molecular graph info and properties. - """ - - # list for all data - data_list = [] - if isinstance(atoms_collection, list): - atoms_iter = atoms_collection - elif isinstance(atoms_collection, ase.db.sqlite.SQLite3Database): - atoms_iter = atoms_collection.select() - elif isinstance( - atoms_collection, ase.io.trajectory.SlicedTrajectory - ) or isinstance(atoms_collection, ase.io.trajectory.TrajectoryReader): - atoms_iter = atoms_collection - else: - raise NotImplementedError - - for atoms in tqdm( - atoms_iter, - desc="converting ASE atoms collection to graphs", - total=len(atoms_collection), - unit=" systems", - disable=disable_tqdm, - ): - # check if atoms is an ASE Atoms object this for the ase.db case - if not isinstance(atoms, ase.atoms.Atoms): - atoms = atoms.toatoms() - data = self.convert(atoms) - data_list.append(data) - - if collate_and_save: - data, slices = collate(data_list) - torch.save((data, slices), processed_file_path) - - return data_list diff --git a/ocpmodels/tasks/__init__.py b/ocpmodels/tasks/__init__.py deleted file mode 100644 index 0e69e1df291f122a2bdf62d2c9010c80760468b5..0000000000000000000000000000000000000000 --- a/ocpmodels/tasks/__init__.py +++ /dev/null @@ -1,8 +0,0 @@ -# Copyright (c) Facebook, Inc. and its affiliates. -# -# This source code is licensed under the MIT license found in the -# LICENSE file in the root directory of this source tree. - -__all__ = ["TrainTask", "PredictTask", "ValidateTask", "RelxationTask"] - -from .task import PredictTask, RelxationTask, TrainTask, ValidateTask diff --git a/ocpmodels/tasks/task.py b/ocpmodels/tasks/task.py deleted file mode 100644 index b8026be03f729f99b85e2fb30dd5b47d1df4cca2..0000000000000000000000000000000000000000 --- a/ocpmodels/tasks/task.py +++ /dev/null @@ -1,96 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import logging -import os - -from ocpmodels.common.registry import registry -from ocpmodels.trainers.forces_trainer import ForcesTrainer - - -class BaseTask: - def __init__(self, config): - self.config = config - - def setup(self, trainer): - self.trainer = trainer - if self.config["checkpoint"] is not None: - self.trainer.load_checkpoint(self.config["checkpoint"]) - - # save checkpoint path to runner state for slurm resubmissions - self.chkpt_path = os.path.join( - self.trainer.config["cmd"]["checkpoint_dir"], "checkpoint.pt" - ) - - def run(self): - raise NotImplementedError - - -@registry.register_task("train") -class TrainTask(BaseTask): - def _process_error(self, e: RuntimeError): - e_str = str(e) - if ( - "find_unused_parameters" in e_str - and "torch.nn.parallel.DistributedDataParallel" in e_str - ): - for name, parameter in self.trainer.model.named_parameters(): - if parameter.requires_grad and parameter.grad is None: - logging.warning( - f"Parameter {name} has no gradient. Consider removing it from the model." - ) - - def run(self): - try: - self.trainer.train( - disable_eval_tqdm=self.config.get("hide_eval_progressbar", False) - ) - except RuntimeError as e: - self._process_error(e) - raise e - - -@registry.register_task("predict") -class PredictTask(BaseTask): - def run(self): - assert ( - self.trainer.test_loader is not None - ), "Test dataset is required for making predictions" - assert self.config["checkpoint"] - results_file = "predictions" - self.trainer.predict( - self.trainer.test_loader, - results_file=results_file, - disable_tqdm=self.config.get("hide_eval_progressbar", False), - ) - - -@registry.register_task("validate") -class ValidateTask(BaseTask): - def run(self): - # Note that the results won't be precise on multi GPUs due to padding of extra images (although the difference should be minor) - assert ( - self.trainer.val_loader is not None - ), "Val dataset is required for making predictions" - assert self.config["checkpoint"] - self.trainer.validate( - split="val", - disable_tqdm=self.config.get("hide_eval_progressbar", False), - ) - - -@registry.register_task("run-relaxations") -class RelxationTask(BaseTask): - def run(self): - assert isinstance( - self.trainer, ForcesTrainer - ), "Relaxations are only possible for ForcesTrainer" - assert ( - self.trainer.relax_dataset is not None - ), "Relax dataset is required for making predictions" - assert self.config["checkpoint"] - self.trainer.run_relaxations() diff --git a/ocpmodels/trainers/__init__.py b/ocpmodels/trainers/__init__.py deleted file mode 100644 index a93fc680bf9d8f9b05ed7862ae2d1435b39a9e8e..0000000000000000000000000000000000000000 --- a/ocpmodels/trainers/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# Copyright (c) Facebook, Inc. and its affiliates. -# -# This source code is licensed under the MIT license found in the -# LICENSE file in the root directory of this source tree. - -__all__ = [ - "BaseTrainer", - "ForcesTrainer", - "EnergyTrainer", -] - -from .base_trainer import BaseTrainer -from .energy_trainer import EnergyTrainer -from .forces_trainer import ForcesTrainer diff --git a/ocpmodels/trainers/base_trainer.py b/ocpmodels/trainers/base_trainer.py deleted file mode 100644 index d2332ad9d00865ce4f5eb6507b197b6d83b80d0e..0000000000000000000000000000000000000000 --- a/ocpmodels/trainers/base_trainer.py +++ /dev/null @@ -1,771 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import datetime -import errno -import logging -import os -import random -import subprocess -from abc import ABC, abstractmethod -from collections import defaultdict - -import numpy as np -import torch -import torch.nn as nn -import torch.optim as optim -import yaml -from torch.nn.parallel.distributed import DistributedDataParallel -from torch.utils.data import DataLoader -from tqdm import tqdm - -import ocpmodels -from ocpmodels.common import distutils, gp_utils -from ocpmodels.common.data_parallel import ( - BalancedBatchSampler, - OCPDataParallel, - ParallelCollater, -) -from ocpmodels.common.registry import registry -from ocpmodels.common.utils import load_state_dict, save_checkpoint -from ocpmodels.modules.evaluator import Evaluator -from ocpmodels.modules.exponential_moving_average import ( - ExponentialMovingAverage, -) -from ocpmodels.modules.loss import AtomwiseL2Loss, DDPLoss, L2MAELoss -from ocpmodels.modules.normalizer import Normalizer -from ocpmodels.modules.scaling.compat import load_scales_compat -from ocpmodels.modules.scaling.util import ensure_fitted -from ocpmodels.modules.scheduler import LRScheduler - - -@registry.register_trainer("base") -class BaseTrainer(ABC): - @property - def _unwrapped_model(self): - module = self.model - while isinstance(module, (OCPDataParallel, DistributedDataParallel)): - module = module.module - return module - - def __init__( - self, - task, - model, - dataset, - optimizer, - identifier, - normalizer=None, - timestamp_id=None, - run_dir=None, - is_debug=False, - is_hpo=False, - print_every=100, - seed=None, - logger="tensorboard", - local_rank=0, - amp=False, - cpu=False, - name="base_trainer", - slurm={}, - noddp=False, - ): - self.name = name - self.cpu = cpu - self.epoch = 0 - self.step = 0 - - if torch.cuda.is_available() and not self.cpu: - self.device = torch.device(f"cuda:{local_rank}") - else: - self.device = torch.device("cpu") - self.cpu = True # handle case when `--cpu` isn't specified - # but there are no gpu devices available - if run_dir is None: - run_dir = os.getcwd() - - if timestamp_id is None: - timestamp = torch.tensor(datetime.datetime.now().timestamp()).to( - self.device - ) - # create directories from master rank only - distutils.broadcast(timestamp, 0) - timestamp = datetime.datetime.fromtimestamp(timestamp.int()).strftime( - "%Y-%m-%d-%H-%M-%S" - ) - if identifier: - self.timestamp_id = f"{timestamp}-{identifier}" - else: - self.timestamp_id = timestamp - else: - self.timestamp_id = timestamp_id - - try: - commit_hash = ( - subprocess.check_output( - [ - "git", - "-C", - ocpmodels.__path__[0], - "describe", - "--always", - ] - ) - .strip() - .decode("ascii") - ) - # catch instances where code is not being run from a git repo - except Exception: - commit_hash = None - - logger_name = logger if isinstance(logger, str) else logger["name"] - self.config = { - "task": task, - "trainer": "forces" if name == "s2ef" else "energy", - "model": model.pop("name"), - "model_attributes": model, - "optim": optimizer, - "logger": logger, - "amp": amp, - "gpus": distutils.get_world_size() if not self.cpu else 0, - "cmd": { - "identifier": identifier, - "print_every": print_every, - "seed": seed, - "timestamp_id": self.timestamp_id, - "commit": commit_hash, - "checkpoint_dir": os.path.join( - run_dir, "checkpoints", self.timestamp_id - ), - "results_dir": os.path.join(run_dir, "results", self.timestamp_id), - "logs_dir": os.path.join( - run_dir, "logs", logger_name, self.timestamp_id - ), - }, - "slurm": slurm, - "noddp": noddp, - } - # AMP Scaler - self.scaler = torch.cuda.amp.GradScaler() if amp else None - - if "SLURM_JOB_ID" in os.environ and "folder" in self.config["slurm"]: - if "SLURM_ARRAY_JOB_ID" in os.environ: - self.config["slurm"]["job_id"] = "%s_%s" % ( - os.environ["SLURM_ARRAY_JOB_ID"], - os.environ["SLURM_ARRAY_TASK_ID"], - ) - else: - self.config["slurm"]["job_id"] = os.environ["SLURM_JOB_ID"] - self.config["slurm"]["folder"] = self.config["slurm"]["folder"].replace( - "%j", self.config["slurm"]["job_id"] - ) - if isinstance(dataset, list): - if len(dataset) > 0: - self.config["dataset"] = dataset[0] - if len(dataset) > 1: - self.config["val_dataset"] = dataset[1] - if len(dataset) > 2: - self.config["test_dataset"] = dataset[2] - elif isinstance(dataset, dict): - self.config["dataset"] = dataset.get("train", None) - self.config["val_dataset"] = dataset.get("val", None) - self.config["test_dataset"] = dataset.get("test", None) - else: - self.config["dataset"] = dataset - - self.normalizer = normalizer - # This supports the legacy way of providing norm parameters in dataset - if self.config.get("dataset", None) is not None and normalizer is None: - self.normalizer = self.config["dataset"] - - if not is_debug and distutils.is_master() and not is_hpo: - os.makedirs(self.config["cmd"]["checkpoint_dir"], exist_ok=True) - os.makedirs(self.config["cmd"]["results_dir"], exist_ok=True) - os.makedirs(self.config["cmd"]["logs_dir"], exist_ok=True) - - self.is_debug = is_debug - self.is_hpo = is_hpo - - if self.is_hpo: - # conditional import is necessary for checkpointing - - # sets the hpo checkpoint frequency - # default is no checkpointing - self.hpo_checkpoint_every = self.config["optim"].get("checkpoint_every", -1) - - if distutils.is_master(): - print(yaml.dump(self.config, default_flow_style=False)) - self.load() - - self.evaluator = Evaluator(task=name) - - def load(self): - self.load_seed_from_config() - self.load_logger() - self.load_datasets() - self.load_task() - self.load_model() - self.load_loss() - self.load_optimizer() - self.load_extras() - - def load_seed_from_config(self): - # https://pytorch.org/docs/stable/notes/randomness.html - seed = self.config["cmd"]["seed"] - if seed is None: - return - - random.seed(seed) - np.random.seed(seed) - torch.manual_seed(seed) - torch.cuda.manual_seed_all(seed) - torch.backends.cudnn.deterministic = True - torch.backends.cudnn.benchmark = False - - def load_logger(self): - self.logger = None - if not self.is_debug and distutils.is_master() and not self.is_hpo: - assert self.config["logger"] is not None, "Specify logger in config" - - logger = self.config["logger"] - logger_name = logger if isinstance(logger, str) else logger["name"] - assert logger_name, "Specify logger name" - - self.logger = registry.get_logger_class(logger_name)(self.config) - - def get_sampler(self, dataset, batch_size, shuffle): - if "load_balancing" in self.config["optim"]: - balancing_mode = self.config["optim"]["load_balancing"] - force_balancing = True - else: - balancing_mode = "atoms" - force_balancing = False - - if gp_utils.initialized(): - num_replicas = gp_utils.get_dp_world_size() - rank = gp_utils.get_dp_rank() - else: - num_replicas = distutils.get_world_size() - rank = distutils.get_rank() - sampler = BalancedBatchSampler( - dataset, - batch_size=batch_size, - num_replicas=num_replicas, - rank=rank, - device=self.device, - mode=balancing_mode, - shuffle=shuffle, - force_balancing=force_balancing, - ) - return sampler - - def get_dataloader(self, dataset, sampler): - loader = DataLoader( - dataset, - collate_fn=self.parallel_collater, - num_workers=self.config["optim"]["num_workers"], - pin_memory=True, - batch_sampler=sampler, - ) - return loader - - def load_datasets(self): - self.parallel_collater = ParallelCollater( - 0 if self.cpu else 1, - self.config["model_attributes"].get("otf_graph", False), - ) - - self.train_loader = self.val_loader = self.test_loader = None - - if self.config.get("dataset", None): - self.train_dataset = registry.get_dataset_class( - self.config["task"]["dataset"] - )(self.config["dataset"]) - self.train_sampler = self.get_sampler( - self.train_dataset, - self.config["optim"]["batch_size"], - shuffle=True, - ) - self.train_loader = self.get_dataloader( - self.train_dataset, - self.train_sampler, - ) - - if self.config.get("val_dataset", None): - self.val_dataset = registry.get_dataset_class( - self.config["task"]["dataset"] - )(self.config["val_dataset"]) - self.val_sampler = self.get_sampler( - self.val_dataset, - self.config["optim"].get( - "eval_batch_size", self.config["optim"]["batch_size"] - ), - shuffle=False, - ) - self.val_loader = self.get_dataloader( - self.val_dataset, - self.val_sampler, - ) - - if self.config.get("test_dataset", None): - self.test_dataset = registry.get_dataset_class( - self.config["task"]["dataset"] - )(self.config["test_dataset"]) - self.test_sampler = self.get_sampler( - self.test_dataset, - self.config["optim"].get( - "eval_batch_size", self.config["optim"]["batch_size"] - ), - shuffle=False, - ) - self.test_loader = self.get_dataloader( - self.test_dataset, - self.test_sampler, - ) - - # Normalizer for the dataset. - # Compute mean, std of training set labels. - self.normalizers = {} - if self.normalizer.get("normalize_labels", False): - if "target_mean" in self.normalizer: - self.normalizers["target"] = Normalizer( - mean=self.normalizer["target_mean"], - std=self.normalizer["target_std"], - device=self.device, - ) - else: - self.normalizers["target"] = Normalizer( - tensor=self.train_loader.dataset.data.y[ - self.train_loader.dataset.__indices__ - ], - device=self.device, - ) - - @abstractmethod - def load_task(self): - """Initialize task-specific information. Derived classes should implement this function.""" - - def load_model(self): - # Build model - if distutils.is_master(): - logging.info(f"Loading model: {self.config['model']}") - - # TODO: depreicated, remove. - bond_feat_dim = None - bond_feat_dim = self.config["model_attributes"].get("num_gaussians", 50) - - loader = self.train_loader or self.val_loader or self.test_loader - self.model = registry.get_model_class(self.config["model"])( - ( - loader.dataset[0].x.shape[-1] - if loader - and hasattr(loader.dataset[0], "x") - and loader.dataset[0].x is not None - else None - ), - bond_feat_dim, - self.num_targets, - **self.config["model_attributes"], - ).to(self.device) - - if distutils.is_master(): - logging.info( - f"Loaded {self.model.__class__.__name__} with " - f"{self.model.num_params} parameters." - ) - - if self.logger is not None: - self.logger.watch(self.model) - - self.model = OCPDataParallel( - self.model, - output_device=self.device, - num_gpus=1 if not self.cpu else 0, - ) - if distutils.initialized() and not self.config["noddp"]: - self.model = DistributedDataParallel(self.model, device_ids=[self.device]) - - def load_checkpoint(self, checkpoint_path): - if not os.path.isfile(checkpoint_path): - raise FileNotFoundError( - errno.ENOENT, "Checkpoint file not found", checkpoint_path - ) - - logging.info(f"Loading checkpoint from: {checkpoint_path}") - map_location = torch.device("cpu") if self.cpu else self.device - checkpoint = torch.load(checkpoint_path, map_location=map_location) - self.epoch = checkpoint.get("epoch", 0) - self.step = checkpoint.get("step", 0) - self.best_val_metric = checkpoint.get("best_val_metric", None) - self.primary_metric = checkpoint.get("primary_metric", None) - - # Match the "module." count in the keys of model and checkpoint state_dict - # DataParallel model has 1 "module.", DistributedDataParallel has 2 "module." - # Not using either of the above two would have no "module." - - ckpt_key_count = next(iter(checkpoint["state_dict"])).count("module") - mod_key_count = next(iter(self.model.state_dict())).count("module") - key_count_diff = mod_key_count - ckpt_key_count - - if key_count_diff > 0: - new_dict = { - key_count_diff * "module." + k: v - for k, v in checkpoint["state_dict"].items() - } - elif key_count_diff < 0: - new_dict = { - k[len("module.") * abs(key_count_diff) :]: v - for k, v in checkpoint["state_dict"].items() - } - else: - new_dict = checkpoint["state_dict"] - - strict = self.config["task"].get("strict_load", True) - load_state_dict(self.model, new_dict, strict=strict) - - if "optimizer" in checkpoint: - self.optimizer.load_state_dict(checkpoint["optimizer"]) - if "scheduler" in checkpoint and checkpoint["scheduler"] is not None: - self.scheduler.scheduler.load_state_dict(checkpoint["scheduler"]) - if "ema" in checkpoint and checkpoint["ema"] is not None: - self.ema.load_state_dict(checkpoint["ema"]) - else: - self.ema = None - - scale_dict = checkpoint.get("scale_dict", None) - if scale_dict: - logging.info( - "Overwriting scaling factors with those loaded from checkpoint. " - "If you're generating predictions with a pretrained checkpoint, this is the correct behavior. " - "To disable this, delete `scale_dict` from the checkpoint. " - ) - load_scales_compat(self._unwrapped_model, scale_dict) - - for key in checkpoint["normalizers"]: - if key in self.normalizers: - self.normalizers[key].load_state_dict(checkpoint["normalizers"][key]) - if self.scaler and checkpoint["amp"]: - self.scaler.load_state_dict(checkpoint["amp"]) - - def load_loss(self): - self.loss_fn = {} - self.loss_fn["energy"] = self.config["optim"].get("loss_energy", "mae") - self.loss_fn["force"] = self.config["optim"].get("loss_force", "mae") - for loss, loss_name in self.loss_fn.items(): - if loss_name in ["l1", "mae"]: - self.loss_fn[loss] = nn.L1Loss() - elif loss_name == "mse": - self.loss_fn[loss] = nn.MSELoss() - elif loss_name == "l2mae": - self.loss_fn[loss] = L2MAELoss() - elif loss_name == "atomwisel2": - self.loss_fn[loss] = AtomwiseL2Loss() - else: - raise NotImplementedError(f"Unknown loss function name: {loss_name}") - self.loss_fn[loss] = DDPLoss(self.loss_fn[loss]) - - def load_optimizer(self): - optimizer = self.config["optim"].get("optimizer", "AdamW") - optimizer = getattr(optim, optimizer) - - if self.config["optim"].get("weight_decay", 0) > 0: - - # Do not regularize bias etc. - params_decay = [] - params_no_decay = [] - for name, param in self.model.named_parameters(): - if param.requires_grad: - if "embedding" in name: - params_no_decay += [param] - elif "frequencies" in name: - params_no_decay += [param] - elif "bias" in name: - params_no_decay += [param] - else: - params_decay += [param] - - self.optimizer = optimizer( - [ - {"params": params_no_decay, "weight_decay": 0}, - { - "params": params_decay, - "weight_decay": self.config["optim"]["weight_decay"], - }, - ], - lr=self.config["optim"]["lr_initial"], - **self.config["optim"].get("optimizer_params", {}), - ) - else: - self.optimizer = optimizer( - params=self.model.parameters(), - lr=self.config["optim"]["lr_initial"], - **self.config["optim"].get("optimizer_params", {}), - ) - - def load_extras(self): - self.scheduler = LRScheduler(self.optimizer, self.config["optim"]) - self.clip_grad_norm = self.config["optim"].get("clip_grad_norm") - self.ema_decay = self.config["optim"].get("ema_decay") - if self.ema_decay: - self.ema = ExponentialMovingAverage( - self.model.parameters(), - self.ema_decay, - ) - else: - self.ema = None - - def save( - self, - metrics=None, - checkpoint_file="checkpoint.pt", - training_state=True, - ): - if not self.is_debug and distutils.is_master(): - if training_state: - return save_checkpoint( - { - "epoch": self.epoch, - "step": self.step, - "state_dict": self.model.state_dict(), - "optimizer": self.optimizer.state_dict(), - "scheduler": ( - self.scheduler.scheduler.state_dict() - if self.scheduler.scheduler_type != "Null" - else None - ), - "normalizers": { - key: value.state_dict() - for key, value in self.normalizers.items() - }, - "config": self.config, - "val_metrics": metrics, - "ema": self.ema.state_dict() if self.ema else None, - "amp": self.scaler.state_dict() if self.scaler else None, - "best_val_metric": self.best_val_metric, - "primary_metric": self.config["task"].get( - "primary_metric", - self.evaluator.task_primary_metric[self.name], - ), - }, - checkpoint_dir=self.config["cmd"]["checkpoint_dir"], - checkpoint_file=checkpoint_file, - ) - else: - if self.ema: - self.ema.store() - self.ema.copy_to() - ckpt_path = save_checkpoint( - { - "state_dict": self.model.state_dict(), - "normalizers": { - key: value.state_dict() - for key, value in self.normalizers.items() - }, - "config": self.config, - "val_metrics": metrics, - "amp": self.scaler.state_dict() if self.scaler else None, - }, - checkpoint_dir=self.config["cmd"]["checkpoint_dir"], - checkpoint_file=checkpoint_file, - ) - if self.ema: - self.ema.restore() - return ckpt_path - return None - - def save_hpo(self, epoch, step, metrics, checkpoint_every): - # default is no checkpointing - # checkpointing frequency can be adjusted by setting checkpoint_every in steps - # to checkpoint every time results are communicated to Ray Tune set checkpoint_every=1 - if checkpoint_every != -1 and step % checkpoint_every == 0: - with tune.checkpoint_dir(step=step) as checkpoint_dir: # noqa: F821 - path = os.path.join(checkpoint_dir, "checkpoint") - torch.save(self.save_state(epoch, step, metrics), path) - - def hpo_update(self, epoch, step, train_metrics, val_metrics, test_metrics=None): - progress = { - "steps": step, - "epochs": epoch, - "act_lr": self.optimizer.param_groups[0]["lr"], - } - # checkpointing must occur before reporter - # default is no checkpointing - self.save_hpo( - epoch, - step, - val_metrics, - self.hpo_checkpoint_every, - ) - # report metrics to tune - tune_reporter( # noqa: F821 - iters=progress, - train_metrics={k: train_metrics[k]["metric"] for k in self.metrics}, - val_metrics={k: val_metrics[k]["metric"] for k in val_metrics}, - test_metrics=test_metrics, - ) - - @abstractmethod - def train(self): - """Derived classes should implement this function.""" - - @torch.no_grad() - def validate(self, split="val", disable_tqdm=False): - ensure_fitted(self._unwrapped_model, warn=True) - - if distutils.is_master(): - logging.info(f"Evaluating on {split}.") - if self.is_hpo: - disable_tqdm = True - - self.model.eval() - if self.ema: - self.ema.store() - self.ema.copy_to() - - evaluator, metrics = Evaluator(task=self.name), {} - rank = distutils.get_rank() - - loader = self.val_loader if split == "val" else self.test_loader - - for i, batch in tqdm( - enumerate(loader), - total=len(loader), - position=rank, - desc="device {}".format(rank), - disable=disable_tqdm, - ): - # Forward. - with torch.amp.autocast("cuda", enabled=self.scaler is not None): - out = self._forward(batch) - loss = self._compute_loss(out, batch) - - # Compute metrics. - metrics = self._compute_metrics(out, batch, evaluator, metrics) - metrics = evaluator.update("loss", loss.item(), metrics) - - aggregated_metrics = {} - for k in metrics: - aggregated_metrics[k] = { - "total": distutils.all_reduce( - metrics[k]["total"], average=False, device=self.device - ), - "numel": distutils.all_reduce( - metrics[k]["numel"], average=False, device=self.device - ), - } - aggregated_metrics[k]["metric"] = ( - aggregated_metrics[k]["total"] / aggregated_metrics[k]["numel"] - ) - metrics = aggregated_metrics - - log_dict = {k: metrics[k]["metric"] for k in metrics} - log_dict.update({"epoch": self.epoch}) - if distutils.is_master(): - log_str = ["{}: {:.4f}".format(k, v) for k, v in log_dict.items()] - logging.info(", ".join(log_str)) - - # Make plots. - if self.logger is not None: - self.logger.log( - log_dict, - step=self.step, - split=split, - ) - - if self.ema: - self.ema.restore() - - return metrics - - @abstractmethod - def _forward(self, batch_list): - """Derived classes should implement this function.""" - - @abstractmethod - def _compute_loss(self, out, batch_list): - """Derived classes should implement this function.""" - - def _backward(self, loss): - self.optimizer.zero_grad() - loss.backward() - # Scale down the gradients of shared parameters - if hasattr(self.model.module, "shared_parameters"): - for p, factor in self.model.module.shared_parameters: - if hasattr(p, "grad") and p.grad is not None: - p.grad.detach().div_(factor) - else: - if not hasattr(self, "warned_shared_param_no_grad"): - self.warned_shared_param_no_grad = True - logging.warning( - "Some shared parameters do not have a gradient. " - "Please check if all shared parameters are used " - "and point to PyTorch parameters." - ) - if self.clip_grad_norm: - if self.scaler: - self.scaler.unscale_(self.optimizer) - grad_norm = torch.nn.utils.clip_grad_norm_( - self.model.parameters(), - max_norm=self.clip_grad_norm, - ) - if self.logger is not None: - self.logger.log({"grad_norm": grad_norm}, step=self.step, split="train") - if self.scaler: - self.scaler.step(self.optimizer) - self.scaler.update() - else: - self.optimizer.step() - if self.ema: - self.ema.update() - - def save_results(self, predictions, results_file, keys): - if results_file is None: - return - - results_file_path = os.path.join( - self.config["cmd"]["results_dir"], - f"{self.name}_{results_file}_{distutils.get_rank()}.npz", - ) - np.savez_compressed( - results_file_path, - ids=predictions["id"], - **{key: predictions[key] for key in keys}, - ) - - distutils.synchronize() - if distutils.is_master(): - gather_results = defaultdict(list) - full_path = os.path.join( - self.config["cmd"]["results_dir"], - f"{self.name}_{results_file}.npz", - ) - - for i in range(distutils.get_world_size()): - rank_path = os.path.join( - self.config["cmd"]["results_dir"], - f"{self.name}_{results_file}_{i}.npz", - ) - rank_results = np.load(rank_path, allow_pickle=True) - gather_results["ids"].extend(rank_results["ids"]) - for key in keys: - gather_results[key].extend(rank_results[key]) - os.remove(rank_path) - - # Because of how distributed sampler works, some system ids - # might be repeated to make no. of samples even across GPUs. - _, idx = np.unique(gather_results["ids"], return_index=True) - gather_results["ids"] = np.array(gather_results["ids"])[idx] - for k in keys: - if k == "forces": - gather_results[k] = np.concatenate(np.array(gather_results[k])[idx]) - elif k == "chunk_idx": - gather_results[k] = np.cumsum(np.array(gather_results[k])[idx])[:-1] - else: - gather_results[k] = np.array(gather_results[k])[idx] - - logging.info(f"Writing results to {full_path}") - np.savez_compressed(full_path, **gather_results) diff --git a/ocpmodels/trainers/energy_trainer.py b/ocpmodels/trainers/energy_trainer.py deleted file mode 100644 index ff76999d32766793b15d6089afcb82cafc5482b3..0000000000000000000000000000000000000000 --- a/ocpmodels/trainers/energy_trainer.py +++ /dev/null @@ -1,319 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import logging - -import torch -import torch_geometric -from tqdm import tqdm - -from ocpmodels.common import distutils -from ocpmodels.common.registry import registry -from ocpmodels.modules.scaling.util import ensure_fitted -from ocpmodels.trainers.base_trainer import BaseTrainer - - -@registry.register_trainer("energy") -class EnergyTrainer(BaseTrainer): - """ - Trainer class for the Initial Structure to Relaxed Energy (IS2RE) task. - - .. note:: - - Examples of configurations for task, model, dataset and optimizer - can be found in `configs/ocp_is2re `_. - - - Args: - task (dict): Task configuration. - model (dict): Model configuration. - dataset (dict): Dataset configuration. The dataset needs to be a SinglePointLMDB dataset. - optimizer (dict): Optimizer configuration. - identifier (str): Experiment identifier that is appended to log directory. - run_dir (str, optional): Path to the run directory where logs are to be saved. - (default: :obj:`None`) - is_debug (bool, optional): Run in debug mode. - (default: :obj:`False`) - is_hpo (bool, optional): Run hyperparameter optimization with Ray Tune. - (default: :obj:`False`) - print_every (int, optional): Frequency of printing logs. - (default: :obj:`100`) - seed (int, optional): Random number seed. - (default: :obj:`None`) - logger (str, optional): Type of logger to be used. - (default: :obj:`tensorboard`) - local_rank (int, optional): Local rank of the process, only applicable for distributed training. - (default: :obj:`0`) - amp (bool, optional): Run using automatic mixed precision. - (default: :obj:`False`) - slurm (dict): Slurm configuration. Currently just for keeping track. - (default: :obj:`{}`) - """ - - def __init__( - self, - task, - model, - dataset, - optimizer, - identifier, - normalizer=None, - timestamp_id=None, - run_dir=None, - is_debug=False, - is_hpo=False, - print_every=100, - seed=None, - logger="tensorboard", - local_rank=0, - amp=False, - cpu=False, - slurm={}, - noddp=False, - ): - super().__init__( - task=task, - model=model, - dataset=dataset, - optimizer=optimizer, - identifier=identifier, - normalizer=normalizer, - timestamp_id=timestamp_id, - run_dir=run_dir, - is_debug=is_debug, - is_hpo=is_hpo, - print_every=print_every, - seed=seed, - logger=logger, - local_rank=local_rank, - amp=amp, - cpu=cpu, - name="is2re", - slurm=slurm, - noddp=noddp, - ) - - def load_task(self): - logging.info(f"Loading dataset: {self.config['task']['dataset']}") - self.num_targets = 1 - - @torch.no_grad() - def predict(self, loader, per_image=True, results_file=None, disable_tqdm=False): - ensure_fitted(self._unwrapped_model) - - if distutils.is_master() and not disable_tqdm: - logging.info("Predicting on test.") - assert isinstance( - loader, - ( - torch.utils.data.dataloader.DataLoader, - torch_geometric.data.Batch, - ), - ) - rank = distutils.get_rank() - - if isinstance(loader, torch_geometric.data.Batch): - loader = [[loader]] - - self.model.eval() - if self.ema: - self.ema.store() - self.ema.copy_to() - - if self.normalizers is not None and "target" in self.normalizers: - self.normalizers["target"].to(self.device) - predictions = {"id": [], "energy": []} - - for i, batch in tqdm( - enumerate(loader), - total=len(loader), - position=rank, - desc="device {}".format(rank), - disable=disable_tqdm, - ): - with torch.amp.autocast("cuda", enabled=self.scaler is not None): - out = self._forward(batch) - - if self.normalizers is not None and "target" in self.normalizers: - out["energy"] = self.normalizers["target"].denorm(out["energy"]) - - if per_image: - predictions["id"].extend([str(i) for i in batch[0].sid.tolist()]) - predictions["energy"].extend(out["energy"].cpu().detach().numpy()) - else: - predictions["energy"] = out["energy"].detach() - return predictions - - self.save_results(predictions, results_file, keys=["energy"]) - - if self.ema: - self.ema.restore() - - return predictions - - def train(self, disable_eval_tqdm=False): - ensure_fitted(self._unwrapped_model, warn=True) - - eval_every = self.config["optim"].get("eval_every", len(self.train_loader)) - primary_metric = self.config["task"].get( - "primary_metric", self.evaluator.task_primary_metric[self.name] - ) - self.best_val_metric = 1e9 - - # Calculate start_epoch from step instead of loading the epoch number - # to prevent inconsistencies due to different batch size in checkpoint. - start_epoch = self.step // len(self.train_loader) - - for epoch_int in range(start_epoch, self.config["optim"]["max_epochs"]): - self.train_sampler.set_epoch(epoch_int) - skip_steps = self.step % len(self.train_loader) - train_loader_iter = iter(self.train_loader) - - for i in range(skip_steps, len(self.train_loader)): - self.epoch = epoch_int + (i + 1) / len(self.train_loader) - self.step = epoch_int * len(self.train_loader) + i + 1 - self.model.train() - - # Get a batch. - batch = next(train_loader_iter) - - # Forward, loss, backward. - with torch.amp.autocast("cuda", enabled=self.scaler is not None): - out = self._forward(batch) - loss = self._compute_loss(out, batch) - loss = self.scaler.scale(loss) if self.scaler else loss - self._backward(loss) - scale = self.scaler.get_scale() if self.scaler else 1.0 - - # Compute metrics. - self.metrics = self._compute_metrics( - out, - batch, - self.evaluator, - metrics={}, - ) - self.metrics = self.evaluator.update( - "loss", loss.item() / scale, self.metrics - ) - - # Log metrics. - log_dict = {k: self.metrics[k]["metric"] for k in self.metrics} - log_dict.update( - { - "lr": self.scheduler.get_lr(), - "epoch": self.epoch, - "step": self.step, - } - ) - if ( - self.step % self.config["cmd"]["print_every"] == 0 - and distutils.is_master() - and not self.is_hpo - ): - log_str = ["{}: {:.2e}".format(k, v) for k, v in log_dict.items()] - print(", ".join(log_str)) - self.metrics = {} - - if self.logger is not None: - self.logger.log( - log_dict, - step=self.step, - split="train", - ) - - # Evaluate on val set after every `eval_every` iterations. - if self.step % eval_every == 0: - self.save(checkpoint_file="checkpoint.pt", training_state=True) - - if self.val_loader is not None: - val_metrics = self.validate( - split="val", - disable_tqdm=disable_eval_tqdm, - ) - if ( - val_metrics[self.evaluator.task_primary_metric[self.name]][ - "metric" - ] - < self.best_val_metric - ): - self.best_val_metric = val_metrics[ - self.evaluator.task_primary_metric[self.name] - ]["metric"] - self.save( - metrics=val_metrics, - checkpoint_file="best_checkpoint.pt", - training_state=False, - ) - if self.test_loader is not None: - self.predict( - self.test_loader, - results_file="predictions", - disable_tqdm=False, - ) - - if self.is_hpo: - self.hpo_update( - self.epoch, - self.step, - self.metrics, - val_metrics, - ) - - if self.scheduler.scheduler_type == "ReduceLROnPlateau": - if self.step % eval_every == 0: - self.scheduler.step( - metrics=val_metrics[primary_metric]["metric"], - ) - else: - self.scheduler.step() - - torch.cuda.empty_cache() - - self.train_dataset.close_db() - if self.config.get("val_dataset", False): - self.val_dataset.close_db() - if self.config.get("test_dataset", False): - self.test_dataset.close_db() - - def _forward(self, batch_list): - output = self.model(batch_list) - - if output.shape[-1] == 1: - output = output.view(-1) - - return { - "energy": output, - } - - def _compute_loss(self, out, batch_list): - energy_target = torch.cat( - [batch.y_relaxed.to(self.device) for batch in batch_list], dim=0 - ) - - if self.normalizer.get("normalize_labels", False): - target_normed = self.normalizers["target"].norm(energy_target) - else: - target_normed = energy_target - - loss = self.loss_fn["energy"](out["energy"], target_normed) - return loss - - def _compute_metrics(self, out, batch_list, evaluator, metrics={}): - energy_target = torch.cat( - [batch.y_relaxed.to(self.device) for batch in batch_list], dim=0 - ) - - if self.normalizer.get("normalize_labels", False): - out["energy"] = self.normalizers["target"].denorm(out["energy"]) - - metrics = evaluator.eval( - out, - {"energy": energy_target}, - prev_metrics=metrics, - ) - - return metrics diff --git a/ocpmodels/trainers/forces_trainer.py b/ocpmodels/trainers/forces_trainer.py deleted file mode 100644 index 045e5de7f262b87b4fe9365e721eec2017936f93..0000000000000000000000000000000000000000 --- a/ocpmodels/trainers/forces_trainer.py +++ /dev/null @@ -1,754 +0,0 @@ -""" -Copyright (c) Facebook, Inc. and its affiliates. - -This source code is licensed under the MIT license found in the -LICENSE file in the root directory of this source tree. -""" - -import logging -import os -import pathlib -from collections import defaultdict -from pathlib import Path - -import numpy as np -import torch -import torch_geometric -from tqdm import tqdm - -from ocpmodels.common import distutils -from ocpmodels.common.registry import registry -from ocpmodels.common.relaxation.ml_relaxation import ml_relax -from ocpmodels.common.utils import check_traj_files -from ocpmodels.modules.evaluator import Evaluator -from ocpmodels.modules.normalizer import Normalizer -from ocpmodels.modules.scaling.util import ensure_fitted -from ocpmodels.trainers.base_trainer import BaseTrainer - - -@registry.register_trainer("forces") -class ForcesTrainer(BaseTrainer): - """ - Trainer class for the Structure to Energy & Force (S2EF) and Initial State to - Relaxed State (IS2RS) tasks. - - .. note:: - - Examples of configurations for task, model, dataset and optimizer - can be found in `configs/ocp_s2ef `_ - and `configs/ocp_is2rs `_. - - Args: - task (dict): Task configuration. - model (dict): Model configuration. - dataset (dict): Dataset configuration. The dataset needs to be a SinglePointLMDB dataset. - optimizer (dict): Optimizer configuration. - identifier (str): Experiment identifier that is appended to log directory. - run_dir (str, optional): Path to the run directory where logs are to be saved. - (default: :obj:`None`) - is_debug (bool, optional): Run in debug mode. - (default: :obj:`False`) - is_hpo (bool, optional): Run hyperparameter optimization with Ray Tune. - (default: :obj:`False`) - print_every (int, optional): Frequency of printing logs. - (default: :obj:`100`) - seed (int, optional): Random number seed. - (default: :obj:`None`) - logger (str, optional): Type of logger to be used. - (default: :obj:`tensorboard`) - local_rank (int, optional): Local rank of the process, only applicable for distributed training. - (default: :obj:`0`) - amp (bool, optional): Run using automatic mixed precision. - (default: :obj:`False`) - slurm (dict): Slurm configuration. Currently just for keeping track. - (default: :obj:`{}`) - """ - - def __init__( - self, - task, - model, - dataset, - optimizer, - identifier, - normalizer=None, - timestamp_id=None, - run_dir=None, - is_debug=False, - is_hpo=False, - print_every=100, - seed=None, - logger="tensorboard", - local_rank=0, - amp=False, - cpu=False, - slurm={}, - noddp=False, - ): - super().__init__( - task=task, - model=model, - dataset=dataset, - optimizer=optimizer, - identifier=identifier, - normalizer=normalizer, - timestamp_id=timestamp_id, - run_dir=run_dir, - is_debug=is_debug, - is_hpo=is_hpo, - print_every=print_every, - seed=seed, - logger=logger, - local_rank=local_rank, - amp=amp, - cpu=cpu, - name="s2ef", - slurm=slurm, - noddp=noddp, - ) - - def load_task(self): - logging.info(f"Loading dataset: {self.config['task']['dataset']}") - - if "relax_dataset" in self.config["task"]: - self.relax_dataset = registry.get_dataset_class("lmdb")( - self.config["task"]["relax_dataset"] - ) - self.relax_sampler = self.get_sampler( - self.relax_dataset, - self.config["optim"].get( - "eval_batch_size", self.config["optim"]["batch_size"] - ), - shuffle=False, - ) - self.relax_loader = self.get_dataloader( - self.relax_dataset, - self.relax_sampler, - ) - - self.num_targets = 1 - - # If we're computing gradients wrt input, set mean of normalizer to 0 -- - # since it is lost when compute dy / dx -- and std to forward target std - if self.config["model_attributes"].get("regress_forces", True): - if self.normalizer.get("normalize_labels", False): - if "grad_target_mean" in self.normalizer: - self.normalizers["grad_target"] = Normalizer( - mean=self.normalizer["grad_target_mean"], - std=self.normalizer["grad_target_std"], - device=self.device, - ) - else: - self.normalizers["grad_target"] = Normalizer( - tensor=self.train_loader.dataset.data.y[ - self.train_loader.dataset.__indices__ - ], - device=self.device, - ) - self.normalizers["grad_target"].mean.fill_(0) - - # Takes in a new data source and generates predictions on it. - @torch.no_grad() - def predict( - self, - data_loader, - per_image=True, - results_file=None, - disable_tqdm=False, - ): - ensure_fitted(self._unwrapped_model, warn=True) - - if distutils.is_master() and not disable_tqdm: - logging.info("Predicting on test.") - assert isinstance( - data_loader, - ( - torch.utils.data.dataloader.DataLoader, - torch_geometric.data.Batch, - ), - ) - rank = distutils.get_rank() - - if isinstance(data_loader, torch_geometric.data.Batch): - data_loader = [[data_loader]] - - self.model.eval() - if self.ema: - self.ema.store() - self.ema.copy_to() - - if self.normalizers is not None and "target" in self.normalizers: - self.normalizers["target"].to(self.device) - self.normalizers["grad_target"].to(self.device) - - predictions = {"id": [], "energy": [], "forces": [], "chunk_idx": []} - - for i, batch_list in tqdm( - enumerate(data_loader), - total=len(data_loader), - position=rank, - desc="device {}".format(rank), - disable=disable_tqdm, - ): - with torch.amp.autocast("cuda", enabled=self.scaler is not None): - out = self._forward(batch_list) - - if self.normalizers is not None and "target" in self.normalizers: - out["energy"] = self.normalizers["target"].denorm(out["energy"]) - out["forces"] = self.normalizers["grad_target"].denorm(out["forces"]) - if per_image: - systemids = [ - str(i) + "_" + str(j) - for i, j in zip( - batch_list[0].sid.tolist(), batch_list[0].fid.tolist() - ) - ] - predictions["id"].extend(systemids) - batch_natoms = torch.cat([batch.natoms for batch in batch_list]) - batch_fixed = torch.cat([batch.fixed for batch in batch_list]) - # total energy target requires predictions to be saved in float32 - # default is float16 - if ( - self.config["task"].get("prediction_dtype", "float16") == "float32" - or self.config["task"]["dataset"] == "oc22_lmdb" - ): - predictions["energy"].extend( - out["energy"].cpu().detach().to(torch.float32).numpy() - ) - forces = out["forces"].cpu().detach().to(torch.float32) - else: - predictions["energy"].extend( - out["energy"].cpu().detach().to(torch.float16).numpy() - ) - forces = out["forces"].cpu().detach().to(torch.float16) - per_image_forces = torch.split(forces, batch_natoms.tolist()) - per_image_forces = [force.numpy() for force in per_image_forces] - # evalAI only requires forces on free atoms - if results_file is not None: - _per_image_fixed = torch.split(batch_fixed, batch_natoms.tolist()) - _per_image_free_forces = [ - force[(fixed == 0).tolist()] - for force, fixed in zip(per_image_forces, _per_image_fixed) - ] - _chunk_idx = np.array( - [free_force.shape[0] for free_force in _per_image_free_forces] - ) - per_image_forces = _per_image_free_forces - predictions["chunk_idx"].extend(_chunk_idx) - predictions["forces"].extend(per_image_forces) - else: - predictions["energy"] = out["energy"].detach() - predictions["forces"] = out["forces"].detach() - if self.ema: - self.ema.restore() - return predictions - - predictions["forces"] = np.array(predictions["forces"]) - predictions["chunk_idx"] = np.array(predictions["chunk_idx"]) - predictions["energy"] = np.array(predictions["energy"]) - predictions["id"] = np.array(predictions["id"]) - self.save_results( - predictions, results_file, keys=["energy", "forces", "chunk_idx"] - ) - - if self.ema: - self.ema.restore() - - return predictions - - def update_best( - self, - primary_metric, - val_metrics, - disable_eval_tqdm=True, - ): - if ( - "mae" in primary_metric - and val_metrics[primary_metric]["metric"] < self.best_val_metric - ) or ( - "mae" not in primary_metric - and val_metrics[primary_metric]["metric"] > self.best_val_metric - ): - self.best_val_metric = val_metrics[primary_metric]["metric"] - self.save( - metrics=val_metrics, - checkpoint_file="best_checkpoint.pt", - training_state=False, - ) - if self.test_loader is not None: - self.predict( - self.test_loader, - results_file="predictions", - disable_tqdm=disable_eval_tqdm, - ) - - def train(self, disable_eval_tqdm=False): - ensure_fitted(self._unwrapped_model, warn=True) - - eval_every = self.config["optim"].get("eval_every", len(self.train_loader)) - checkpoint_every = self.config["optim"].get("checkpoint_every", eval_every) - primary_metric = self.config["task"].get( - "primary_metric", self.evaluator.task_primary_metric[self.name] - ) - if not hasattr(self, "primary_metric") or self.primary_metric != primary_metric: - self.best_val_metric = 1e9 if "mae" in primary_metric else -1.0 - else: - primary_metric = self.primary_metric - self.metrics = {} - - # Calculate start_epoch from step instead of loading the epoch number - # to prevent inconsistencies due to different batch size in checkpoint. - start_epoch = self.step // len(self.train_loader) - - for epoch_int in range(start_epoch, self.config["optim"]["max_epochs"]): - self.train_sampler.set_epoch(epoch_int) - skip_steps = self.step % len(self.train_loader) - train_loader_iter = iter(self.train_loader) - - for i in range(skip_steps, len(self.train_loader)): - self.epoch = epoch_int + (i + 1) / len(self.train_loader) - self.step = epoch_int * len(self.train_loader) + i + 1 - self.model.train() - - # Get a batch. - batch = next(train_loader_iter) - - # Forward, loss, backward. - with torch.amp.autocast("cuda", enabled=self.scaler is not None): - out = self._forward(batch) - loss = self._compute_loss(out, batch) - loss = self.scaler.scale(loss) if self.scaler else loss - self._backward(loss) - scale = self.scaler.get_scale() if self.scaler else 1.0 - - # Compute metrics. - self.metrics = self._compute_metrics( - out, - batch, - self.evaluator, - self.metrics, - ) - self.metrics = self.evaluator.update( - "loss", loss.item() / scale, self.metrics - ) - - # Log metrics. - log_dict = {k: self.metrics[k]["metric"] for k in self.metrics} - log_dict.update( - { - "lr": self.scheduler.get_lr(), - "epoch": self.epoch, - "step": self.step, - } - ) - if ( - self.step % self.config["cmd"]["print_every"] == 0 - and distutils.is_master() - and not self.is_hpo - ): - log_str = ["{}: {:.2e}".format(k, v) for k, v in log_dict.items()] - logging.info(", ".join(log_str)) - self.metrics = {} - - if self.logger is not None: - self.logger.log( - log_dict, - step=self.step, - split="train", - ) - - if checkpoint_every != -1 and self.step % checkpoint_every == 0: - self.save(checkpoint_file="checkpoint.pt", training_state=True) - - # Evaluate on val set every `eval_every` iterations. - if self.step % eval_every == 0: - if self.val_loader is not None: - val_metrics = self.validate( - split="val", - disable_tqdm=disable_eval_tqdm, - ) - self.update_best( - primary_metric, - val_metrics, - disable_eval_tqdm=disable_eval_tqdm, - ) - if self.is_hpo: - self.hpo_update( - self.epoch, - self.step, - self.metrics, - val_metrics, - ) - - if self.config["task"].get("eval_relaxations", False): - if "relax_dataset" not in self.config["task"]: - logging.warning( - "Cannot evaluate relaxations, relax_dataset not specified" - ) - else: - self.run_relaxations() - - if self.scheduler.scheduler_type == "ReduceLROnPlateau": - if self.step % eval_every == 0: - self.scheduler.step( - metrics=val_metrics[primary_metric]["metric"], - ) - else: - self.scheduler.step() - - torch.cuda.empty_cache() - - if checkpoint_every == -1: - self.save(checkpoint_file="checkpoint.pt", training_state=True) - - self.train_dataset.close_db() - if self.config.get("val_dataset", False): - self.val_dataset.close_db() - if self.config.get("test_dataset", False): - self.test_dataset.close_db() - - def _forward(self, batch_list): - # forward pass. - if self.config["model_attributes"].get("regress_forces", True): - out_energy, out_forces = self.model(batch_list) - else: - out_energy = self.model(batch_list) - - if out_energy.shape[-1] == 1: - out_energy = out_energy.view(-1) - - out = { - "energy": out_energy, - } - - if self.config["model_attributes"].get("regress_forces", True): - out["forces"] = out_forces - - return out - - def _compute_loss(self, out, batch_list): - loss = [] - - # Energy loss. - energy_target = torch.cat( - [batch.y.to(self.device) for batch in batch_list], dim=0 - ) - if self.normalizer.get("normalize_labels", False): - energy_target = self.normalizers["target"].norm(energy_target) - energy_mult = self.config["optim"].get("energy_coefficient", 1) - loss.append(energy_mult * self.loss_fn["energy"](out["energy"], energy_target)) - - # Force loss. - if self.config["model_attributes"].get("regress_forces", True): - force_target = torch.cat( - [batch.force.to(self.device) for batch in batch_list], dim=0 - ) - if self.normalizer.get("normalize_labels", False): - force_target = self.normalizers["grad_target"].norm(force_target) - - tag_specific_weights = self.config["task"].get("tag_specific_weights", []) - if tag_specific_weights != []: - # handle tag specific weights as introduced in forcenet - assert len(tag_specific_weights) == 3 - - batch_tags = torch.cat( - [batch.tags.float().to(self.device) for batch in batch_list], - dim=0, - ) - weight = torch.zeros_like(batch_tags) - weight[batch_tags == 0] = tag_specific_weights[0] - weight[batch_tags == 1] = tag_specific_weights[1] - weight[batch_tags == 2] = tag_specific_weights[2] - - if self.config["optim"].get("loss_force", "l2mae") == "l2mae": - # zero out nans, if any - found_nans_or_infs = not torch.all(out["forces"].isfinite()) - if found_nans_or_infs is True: - logging.warning("Found nans while computing loss") - out["forces"] = torch.nan_to_num(out["forces"], nan=0.0) - - dists = torch.norm(out["forces"] - force_target, p=2, dim=-1) - weighted_dists_sum = (dists * weight).sum() - - num_samples = out["forces"].shape[0] - num_samples = distutils.all_reduce(num_samples, device=self.device) - weighted_dists_sum = ( - weighted_dists_sum * distutils.get_world_size() / num_samples - ) - - force_mult = self.config["optim"].get("force_coefficient", 30) - loss.append(force_mult * weighted_dists_sum) - else: - raise NotImplementedError - else: - # Force coefficient = 30 has been working well for us. - force_mult = self.config["optim"].get("force_coefficient", 30) - if self.config["task"].get("train_on_free_atoms", False): - fixed = torch.cat( - [batch.fixed.to(self.device) for batch in batch_list] - ) - mask = fixed == 0 - if ( - self.config["optim"] - .get("loss_force", "mae") - .startswith("atomwise") - ): - force_mult = self.config["optim"].get("force_coefficient", 1) - natoms = torch.cat( - [batch.natoms.to(self.device) for batch in batch_list] - ) - natoms = torch.repeat_interleave(natoms, natoms) - force_loss = force_mult * self.loss_fn["force"]( - out["forces"][mask], - force_target[mask], - natoms=natoms[mask], - batch_size=batch_list[0].natoms.shape[0], - ) - loss.append(force_loss) - else: - loss.append( - force_mult - * self.loss_fn["force"]( - out["forces"][mask], force_target[mask] - ) - ) - else: - loss.append( - force_mult * self.loss_fn["force"](out["forces"], force_target) - ) - - # Sanity check to make sure the compute graph is correct. - for lc in loss: - assert hasattr(lc, "grad_fn") - - loss = sum(loss) - return loss - - def _compute_metrics(self, out, batch_list, evaluator, metrics={}): - natoms = torch.cat( - [batch.natoms.to(self.device) for batch in batch_list], dim=0 - ) - - target = { - "energy": torch.cat( - [batch.y.to(self.device) for batch in batch_list], dim=0 - ), - "forces": torch.cat( - [batch.force.to(self.device) for batch in batch_list], dim=0 - ), - "natoms": natoms, - } - - out["natoms"] = natoms - - if self.config["task"].get("eval_on_free_atoms", True): - fixed = torch.cat([batch.fixed.to(self.device) for batch in batch_list]) - mask = fixed == 0 - out["forces"] = out["forces"][mask] - target["forces"] = target["forces"][mask] - - s_idx = 0 - natoms_free = [] - for natoms in target["natoms"]: - natoms_free.append(torch.sum(mask[s_idx : s_idx + natoms]).item()) - s_idx += natoms - target["natoms"] = torch.LongTensor(natoms_free).to(self.device) - out["natoms"] = torch.LongTensor(natoms_free).to(self.device) - - if self.normalizer.get("normalize_labels", False): - out["energy"] = self.normalizers["target"].denorm(out["energy"]) - out["forces"] = self.normalizers["grad_target"].denorm(out["forces"]) - - metrics = evaluator.eval(out, target, prev_metrics=metrics) - return metrics - - def run_relaxations(self, split="val"): - ensure_fitted(self._unwrapped_model) - - # When set to true, uses deterministic CUDA scatter ops, if available. - # https://pytorch.org/docs/stable/generated/torch.use_deterministic_algorithms.html#torch.use_deterministic_algorithms - # Only implemented for GemNet-OC currently. - registry.register( - "set_deterministic_scatter", - self.config["task"].get("set_deterministic_scatter", False), - ) - - logging.info("Running ML-relaxations") - self.model.eval() - if self.ema: - self.ema.store() - self.ema.copy_to() - - evaluator_is2rs, metrics_is2rs = Evaluator(task="is2rs"), {} - evaluator_is2re, metrics_is2re = Evaluator(task="is2re"), {} - - # Need both `pos_relaxed` and `y_relaxed` to compute val IS2R* metrics. - # Else just generate predictions. - if ( - hasattr(self.relax_dataset[0], "pos_relaxed") - and self.relax_dataset[0].pos_relaxed is not None - ) and ( - hasattr(self.relax_dataset[0], "y_relaxed") - and self.relax_dataset[0].y_relaxed is not None - ): - split = "val" - else: - split = "test" - - ids = [] - relaxed_positions = [] - chunk_idx = [] - for i, batch in tqdm( - enumerate(self.relax_loader), total=len(self.relax_loader) - ): - if i >= self.config["task"].get("num_relaxation_batches", 1e9): - break - - # If all traj files already exist, then skip this batch - if check_traj_files( - batch, self.config["task"]["relax_opt"].get("traj_dir", None) - ): - logging.info(f"Skipping batch: {batch[0].sid.tolist()}") - continue - - relaxed_batch = ml_relax( - batch=batch, - model=self, - steps=self.config["task"].get("relaxation_steps", 200), - fmax=self.config["task"].get("relaxation_fmax", 0.0), - relax_opt=self.config["task"]["relax_opt"], - save_full_traj=self.config["task"].get("save_full_traj", True), - device=self.device, - transform=None, - ) - - if self.config["task"].get("write_pos", False): - systemids = [str(i) for i in relaxed_batch.sid.tolist()] - natoms = relaxed_batch.natoms.tolist() - positions = torch.split(relaxed_batch.pos, natoms) - batch_relaxed_positions = [pos.tolist() for pos in positions] - - relaxed_positions += batch_relaxed_positions - chunk_idx += natoms - ids += systemids - - if split == "val": - mask = relaxed_batch.fixed == 0 - s_idx = 0 - natoms_free = [] - for natoms in relaxed_batch.natoms: - natoms_free.append(torch.sum(mask[s_idx : s_idx + natoms]).item()) - s_idx += natoms - - target = { - "energy": relaxed_batch.y_relaxed, - "positions": relaxed_batch.pos_relaxed[mask], - "cell": relaxed_batch.cell, - "pbc": torch.tensor([True, True, True]), - "natoms": torch.LongTensor(natoms_free), - } - - prediction = { - "energy": relaxed_batch.y, - "positions": relaxed_batch.pos[mask], - "cell": relaxed_batch.cell, - "pbc": torch.tensor([True, True, True]), - "natoms": torch.LongTensor(natoms_free), - } - - metrics_is2rs = evaluator_is2rs.eval( - prediction, - target, - metrics_is2rs, - ) - metrics_is2re = evaluator_is2re.eval( - {"energy": prediction["energy"]}, - {"energy": target["energy"]}, - metrics_is2re, - ) - - if self.config["task"].get("write_pos", False): - rank = distutils.get_rank() - pos_filename = os.path.join( - self.config["cmd"]["results_dir"], f"relaxed_pos_{rank}.npz" - ) - np.savez_compressed( - pos_filename, - ids=ids, - pos=np.array(relaxed_positions, dtype=object), - chunk_idx=chunk_idx, - ) - - distutils.synchronize() - if distutils.is_master(): - gather_results = defaultdict(list) - full_path = os.path.join( - self.config["cmd"]["results_dir"], - "relaxed_positions.npz", - ) - - for i in range(distutils.get_world_size()): - rank_path = os.path.join( - self.config["cmd"]["results_dir"], - f"relaxed_pos_{i}.npz", - ) - rank_results = np.load(rank_path, allow_pickle=True) - gather_results["ids"].extend(rank_results["ids"]) - gather_results["pos"].extend(rank_results["pos"]) - gather_results["chunk_idx"].extend(rank_results["chunk_idx"]) - os.remove(rank_path) - - # Because of how distributed sampler works, some system ids - # might be repeated to make no. of samples even across GPUs. - _, idx = np.unique(gather_results["ids"], return_index=True) - gather_results["ids"] = np.array(gather_results["ids"])[idx] - gather_results["pos"] = np.concatenate( - np.array(gather_results["pos"])[idx] - ) - gather_results["chunk_idx"] = np.cumsum( - np.array(gather_results["chunk_idx"])[idx] - )[ - :-1 - ] # np.split does not need last idx, assumes n-1:end - - logging.info(f"Writing results to {full_path}") - np.savez_compressed(full_path, **gather_results) - - if split == "val": - for task in ["is2rs", "is2re"]: - metrics = eval(f"metrics_{task}") - aggregated_metrics = {} - for k in metrics: - aggregated_metrics[k] = { - "total": distutils.all_reduce( - metrics[k]["total"], - average=False, - device=self.device, - ), - "numel": distutils.all_reduce( - metrics[k]["numel"], - average=False, - device=self.device, - ), - } - aggregated_metrics[k]["metric"] = ( - aggregated_metrics[k]["total"] / aggregated_metrics[k]["numel"] - ) - metrics = aggregated_metrics - - # Make plots. - log_dict = {f"{task}_{k}": metrics[k]["metric"] for k in metrics} - if self.logger is not None: - self.logger.log( - log_dict, - step=self.step, - split=split, - ) - - if distutils.is_master(): - logging.info(metrics) - - if self.ema: - self.ema.restore() - - registry.unregister("set_deterministic_scatter") diff --git a/ocpmodels/units.py b/ocpmodels/units.py deleted file mode 100644 index dda0a0e00fa0022ea0aeea525f7bdac268671e17..0000000000000000000000000000000000000000 --- a/ocpmodels/units.py +++ /dev/null @@ -1,142 +0,0 @@ -# unit conversions based on NIST values - -# HORM used before: -# hartree_to_ev = 27.2114 -# bohr_to_angstrom = 0.529177 -# ev_angstrom_2_to_hartree_bohr_2 = (bohr_to_angstrom**2) / hartree_to_ev - -# https://physics.nist.gov/cgi-bin/cuu/Value?evhr -# 1 eV = 3.674 932 217 5665 x 10-2 Eh -ev_to_hartree = 3.6749322175665e-2 -hartree_to_ev = 1 / ev_to_hartree - -# https://physics.nist.gov/cgi-bin/cuu/Value?bohrrada0 -# 1 bohr = 5.291 772 105 44 x 10-11 m -bohr_to_angstrom = 0.529177210544 -angstrom_to_bohr = 1 / bohr_to_angstrom - -ev_angstrom_2_to_hartree_bohr_2 = (bohr_to_angstrom**2) / hartree_to_ev - - -ELEMENT_TO_ATOMIC_NUMBER = { - "H": 1, - "He": 2, - "Li": 3, - "Be": 4, - "B": 5, - "C": 6, - "N": 7, - "O": 8, - "F": 9, - "Ne": 10, - "Na": 11, - "Mg": 12, - "Al": 13, - "Si": 14, - "P": 15, - "S": 16, - "Cl": 17, - "Ar": 18, - "K": 19, - "Ca": 20, - "Sc": 21, - "Ti": 22, - "V": 23, - "Cr": 24, - "Mn": 25, - "Fe": 26, - "Co": 27, - "Ni": 28, - "Cu": 29, - "Zn": 30, - "Ga": 31, - "Ge": 32, - "As": 33, - "Se": 34, - "Br": 35, - "Kr": 36, - "Rb": 37, - "Sr": 38, - "Y": 39, - "Zr": 40, - "Nb": 41, - "Mo": 42, - "Tc": 43, - "Ru": 44, - "Rh": 45, - "Pd": 46, - "Ag": 47, - "Cd": 48, - "In": 49, - "Sn": 50, - "Sb": 51, - "Te": 52, - "I": 53, - "Xe": 54, - "Cs": 55, - "Ba": 56, - "La": 57, - "Ce": 58, - "Pr": 59, - "Nd": 60, - "Pm": 61, - "Sm": 62, - "Eu": 63, - "Gd": 64, - "Tb": 65, - "Dy": 66, - "Ho": 67, - "Er": 68, - "Tm": 69, - "Yb": 70, - "Lu": 71, - "Hf": 72, - "Ta": 73, - "W": 74, - "Re": 75, - "Os": 76, - "Ir": 77, - "Pt": 78, - "Au": 79, - "Hg": 80, - "Tl": 81, - "Pb": 82, - "Bi": 83, - "Po": 84, - "At": 85, - "Rn": 86, - "Fr": 87, - "Ra": 88, - "Ac": 89, - "Th": 90, - "Pa": 91, - "U": 92, - "Np": 93, - "Pu": 94, - "Am": 95, - "Cm": 96, - "Bk": 97, - "Cf": 98, - "Es": 99, - "Fm": 100, - "Md": 101, - "No": 102, - "Lr": 103, - "Rf": 104, - "Db": 105, - "Sg": 106, - "Bh": 107, - "Hs": 108, - "Mt": 109, - "Ds": 110, - "Rg": 111, - "Cn": 112, - "Nh": 113, - "Fl": 114, - "Mc": 115, - "Lv": 116, - "Ts": 117, - "Og": 118, -} - -ATOMIC_NUMBER_TO_ELEMENT = {v: k for k, v in ELEMENT_TO_ATOMIC_NUMBER.items()}