Skip to content

LCAONet - MPNN including electronic structure and orbital information, physically motivatied by the LCAO method.

License

Notifications You must be signed in to change notification settings

nmdl-mizo/lcaonet

Repository files navigation

LCAONet

python pytorch Code style: black isort pre-commit license

model

Installation

Requirements

Note: Using a GPU is recommended.

Install from source

First, clone this repository.

git clone https://github.com/nmdl-mizo/lcaonet.git

It is possible to build a virtual environment using conda, venv, or docker.

Using conda

You can create a new virtual environment with conda by running below commands:

conda create -n lcaonet python=3.10
conda activate lcaonet

Install dependencies in your environment:

conda install pytorch=2.0.0 -c pytorch
conda install pyg pytorch-scatter pytorch-sparse -c pyg
conda install numpy=1.* scipy=1.* sympy=1.* ase=3.* -c anaconda -c conda-forge

Install LCAONet:

cd lcaonet/conda
chmod +x build_conda.sh
./build_conda.sh

Using venv

You can create a new virtual environments with venv by running below commands:

python3 -m venv lcaonet-venv
source lcaonet-venv/bin/activate

Install dependencies in your environment:

cd lcaonet
pip install -r requirements.txt

Install LCAONet:

pip install .

Using docker

You can use the docker image of base environment from here.

Usage

You can train LCAONet with custom data in the following three steps.

  1. Prepare data

    First, prepare a list of ase.Atoms objects and a dict of physical property values to be labels. Then, convert them to lcaonet.data.List2GraphDataset object which inherits the torch_geoemtric.data.Dataset class.

    from numpy import ndarray
    from torch import Tensor
    from ase import Atoms
    
    from lcaonet.data import List2GraphDataset
    
    # Prepare a list of ase.Atoms objects
    data_list: list[Atoms] = ...
    # Prepare a dict of physical property values(Key: label name, Value: array of label values).
    label_list: dict[str, list[float] | ndarray | Tensor] = ...
    
    # Convert to List2GraphDataset object
    dataset = List2GraphDataset(data_list, y_values=label_list, cutoff=5.0)
  2. Define model

    Define LCAONet with any hyperparameters.

    from lcaonet.nn.cutoff import BaseCutoff
    from lcaonet.model import LCAONet
    
    # Define LCAONet
    model = LCAONet(
        hidden_dim: int = 128,
        coeffs_dim: int = 128,
        conv_dim: int = 128,
        out_dim: int = 1,
        n_interaction: int = 3,
        n_per_orb: int = 1,
        cutoff: float | None = None,
        rbf_type: str | type[BaseRadialBasis] = "hydrogen",
        cutoff_net: str | type[BaseCutoff] | None = "polynomial",
        max_z: int = 36,
        max_orb: str | None = None,
        elec_to_node: bool = True,
        add_valence: bool = False,
        extend_orb: bool = False,
        is_extensive: bool = True,
        activation: str = "SiLU",
        weight_init: str | None = "glorotorthogonal",
    )
  3. Train model

    Train with the interface of your choice (either plain PyTorch or PytorchLighting).

    import torch
    from torch_geometric.loader import DataLoader
    
    # Prepare DataLoader
    loader = DataLoader(dataset, batch_size=32, shuffle=True)
    
    for _ in range(epochs):
        for batch in loader:
            # Forward
            y_pred = model(batch)
            # Calculate loss
            loss = ...
            loss.backward()
            ...

References

  1. K. Nishio, K. Shibata, T. Mizoguchi. LCAONet: Message passing with physically optimized atomic basis functions. (2023) Paper