work on kgcnn.io.file updated graph.base and crystal.base to match indices. Added schedule to training. Added coGN to hyper_mp_e_form.py

PatReis · PatReis · commit 245c1efb2caa · 2023-05-12T17:59:00.000+02:00
diff --git a/changelog.md b/changelog.md
@@ -3,7 +3,7 @@ v3.0.1
 * Removed deprecated molecules.
 * Fix error in ``kgcnn.data.transform.scaler.serial``
 * Fix error in ``QMDataset`` for molecular features.
-* 
+* Fix error in ``kgcnn.layers.conv.GraphSageNodeLayer`` .
 
 
 v3.0.0
diff --git a/kgcnn/crystal/base.py b/kgcnn/crystal/base.py
@@ -55,7 +55,7 @@ def __call__(self, structure: Structure) -> Union[MultiDiGraph, GraphDict]:
             g = GraphDict()
             g.from_networkx(
                 nxg, node_attributes=self.node_attributes, edge_attributes=self.edge_attributes,
-                graph_attributes=self.graph_attributes)
+                graph_attributes=self.graph_attributes, reverse_edge_indices=True)
             return g
         return self.call(structure)
 
diff --git a/kgcnn/graph/base.py b/kgcnn/graph/base.py
@@ -179,7 +179,8 @@ def from_networkx(self, graph: nx.Graph,
                       node_attributes: Union[str, List[str]] = None,
                       edge_attributes: Union[str, List[str]] = None,
                       graph_attributes: Union[str, List[str]] = None,
-                      node_labels: str = None):
+                      node_labels: str = None,
+                      reverse_edge_indices: bool = False):
         r"""Convert a networkx graph instance into a dictionary of graph-tensors. The networkx graph is always converted
         into integer node labels. The former node IDs can be hold in :obj:`node_labels`. Furthermore, node or edge
         data can be cast into attributes via :obj:`node_attributes` and :obj:`edge_attributes`.
@@ -195,6 +196,7 @@ def from_networkx(self, graph: nx.Graph,
             graph_attributes (str, list): Name of graph attributes to add from graph data. Can also be a list of names.
                 Default is None.
             node_labels (str): Name of the labels of nodes to store former node IDs into. Default is None.
+            reverse_edge_indices (bool): Whether to reverse edge indices for notation '(ij, i<-j)'. Default is False.
 
         Returns:
             self.
@@ -230,7 +232,10 @@ def _attr_to_list(attr):
         edges_attr = _attr_to_list(edge_attributes)
         edges_attr_dict = {x: [None] * graph_int.number_of_edges() for x in edges_attr}
         for i, x in enumerate(graph_int.edges.data()):
-            edge_id.append(x[:2])
+            if reverse_edge_indices:
+                edge_id.append([x[1], x[0]])
+            else:
+                edge_id.append([x[0], x[1]])
             for d in edges_attr:
                 if d not in x[2]:
                     raise KeyError("Edge does not have property '%s'." % d)
diff --git a/kgcnn/io/file.py b/kgcnn/io/file.py
@@ -1,5 +1,6 @@
 import numpy as np
 import h5py
+from typing import List, Union
 
 
 class RaggedArrayNumpyFile:
@@ -8,13 +9,12 @@ def __init__(self, file_path: str, compressed: bool = False):
         self.file_path = file_path
         self.compressed = compressed
 
-    def write(self, ragged_array: list):
-        # Only support ragged one for the moment.
-        leading_shape = ragged_array[0].shape
-        assert all([leading_shape[1:] == x.shape[1:] for x in ragged_array]), "Only support ragged rank == 1."
+    def write(self, ragged_array: List[np.ndarray]):
+        inner_shape = ragged_array[0].shape
         values = np.concatenate([x for x in ragged_array], axis=0)
         row_splits = np.cumsum(np.array([len(x) for x in ragged_array], dtype="int64"), dtype="int64")
-        out = {"values": values, "row_splits": row_splits}
+        row_splits = np.pad(row_splits, [1, 0])
+        out = {"values": values, "row_splits": row_splits, "shape": np.array([])}
         if self.compressed:
             np.savez_compressed(self.file_path, **out)
         else:
@@ -24,16 +24,19 @@ def read(self):
         data = np.load(self.file_path)
         values = data.get("values")
         row_splits = data.get("row_splits")
-        return np.split(values, row_splits[:-1])
+        return np.split(values, row_splits[1:-1])
+
+    def __getitem__(self, item):
+        raise NotImplementedError("Not implemented for file reference load.")
 
 
 class RaggedArrayHDFile:
 
-    def __init__(self, file_path: str, compressed: bool = False):
+    def __init__(self, file_path: str, compressed: bool = None):
         self.file_path = file_path
         self.compressed = compressed
 
-    def write(self, ragged_array: list):
+    def write(self, ragged_array: List[np.ndarray]):
         """Write ragged array to file.
 
         .. code-block:: python
