This is the official repository of the paper "Explainable AI in Drug Discovery: Self-interpretable Graph Neural Network for molecular property prediction using Concept Whitening. Michela Proietti,Alessio Ragno Biagio La Rosa, Rino Ragno, and Roberto Capobianco. Machine Learning (2024)".

This repository implements the first concept-based explanability method for graph neural networks, obtained by adapting the code for (https://github.com/zhiCHEN96/ConceptWhitening "Concept Whitening"). It consists of a module which can be inserted straight after a convolutional layer in a GNN, in place of batch normalization layers. The module aligns the axes of the latent space with known concept of interest, in our case molecular properties that are known to influence the bioactivity of molecules.

In the config.yaml file in the config folder, you need to change the paths in the base_dir, statistics_dir, and concept_dir items. In order to use concept whitening, you will need to extract the molecular properties from the dataset's molecules. This can be done by executing the command:

python extract_molecular_concepts.py

To train a new model or test a pre-trained one, execute the command

python train_gnns.py

after setting the desired parameters by changing the files in the config folder.

@article{proietti2024explainable,
  title={Explainable AI in drug discovery: self-interpretable graph neural network for molecular property prediction using concept whitening},
  author={Proietti, Michela and Ragno, Alessio and Rosa, Biagio La and Ragno, Rino and Capobianco, Roberto},
  journal={Machine Learning},
  volume={113},
  number={4},
  pages={2013--2044},
  year={2024},
  publisher={Springer}
}

If you have any question, do not hesitate to contact us at mproietti[at]diag[dot]uniroma1[dot]it.