Skip to content
BranchesTags

Repository files navigation

python black Tests Build Docs Coverage License: MIT

UDM logo

UDM - Unified Data Module

The unified datamodule is a project designed to provide a unified way of accessing and bundling data from different UKBB applications/modalities to use across projects at the BIH.

The documentation is available at luisherrmann.github.io/udm.

Getting started

Installation from GitHub

You can install the package with pip directly from GitHub:

pip install "udm @ git+https://github.com/luisherrmann/udm.git" --extra-index-url https://download.pytorch.org/whl/cu116

Note that the --extra-index-url parameter is required to install the correct PyTorch dependencies.

If you are using poetry for dependency management, you can also install the package by adding these lines

pyproject.toml

[[tool.poetry.source]]
name = "torch"
url = "https://download.pytorch.org/whl/cu116"
secondary = true

[tool.poetry.dependencies]
udm = {git = "https://github.com/luisherrmann/udm.git"}

to the pyproject.toml of your project, and run either poetry install or poetry update in your project.

Installation from local source

If you want to install the UDM from a local source, activate the Python environment where you wish to install the UDM (e.g. if you are managing your environments with conda, use conda activate <ENV> to activate the environment) and install the UDM using

pip install -e <PATH>

It is highly recommended that you install in development mode (i.e. providing the -e editable flag), since it may be necessary to introduce additions or modifications for your own project, requiring the package to be editable.

If you are using poetry for dependency and environment management, include the following line

pyproject.toml

[tool.poetry.dependencies]
foo-package = { file = "relative/path/to/distribution" }

in your pyproject.toml and install or update the dependencies.

Post installation

If your project is already using Hydra configs, it is recommended that you use the Hydra configs in the config folder as a starting point for building your own Hydra configs by cloning the config folder into your respective config subfolder. For instance, in the OGM/Umbrella project, the config templates of the UDM were included through

cp -r config umbrella/config/run/datamodule

so the config of the UDM becomes a datamodule subconfig of the run config.

Overview: How it works

UDM Overview

The UDM relies on the following three classes

  1. DataPlugin
  2. GeneralDataset
  3. GeneralDatamodule

Additional helper classes are provided to enable

  1. Transforms (PluginTransform, DatasetTransform)
  2. Filtering (PluginRowFilter, PluginColFilter and more)

1. DataPlugin

The DataPlugin class is an abstract class that defines the general interface for interacting with different data modalities. The idea is that for each data modality or each way of interacting with the data, there should be a class extending from DataPlugin. For instance, there could be a GeneticPlugin, CovariatePlugin, ... class implementing the interface of DataPlugin and providing additional methods specific to the respective data modality. There is also a pre-existing generic class TabularPlugin which can be used for reading generic tabular data from a .feather file.

The __getitem__() method is called with an eid of the DataPlugin instance and an optional dictionary of feature selections and returns the respective data sample corresponding to that eid for the features provided.

An important aspect to keep in mind is that the eids between different applications of the UKBB may be different, e.g. the eids used by the data in the CovariatesPlugin might differ from those of the GeneticPlugin. Thus, there is a distinction between native eids and master eids of the samples controlled by a data plugin.

  1. native eids are ids that are native to the application used by the DataPlugin.
  2. master eids are the ids that are used to retrieve elements from the DataPlugin.

The user needs to ensure that all DataPlugins use the same master eids. If they use different native eids, one set of eids is taken as the master eids and the DataPlugins using different native eids need to be provided with a .csv mapping file to map native eids to master eids. This can be done using the eid_map_path option to specify the mapping file, as well as eid_map_from and eid_map_to to specify the columns containing the native and master eids, respectively.

If the user needs to retrieve metadata from the DataPlugin, this can be done using the get_metadata() method. This method should provide at least the following info:

  1. 'features': A list of the feature names of the data controlled by the DataPlugin.
  2. 'feature_types': A list of the data types used by the respective feature names.
  3. 'eids': A list of all the master eids controlled by this DataPlugin.
  4. 'tags': A list of strings that can be used to tag the DataPlugin for later identification and metadata aggregation across multiple DataPlugin instances. These are passed to the DataPlugin during initialization.

