# -*- coding: utf-8 -*-
"""Clinical Knowledge Graph."""
import pathlib
import tarfile
from pathlib import Path
from typing import Iterable, Optional
from urllib.request import urlretrieve
import click
import pandas as pd
from docdata import parse_docdata
from more_click import verbose_option
from .base import TabbedDataset
from ..typing import TorchRandomHint
__all__ = [
'CKG',
]
URL = 'https://md-datasets-public-files-prod.s3.eu-west-1.amazonaws.com/d1e8d3df-2342-468a-91a9-97a981a479ad'
COLUMNS = ['START_ID', 'TYPE', 'END_ID']
[docs]@parse_docdata
class CKG(TabbedDataset):
"""The Clinical Knowledge Graph (CKG) dataset from [santos2020]_.
---
name: Clinical Knowledge Graph
citation:
author: Santos
year: 2020
link: https://doi.org/10.1101/2020.05.09.084897
github: MannLabs/CKG
single: true
statistics:
entities: 7617419
relations: 11
triples: 26691525
"""
def __init__(
self,
create_inverse_triples: bool = False,
random_state: TorchRandomHint = 0,
**kwargs,
):
"""Initialize the `CKG <https://github.com/MannLabs/CKG>`_ dataset from [santos2020]_.
:param create_inverse_triples: Should inverse triples be created? Defaults to false.
:param random_state: The random seed to use in splitting the dataset. Defaults to 0.
:param kwargs: keyword arguments passed to :class:`pykeen.datasets.base.TabbedDataset`.
"""
super().__init__(
create_inverse_triples=create_inverse_triples,
random_state=random_state,
**kwargs,
)
self.preloaded_path = self.cache_root.joinpath('preloaded.tsv.gz')
def _get_path(self) -> Optional[pathlib.Path]:
return self.preloaded_path
def _get_df(self) -> pd.DataFrame:
if self.preloaded_path.exists():
return pd.read_csv(self.preloaded_path, sep='\t', dtype=str)
df = pd.concat(self._iterate_dataframes())
df.to_csv(self.preloaded_path, sep='\t', index=False)
return df
def _iterate_dataframes(self) -> Iterable[pd.DataFrame]:
archive_path = self.cache_root / 'data.tar.gz'
if not archive_path.exists():
urlretrieve(URL, archive_path) # noqa:S310
with tarfile.TarFile.open(archive_path) as tar_file:
if tar_file is None:
raise ValueError
for tarinfo in tar_file:
if not tarinfo.name.startswith('data/imports/') or not tarinfo.name.endswith('.tsv'):
continue
path = Path(tarinfo.name)
if path.name.startswith('.'):
continue
_inner_file = tar_file.extractfile(tarinfo)
if _inner_file is None:
raise ValueError(f'Unable to open inner file: {tarinfo}')
with _inner_file as file:
df = pd.read_csv(file, usecols=COLUMNS, sep='\t', dtype=str)
df = df[COLUMNS]
yield df
@click.command()
@verbose_option
def _main():
d = CKG()
d.summarize()
if __name__ == '__main__':
_main()