@@ -44,30 +47,57 @@ def write(self, ragged_array: list):
             f = RaggedArrayHDFile("test.hdf5")
             f.write(data)
             print(f.read())
-            print(f[1])
 
         Args:
-            ragged_array (list): List of numpy arrays.
+            ragged_array (list, tf.RaggedTensor): List or list of numpy arrays.
 
         Returns:
             None.
         """
+        inner_shape = ragged_array[0].shape
         values = np.concatenate([x for x in ragged_array], axis=0)
         row_splits = np.cumsum(np.array([len(x) for x in ragged_array], dtype="int64"), dtype="int64")
-        with h5py.File(self.file_path, "w") as f:
-            f.create_dataset("values", data=values)
-            f.create_dataset("row_splits", data=row_splits)
+        row_splits = np.pad(row_splits, [1, 0])
+        with h5py.File(self.file_path, "w") as file:
+            file.create_dataset("values", data=values, maxshape=[None] + list(inner_shape)[1:])
+            file.create_dataset("row_splits", data=row_splits, maxshape=(None, ))
+            file.create_dataset("shape", data=np.array([]))
 
     def read(self):
-        file = h5py.File(self.file_path)
-        data = np.split(file["values"][()], file["row_splits"][:1])
-        file.close()
+        with h5py.File(self.file_path, "r") as file:
+            data = np.split(file["values"][()], file["row_splits"][1:-1])
         return data
 
-    def __getitem__(self, item):
-        file = h5py.File(self.file_path)
-        row_splits = file["row_splits"]
-        row_splits = np.pad(row_splits, [1, 0])
-        out_data = file["values"][row_splits[item]:row_splits[item+1]]
-        file.close()
+    def __getitem__(self, item: int):
+        with h5py.File(self.file_path, "r") as file:
+            row_splits = file["row_splits"]
+            out_data = file["values"][row_splits[item]:row_splits[item+1]]
         return out_data
+
+    def append(self, item):
+        with h5py.File(self.file_path, "r+") as file:
+            file["values"].resize(
+                file["values"].shape[0] + len(item), axis=0
+            )
+            split_last = file["row_splits"][-1]
+            file["row_splits"].resize(
+                file["row_splits"].shape[0] + 1, axis=0
+            )
+            len_last = len(item)
+            file["row_splits"][-1] = split_last + len_last
+            file["values"][split_last:split_last+len_last] = item
+
+    def append_multiple(self, items: list):
+        new_values = np.concatenate(items, axis=0)
+        new_len = len(items)
+        new_splits = np.cumsum([len(x) for x in items])
+        with h5py.File(self.file_path, "r+") as file:
+            file["values"].resize(
+                file["values"].shape[0] + new_values.shape[0], axis=0
+            )
+            split_last = file["row_splits"][-1]
+            file["row_splits"].resize(
+                file["row_splits"].shape[0] + new_len, axis=0
+            )
+            file["row_splits"][-new_len:] = split_last + new_splits
+            file["values"][split_last:+split_last+new_splits[-1]] = new_values
diff --git a/kgcnn/training/schedule.py b/kgcnn/training/schedule.py
@@ -1,8 +1,10 @@
 import numpy as np
 import tensorflow as tf
 
+ks = tf.keras
 
-@tf.keras.utils.register_keras_serializable(package='kgcnn', name='LinearWarmupExponentialDecay')
+
+@ks.utils.register_keras_serializable(package='kgcnn', name='LinearWarmupExponentialDecay')
 class LinearWarmupExponentialDecay(tf.optimizers.schedules.LearningRateSchedule):
     r"""This schedule combines a linear warmup with an exponential decay.
     Combines :obj:` tf.optimizers.schedules.PolynomialDecay` with an actual increase during warmup
@@ -63,3 +65,31 @@ def get_config(self):
         config = {}
         config.update(self._input_config_settings)
         return config