2. GeneralDataset

The GeneralDataset class extends from the PyTorch Dataset (torch.utils.data) and allows for the creation of a dataset from which to build a DataLoader. It expects a list of (optionally named) DataPlugin instances and a list of eids to use from those plugins. For instance, given

>>> plugins = {
...     "geno": GeneticPlugin(...)
...     "cov": CovariatesPlugin(...)
... }

containing data for hypothetical eids [1, 2, ..., 100], one could define datasets representing train, validation, and test splits through GeneralDataset instances:

>>> train_ds = GeneralDataset(plugins, eids=[0, ..., 80], ...)
>>> valid_ds = GeneralDataset(plugins, eids=[80, ..., 90], ...)
>>> test_ds = GeneralDataset(plugins, eids=[90, ..., 100], ...)

Contrary to the DataPlugin, the __getitem__() method of the GeneralDataset returns elements of the dataset by index, where the index is a number between 0 and the number of valid eids in use by the dataset. The return value is a dictionary with plugin names mapping to the sample obtained from the respective plugin at that index, as well as a list of master eids of the samples extracted. For example, for the valid_ds using the aforementioned plugins, retrieving valid_ds[0] would return

>>> valid_ds[0]
{
    'geno': {
        'genetic': torch.Tensor(...)
    },
    'cov': {
        'covariates': torch.Tensor(...)
    },
    'eids': [80]
}

The get_metadata() method returns metadata of the respective DataPlugins as a nested dictionary, i.e.

>>> valid_ds.get_metadata()
{
    'geno': {
        'features': ...,
        'feature_types': ...,
        ...
    }
    'cov': {
        'features': ...,
        'feature_types': ...,
        ...
    }
}

3. GeneralDatamodule

The GeneralDatamodule extends from the LightningDatamodule of PyTorch Lightning. It is initialized with a config of plugins, as well as a mapping of eids denoting the respective splits. For instance, a GeneralDatamodule using the CovariatePlugin and the GeneticPlugin might be initialized by something like

datamodule = GeneralDatamodule(
    plugins=[
        {
            'name': 'GeneticPlugin',
            '__init__': '__init__',
            ...
        },
        {
            'name': 'CovariatePlugin',
            '__init__': '__init__',
            ...
        }
    ],
    splits={
        'train': [0, ..., 80],
        'valid': [80, ..., 90],
        'test': [90, ..., 100]
    }
)

Ideally, the eids of all the plugins should align. However, when they do not, eids available from the DataPlugin will be obtained through an intersect or union operation of the respective eid sets (the behaviour is controlled by the combine_eids_as parameter).

Another thing to take into account is that the GeneralDatamodule supports passing of multiple splits. The respective datasets will all be set during the setup of the module.

The splits to be used for training, validation and testing can be reassigned at any point during the lifecycle of the UDM. However, you will have to rerun prepare_data() and setup() for the new splits to be used rather than the old ones.

4. Transforms

All DataPlugin and GeneralDataset subclasses can be instantiated with additional transformations that are applied on top of the data of the DataPlugin or GeneralDataset every time a single element (or batch thereof) is sampled from the respective instance. The DataPlugin needs to be provided with an instance of PluginTransform, while a GeneralDataset must be provided with a DatasetTransform at initialization.

These transforms can be understood as follows:

  • TensorTransform A function that takes a tensor as input and returns a tensor as output torch.Tensor -> torch.Tensor. You can use the get_transform() function from the transforms.tensor_transforms.factory package to instantiate transformations from Hydra configurations. By default, all transforms from torchvision.transforms and torch.nn.functional are also supported. For example:
from omegaconf import DictConfig
config = DictConfig({
    "name": "pad",
    "pad": (1, 1),
    "mode": "constant",
    "value": 0
})
transform = get_transform(config)
data = torch.tensor([1.0, 2.0, 3.0])
data_ = transform(data)
# data_ == torch.tensor([0.0, 1.0, 2.0, 3.0, 0.0]),
  • PluginTransform A mapping of component names to TensorTransform instances to be applied to the respective component. For instance, the EHRPlugin could be equipped with a single transformation {records: cnormalize}, where cnormalize = lambda x: normalize(x, mean=(0.0, 0.5), std=(0.278, 1.023)). Every time an element gets sampled from the DataPlugin, the tensor of component records is modified by the cnormalize function. For example:
data = {
    "x": torch.tensor([[1., 2., 3.]]),
    "y": torch.tensor([[2., 3., 4.]])
}
transform = PluginTransform(
    "x": ScaleTrafo(2.0)
)
data_ = transform(data)
# data = {
#     "x": torch.tensor([[2., 4., 6.]]),
#     "y": torch.tensor([[2., 3., 4.]])
# }
  • DatasetTransform A mapping of plugin names to PluginTransform instances to be applied to the respective DataPlugins when sampling. Every time an element gets sampled from the dataset, the tensor of every plugin is modified by the respective PluginTransform. Note that if a DataPlugin plugin1, controlled by the dataset ds, was initialized with a PluginTransform plugin1_trafo, then the transformation will be applied in any dataset controlling plugin1. If a dataset ds has its own plugin transform ds_plugin1_trafo, then sampling an original element x from ds will lead to transformations x[plugin1] -> plugin1_trafo -> ds_plugin1_trafo of the original output.

5. Filtering

All DataPlugin instances can be provided with an instance of PluginRowFilter or PluginColFilter to filter rows or columns of the data, respectively. The filter will be applied once to the entire dataset during the setup() of the respective plugin, and the dataset will only keep rows and columns which satisfy the filtering condition (assuming a non-empty row_filter or col_filter argument is provided). Subsequent sampling of filtered eids will cause a KeyError, as will the selection of filtered columns. The hierarchy of filters works as follows:

  • Filter A filter is essentially a function instantiated with certain parameters which can be called with a 2D tensor $X$ of shape $M \times N$, and which returns a boolean tensor mask $\mu$ of length $M$, where $\mu_i$ indicates whether row $X_{i,:}$ should be kept or not. Column filtering for a 2D tensor $X$ can be handled analogously by calling a filter with $X^T$. For example:
v_filter = AnyNan()
data = torch.tensor([[1., 0.], [np.nan, 2.]])
mask = v_filter(data)
# mask == torch.tensor([False, True])
data = data[mask, :] # torch.tensor([[np.nan, 2.]])
  • KeyFilter A special case of Filter to be applied to keys rather than values, i.e. to a sequence of hashable values of shape $M$ to mark rows of data to be discarded by eid or feature name. For example:
k_filter = IsIn([2, 4])
eids = [0, 1, 2, 3, 4, 5]
mask = k_filter(eids)
# mask == torch.tensor([False, False, True, False, True, False])
eids = eids[mask] # [2, 4]
  • ComposedFilter More complex filters can be built by aggregating existing filters into ComposedFilters. A ComposedFilter is always initialized with multiple filters and combines the results of the individual filters in some way to produce a single output mask. For example, we could mark any rows that have a torch.nan value for discarding by the DataPlugin as follows:
  c_filter = Not(AnyNan())
  data = torch.tensor([[1., 2.],[3., torch.nan]])
  mask = c_filter(data)
  # mask == torch.tensor([True, False])
  data = data[mask, :] # torch.tensor([[1., 2.]])
  • PluginRowFilter Instantiated with an optional KeyFilter and a mapping of component names to Filter instances to implement a plugin that selects rows to be filtered according to eids and the data from each component. Given $M$ keys, the filter returns a boolean mask of size $M$, obtained through logical anding of all individual masks. I.e. only rows satisfying all filters of the PluginRowFilter are marked for preservation. For example, consider the following case:
