292 lines
11 KiB
Plaintext
292 lines
11 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"# sqlmodel-pg-kit — Cheminformatics Quickstart\n\n",
|
|
"This notebook demonstrates how to use `sqlmodel-pg-kit` for day-to-day data work, including:\n",
|
|
"- Environment setup (SQLite smoke vs PostgreSQL)\n",
|
|
"- Core CRUD with SQLModel sessions\n",
|
|
"- Multi-table schema and joins for molecules and datasets\n",
|
|
"- Optional RDKit + Mordred descriptor computation and storage\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## 0. Install (in Jupyter)\n",
|
|
"Uncomment as needed to install the package and optional chem deps."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# %pip install -e . pytest\n",
|
|
"# Optional: RDKit + Mordred\n",
|
|
"# %pip install rdkit-pypi mordred\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## 1. Basic kit usage (SQLite smoke)\n",
|
|
"Use the built-in smoke test pattern with an in-memory SQLite DB to verify CRUD paths."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"import sys\n",
|
|
"import os\n",
|
|
"\n",
|
|
"# Add the src directory to the path so we can import sqlmodel_pg_kit\n",
|
|
"sys.path.insert(0, os.path.join(os.getcwd(), '..', 'src'))\n",
|
|
"\n",
|
|
"from sqlmodel_pg_kit import db, create_all, Repository\n",
|
|
"\n",
|
|
"# Override to SQLite in-memory for quick check\n",
|
|
"db.cfg = db.DatabaseConfig(host='', port=0, user='', password='', database=':memory:', sslmode='disable')\n",
|
|
"db.engine = db.create_engine('sqlite:///:memory:', echo=False)\n",
|
|
"\n",
|
|
"from typing import Optional\n",
|
|
"from sqlmodel import SQLModel, Field\n",
|
|
"\n",
|
|
"class Hero(SQLModel, table=True):\n",
|
|
" id: Optional[int] = Field(default=None, primary_key=True)\n",
|
|
" name: str = Field(index=True)\n",
|
|
" age: Optional[int] = None\n",
|
|
"\n",
|
|
"create_all()\n",
|
|
"repo = Repository(Hero)\n",
|
|
"from sqlmodel_pg_kit.db import get_session\n",
|
|
"with get_session() as s:\n",
|
|
" repo.create(s, {'name': 'Iron Man', 'age': 45})\n",
|
|
" [h.name for h in repo.list(s)]\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## 2. Connect to PostgreSQL (optional)\n",
|
|
"Export `SQL_*` or `PG*` variables in your shell before starting Jupyter or set them here."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"import os\n",
|
|
"# Example (adjust to your environment, uncomment to set in-process):\n",
|
|
"# os.environ['SQL_HOST'] = '127.0.0.1'\n",
|
|
"# os.environ['SQL_PORT'] = '5432'\n",
|
|
"# os.environ['SQL_USER'] = 'postgres'\n",
|
|
"# os.environ['SQL_PASSWORD'] = 'change-me-strong'\n",
|
|
"# os.environ['SQL_DATABASE'] = 'appdb'\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## 3. Cheminformatics schema (Molecule, Dataset, MoleculeDataset)\n",
|
|
"Define models inline (you can also place them in your own package)."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from __future__ import annotations\n",
|
|
"from dataclasses import dataclass\n",
|
|
"from datetime import datetime\n",
|
|
"from typing import List, Optional\n",
|
|
"\n",
|
|
"from sqlmodel_pg_kit.db import get_session, engine\n",
|
|
"from sqlmodel import SQLModel, Field, Relationship, select\n",
|
|
"from sqlalchemy.orm import Mapped\n",
|
|
"\n",
|
|
"class MoleculeDataset(SQLModel, table=True):\n",
|
|
" molecule_id: int = Field(foreign_key='molecule.id', primary_key=True)\n",
|
|
" dataset_id: int = Field(foreign_key='dataset.id', primary_key=True)\n",
|
|
" added_at: datetime = Field(default_factory=datetime.utcnow)\n",
|
|
"\n",
|
|
"class Molecule(SQLModel, table=True):\n",
|
|
" id: Optional[int] = Field(default=None, primary_key=True)\n",
|
|
" smiles: str = Field(index=True)\n",
|
|
" selfies: Optional[str] = Field(default=None)\n",
|
|
" qed: Optional[float] = Field(default=None, index=True)\n",
|
|
" sa_score: Optional[float] = Field(default=None, index=True)\n",
|
|
" created_at: datetime = Field(default_factory=datetime.utcnow)\n",
|
|
" updated_at: datetime = Field(default_factory=datetime.utcnow)\n",
|
|
" datasets: Mapped[List[\"Dataset\"]] = Relationship(back_populates=\"molecules\", link_model=MoleculeDataset)\n",
|
|
"\n",
|
|
"class Dataset(SQLModel, table=True):\n",
|
|
" id: Optional[int] = Field(default=None, primary_key=True)\n",
|
|
" name: str = Field(index=True)\n",
|
|
" molecules: Mapped[List[\"Molecule\"]] = Relationship(back_populates=\"datasets\", link_model=MoleculeDataset)\n",
|
|
"\n",
|
|
"@dataclass\n",
|
|
"class MoleculeDTO:\n",
|
|
" smiles: str\n",
|
|
" selfies: Optional[str] = None\n",
|
|
" qed: Optional[float] = None\n",
|
|
" sa_score: Optional[float] = None\n",
|
|
" def to_model(self) -> Molecule:\n",
|
|
" return Molecule(**vars(self))\n",
|
|
"\n",
|
|
"# Create tables for these models\n",
|
|
"SQLModel.metadata.create_all(engine)\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## 4. CRUD and common queries\n",
|
|
"Insert molecules/datasets, link them, filter and join."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# Clean\n",
|
|
"with get_session() as s:\n",
|
|
" s.execute(MoleculeDataset.__table__.delete())\n",
|
|
" s.execute(Molecule.__table__.delete())\n",
|
|
" s.execute(Dataset.__table__.delete())\n",
|
|
" s.commit()\n",
|
|
"\n",
|
|
"# Insert molecules via DTO\n",
|
|
"mols = [\n",
|
|
" MoleculeDTO(smiles='CCO', qed=0.45, sa_score=2.1),\n",
|
|
" MoleculeDTO(smiles='c1ccccc1', qed=0.76, sa_score=3.5),\n",
|
|
" MoleculeDTO(smiles='CCN(CC)CC', qed=0.62, sa_score=2.8),\n",
|
|
"]\n",
|
|
"with get_session() as s:\n",
|
|
" s.add_all([dto.to_model() for dto in mols])\n",
|
|
" s.commit()\n",
|
|
"\n",
|
|
"# Datasets and linking\n",
|
|
"with get_session() as s:\n",
|
|
" ds_train = Dataset(name='train'); ds_holdout = Dataset(name='holdout')\n",
|
|
" s.add_all([ds_train, ds_holdout]); s.commit(); s.refresh(ds_train); s.refresh(ds_holdout)\n",
|
|
" mol_list = s.exec(select(Molecule).order_by(Molecule.id.asc())).all()\n",
|
|
" links = [\n",
|
|
" MoleculeDataset(molecule_id=mol_list[0].id, dataset_id=ds_train.id),\n",
|
|
" MoleculeDataset(molecule_id=mol_list[1].id, dataset_id=ds_train.id),\n",
|
|
" MoleculeDataset(molecule_id=mol_list[2].id, dataset_id=ds_holdout.id),\n",
|
|
" ]\n",
|
|
" s.add_all(links); s.commit()\n",
|
|
"\n",
|
|
"# Update one\n",
|
|
"with get_session() as s:\n",
|
|
" mol = s.exec(select(Molecule).where(Molecule.smiles=='CCO')).one()\n",
|
|
" mol.qed = 0.50; s.add(mol); s.commit(); s.refresh(mol)\n",
|
|
" mol.qed\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# Filters and joins\n",
|
|
"from sqlalchemy.orm import selectinload\n",
|
|
"\n",
|
|
"with get_session() as s:\n",
|
|
" hi_qed = s.exec(select(Molecule).where(Molecule.qed>=0.6).order_by(Molecule.sa_score.asc())).all()\n",
|
|
" hi_qed_view = [(m.smiles, m.qed, m.sa_score) for m in hi_qed]\n",
|
|
"\n",
|
|
"with get_session() as s:\n",
|
|
" stmt = select(Molecule).options(selectinload(Molecule.datasets)).order_by(Molecule.id.asc())\n",
|
|
" mols_with_ds = s.exec(stmt).all()\n",
|
|
" mols_with_ds_view = [(m.smiles, [d.name for d in m.datasets]) for m in mols_with_ds]\n",
|
|
"\n",
|
|
"with get_session() as s:\n",
|
|
" stmt = (select(Molecule)\n",
|
|
" .join(MoleculeDataset, Molecule.id==MoleculeDataset.molecule_id)\n",
|
|
" .join(Dataset, Dataset.id==MoleculeDataset.dataset_id)\n",
|
|
" .where(Dataset.name=='train')\n",
|
|
" .order_by(Molecule.id.asc()))\n",
|
|
" train_mols = s.exec(stmt).all()\n",
|
|
" train_mols_view = [m.smiles for m in train_mols]\n",
|
|
"\n",
|
|
"hi_qed_view, mols_with_ds_view, train_mols_view\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## 5. Optional: RDKit + Mordred integration\n",
|
|
"Compute descriptors and store them. If you prefer flexible storage, use a JSONB column or a normalized EAV table."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"try:\n",
|
|
" from rdkit import Chem\n",
|
|
" from rdkit.Chem import QED\n",
|
|
" from mordred import Calculator, descriptors\n",
|
|
" rdkit_ok = True\n",
|
|
"except Exception as e:\n",
|
|
" rdkit_ok = False\n",
|
|
" print('RDKit/Mordred not available in this environment. Skipping demo.\n', e)\n",
|
|
"\n",
|
|
"if rdkit_ok:\n",
|
|
" mol = Chem.MolFromSmiles('c1ccccc1O')\n",
|
|
" qed = float(QED.qed(mol))\n",
|
|
" calc = Calculator(descriptors, ignore_3D=True)\n",
|
|
" md = calc(mol)\n",
|
|
" # keep numeric descriptors only\n",
|
|
" desc = {k: float(v) for k, v in md.items() if v is not None and isinstance(v, (int, float))}\n",
|
|
" print('qed:', qed, 'num_desc:', len(desc))\n",
|
|
"\n",
|
|
" # Upsert molecule with refined qed as a column; optionally also persist `desc` via JSONB/EAV patterns.\n",
|
|
" with get_session() as s:\n",
|
|
" m = s.exec(select(Molecule).where(Molecule.smiles=='c1ccccc1O')).first()\n",
|
|
" if m is None:\n",
|
|
" m = Molecule(smiles='c1ccccc1O', qed=qed)\n",
|
|
" else:\n",
|
|
" m.qed = qed\n",
|
|
" s.add(m); s.commit(); s.refresh(m)\n",
|
|
" print('Stored molecule id:', m.id)\n"
|
|
]
|
|
}
|
|
],
|
|
"metadata": {
|
|
"kernelspec": {
|
|
"display_name": "Python 3",
|
|
"language": "python",
|
|
"name": "python3"
|
|
},
|
|
"language_info": {
|
|
"name": "python",
|
|
"pygments_lexer": "ipython3"
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 5
|
|
} |