diff --git a/README.md b/README.md
index 73dac3e7..9f0a9f6f 100644
--- a/README.md
+++ b/README.md
@@ -16,11 +16,8 @@
A set of layers for graph convolutions in Keras.
> [!IMPORTANT]
-> The versions of kgcnn<=3.1.0 were focused on ragged tensors of tensorflow. With keras as multi backend framework,
-> we plan to port kgcnn to keras-core for version kgcnn==4.0.0.
-> This means kgcnn models and layers can be used with tensorflow, jax and pytorch as backend in the future.
-> This means however, that layers and tensor representation will change. However, the workflow, weights and frontend behaviour should remain the same.
-> We will try to make kgcnn also more compatible with e.g. PytorchGeometric.
+> The versions of kgcnn<=3.1.0 were focused on ragged tensors of tensorflow.
+> The current main branch is developped for Keras 3.0 . Please use last release of 3.1.0 for previous version of kgcnn.
[General](#general) | [Requirements](#requirements) | [Installation](#installation) | [Documentation](#documentation) | [Implementation details](#implementation-details)
| [Literature](#literature) | [Data](#data) | [Datasets](#datasets) | [Training](#training) | [Issues](#issues) | [Citing](#citing) | [References](#references)
@@ -63,85 +60,20 @@ Auto-documentation is generated at https://kgcnn.readthedocs.io/en/latest/index.
# Implementation details
### Representation
-The most frequent usage for graph convolutions is either node or graph classification. As for their size, either a single large graph, e.g. citation network or small (batched) graphs like molecules have to be considered.
-Graphs can be represented by an index list of connections plus feature information. Typical quantities in tensor format to describe a graph are listed below.
-* `nodes`: Node-list of shape `(batch, [N], F)` where `N` is the number of nodes and `F` is the node feature dimension.
-* `edges`: Edge-list of shape `(batch, [M], F)` where `M` is the number of edges and `F` is the edge feature dimension.
-* `indices`: Connection-list of shape `(batch, [M], 2)` where `M` is the number of edges. The indices denote a connection of incoming or receiving node `i` and outgoing or sending node `j` as `(i, j)`.
-* `state`: Graph state information of shape `(batch, F)` where `F` denotes the feature dimension.
-
-A major issue for graphs is their flexible size and shape, when using mini-batches. Here, for a graph implementation in the spirit of keras, the batch dimension should be kept also in between layers. This is realized by using `RaggedTensor`s.
+TODO
### Input
-Graph tensors for edge-indices or attributes for multiple graphs is passed to the model in form of ragged tensors
-of shape `(batch, None, Dim)` where `Dim` denotes a fixed feature or index dimension.
-Such a ragged tensor has `ragged_rank=1` with one ragged dimension indicated by `None` and is build from a value plus partition tensor.
-For example, the graph structure is represented by an index-list of shape `(batch, None, 2)` with index of incoming or receiving node `i` and outgoing or sending node `j` as `(i, j)`.
-Note, an additional edge with `(j, i)` is required for undirected graphs.
-A ragged constant can be easily created and passed to a model:
-
-```python
-import tensorflow as tf
-import numpy as np
-idx = [[[0, 1], [1, 0]], [[0, 1], [1, 2], [2, 0]], [[0, 0]]] # batch_size=3
-# Get ragged tensor of shape (3, None, 2)
-print(tf.ragged.constant(idx, ragged_rank=1, inner_shape=(2, )).shape)
-print(tf.RaggedTensor.from_row_lengths(np.concatenate(idx), [len(i) for i in idx]).shape)
-```
+TODO
### Model
Models can be set up in a functional way. Example message passing from fundamental operations:
-```python
-import tensorflow as tf
-from kgcnn.layers.gather import GatherNodes
-from kgcnn.layers.modules import Dense, LazyConcatenate # ragged support
-from kgcnn.layers.aggr import AggregateLocalMessages
-from kgcnn.layers.pooling import PoolingNodes
-
-ks = tf.keras
-
-n = ks.layers.Input(shape=(None, 3), name='node_input', dtype="float32", ragged=True)
-ei = ks.layers.Input(shape=(None, 2), name='edge_index_input', dtype="int64", ragged=True)
-
-n_in_out = GatherNodes()([n, ei])
-node_messages = Dense(10, activation='relu')(n_in_out)
-node_updates = AggregateLocalMessages(is_sorted=False)([n, node_messages, ei])
-n_node_updates = LazyConcatenate(axis=-1)([n, node_updates])
-n_embedding = Dense(1)(n_node_updates)
-g_embedding = PoolingNodes()(n_embedding)
-
-message_passing = ks.models.Model(inputs=[n, ei], outputs=g_embedding)
-```
-
-or via sub-classing of the message passing base layer. Where only `message_function` and `update_nodes` must be implemented:
-
-```python
-
-from kgcnn.layers.message import MessagePassingBase
-from kgcnn.layers.modules import Dense, LazyConcatenate
-
-
-class MyMessageNN(MessagePassingBase):
-
- def __init__(self, units, **kwargs):
- super(MyMessageNN, self).__init__(**kwargs)
- self.dense = Dense(units)
- self.add = LazyConcatenate(axis=-1)
-
- def message_function(self, inputs, **kwargs):
- n_in, n_out, edges = inputs
- return self.dense(n_out)
-
- def update_nodes(self, inputs, **kwargs):
- nodes, nodes_update = inputs
- return self.add([nodes, nodes_update])
-```
+TODO
# Literature
@@ -153,42 +85,20 @@ original implementations (with proper licencing).
* **[Schnet](kgcnn/literature/Schnet)**: [SchNet – A deep learning architecture for molecules and materials ](https://aip.scitation.org/doi/10.1063/1.5019779) by Schütt et al. (2017)
* **[GAT](kgcnn/literature/GAT)**: [Graph Attention Networks](https://arxiv.org/abs/1710.10903) by Veličković et al. (2018)
* **[GraphSAGE](kgcnn/literature/GraphSAGE)**: [Inductive Representation Learning on Large Graphs](http://arxiv.org/abs/1706.02216) by Hamilton et al. (2017)
-* **[DimeNetPP](kgcnn/literature/DimeNetPP)**: [Fast and Uncertainty-Aware Directional Message Passing for Non-Equilibrium Molecules](https://arxiv.org/abs/2011.14115) by Klicpera et al. (2020)
-* **[GNNExplainer](kgcnn/literature/GNNExplain)**: [GNNExplainer: Generating Explanations for Graph Neural Networks](https://arxiv.org/abs/1903.03894) by Ying et al. (2019)
-* **[AttentiveFP](kgcnn/literature/AttentiveFP)**: [Pushing the Boundaries of Molecular Representation for Drug Discovery with the Graph Attention Mechanism](https://pubs.acs.org/doi/10.1021/acs.jmedchem.9b00959) by Xiong et al. (2019)
+* **[GIN](kgcnn/literature/GIN)**: [How Powerful are Graph Neural Networks?](https://arxiv.org/abs/1810.00826) by Xu et al. (2019)
... and many more (click to expand).
-* **[INorp](kgcnn/literature/INorp)**: [Interaction Networks for Learning about Objects,Relations and Physics](https://arxiv.org/abs/1612.00222) by Battaglia et al. (2016)
-* **[Megnet](kgcnn/literature/Megnet)**: [Graph Networks as a Universal Machine Learning Framework for Molecules and Crystals](https://doi.org/10.1021/acs.chemmater.9b01294) by Chen et al. (2019)
-* **[NMPN](kgcnn/literature/NMPN)**: [Neural Message Passing for Quantum Chemistry](http://arxiv.org/abs/1704.01212) by Gilmer et al. (2017)
-* **[Unet](kgcnn/literature/Unet)**: [Graph U-Nets](http://proceedings.mlr.press/v97/gao19a/gao19a.pdf) by H. Gao and S. Ji (2019)
* **[GATv2](kgcnn/literature/GATv2)**: [How Attentive are Graph Attention Networks?](https://arxiv.org/abs/2105.14491) by Brody et al. (2021)
-* **[GIN](kgcnn/literature/GIN)**: [How Powerful are Graph Neural Networks?](https://arxiv.org/abs/1810.00826) by Xu et al. (2019)
-* **[PAiNN](kgcnn/literature/PAiNN)**: [Equivariant message passing for the prediction of tensorial properties and molecular spectra](https://arxiv.org/pdf/2102.03150.pdf) by Schütt et al. (2020)
-* **[DMPNN](kgcnn/literature/DMPNN)**: [Analyzing Learned Molecular Representations for Property Prediction](https://pubs.acs.org/doi/abs/10.1021/acs.jcim.9b00237) by Yang et al. (2019)
-* **[HamNet](kgcnn/literature/HamNet)**: [HamNet: Conformation-Guided Molecular Representation with Hamiltonian Neural Networks](https://arxiv.org/abs/2105.03688) by Li et al. (2021)
-* **[CGCNN](kgcnn/literature/CGCNN)**: [Crystal Graph Convolutional Neural Networks for an Accurate and Interpretable Prediction of Material Properties](https://journals.aps.org/prl/abstract/10.1103/PhysRevLett.120.145301) by Xie et al. (2018)
-* **[CMPNN](kgcnn/literature/CMPNN)**: [Communicative Representation Learning on Attributed Molecular Graphs](https://www.ijcai.org/proceedings/2020/0392.pdf) by Song et al. (2020)
-* **[EGNN](kgcnn/literature/EGNN)**: [E(n) Equivariant Graph Neural Networks](https://arxiv.org/abs/2102.09844) by Satorras et al. (2021)
-* **[MAT](kgcnn/literature/MAT)**: [Molecule Attention Transformer](https://arxiv.org/abs/2002.08264) by Maziarka et al. (2020)
-* **[MXMNet](kgcnn/literature/MXMNet)**: [Molecular Mechanics-Driven Graph Neural Network with Multiplex Graph for Molecular Structures](https://arxiv.org/abs/2011.07457) by Zhang et al. (2020)
-* **[RGCN](kgcnn/literature/RGCN)**: [Modeling Relational Data with Graph Convolutional Networks](https://arxiv.org/abs/1703.06103) by Schlichtkrull et al. (2017)
-* **[GNNFilm](kgcnn/literature/GNNFilm)**: [GNN-FiLM: Graph Neural Networks with Feature-wise Linear Modulation](https://arxiv.org/abs/1906.12192) by Marc Brockschmidt (2020)
-* **[HDNNP2nd](kgcnn/literature/HDNNP2nd)**: [Atom-centered symmetry functions for constructing high-dimensional neural network potentials](https://aip.scitation.org/doi/abs/10.1063/1.3553717) by Jörg Behler (2011)
-* **[HDNNP4th](kgcnn/literature/HDNNP4th)**: [A fourth-generation high-dimensional neural network potential with accurate electrostatics including non-local charge transfer](https://www.nature.com/articles/s41467-020-20427-2) by Ko et al. (2021)
-* **[DGIN](kgcnn/literature/DGIN)**: [Improved Lipophilicity and Aqueous Solubility Prediction with Composite Graph Neural Networks ](https://pubmed.ncbi.nlm.nih.gov/34684766/) by Wieder et al. (2021)
-* **[MoGAT](kgcnn/literature/MoGAT)**: [Multi-order graph attention network for water solubility prediction and interpretation](https://www.nature.com/articles/s41598-022-25701-5) by Lee et al. (2023)
-* **[rGIN](kgcnn/literature/rGIN)** [Random Features Strengthen Graph Neural Networks](https://arxiv.org/abs/2002.03155) by Sato et al. (2020)
+
# Data
-How to construct ragged tensors is shown [above](#implementation-details).
-Moreover, some data handling classes are given in `kgcnn.data`.
+Data handling classes are given in `kgcnn.data` which stores graphs as `List[Dict]` .
#### Graph dictionary
@@ -251,7 +161,7 @@ dataset = MemoryGraphDataset(data_directory="ExampleDir/",
dataset.save() # opposite is load().
```
-The subclasses `QMDataset`, `MoleculeNetDataset`, `CrystalDataset`, `VisualGraphDataset` and `GraphTUDataset` further have functions required for the specific dataset type to convert and process files such as '.txt', '.sdf', '.xyz' etc.
+The subclasses `QMDataset`, `ForceDataset`, `MoleculeNetDataset`, `CrystalDataset`, `VisualGraphDataset` and `GraphTUDataset` further have functions required for the specific dataset type to convert and process files such as '.txt', '.sdf', '.xyz' etc.
Most subclasses implement `prepare_data()` and `read_in_memory()` with dataset dependent arguments.
An example for `MoleculeNetDataset` is shown below.
For more details find tutorials in [notebooks](notebooks).
@@ -290,9 +200,7 @@ You can find a [table](training/results/README.md) of common benchmark datasets
# Issues
Some known issues to be aware of, if using and making new models or layers with `kgcnn`.
-* RaggedTensor can not yet be used as a keras model output [(issue)](https://github.com/tensorflow/tensorflow/issues/42320), which has been mostly resolved in TF 2.8.
-* Using `RaggedTensor`'s for arbitrary ragged rank apart from `kgcnn.layers.modules` can cause significant performance decrease. This is due to shape check during add, multiply or concatenate (we think).
- We therefore use lazy add and concat in the `kgcnn.layers.modules` layers or directly operate on the value tensor for possible rank.
+* TODO
# Citing