eids = [1, 2, 3]
data = {
    "x": torch.tensor([
        [1., 2.],
        [3., torch.nan],
        [5., 6.]]
    ),
    "y": torch.tensor([
        [1., 4.],
        [2., 5.],
        [3., 6.]]
    )
}
row_filter = PluginRowFilter(
    key_filter = IsIn([1, 2]),
    val_filters = {
        "x": Not(AnyNan()),
    }
)
mask = row_filter(eids, data)
# mask == torch.tensor([True, False, False])
# == torch.tensor([True, True, False])
# && torch.tensor([True, False, True])
#
# Corresponds to:
# data == {
#     "x": torch.tensor([[1., 2.]])
#     "y": torch.tensor([[1., 4.]])
# }
  • PluginColFilter Instantiated with an optional mapping of component name to KeyFilter, and an optional mapping of component name to Filter. The key filters and value filters are applied to each component of the provided data separately and the filter returns a dictionary mapping components to masks to be applied to each component separately. Only columns satisfying their respective KeyFilter and Filter are marked for preservation. Consider the following example:
features = {
    "categorical": ["sex", "eye_color"],
    "continuous": ["height", "weight", "age"]
}
data = {
    "categorical": [
        [1, 0],
        [0, 3],
        [1, 1]
    ],
    "continuous": [
        [180, 89, 50],
        [160, 62, 42],
        [178, torch.nan, 30]
    ]
}
col_filter = PluginColFilter(
    key_filters = {
        "categorical": Isin(["sex"])
        "continuous": Not(IsIn(["age"]))
    },
    val_filters = {
        "continuous": Not(AnyNan())
    }
)
masks = col_filters(data)
# masks == {
#     "categorical": torch.tensor([True, False]),
#     "continuous": torch.tensor([True, False, False])
# }
# Corresponds to:
# data == {
#     "categorical": torch.tensor([[1], [0], [1]]),
#     "continuous": torch.tensor([[180], [160], [178]])
# }

NOTE: Some of the plugins, such as the H5adPlugin, had their own systems for filtering columns in place (e.g. usecols and dropcols arguments). These individual filtering systems are still in place for backward compatibility reasons, but will eventually be removed.

Extending the UDM

1. Including the DataPlugin class

You can extend the UDM by providing new classes that extend from DataPlugin (or from a pre-existing subclass of DataPlugin), and putting them in a submodule in the udm/plugins package. For instance, say you have implemented a class

proteomics_plugin.py

class ProteomicsPlugin(DataPlugin):

    def __init__(self, src_path, memmap=False, **kwargs):
        ...

    @classmethod
    def from_db(self, db_user, db_pass, db_table):
        ...

    ...

which extends from DataPlugin, enables the use of proteomics data, and has its code in a file called proteomics_plugin.py. Following Python convention, the source code for each DataPlugin subclass should be in a file containing no other classes than the subclass itself, and the class name should be written in camel case (e.g. ProteomicsPlugin), while the file name should be written in snake case (e.g. proteomics_plugin). NOTE: Please make sure the name of the plugin does not match the name of any other pre-existing DataPlugin in the project!

To include this DataPlugin in the repository, you could put the file proteomics_plugin.py in a subdirectory of plugins like so:

.
├── config
├── plugins
│   ├── genetics
│   │   └── ...
│   ├── proteomics_plugin.py
│   └── ...
└── ...

Or better yet:

.
├── config
├── plugins
│   ├── genetics
│   │   └── ...
│   ├── proteomics
│   │   ├── __init__.py
│   │   └── proteomics_plugin.py
│   └── ...
└── ...

Arbitrary levels of nesting are possible, as the GeneralDatamodule will automatically discover all subclasses of DataPlugin within the plugins package. However, we encourage you to use a low amount of nesting to keep a clean directory structure.

