BUCToolkit 1.0b0__tar.gz → 1.0b2__tar.gz

buctoolkit-1.0b2/BUCToolkit/BatchMD/constrained_md/README.md

@@ -0,0 +1,4 @@
+### Here algorithms require all input batch are regular, i.e., in one batch, all structures have the same atom numbers.
+* Typically, they are the same chemical systems with different configurations.
+* It is useful for independent repeated MD, blue-moon ensembles, or application for Jarzynski identity.
+* It is more efficiency than the irregular batch version algorithms.

{buctoolkit-1.0b0 → buctoolkit-1.0b2}/BUCToolkit/BatchStructures/BatchStructuresBase.py

@@ -423,13 +423,15 @@ class BatchStructures(object):
         _STANDARD_FILE_LIST = set(list(self._ATTR_NAMES) + ['head', ])
         if mode not in {'w', 'a'}:
             raise ValueError(f'Unknown `mode` value {mode}. It must be "w" or "a".')
-        if not os.path.isdir(path):  # path do not exist
+        if not os.path.exists(path):  # path do not exist
             mode = 'w'
             os.makedirs(path, )
+        elif os.path.isfile(path):
+            raise OSError(f"The file `{path}` already exists.")
         elif len(os.listdir(path)) > 0:
             if mode == 'a':
                 file_set = set(os.listdir(path))
-                assert file_set == _STANDARD_FILE_LIST, 'Existing files do not match the saving data.'
+                if file_set != _STANDARD_FILE_LIST: raise RuntimeError('Existing files do not match the saving data.')
             else:
                 for ff in os.listdir(path):
                     if ff not in _STANDARD_FILE_LIST:
@@ -724,8 +726,10 @@ class BatchStructures(object):
             is_convert = True
         else:
             is_convert = False
-        if not os.path.isdir(output_path):
+        if not os.path.exists(output_path):
             os.makedirs(output_path)
+        elif os.path.isfile(output_path):
+            raise OSError(f"There is already a file named '{output_path}'.")
 
         try:
             file_name_list = [f'{_}' for _ in file_name_list]

buctoolkit-1.0b2/BUCToolkit/BatchStructures/README.md

