File size: 5,174 Bytes
3e02ab8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 | # This file is a part of the `nequip` package. Please see LICENSE and README at the root for information on using it.
from typing import List, Dict, Union
import torch
from onescience.utils.nequip.internal.global_dtype import _GLOBAL_DTYPE
# conversion flow: partial_dict -> full_dict -> tensor -> str
# |
# v
# full_dict
def cutoff_partialdict_to_fulldict(
partial_dict: Dict[str, Union[float, Dict[str, float]]],
type_names: List[str],
r_max: float,
) -> Dict[str, Dict[str, float]]:
"""Convert partial cutoff dict to full dict with all entries.
Fills missing entries with ``r_max``.
Args:
partial_dict: partial specification from config,
e.g. ``{"H": 2.0, "C": {"H": 4.0, "C": 3.5}}``
type_names: list of atom type names
r_max: global cutoff radius (default for missing entries)
Returns:
full dict with all source -> target pairs specified,
e.g. ``{"H": {"H": 2.0, "C": 2.0}, "C": {"H": 4.0, "C": 3.5}}``
"""
full_dict = {}
for source_type in type_names:
full_dict[source_type] = {}
if source_type in partial_dict:
entry = partial_dict[source_type]
if isinstance(entry, float):
# uniform cutoff for this source type
for target_type in type_names:
full_dict[source_type][target_type] = entry
else:
# per-target specification
for target_type in type_names:
if target_type in entry:
full_dict[source_type][target_type] = entry[target_type]
else:
# missing target defaults to r_max
full_dict[source_type][target_type] = r_max
else:
# missing source defaults to r_max for all targets
for target_type in type_names:
full_dict[source_type][target_type] = r_max
return full_dict
def cutoff_fulldict_to_tensor(
full_dict: Dict[str, Dict[str, float]],
type_names: List[str],
) -> torch.Tensor:
"""Convert full cutoff dict to tensor.
Args:
full_dict: full specification with all source -> target pairs
type_names: list of atom type names
Returns:
tensor of shape ``(num_types, num_types)`` with per-edge-type cutoffs
"""
num_types = len(type_names)
cutoff_list = []
for source_type in type_names:
row = []
for target_type in type_names:
row.append(full_dict[source_type][target_type])
cutoff_list.append(row)
cutoff_tensor = torch.as_tensor(cutoff_list, dtype=_GLOBAL_DTYPE).contiguous()
assert cutoff_tensor.shape == (num_types, num_types)
assert torch.all(cutoff_tensor > 0)
return cutoff_tensor
def cutoff_tensor_to_str(cutoff_tensor: torch.Tensor) -> str:
"""Convert tensor to metadata string format.
Args:
cutoff_tensor: cutoff values as tensor (any shape, will be flattened)
Returns:
space-separated string of cutoff values in row-major order
"""
return " ".join(str(r.item()) for r in cutoff_tensor.reshape(-1))
def cutoff_str_to_fulldict(
cutoff_str: str,
type_names: List[str],
) -> Dict[str, Dict[str, float]]:
"""Convert metadata string to full dict format.
Args:
cutoff_str: space-separated string of cutoff values
type_names: list of atom type names
Returns:
full dict with all source -> target pairs specified
"""
if cutoff_str in ("", None):
return None
cutoff_values = [float(x) for x in cutoff_str.split()]
num_types = len(type_names)
assert len(cutoff_values) == num_types * num_types, (
f"Expected {num_types * num_types} cutoff values, got {len(cutoff_values)}"
)
full_dict = {}
for i, source_type in enumerate(type_names):
full_dict[source_type] = {}
for j, target_type in enumerate(type_names):
full_dict[source_type][target_type] = cutoff_values[i * num_types + j]
return full_dict
def cutoff_partialdict_to_tensor(
partial_dict: Dict[str, Union[float, Dict[str, float]]],
type_names: List[str],
r_max: float,
) -> torch.Tensor:
"""Composes ``cutoff_partialdict_to_fulldict`` and ``cutoff_fulldict_to_tensor``."""
full_dict = cutoff_partialdict_to_fulldict(partial_dict, type_names, r_max)
cutoff_tensor = cutoff_fulldict_to_tensor(full_dict, type_names)
assert torch.all(cutoff_tensor <= r_max)
return cutoff_tensor
def cutoff_partialdict_to_str(
partial_dict: Dict[str, Union[float, Dict[str, float]]],
type_names: List[str],
r_max: float,
) -> str:
"""Composes ``cutoff_partialdict_to_fulldict``, ``cutoff_fulldict_to_tensor``, and ``cutoff_tensor_to_str``."""
full_dict = cutoff_partialdict_to_fulldict(partial_dict, type_names, r_max)
tensor = cutoff_fulldict_to_tensor(full_dict, type_names)
return cutoff_tensor_to_str(tensor)
|