+
+
+@ks.utils.register_keras_serializable(package='kgcnn', name='KerasPolynomialDecaySchedule')
+class KerasPolynomialDecaySchedule(ks.optimizers.schedules.PolynomialDecay):
+    r"""This schedule extends :obj:` tf.optimizers.schedules.PolynomialDecay` ."""
+
+    def __init__(self, dataset_size: int, batch_size: int, epochs: int, lr_start: float = 0.0005,
+                 lr_stop: float = 1e-5):
+        """Initialize class.
+
+        Args:
+            dataset_size (int): Size of the dataset.
+            batch_size (int): Batch size for training.
+            epochs (int): Total epochs to run schedule on.
+            lr_start (int): Learning rate at the start.
+            lr_stop (int): Final learning rate.
+        """
+        self._input_config_settings = {"lr_start": lr_start, "lr_stop": lr_stop,
+                                       "epochs": epochs, "batch_size": batch_size, "dataset_size": dataset_size}
+        steps_per_epoch = dataset_size / batch_size
+        num_steps = epochs * steps_per_epoch
+        super().__init__(initial_learning_rate=lr_start, decay_steps=num_steps, end_learning_rate=lr_stop)
+
+    def get_config(self):
+        """Get config for this class."""
+        config = {}
+        config.update(self._input_config_settings)
+        return config
diff --git a/requirements.txt b/requirements.txt
@@ -16,4 +16,5 @@ ase>=3.22.1
 click>=7.1.2
 visual_graph_datasets>=0.7.1
 brotli>=1.0.9
-pyxtal>=0.5.5
+pyxtal>=0.5.5
+h5py
diff --git a/training/hyper/hyper_mp_e_form.py b/training/hyper/hyper_mp_e_form.py
@@ -419,4 +419,82 @@
             "kgcnn_version": "2.1.1"
         }
     },
+    "coGN": {
+        "model": {
+            "module_name": "kgcnn.literature.coGN",
+            "class_name": "make_model",
+            "config": {
+                "name": "coGN",
+                "inputs": {
+                    "offset": {"shape": (None, 3), "name": "offset", "dtype": "float32", "ragged": True},
+                    "cell_translation": None,
+                    "affine_matrix": None,
+                    "voronoi_ridge_area": None,
+                    "atomic_number": {"shape": (None,), "name": "atomic_number", "dtype": "int32", "ragged": True},
+                    "frac_coords": None,
+                    "coords": None,
+                    "multiplicity": {"shape": (None,), "name": "multiplicity", "dtype": "int32", "ragged": True},
+                    "lattice_matrix": None,
+                    "edge_indices": {"shape": (None, 2), "name": "edge_indices", "dtype": "int32", "ragged": True},
+                    "line_graph_edge_indices": None,
+                },
+                # All default.
+            }
+        },
+        "training": {
+            "cross_validation": {"class_name": "KFold",
+                                 "config": {"n_splits": 5, "random_state": 42, "shuffle": True}},
+            "fit": {
+                "batch_size": 64, "epochs": 800, "validation_freq": 10, "verbose": 2,
+                "callbacks": [
+                    # {"class_name": "kgcnn>LinearLearningRateScheduler", "config": {
+                    #     "learning_rate_start": 0.0005, "learning_rate_stop": 0.5e-05, "epo_min": 0, "epo": 800,
+                    #     "verbose": 0}
+                    #  }
+                ]
+            },
+            "compile": {
+                "optimizer": {
+                    "class_name": "Adam",
+                        "config": {
+                            "learning_rate": {
+                                "class_name": "kgcnn>KerasPolynomialDecaySchedule",
+                                "config": {
+                                    "dataset_size": 106.201, "batch_size": 64, "epochs": 800,
+                                    "lr_start": 0.0005, "lr_stop": 1.0e-05
+                                }
+                            }
+                        }
+                },
+                "loss": "mean_absolute_error"
+            },
+            "scaler": {
+                "class_name": "StandardScaler",
+                "module_name": "kgcnn.data.transform.scaler.standard",
+                "config": {"with_std": True, "with_mean": True, "copy": True}
+            },
+            "multi_target_indices": None
+        },
+        "data": {
+            "dataset": {
+                "class_name": "MatProjectEFormDataset",
+                "module_name": "kgcnn.data.datasets.MatProjectEFormDataset",
+                "config": {},
+                "methods": [
+                    {"set_representation": {
+                        "pre_processor": {"class_name": "KNNAsymmetricUnitCell",
+                                          "module_name": "kgcnn.crystal.preprocessor",
+                                          "config": {"k": 24}
+                                          },
+                        "reset_graphs": False}}
+                ]
+            },
+            "data_unit": "eV/atom"
+        },
+        "info": {
+            "postfix": "",
+            "postfix_file": "",
+            "kgcnn_version": "3.0.1"
+        }
+    },
 }