@@ -0,0 +1,14 @@
+# Here `data.py` and `batch.py` are a subset from `pytorch_geometric`
+
+BUCToolkit uses the most basic graph data structures of
+[`pytorch_geometric`](https://pytorch-geometric.readthedocs.io/en/latest/) [1, 2] framework as a 
+basic standard input data format.
+
+Thus, one needs not install the whole `pytorch_geometric` with heavy dependencies to run our methods.
+
+We follow the same approach to NequIP (https://github.com/mir-group/nequip/tree/main/nequip) and 
+mace (https://github.com/ACEsuit/mace), and copy their code of `data.py` and `batch.py` here.
+
+
+[1]  Fey, M., & Lenssen, J. E. (2019). Fast Graph Representation Learning with PyTorch Geometric (Version 2.0.1) [Computer software]. https://github.com/pyg-team/pytorch_geometric <br>
+[2]  https://arxiv.org/abs/1903.02428

{buctoolkit-1.0b0 → buctoolkit-1.0b2}/BUCToolkit/BatchStructures/__init__.py

@@ -4,7 +4,11 @@
 #  File: __init__.py
 #  Environment: Python 3.12
 
+from BUCToolkit.BatchStructures.BatchStructuresBase import BatchStructures
 from BUCToolkit.BatchStructures.StructuresIO import read_opt_structures, read_md_traj, read_mc_traj, ArrayDumper, ArrayDumpReader
+from BUCToolkit.BatchStructures.batch import Batch
+from BUCToolkit.BatchStructures.data import Data
+
 
 __all__ = [
     "read_opt_structures",
@@ -12,4 +16,7 @@ __all__ = [
     "read_mc_traj",
     "ArrayDumper",
     "ArrayDumpReader",
+    'Batch',
+    'Data',
+    'BatchStructures'
 ]

buctoolkit-1.0b2/BUCToolkit/BatchStructures/batch.py

@@ -0,0 +1,256 @@
+from collections.abc import Sequence
+from typing import List
+
+import numpy as np
+import torch
+from torch import Tensor
+
+from .data import Data
+
+
+class Batch(Data):
+    r"""A plain old python object modeling a batch of graphs as one big
+    (disconnected) graph. With :class:`torch_geometric.data.Data` being the
+    base class, all its methods can also be used here.
+    In addition, single graphs can be reconstructed via the assignment vector
+    :obj:`batch`, which maps each node to its respective graph identifier.
+    """
+
+    def __init__(self, batch=None, ptr=None, **kwargs):
+        super(Batch, self).__init__(**kwargs)
+
+        for key, item in kwargs.items():
+            if key == "num_nodes":
+                self.__num_nodes__ = item
+            else:
+                self[key] = item
+
+        self.batch = batch
+        self.ptr = ptr
+        self.__data_class__ = Data
+        self.__slices__ = None
+        self.__cumsum__ = None
+        self.__cat_dims__ = None
+        self.__num_nodes_list__ = None
+        self.__num_graphs__ = None
+
+    @classmethod
+    def from_data_list(cls, data_list, follow_batch=[], exclude_keys=[]):
+        r"""Constructs a batch object from a python list holding
+        :class:`torch_geometric.data.Data` objects.
+        The assignment vector :obj:`batch` is created on the fly.
+        Additionally, creates assignment batch vectors for each key in
+        :obj:`follow_batch`.
+        Will exclude any keys given in :obj:`exclude_keys`."""
+
+        keys = list(set(data_list[0].keys) - set(exclude_keys))
+        assert "batch" not in keys and "ptr" not in keys
+
+        batch = cls()
+        for key in data_list[0].__dict__.keys():
+            if key[:2] != "__" and key[-2:] != "__":
+                batch[key] = None
+
+        batch.__num_graphs__ = len(data_list)
+        batch.__data_class__ = data_list[0].__class__
+        for key in keys + ["batch"]:
+            batch[key] = []
+        batch["ptr"] = [0]
+
+        device = None
+        slices = {key: [0] for key in keys}
+        cumsum = {key: [0] for key in keys}
+        cat_dims = {}
+        num_nodes_list = []
+        for i, data in enumerate(data_list):
+            for key in keys:
+                item = data[key]
+
+                # Increase values by `cumsum` value.
+                cum = cumsum[key][-1]
+                if isinstance(item, Tensor) and item.dtype != torch.bool:
+                    if not isinstance(cum, int) or cum != 0:
+                        item = item + cum
+                elif isinstance(item, (int, float)):
+                    item = item + cum
+
+                # Gather the size of the `cat` dimension.
+                size = 1
+                cat_dim = data.__cat_dim__(key, data[key])
+                # 0-dimensional tensors have no dimension along which to
+                # concatenate, so we set `cat_dim` to `None`.
+                if isinstance(item, Tensor) and item.dim() == 0:
+                    cat_dim = None
+                cat_dims[key] = cat_dim
+
+                # Add a batch dimension to items whose `cat_dim` is `None`:
+                if isinstance(item, Tensor) and cat_dim is None:
+                    cat_dim = 0  # Concatenate along this new batch dimension.
+                    item = item.unsqueeze(0)
+                    device = item.device
+                elif isinstance(item, Tensor):
+                    size = item.size(cat_dim)
+                    device = item.device
+
+                batch[key].append(item)  # Append item to the attribute list.
+
+                slices[key].append(size + slices[key][-1])
+                inc = data.__inc__(key, item)
+                if isinstance(inc, (tuple, list)):
+                    inc = torch.tensor(inc)
+                cumsum[key].append(inc + cumsum[key][-1])
+
+                if key in follow_batch:
+                    if isinstance(size, Tensor):
+                        for j, size in enumerate(size.tolist()):
+                            tmp = f"{key}_{j}_batch"
+                            batch[tmp] = [] if i == 0 else batch[tmp]
+                            batch[tmp].append(
+                                torch.full((size,), i, dtype=torch.long, device=device)
+                            )
+                    else:
+                        tmp = f"{key}_batch"
+                        batch[tmp] = [] if i == 0 else batch[tmp]
+                        batch[tmp].append(
+                            torch.full((size,), i, dtype=torch.long, device=device)
+                        )
+
+            if hasattr(data, "__num_nodes__"):
+                num_nodes_list.append(data.__num_nodes__)
+            else:
+                num_nodes_list.append(None)
+
+            num_nodes = data.num_nodes
+            if num_nodes is not None:
+                item = torch.full((num_nodes,), i, dtype=torch.long, device=device)
+                batch.batch.append(item)
+                batch.ptr.append(batch.ptr[-1] + num_nodes)
+
+        batch.batch = None if len(batch.batch) == 0 else batch.batch
+        batch.ptr = None if len(batch.ptr) == 1 else batch.ptr
+        batch.__slices__ = slices
+        batch.__cumsum__ = cumsum
+        batch.__cat_dims__ = cat_dims
+        batch.__num_nodes_list__ = num_nodes_list
+
+        ref_data = data_list[0]
+        for key in batch.keys:
+            items = batch[key]
+            item = items[0]
+            cat_dim = ref_data.__cat_dim__(key, item)
+            cat_dim = 0 if cat_dim is None else cat_dim
+            if isinstance(item, Tensor):
+                batch[key] = torch.cat(items, cat_dim)
+            elif isinstance(item, (int, float)):
+                batch[key] = torch.tensor(items)
+
+        # if torch_geometric.is_debug_enabled():
+        #     batch.debug()
+
+        return batch.contiguous()
+
+    def get_example(self, idx: int) -> Data:
+        r"""Reconstructs the :class:`torch_geometric.data.Data` object at index
+        :obj:`idx` from the batch object.
+        The batch object must have been created via :meth:`from_data_list` in
+        order to be able to reconstruct the initial objects."""
+
+        if self.__slices__ is None:
+            raise RuntimeError(
+                (
+                    "Cannot reconstruct data list from batch because the batch "
+                    "object was not created using `Batch.from_data_list()`."
+                )
+            )
+
+        data = self.__data_class__()
+        idx = self.num_graphs + idx if idx < 0 else idx
+
+        for key in self.__slices__.keys():
+            item = self[key]
+            if self.__cat_dims__[key] is None:
+                # The item was concatenated along a new batch dimension,
+                # so just index in that dimension:
+                item = item[idx]
+            else:
+                # Narrow the item based on the values in `__slices__`.
+                if isinstance(item, Tensor):
+                    dim = self.__cat_dims__[key]
+                    start = self.__slices__[key][idx]
+                    end = self.__slices__[key][idx + 1]
+                    item = item.narrow(dim, start, end - start)
+                else:
+                    start = self.__slices__[key][idx]
+                    end = self.__slices__[key][idx + 1]
+                    item = item[start:end]
+                    item = item[0] if len(item) == 1 else item
+
+            # Decrease its value by `cumsum` value:
+            cum = self.__cumsum__[key][idx]
+            if isinstance(item, Tensor):
+                if not isinstance(cum, int) or cum != 0:
+                    item = item - cum
+            elif isinstance(item, (int, float)):
+                item = item - cum
+
+            data[key] = item
+
+        if self.__num_nodes_list__[idx] is not None:
+            data.num_nodes = self.__num_nodes_list__[idx]
+
+        return data
+
+    def index_select(self, idx: List) -> List[Data]:
+        if isinstance(idx, slice):
+            idx = list(range(self.num_graphs)[idx])
+
+        elif isinstance(idx, Tensor) and idx.dtype == torch.long:
+            idx = idx.flatten().tolist()
+
+        elif isinstance(idx, Tensor) and idx.dtype == torch.bool:
+            idx = idx.flatten().nonzero(as_tuple=False).flatten().tolist()
+
+        elif isinstance(idx, np.ndarray) and idx.dtype == np.int64:
+            idx = idx.flatten().tolist()
+
+        elif isinstance(idx, np.ndarray) and idx.dtype == np.bool:
+            idx = idx.flatten().nonzero()[0].flatten().tolist()
+
+        elif isinstance(idx, Sequence) and not isinstance(idx, str):
+            pass
+
+        else:
+            raise IndexError(
+                f"Only integers, slices (':'), list, tuples, torch.tensor and "
+                f"np.ndarray of dtype long or bool are valid indices (got "
+                f"'{type(idx).__name__}')"
+            )
+
+        return [self.get_example(i) for i in idx]
+
+    def __getitem__(self, idx):
+        if isinstance(idx, str):
+            return super(Batch, self).__getitem__(idx)
+        elif isinstance(idx, (int, np.integer)):
+            return self.get_example(idx)
+        else:
+            return self.index_select(idx)
+
+    def to_data_list(self) -> List[Data]:
+        r"""Reconstructs the list of :class:`torch_geometric.data.Data` objects
+        from the batch object.
+        The batch object must have been created via :meth:`from_data_list` in
+        order to be able to reconstruct the initial objects."""
+        return [self.get_example(i) for i in range(self.num_graphs)]
+
+    @property
+    def num_graphs(self) -> int:
+        """Returns the number of graphs in the batch."""
+        if self.__num_graphs__ is not None:
+            return self.__num_graphs__
+        elif self.ptr is not None:
+            return self.ptr.numel() - 1
+        elif self.batch is not None:
+            return int(self.batch.max()) + 1
+        else:
+            raise ValueError