{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "98a92d0a",
      "metadata": {},
      "source": [
        "# Defining a Custom BioMol Class\n",
        "\n",
        "```{eval-rst}\n",
        ".. currentmodule:: biomol\n",
        "```\n",
        "\n",
        "In this tutorial, you will\n",
        "\n",
        "- Learn how to create a custom class that inherits from the {py:class}`BioMol` class.\n",
        "- Define a custom {py:class}`View <core.View>` for feature type-hinting.\n",
        "- Understand the benefits of using type hints in your code."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "1ed1f4da",
      "metadata": {},
      "source": [
        "## The Problem: Lack of Autocompletion\n",
        "\n",
        "Before we begin, let's discuss why type-hinting is important. \n",
        "Because the {py:class}`BioMol` class **doesn't predefine** specific features, it's hard to know what features are available in a {py:class}`BioMol` instance. \n",
        "For example, autocompletion for `mol.atoms.positions`  doesn't work.\n",
        "\n",
        "To address this limitation, you will:\n",
        "\n",
        "1. Define a custom {py:class}`View <core.View>` for feature type-hinting.\n",
        "2. Create a custom class that inherits from {py:class}`BioMol` and uses the defined {py:class}`View <core.View>`.\n",
        "\n",
        "Let's start by importing the necessary modules."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "id": "b81f4526",
      "metadata": {},
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "\n",
        "from biomol import BioMol\n",
        "from biomol.core import FeatureContainer, IndexTable, NodeFeature, View"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "cf7e531c",
      "metadata": {},
      "source": [
        "## Preparing container and index table\n",
        "\n",
        "As in the {doc}`previous tutorial <creating_biomol>`, prepare a {py:class}`FeatureContainer <core.FeatureContainer>` and an {py:class}`IndexTable <core.IndexTable>` to hold features and establish hierarchical relationships."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "id": "e22da61e",
      "metadata": {},
      "outputs": [],
      "source": [
        "atom_positions = NodeFeature(\n",
        "    value=np.array(\n",
        "        [\n",
        "            [0.0, 0.0, 0.0],\n",
        "            [1.4, 0.0, 0.0],\n",
        "            [1.4, 1.4, 0.0],\n",
        "            [0.0, 1.4, 0.0],  # ALA-1\n",
        "            [2.8, 0.0, 0.0],\n",
        "            [4.2, 0.0, 0.0],\n",
        "            [4.2, 1.4, 0.0],\n",
        "            [2.8, 1.4, 0.0],  # GLY-2\n",
        "            [5.6, 0.0, 0.0],\n",
        "            [7.0, 0.0, 0.0],\n",
        "            [7.0, 1.4, 0.0],\n",
        "            [5.6, 1.4, 0.0],  # ALA-3\n",
        "        ],\n",
        "    ),\n",
        ")\n",
        "atom_names = NodeFeature(value=np.array([\"N\", \"CA\", \"C\", \"O\"] * 3))\n",
        "\n",
        "residue_ids = NodeFeature(value=np.array([1, 2, 3]))\n",
        "residue_names = NodeFeature(value=np.array([\"ALA\", \"GLY\", \"ALA\"]))\n",
        "\n",
        "chain_ids = NodeFeature(value=np.array([\"A\"]))\n",
        "chain_entities = NodeFeature(value=np.array([\"PROTEIN\"]))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "id": "64940473",
      "metadata": {},
      "outputs": [],
      "source": [
        "atom_container = FeatureContainer(\n",
        "    {\"positions\": atom_positions, \"name\": atom_names},\n",
        ")\n",
        "\n",
        "residue_container = FeatureContainer(\n",
        "    {\"id\": residue_ids, \"name\": residue_names},\n",
        ")\n",
        "\n",
        "chain_container = FeatureContainer(\n",
        "    {\"id\": chain_ids, \"entity\": chain_entities},\n",
        ")\n",
        "\n",
        "index_table = IndexTable.from_parents(\n",
        "    atom_to_res=np.array([0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2]),\n",
        "    res_to_chain=np.array([0, 0, 0]),\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "18e7c0fa",
      "metadata": {},
      "source": [
        "## Defining a Custom BioMol\n",
        "\n",
        "Next, define a custom {py:class}`View <core.View>` and a custom class that inherits from  {py:class}`BioMol` class.\n",
        "The {py:class}`View <core.View>` specifies which features are exposed on the {py:class}`FeatureContainer <core.FeatureContainer>` via property methods.\n",
        "\n",
        "```{note}\n",
        "{py:class}`View <core.View>` is used solely for static type-hinting. \n",
        "It does **not** affect runtime behavior and therefore cannot guarantee that features actually exist at runtime.\n",
        "```"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "id": "80a3f51e",
      "metadata": {},
      "outputs": [],
      "source": [
        "class MyAtomView(View[\"MyAtomView\", \"MyResidueView\", \"MyChainView\", \"MyBioMol\"]):\n",
        "    \"\"\"Custom atom-level View.\"\"\"\n",
        "\n",
        "    @property\n",
        "    def positions(self) -> NodeFeature:\n",
        "        \"\"\"XYZ coordinates of atoms.\"\"\"\n",
        "\n",
        "    @property\n",
        "    def name(self) -> NodeFeature:\n",
        "        \"\"\"Atom name.\"\"\"\n",
        "\n",
        "\n",
        "class MyResidueView(View[\"MyAtomView\", \"MyResidueView\", \"MyChainView\", \"MyBioMol\"]):\n",
        "    \"\"\"Custom residue-level View.\"\"\"\n",
        "\n",
        "    @property\n",
        "    def id(self) -> NodeFeature:\n",
        "        \"\"\"Residue ID.\"\"\"\n",
        "\n",
        "    @property\n",
        "    def name(self) -> NodeFeature:\n",
        "        \"\"\"Residue name.\"\"\"\n",
        "\n",
        "\n",
        "class MyChainView(View[\"MyAtomView\", \"MyResidueView\", \"MyChainView\", \"MyBioMol\"]):\n",
        "    \"\"\"Custom chain-level View.\"\"\"\n",
        "\n",
        "    @property\n",
        "    def id(self) -> NodeFeature:\n",
        "        \"\"\"Chain ID.\"\"\"\n",
        "\n",
        "    @property\n",
        "    def entity(self) -> NodeFeature:\n",
        "        \"\"\"Chain entity type.\"\"\"\n",
        "\n",
        "\n",
        "class MyBioMol(BioMol[\"MyAtomView\", \"MyResidueView\", \"MyChainView\"]):\n",
        "    \"\"\"Custom BioMol class.\"\"\""
      ]
    },
    {
      "cell_type": "markdown",
      "id": "673f4c9b",
      "metadata": {},
      "source": [
        "Now, you can create an instance of your custom {py:class}`BioMol` class and access its features with type-hinting support."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "id": "9c0147ac",
      "metadata": {},
      "outputs": [],
      "source": [
        "mol = MyBioMol(\n",
        "    atom_container=atom_container,\n",
        "    residue_container=residue_container,\n",
        "    chain_container=chain_container,\n",
        "    index_table=index_table,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "id": "399888c9",
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/plain": [
              "NodeFeature(value=array(['PROTEIN'], dtype='<U7'), description=None)"
            ]
          },
          "execution_count": 6,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# ruff: noqa: B018\n",
        "\n",
        "mol.atoms.positions\n",
        "mol.atoms.name\n",
        "\n",
        "mol.residues.id\n",
        "mol.residues.name\n",
        "\n",
        "mol.chains.id\n",
        "mol.chains.entity"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "705c0d37",
      "metadata": {},
      "source": [
        "Type hints remain effective even in complex expressions that involve multiple hierarchical traversals."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "id": "5e7978ba",
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/plain": [
              "NodeFeature(value=array([1, 2, 3]), description=None)"
            ]
          },
          "execution_count": 7,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "mol.atoms.residues.chains.id\n",
        "mol.atoms.chains.residues.atoms.name\n",
        "mol.chains.atoms.chains.residues.id"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "default",
      "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.18"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}