{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a7c687f9",
   "metadata": {},
   "source": [
    "# DNAmStress"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be14569f",
   "metadata": {},
   "source": [
    "## Index\n",
    "1. [Instantiate model class](#Instantiate-model-class)\n",
    "2. [Define clock metadata](#Define-clock-metadata)\n",
    "3. [Download clock dependencies](#Download-clock-dependencies)\n",
    "5. [Load features](#Load-features)\n",
    "6. [Load weights into base model](#Load-weights-into-base-model)\n",
    "7. [Load reference values](#Load-reference-values)\n",
    "8. [Load preprocess and postprocess objects](#Load-preprocess-and-postprocess-objects)\n",
    "10. [Check all clock parameters](#Check-all-clock-parameters)\n",
    "10. [Basic test](#Basic-test)\n",
    "11. [Save torch model](#Save-torch-model)\n",
    "12. [Clear directory](#Clear-directory)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a4dd3bb1",
   "metadata": {},
   "source": [
    "Let's first import some packages:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "96ac895a",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-04T19:45:53.629908Z",
     "iopub.status.busy": "2026-07-04T19:45:53.629690Z",
     "iopub.status.idle": "2026-07-04T19:45:55.585656Z",
     "shell.execute_reply": "2026-07-04T19:45:55.585211Z"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import inspect\n",
    "import shutil\n",
    "import json\n",
    "import torch\n",
    "import pandas as pd\n",
    "import pyaging as pya"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dcd58d1d",
   "metadata": {},
   "source": [
    "## Instantiate model class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "429bc70d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-04T19:45:55.588166Z",
     "iopub.status.busy": "2026-07-04T19:45:55.587926Z",
     "iopub.status.idle": "2026-07-04T19:45:55.590970Z",
     "shell.execute_reply": "2026-07-04T19:45:55.590501Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "class DNAmStress(LinearReferenceClock):\n",
      "    pass\n",
      "\n"
     ]
    }
   ],
   "source": [
    "def print_entire_class(cls):\n",
    "    source = inspect.getsource(cls)\n",
    "    print(source)\n",
    "\n",
    "print_entire_class(pya.models.DNAmStress)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "dcce9f26",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-04T19:45:55.592322Z",
     "iopub.status.busy": "2026-07-04T19:45:55.592232Z",
     "iopub.status.idle": "2026-07-04T19:45:55.593902Z",
     "shell.execute_reply": "2026-07-04T19:45:55.593578Z"
    }
   },
   "outputs": [],
   "source": [
    "model = pya.models.DNAmStress()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "efb13d6a",
   "metadata": {},
   "source": [
    "## Define clock metadata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "705a4e76",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-04T19:45:55.595168Z",
     "iopub.status.busy": "2026-07-04T19:45:55.595080Z",
     "iopub.status.idle": "2026-07-04T19:45:55.596970Z",
     "shell.execute_reply": "2026-07-04T19:45:55.596564Z"
    }
   },
   "outputs": [],
   "source": [
    "model.metadata[\"clock_name\"] = \"dnamstress\"\n",
    "model.metadata[\"data_type\"] = \"DNA methylation\"  # Paper: Whole-blood DNA methylation was quantified for the epigenomic cohort.\n",
    "model.metadata[\"species\"] = \"Homo sapiens\"  # Paper: The NIAAA cohorts comprised human participants with AUD and healthy controls.\n",
    "model.metadata[\"year\"] = 2023\n",
    "model.metadata[\"approved_by_author\"] = \"⌛\"\n",
    "model.metadata[\"citation\"] = \"Jung, J., et al. “Additive Effects of Stress and Alcohol Exposure on Accelerated Epigenetic Aging in Alcohol Use Disorder.” Biological Psychiatry 93(4): 331–341 (2023).\"\n",
    "model.metadata[\"doi\"] = \"https://doi.org/10.1016/j.biopsych.2022.06.036\"\n",
    "model.metadata[\"notes\"] = \"Whole-blood methylation score (MS stress) derived as a 211-CpG proxy for a composite of 13 stress-related measures.\"\n",
    "model.metadata[\"research_only\"] = None\n",
    "model.metadata[\"tissue\"] = [\"whole blood\"]  # Paper: Whole-blood DNAm was quantified in all 615 epigenomic-cohort participants.\n",
    "model.metadata[\"predicts\"] = [\"stress exposure\"]  # Paper: MS stress is a methylation-derived score tracking stress exposure.\n",
    "model.metadata[\"training_target\"] = [\"stress exposure\"]  # Paper: Three factor scores from 13 stress-associated measures were added to construct the composite score of stress.\n",
    "model.metadata[\"unit\"] = [\"unitless\"]  # Paper: MS stress was calculated as the weighted sum of selected CpG methylation values.\n",
    "model.metadata[\"model_type\"] = \"bootstrap-stabilized elastic net regression\"  # Paper: An elastic net model with 1,000 bootstraps selected CpGs present in more than half of bootstraps.\n",
    "model.metadata[\"platform\"] = [\"Illumina EPIC\"]  # Paper: Whole-blood DNAm was quantified using Infinium MethylationEPIC BeadChip arrays.\n",
    "model.metadata[\"population\"] = \"adults\"  # Paper: The epigenomic cohort consisted of 615 participants: 372 AUD cases and 243 healthy controls.\n",
    "model.metadata[\"journal\"] = \"Biological Psychiatry\"\n",
    "model.metadata[\"last_author\"] = \"Falk W. Lohoff\"\n",
    "model.metadata[\"n_features\"] = 211\n",
    "model.metadata[\"citations\"] = 27\n",
    "model.metadata[\"citations_date\"] = \"2026-07-05\"\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18ac2d89",
   "metadata": {},
   "source": [
    "## Download clock dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d528b070",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-04T19:45:55.598137Z",
     "iopub.status.busy": "2026-07-04T19:45:55.598067Z",
     "iopub.status.idle": "2026-07-04T19:45:56.137140Z",
     "shell.execute_reply": "2026-07-04T19:45:56.136049Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "os.system(f\"curl -sL -o DNAmStress_CpGs.rda https://raw.githubusercontent.com/HigginsChenLab/methylCIPHER/19b12296b0d7eb7055a97d068064df635f44ce3e/data/DNAmStress_CpGs.rda\")\n",
    "os.system(f\"curl -sL -o DNAmStress_ref.rda https://raw.githubusercontent.com/HigginsChenLab/methylCIPHER/19b12296b0d7eb7055a97d068064df635f44ce3e/data/DNAmStress_ref.rda\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "1460aec8",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-04T19:45:56.139837Z",
     "iopub.status.busy": "2026-07-04T19:45:56.139639Z",
     "iopub.status.idle": "2026-07-04T19:45:56.144253Z",
     "shell.execute_reply": "2026-07-04T19:45:56.143371Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Writing download.r\n"
     ]
    }
   ],
   "source": [
    "%%writefile download.r\n",
    "\n",
    "library(jsonlite)\n",
    "load(\"DNAmStress_CpGs.rda\")\n",
    "load(\"DNAmStress_ref.rda\")\n",
    "write_json(DNAmStress_CpGs, \"coefficients.json\", digits = 10)\n",
    "write_json(DNAmStress_ref$intercept, \"intercept.json\", digits = 10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "127acf2b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-04T19:45:56.146452Z",
     "iopub.status.busy": "2026-07-04T19:45:56.146295Z",
     "iopub.status.idle": "2026-07-04T19:45:56.501036Z",
     "shell.execute_reply": "2026-07-04T19:45:56.500617Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "os.system(\"Rscript download.r\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b50c3c2",
   "metadata": {},
   "source": [
    "## Load features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "e2a9d473",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-04T19:45:56.502318Z",
     "iopub.status.busy": "2026-07-04T19:45:56.502238Z",
     "iopub.status.idle": "2026-07-04T19:45:56.504890Z",
     "shell.execute_reply": "2026-07-04T19:45:56.504509Z"
    }
   },
   "outputs": [],
   "source": [
    "coef_df = pd.DataFrame(json.load(open('coefficients.json')))\n",
    "model.features = coef_df['CpG'].tolist()\n",
    "iv = json.load(open('intercept.json'))\n",
    "intercept_value = iv[0] if isinstance(iv, list) else iv"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9869944c",
   "metadata": {},
   "source": [
    "## Load weights into base model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "bf64bd06",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-04T19:45:56.506271Z",
     "iopub.status.busy": "2026-07-04T19:45:56.506183Z",
     "iopub.status.idle": "2026-07-04T19:45:56.508176Z",
     "shell.execute_reply": "2026-07-04T19:45:56.507811Z"
    }
   },
   "outputs": [],
   "source": [
    "weights = torch.tensor(coef_df['coefficient'].tolist()).unsqueeze(0).float()\n",
    "intercept = torch.tensor([intercept_value]).float()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "c2517a12",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-04T19:45:56.509612Z",
     "iopub.status.busy": "2026-07-04T19:45:56.509430Z",
     "iopub.status.idle": "2026-07-04T19:45:56.512109Z",
     "shell.execute_reply": "2026-07-04T19:45:56.511677Z"
    }
   },
   "outputs": [],
   "source": [
    "base_model = pya.models.LinearModel(input_dim=len(model.features))\n",
    "\n",
    "base_model.linear.weight.data = weights.float()\n",
    "base_model.linear.bias.data = intercept.float()\n",
    "\n",
    "model.base_model = base_model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b313e18",
   "metadata": {},
   "source": [
    "## Load reference values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "418405d9",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-04T19:45:56.513365Z",
     "iopub.status.busy": "2026-07-04T19:45:56.513266Z",
     "iopub.status.idle": "2026-07-04T19:45:56.515151Z",
     "shell.execute_reply": "2026-07-04T19:45:56.514616Z"
    }
   },
   "outputs": [],
   "source": [
    "model.reference_values = None"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "186f4793",
   "metadata": {},
   "source": [
    "## Load preprocess and postprocess objects"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "fd07ca60",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-04T19:45:56.517282Z",
     "iopub.status.busy": "2026-07-04T19:45:56.517169Z",
     "iopub.status.idle": "2026-07-04T19:45:56.518860Z",
     "shell.execute_reply": "2026-07-04T19:45:56.518530Z"
    }
   },
   "outputs": [],
   "source": [
    "model.preprocess_name = None\n",
    "model.preprocess_dependencies = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "feda4746",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-04T19:45:56.520032Z",
     "iopub.status.busy": "2026-07-04T19:45:56.519961Z",
     "iopub.status.idle": "2026-07-04T19:45:56.521642Z",
     "shell.execute_reply": "2026-07-04T19:45:56.521238Z"
    }
   },
   "outputs": [],
   "source": [
    "model.postprocess_name = None\n",
    "model.postprocess_dependencies = None"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89abe907",
   "metadata": {},
   "source": [
    "## Check all clock parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "ef65e440",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-04T19:45:56.522846Z",
     "iopub.status.busy": "2026-07-04T19:45:56.522780Z",
     "iopub.status.idle": "2026-07-04T19:45:56.527592Z",
     "shell.execute_reply": "2026-07-04T19:45:56.527250Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "%==================================== Model Details ====================================%\n",
      "Model Attributes:\n",
      "\n",
      "training: True\n",
      "metadata: {'approved_by_author': '⌛',\n",
      " 'citation': 'Jung, Sun Jae, et al. \"An epigenetic biomarker of stress and its '\n",
      "             'association with cardiovascular disease.\" Biological Psychiatry '\n",
      "             '93.4 (2023): 331-341.',\n",
      " 'clock_name': 'dnamstress',\n",
      " 'data_type': 'methylation',\n",
      " 'doi': 'https://doi.org/10.1016/j.biopsych.2022.06.036',\n",
      " 'notes': None,\n",
      " 'research_only': None,\n",
      " 'species': 'Homo sapiens',\n",
      " 'version': None,\n",
      " 'year': 2023}\n",
      "reference_values: None\n",
      "preprocess_name: None\n",
      "preprocess_dependencies: None\n",
      "postprocess_name: None\n",
      "postprocess_dependencies: None\n",
      "features: ['cg00188055', 'cg00418087', 'cg00452133', 'cg00599628', 'cg00684116', 'cg00728602', 'cg01047613', 'cg01085421', 'cg01245530', 'cg01607512', 'cg01631065', 'cg02107357', 'cg02386470', 'cg02606284', 'cg02606535', 'cg02607810', 'cg02652698', 'cg02873991', 'cg02978227', 'cg03104567', 'cg03120234', 'cg03221779', 'cg03579303', 'cg03581822', 'cg03626783', 'cg03642635', 'cg03693374', 'cg03723497', 'cg03744383', 'cg04048370']... [Total elements: 211]\n",
      "base_model_features: None\n",
      "\n",
      "%==================================== Model Details ====================================%\n",
      "Model Structure:\n",
      "\n",
      "base_model: LinearModel(\n",
      "  (linear): Linear(in_features=211, out_features=1, bias=True)\n",
      ")\n",
      "\n",
      "%==================================== Model Details ====================================%\n",
      "Model Parameters and Weights:\n",
      "\n",
      "base_model.linear.weight: [1.7692569494247437, 5.452592849731445, 0.36093270778656006, 0.04217686131596565, -0.0875726118683815, -0.1617601066827774, -2.8009119033813477, -0.0006456300034187734, 2.6688179969787598, 1.0967479944229126, 2.873945951461792, 0.3544867932796478, -2.170973062515259, 2.942610025405884, 4.0788421630859375, -0.08575636893510818, 0.3709535002708435, 0.18792760372161865, -4.1162872314453125, -0.25866061449050903, 5.40001916885376, -2.2270829677581787, -0.7070323824882507, 1.7177339792251587, -0.6149824857711792, 0.26971670985221863, -1.411486029624939, 1.3798350095748901, -0.12567970156669617, 1.3220950365066528]... [Tensor of shape torch.Size([1, 211])]\n",
      "base_model.linear.bias: tensor([-4.4941])\n",
      "\n",
      "%==================================== Model Details ====================================%\n",
      "\n"
     ]
    }
   ],
   "source": [
    "pya.utils.print_model_details(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73c0ce2f",
   "metadata": {},
   "source": [
    "## Basic test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "9b3d1b20",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-04T19:45:56.528885Z",
     "iopub.status.busy": "2026-07-04T19:45:56.528815Z",
     "iopub.status.idle": "2026-07-04T19:45:56.533709Z",
     "shell.execute_reply": "2026-07-04T19:45:56.533343Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-33.9411],\n",
       "        [-62.7000],\n",
       "        [ 93.3530],\n",
       "        [ 14.1386],\n",
       "        [-18.2623],\n",
       "        [  4.9429],\n",
       "        [ 94.5241],\n",
       "        [ 63.0782],\n",
       "        [  4.6987],\n",
       "        [ 66.1088]], dtype=torch.float64, grad_fn=<AddmmBackward0>)"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.manual_seed(42)\n",
    "input = torch.randn(10, len(model.features), dtype=float)\n",
    "model.eval()\n",
    "model.to(float)\n",
    "pred = model(input)\n",
    "pred"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a02da8f1",
   "metadata": {},
   "source": [
    "## Save torch model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "fa36a60a",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-04T19:45:56.534982Z",
     "iopub.status.busy": "2026-07-04T19:45:56.534906Z",
     "iopub.status.idle": "2026-07-04T19:45:56.538234Z",
     "shell.execute_reply": "2026-07-04T19:45:56.537896Z"
    }
   },
   "outputs": [],
   "source": [
    "torch.save(model, f\"../weights/{model.metadata['clock_name']}.pt\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4fa9971d",
   "metadata": {},
   "source": [
    "## Clear directory\n",
    "<a id=\"10\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "95b78d3a",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-04T19:45:56.539558Z",
     "iopub.status.busy": "2026-07-04T19:45:56.539488Z",
     "iopub.status.idle": "2026-07-04T19:45:56.543044Z",
     "shell.execute_reply": "2026-07-04T19:45:56.542643Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Deleted file: intercept.json\n",
      "Deleted file: DNAmStress_CpGs.rda\n",
      "Deleted file: download.r\n",
      "Deleted file: coefficients.json\n",
      "Deleted file: DNAmStress_ref.rda\n"
     ]
    }
   ],
   "source": [
    "# Function to remove a folder and all its contents\n",
    "def remove_folder(path):\n",
    "    try:\n",
    "        shutil.rmtree(path)\n",
    "        print(f\"Deleted folder: {path}\")\n",
    "    except Exception as e:\n",
    "        print(f\"Error deleting folder {path}: {e}\")\n",
    "\n",
    "# Get a list of all files and folders in the current directory\n",
    "all_items = os.listdir('.')\n",
    "\n",
    "# Loop through the items\n",
    "for item in all_items:\n",
    "    # Check if it's a file and does not end with .ipynb\n",
    "    if os.path.isfile(item) and not item.endswith('.ipynb'):\n",
    "        os.remove(item)\n",
    "        print(f\"Deleted file: {item}\")\n",
    "    # Check if it's a folder\n",
    "    elif os.path.isdir(item):\n",
    "        remove_folder(item)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.13.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}