{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# COVID-19 - unvaccinated high risk patients\n",
    "\n",
    "Analysis of FHIR source data using a Pathling FHIR server.\n",
    "\n",
    "This query counts patients, grouped by whether they have received a COVID-19 vaccination and whether they are high-risk based upon a number of factors (CKD, heart disease, BMI).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "22/08/10 01:48:31 WARN NativeCodeLoader: Unable to load native-hadoop library for your platform... using builtin-java classes where applicable\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting default log level to \"WARN\".\n",
      "To adjust logging level use sc.setLogLevel(newLevel). For SparkR, use setLogLevel(newLevel).\n"
     ]
    }
   ],
   "source": [
    "from time import time\n",
    "\n",
    "from pathling import PathlingContext\n",
    "from pathling.coding import Coding\n",
    "from pathling.etc import find_jar\n",
    "from pyspark.sql import SparkSession, functions\n",
    "\n",
    "start = time()\n",
    "spark = SparkSession.builder.config('spark.jars', find_jar()).config('spark.executor.memory',\n",
    "                                                                     '8g').getOrCreate()\n",
    "pc = PathlingContext.create(spark, terminology_server_url='http://localhost:8081/fhir')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "from time import time\n",
    "\n",
    "from pathling import PathlingContext\n",
    "from pathling.coding import Coding\n",
    "from pathling.etc import find_jar\n",
    "from pyspark.sql import SparkSession, functions\n",
    "\n",
    "start = time()\n",
    "spark = SparkSession.builder.config('spark.jars', find_jar()).config('spark.executor.memory',\n",
    "                                                                     '8g').getOrCreate()\n",
    "pc = PathlingContext.create(spark, terminology_server_url='http://localhost:8081/fhir')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Load data from NDJSON\n",
    "\n",
    "First we need to load the data from FHIR NDJSON files into Spark dataframes that we can work with."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def load_resource(path, resource_type):\n",
    "    json_resources = pc.spark.read.text(path)\n",
    "    return pc.encode(json_resources, resource_type)\n",
    "\n",
    "\n",
    "resources = ['Patient', 'Immunization', 'Condition', 'Observation']\n",
    "resource_data = {\n",
    "    resource_type: load_resource(\n",
    "        f'file:///Users/gri306/Library/CloudStorage/OneDrive-CSIRO/Data/synthea/paper-md/fhir/{resource_type}.ndjson',\n",
    "        resource_type)\n",
    "    for resource_type in resources\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Preparation\n",
    "\n",
    "Filter the Patient resources to only those that meet our cohort criteria (ages 18-60)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "patients = resource_data['Patient']\n",
    "filtered_patients = patients.filter(\n",
    "    (patients.birthDate < '2004-07-30') & (patients.birthDate > '1962-07-30'))\n",
    "# filtered_patients.toPandas()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Join the Patient resources to their associated Immunization resources, and explode out the codings ready for later terminology functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "immunizations = resource_data['Immunization']\n",
    "patients_immunizations = filtered_patients.join(\n",
    "    immunizations,\n",
    "    filtered_patients.id == functions.element_at(\n",
    "        functions.split(immunizations.patient.reference, '/'), 2),\n",
    "    'left_outer'\n",
    ")\n",
    "patients_immunizations = patients_immunizations.select(\n",
    "    filtered_patients.id,\n",
    "    functions.explode_outer(immunizations.vaccineCode.getField('coding')).alias('codings')\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Run the immunization codes through the `memberOf` function, testing whether they are within our COVID-19 vaccination value set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "covid_19 = pc.member_of(patients_immunizations, patients_immunizations.codings,\n",
    "                        'https://aehrc.csiro.au/fhir/ValueSet/covid-19-vaccines', 'covid_19')\n",
    "covid_19 = covid_19.groupBy(filtered_patients.id).agg(\n",
    "    functions.when(functions.max(covid_19.covid_19), True).otherwise(False).alias('covid_19'))\n",
    "# covid_19.toPandas()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Join the Patient resources to their associated Condition resources, and explode out the codings ready for later terminology functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "conditions = resource_data['Condition']\n",
    "patients_conditions = filtered_patients.join(\n",
    "    conditions,\n",
    "    filtered_patients.id == functions.element_at(functions.split(conditions.subject.reference, '/'),\n",
    "                                                 2),\n",
    "    'left_outer'\n",
    ")\n",
    "patients_conditions = patients_conditions.select(\n",
    "    filtered_patients.id,\n",
    "    functions.explode_outer(conditions.code.getField('coding')).alias('codings')\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Run the condition codes through the `subsumes` function, testing if they are sub-types of chronic kidney disease."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "ckd = pc.subsumes(patients_conditions,\n",
    "                  output_column_name='ckd',\n",
    "                  left_coding=Coding(system='http://snomed.info/sct', code='709044004'),\n",
    "                  right_coding_column=patients_conditions.codings)\n",
    "ckd = ckd.groupBy(filtered_patients.id).agg(\n",
    "    functions.when(functions.max(ckd.ckd), True).otherwise(False).alias('ckd'))\n",
    "# ckd.toPandas()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Run the condition codes through the `memberOf` function, testing whether they match an ECL expression that describes heart disease."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# << 49601007|Cardiovascular disease| : << 363698007|Finding site| = << 80891009|Structure of heart|\n",
    "heart_disease = pc.member_of(patients_conditions, patients_conditions.codings,\n",
    "                             'http://snomed.info/sct?fhir_vs=ecl/%3C%3C%2049601007%20%3A%20%3C%3C%20363698007%20'\n",
    "                             '%3D%20%3C%3C%2080891009%20', 'heart_disease')\n",
    "heart_disease = heart_disease.groupBy(filtered_patients.id).agg(\n",
    "    functions.when(functions.max(heart_disease.heart_disease), True).otherwise(False).alias(\n",
    "        'heart_disease'))\n",
    "# heart_disease.toPandas()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Join the Patient resources to their associated Observation resources."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "observations = resource_data['Observation']\n",
    "patients_observations = filtered_patients.join(\n",
    "    observations,\n",
    "    filtered_patients.id == functions.element_at(\n",
    "        functions.split(observations.subject.reference, '/'), 2),\n",
    "    'left_outer'\n",
    ")\n",
    "patients_observations = patients_observations.select(\n",
    "    filtered_patients.id,\n",
    "    functions.explode_outer(observations.code.getField('coding')).alias('codings'),\n",
    "    observations.valueQuantity\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Test whether the Quantity values within the BMI observations are greater than 30."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "bmi_over_30 = patients_observations.select(\n",
    "    filtered_patients.id,\n",
    "    functions.when(\n",
    "        (patients_observations.codings.system == 'http://loinc.org') &\n",
    "        (patients_observations.codings.code == '39156-5') &\n",
    "        observations.valueQuantity.isNotNull() &\n",
    "        (observations.valueQuantity.system == 'http://unitsofmeasure.org') &\n",
    "        (observations.valueQuantity.code == 'kg/m2'),\n",
    "        observations.valueQuantity.value > 30).otherwise(False).alias('bmi_over_30')\n",
    ")\n",
    "bmi_over_30 = bmi_over_30.groupBy(filtered_patients.id).agg(\n",
    "    functions.max(bmi_over_30.bmi_over_30).alias('bmi_over_30'))\n",
    "# bmi_over_30.toPandas()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Final aggregation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Join each of the risk factor dataframes together, and create an aggregate column `high_risk` which indicates if the patient exhibits any one of the risk factors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "risk_factors = covid_19.withColumn('patient_id', covid_19.id).alias('covid_19').join(\n",
    "    ckd.alias('ckd'),\n",
    "    functions.col('covid_19.id') == functions.col(\n",
    "        'ckd.id'), 'left_outer').join(\n",
    "    heart_disease.alias('heart_disease'),\n",
    "    functions.col('covid_19.id') == functions.col('heart_disease.id'), 'left_outer').join(\n",
    "    bmi_over_30.alias('bmi_over_30'),\n",
    "    functions.col('covid_19.id') == functions.col('bmi_over_30.id'), 'left_outer').withColumn(\n",
    "    'high_risk', ckd.ckd | heart_disease.heart_disease | bmi_over_30.bmi_over_30)\n",
    "# risk_factors.toPandas()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Count the patients, grouped by whether they have had a COVID-19 vaccination and whether they are high-risk."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "22/08/10 01:48:36 WARN package: Truncated the string representation of a plan since it was too large. This behavior can be adjusted by setting 'spark.sql.debug.maxToStringFields'.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "22/08/10 01:48:50 WARN DAGScheduler: Broadcasting large task binary with size 1306.4 KiB\n",
      "22/08/10 01:48:50 WARN DAGScheduler: Broadcasting large task binary with size 1287.9 KiB\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>covid_19</th>\n",
       "      <th>high_risk</th>\n",
       "      <th>count(id)</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>True</td>\n",
       "      <td>False</td>\n",
       "      <td>266</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>True</td>\n",
       "      <td>True</td>\n",
       "      <td>215</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "      <td>81</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>False</td>\n",
       "      <td>True</td>\n",
       "      <td>85</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   covid_19  high_risk  count(id)\n",
       "0      True      False        266\n",
       "1      True       True        215\n",
       "2     False      False         81\n",
       "3     False       True         85"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "aggregate = risk_factors.groupBy(covid_19.covid_19, risk_factors.high_risk).agg(\n",
    "    functions.countDistinct(covid_19.id))\n",
    "aggregate.toPandas()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## List the high risk unvaccinated patients\n",
    "\n",
    "Finally, we list out the patients that are high-risk and have not been vaccinated, along with the specific risk factors that were identified for each patient."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>given_name</th>\n",
       "      <th>family_name</th>\n",
       "      <th>telephone</th>\n",
       "      <th>ckd</th>\n",
       "      <th>heart_disease</th>\n",
       "      <th>bmi_over_30</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Jerald662</td>\n",
       "      <td>Ratke343</td>\n",
       "      <td>555-323-4093</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Vernon254</td>\n",
       "      <td>Harvey63</td>\n",
       "      <td>555-555-1286</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Lilly908</td>\n",
       "      <td>Volkman526</td>\n",
       "      <td>555-667-5077</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Renda520</td>\n",
       "      <td>Halvorson124</td>\n",
       "      <td>555-786-9558</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Ines699</td>\n",
       "      <td>Muller251</td>\n",
       "      <td>555-440-3532</td>\n",
       "      <td>False</td>\n",
       "      <td>True</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>80</th>\n",
       "      <td>Benito209</td>\n",
       "      <td>Windler79</td>\n",
       "      <td>555-902-5964</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>81</th>\n",
       "      <td>Tisha655</td>\n",
       "      <td>Kirlin939</td>\n",
       "      <td>555-463-7777</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>82</th>\n",
       "      <td>Lloyd546</td>\n",
       "      <td>Dickens475</td>\n",
       "      <td>555-475-1711</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>83</th>\n",
       "      <td>Domenic627</td>\n",
       "      <td>Wunsch504</td>\n",
       "      <td>555-973-4643</td>\n",
       "      <td>False</td>\n",
       "      <td>True</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>84</th>\n",
       "      <td>Tomi599</td>\n",
       "      <td>Hintz995</td>\n",
       "      <td>555-726-7932</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>85 rows × 6 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "    given_name   family_name     telephone    ckd  heart_disease  bmi_over_30\n",
       "0    Jerald662      Ratke343  555-323-4093  False          False         True\n",
       "1    Vernon254      Harvey63  555-555-1286  False          False         True\n",
       "2     Lilly908    Volkman526  555-667-5077  False          False         True\n",
       "3     Renda520  Halvorson124  555-786-9558  False          False         True\n",
       "4      Ines699     Muller251  555-440-3532  False           True        False\n",
       "..         ...           ...           ...    ...            ...          ...\n",
       "80   Benito209     Windler79  555-902-5964  False          False         True\n",
       "81    Tisha655     Kirlin939  555-463-7777  False          False         True\n",
       "82    Lloyd546    Dickens475  555-475-1711  False          False         True\n",
       "83  Domenic627     Wunsch504  555-973-4643  False           True         True\n",
       "84     Tomi599      Hintz995  555-726-7932  False          False         True\n",
       "\n",
       "[85 rows x 6 columns]"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "patients_risk_factors = patients.alias('patients').join(\n",
    "    risk_factors.alias('risk_factors'),\n",
    "    functions.col('patients.id') == functions.col('risk_factors.patient_id'),\n",
    "    'left_outer'\n",
    ").filter((risk_factors.high_risk & ~risk_factors.covid_19))\n",
    "patient_list = patients_risk_factors.select(\n",
    "    functions.element_at(functions.element_at(patients.name, 1)['given'], 1).alias('given_name'),\n",
    "    functions.element_at(patients.name, 1)['family'].alias('family_name'),\n",
    "    functions.explode_outer(patients.telecom).alias('telecom'),\n",
    "    risk_factors.ckd,\n",
    "    risk_factors.heart_disease,\n",
    "    risk_factors.bmi_over_30\n",
    ")\n",
    "patient_list = patient_list.filter(functions.col('telecom').system == 'phone').select(\n",
    "    functions.col('given_name'),\n",
    "    functions.col('family_name'),\n",
    "    functions.col('telecom').value.alias('telephone'),\n",
    "    functions.col('ckd'),\n",
    "    functions.col('heart_disease'),\n",
    "    functions.col('bmi_over_30')\n",
    ")\n",
    "patient_list.toPandas()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total execution time: 27.769\n"
     ]
    }
   ],
   "source": [
    "print(f\"Total execution time: {time() - start:.3f}\")\n"
   ]
  }
 ],
 "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.10.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}