2. Adding config files

In addition to adding the source code for the plugin, you should also add a default config file that can be used by others to create datamodule configurations using your DataPlugin subclass. The config file should be a .yaml file containing the fields

default.yaml

name: ProteomicsPlugin  # mandatory
__init__: __init__      # mandatory
src_path: <PATH>
memmap: false

i.e. the field name gives the name of the class, the field __init__ gives the method by which to initialize an instance of the class, and the remaining fields give default values for the arguments to be passed to the init function of the class. The first two values are mandatory, because they are required by the DataModule to know which DataPlugins to prepare, and what method to use for the initialization. By default, the initialization method will be the regular __init__ method, but in some cases it may be useful to define different __init__ methods for different ways of initializing the DataPlugin for interfacing with the respective data.

For example, let's say the proteomics data to be accessed through the ProteomicsPlugin can also be retrieved from a database. Then, a good pattern would be to enable the ProteomicsPlugin to be initialized through another method from_db(), where database connection arguments are provided. It would be recommended to have a separate default configuration for this case, e.g.:

from_db.yaml

name: ProteomicsPlugin  # mandatory
__init__: from_db       # mandatory
db_user: sher
db_pass: locked
db_table: ukbb_processed

Putting everything together, these two config files should be placed in a subdirectory of configs, preferably mirroring the directory structure of udm/plugins, like so:

.
├── config
│   ├── genetics
│   │   └── ...
│   ├── proteomics
│   │   ├── default.yaml
│   │   └── from_db.yaml
├── udm
│   ├── plugins
│   │   ├── genetics
│   │   │   └── ...
│   │   ├── proteomics
│   │   │   └── proteomics_plugin.py
│   │   └── ...
│   └── ...
└── ...

And that's it, you can now create your own DataModule configurations.

3. Adding tests

In order to ensure that your DataPlugin subclass works properly, it is highly encouraged that you write unit tests to check that your plugin works as intended on small test datasets.

To do so, we recommend you use the default Python practice of mirroring the main source package structure. For example, a unit test for the ProteomicsPlugin class ProteomicsPluginTest would be included in the project as follows:

.
├── config
│   ├── genetics
│   │   └── ...
│   ├── proteomics
│   │   ├── default.yaml
│   │   └── from_db.yaml
├── udm
│   ├── plugins
│   │   ├── genetics
│   │   │   └── ...
│   │   ├── proteomics
│   │   │   └── proteomics_plugin.py
│   │   └── ...
│   └── ...
├── test
│   ├── plugins
│   │   ├── genetics
│   │   │   └── ...
│   │   ├── proteomics
│   │   │   └── proteomics_plugin_test.py
│   │   └── ...
└── ...

We also recommend writing integration tests by extending the tests in the datamodule_test.py to include test scenarios where the GeneralDatamodule is initialized with your custom DataPlugin.

Small datasets for testing may be included in the repository through the res/ directory. However, to ensure compliance with data protection guidelines YOU MAY NOT INCLUDE DATASETS CONTAINING ANY ACTUAL UKBB DATA if you wish to push the debug dataset to a remote repository later on!

4. Contributing to the repo

Before you push anything to the repo, please make sure you have installed the pre-commit hooks by running

pre-commit install

so your code changes can be cleaned up beforehand.

To have your changes added to the main UDM project repo, apply for a collaborator status on the main repo and directly open a pull request.

License

This project is licensed under the MIT License. See LICENSE for details.

About

Unified Data Module for multimodal training applications across BIH projects

luisherrmann.github.io/udm/

Topics

clinical-datadata-preprocessingdata-transformationdatamodulesmultimodal-learningpytorchpytorch-lightninguk-biobankukbb

Resources

Readme
MIT license
Activity

Stars

2 stars

Watchers

0 watching

Forks

0 forks
Report repository

Releases

Packages

Contributors

Languages