docs/tutorials/add_your_own.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Define your own calculator\n",
    "Remember that a calculator is simply a `callable` that takes a molecule as input (either a RDKit `Chem.Mol` object or SMILES string) and returns a dictionary of features.\n",
    "We can thus easily define our own calculator!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[array([13.        , 13.        ,  3.        ,  0.63696169]),\n",
       " array([5.        , 4.        , 4.        , 0.63696169]),\n",
       " array([5.        , 4.        , 0.        , 0.63696169])]"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import numpy as np\n",
    "import datamol as dm\n",
    "\n",
    "from molfeat.trans import MoleculeTransformer\n",
    "from rdkit.Chem.rdMolDescriptors import CalcNumHeteroatoms\n",
    "\n",
    "smiles = dm.freesolv()[\"smiles\"].iloc[:3]\n",
    "def my_calculator(mol):\n",
    "    \"\"\"My custom featurizer\"\"\"\n",
    "    mol = dm.to_mol(mol)\n",
    "    rng = np.random.default_rng(0)\n",
    "    return [mol.GetNumAtoms(), mol.GetNumBonds(), CalcNumHeteroatoms(mol), rng.random()]\n",
    "\n",
    "# This directly works with the MoleculeTransformer\n",
    "mol_transf = MoleculeTransformer(my_calculator)\n",
    "mol_transf(smiles)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "If such functions get more complex, it might instead be easier to wrap it in a class.\n",
    "This also ensures the calculator remains serializable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[array([13.        , 13.        ,  3.        ,  0.63696169]),\n",
       " array([5.        , 4.        , 4.        , 0.63696169]),\n",
       " array([5.        , 4.        , 0.        , 0.63696169])]"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from molfeat.calc import SerializableCalculator\n",
    "\n",
    "\n",
    "class MyCalculator(SerializableCalculator):\n",
    "    def __call__(self, mol):\n",
    "        mol = dm.to_mol(mol)\n",
    "        rng = np.random.default_rng(0)\n",
    "        return [mol.GetNumAtoms(), mol.GetNumBonds(), CalcNumHeteroatoms(mol), rng.random()]\n",
    "\n",
    "\n",
    "mol_transf = MoleculeTransformer(MyCalculator())\n",
    "mol_transf(smiles)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If your calculator can perform featurization of a batch of molecules in an efficient way, then you should implement the optional `batch_compute` method which will then be used by `MoleculeTransformer` instead of the default sequential or parallelization process. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "We are in batch mode!\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([[0.77395605, 0.43887844, 0.85859792, 0.69736803, 0.09417735,\n",
       "        0.97562235, 0.7611397 , 0.78606431, 0.12811363, 0.45038594],\n",
       "       [0.37079802, 0.92676499, 0.64386512, 0.82276161, 0.4434142 ,\n",
       "        0.22723872, 0.55458479, 0.06381726, 0.82763117, 0.6316644 ],\n",
       "       [0.75808774, 0.35452597, 0.97069802, 0.89312112, 0.7783835 ,\n",
       "        0.19463871, 0.466721  , 0.04380377, 0.15428949, 0.68304895]])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from molfeat.calc import SerializableCalculator\n",
    "\n",
    "class MyBatchableCalculator(SerializableCalculator):\n",
    "    def __init__(self, random_seed=42, length=10):\n",
    "        self.random_seed = random_seed\n",
    "        self.length = length\n",
    "        self.rng = np.random.default_rng(self.random_seed)\n",
    "\n",
    "    def __call__(self, mol):\n",
    "        print(\"We are in single compute mode!\")\n",
    "        return self.rng.random(self.length)\n",
    "    \n",
    "    def __len__(self):\n",
    "        return self.length\n",
    "        \n",
    "    def batch_compute(self, mols, **kwargs):\n",
    "        # note that dm.parallelized information is passed along with the molecules list\n",
    "        print(\"We are in batch mode!\")\n",
    "        return self.rng.random((len(mols), self.length))\n",
    "\n",
    "mol_transf = MoleculeTransformer(MyBatchableCalculator())\n",
    "mol_transf(smiles)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Define your own transformer\n",
    "The above example shows that in many cases, there's no direct need to create your own transformer class. You can simply use the `MoleculeTransformer` base class.\n",
    "In more complex cases, such as with pretrained models where batching would be advantageous, it is instead preferable to create your own subclass. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from sklearn.ensemble import RandomForestRegressor\n",
    "from molfeat.trans.pretrained import PretrainedMolTransformer\n",
    "\n",
    "\n",
    "class MyFoundationModel(PretrainedMolTransformer):\n",
    "    \"\"\"\n",
    "    In this dummy example, we train a RF model to predict the cLogP\n",
    "    then use the feature importance of the RF model as the embedding.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self):\n",
    "        super().__init__(dtype=np.float32)\n",
    "        self._featurizer = MoleculeTransformer(\"maccs\", dtype=np.float32)\n",
    "        self._model = RandomForestRegressor()\n",
    "        self.train_dummy_model()\n",
    "\n",
    "    def train_dummy_model(self):\n",
    "        \"\"\"\n",
    "        Load the pretrained model.\n",
    "        In this dummy example, we train a RF model to predict the cLogP\n",
    "        \"\"\"\n",
    "        data = dm.data.freesolv().smiles.values\n",
    "        X = self._featurizer(data)\n",
    "        y = np.array([dm.descriptors.clogp(dm.to_mol(smi)) for smi in data])\n",
    "        self._model.fit(X, y)\n",
    "\n",
    "    def _convert(self, inputs: list, **kwargs):\n",
    "        \"\"\"Convert the molecule to a format that the model expects\"\"\"\n",
    "        return self._featurizer(inputs)\n",
    "\n",
    "    def _embed(self,  mols: list, **kwargs):\n",
    "        \"\"\"\n",
    "        Embed the molecules using the pretrained model\n",
    "        In this dummy example, we simply multiply the features by the importance of the feature\n",
    "        \"\"\"\n",
    "        return [feats * self._model.feature_importances_ for feats in mols]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1, 167)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mol_transf = MyFoundationModel()\n",
    "mol_transf([\"CN1C=NC2=C1C(=O)N(C(=O)N2C)C\"]).shape"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is another example that shows how to extend Molfeat with an existing embedding language model for astrochemistry. \n",
    "\n",
    "```bash\n",
    "pip install astrochem_embedding\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([10, 32])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "import datamol as dm\n",
    "\n",
    "from astrochem_embedding import VICGAE\n",
    "from molfeat.trans.pretrained import PretrainedMolTransformer\n",
    "\n",
    "class MyAstroChemFeaturizer(PretrainedMolTransformer):\n",
    "    \"\"\"\n",
    "    In this more practical example, we use embeddings from VICGAE a variance-invariance-covariance \n",
    "    regularized GRU autoencoder trained on SELFIES strings.\n",
    "    \"\"\"\n",
    "    def __init__(self, *args, **kwargs):\n",
    "        super().__init__(*args, **kwargs)        \n",
    "        self.featurizer = VICGAE.from_pretrained()\n",
    "    \n",
    "    def _embed(self, smiles, **kwargs):\n",
    "        return [self.featurizer.embed_smiles(x) for x in smiles]\n",
    "\n",
    "transformer = MyAstroChemFeaturizer(dtype=torch.float)\n",
    "transformer(dm.freesolv()[\"smiles\"][:10]).shape\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Add it to your Model Store\n",
    "Molfeat has a Model Store to publish your models in a centralized location.\n",
    "The default is a read-only GCP bucket but you can replace this with your own file storage. This can, for example, be useful to share private featurizers with your team."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import platformdirs\n",
    "from molfeat.store.modelstore import ModelStore\n",
    "from molfeat.store import ModelInfo\n",
    "\n",
    "path = dm.fs.join(platformdirs.user_cache_dir(\"molfeat\"), \"custom_model_store\")\n",
    "store = ModelStore(model_store_bucket=path)\n",
    "len(store.available_models)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "805c88dfabe241bdb8d04325fca22448",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0.00/4.00 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-04-07 18:59:03.363 | INFO     | molfeat.store.modelstore:register:124 - Successfuly registered model my_foundation_model !\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[ModelInfo(name='my_foundation_model', inputs='smiles', type='pretrained', version=0, group='my_group', submitter='Datamol', description='Solves chemistry!', representation='vector', require_3D=False, tags=['foundation_model', 'random_forest'], authors=['Datamol'], reference='/fake/ref', created_at=datetime.datetime(2023, 4, 7, 18, 59, 3, 312234), sha256sum='9c298d589a2158eb513cb52191144518a2acab2cb0c04f1df14fca0f712fa4a1')]"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Let's define our model's info\n",
    "info = ModelInfo(\n",
    "    name = \"my_foundation_model\",\n",
    "    inputs = \"smiles\",\n",
    "    type=\"pretrained\",\n",
    "    group=\"my_group\",\n",
    "    version=0,\n",
    "    submitter=\"Datamol\",\n",
    "    description=\"Solves chemistry!\",\n",
    "    representation=\"vector\",\n",
    "    require_3D=False,\n",
    "    tags = [\"foundation_model\", \"random_forest\"],\n",
    "    authors= [\"Datamol\"],\n",
    "    reference = \"/fake/ref\"\n",
    ")\n",
    "\n",
    "store.register(info)\n",
    "store.available_models"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Share with the community\n",
    "We invite you to share your featurizers with the community to help progress the field.\n",
    "To learn more, visit [the developer documentation](../developers/create-plugin.html).\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}