POME is a graph-based representation-learning method for heterogeneous datasets that incorporates missingness structures into the computation of low-dimensional sample and variable embeddings. It is applicable to any tabular datasets consisting of both numeric- and categorical-type features, where missing data patterns are supposed to be taken into account.
POME is implemented as a Python package and is easily installable from PyPI by running
pip install pome-py
or locally from this repository by running
pip install .
POME expects input data to be given in the form of a pandas dataframe object, with rows representing variables/features and columns representing samples. Missing data needs to be encoded by a unique numerical value. Furthermore, POME expects one column storing datatypes of the respective variables. An example dataset could have the following structure, with e.g. value -99 encoding missing data:
| Sample1 | Sample2 | Sample3 | Type | |
|---|---|---|---|---|
| VariableA | 0 | 1 | -99 | cat |
| VariableB | 3.14 | -0.1 | 2.5 | numerical |
| VariableC | 0.3 | 1.2 | -99 | numerical |
| VariableD | 1 | 0 | 2 | cat |
POME's core functionality is integrated into its Embedder class, which handles input transformation, training and output generation. A typical such workflow looks as follows:
import pandas as pd
from pome import Embedder
if __name__ == "__main__":
# Load data and set parameters.
example_df = pd.read_csv("example.csv", index_col=0)
NA_ENCODING = -99.0
DIMENSION = 16
DEVICE = "cpu"
# Initialize embedding object with parameters.
embedder = Embedder(epochs=100,
na_encoding=NA_ENCODING,
embedding_dimension=DIMENSION,
device=DEVICE,
enable_imputation=True)
# Fit embedding object to dataset.
embedder.fit(example_df)
# Output stores low-dimensional embeddings for samples and variables.
sample_embeddings, variable_embeddings, _ , _ = embedder.get_embeddings()
print("Computed sample embeddings: \n", sample_embeddings)
imputed_df = embedder.impute_all(na_value=NA_ENCODING)
print("Imputed data: \n", imputed_df)Once fitted, POME can embed new samples that were not part of the training data using the frozen trained encoder, without any retraining. When new samples are the goal, inductive=True additionally replaces the fixed epoch budget by a cross-validated one that is selected to generalize to unseen samples:
from pome import Embedder, make_deterministic
embedder = Embedder(epochs=500, na_encoding=-99.0, embedding_dimension=16, inductive=True)
embedder.fit(train_df) # CV-tuned epoch count, stored in _optimal_epochs
new_embeddings = embedder.transform(test_df) # frozen encoder, no retrainingPOME's Embedder class allows for the specification of the following parameters:
embedding_dimension : int = 32: Specifies the number of dimensions of the sample & variable embeddings learned by POME.epochs : int = 500: Sets the number of epochs that POME is supposed to be trained.device : str = "cpu": Specifies whether to train on CPU ("cpu") or GPU ("cuda").
type_column : str = "type": Name of the column storing the variable types.na_encoding : float = -99.0: The float encoding value of missing data. ActualNaNentries are not supported and raise.discretization_type : str = "z": How continuous variables are binned -"z"(z-score bins) or"nonlinear"(signed-power bins).bins_per_continuous : int = 15: Number of bins per continuous variable. With"z"discretization, only 3, 7, 11 and 15 are supported; other values raise.
enable_imputation : bool = False: Set this to true if you want to use POME for imputation after training. It has to be set at construction time.
With inductive=True, fit() does not train for a fixed number of epochs. It first runs a cross-validation that holds out whole samples and picks the epoch count that best generalizes to unseen ones, using a label-free train-vs-held-out effective-rank gap as the stopping signal. The selected value is stored in _optimal_epochs, while epochs serves as the upper cap and is restored after fitting.
inductive : bool = False: Enables the CV-based epoch selection.
All of these are disabled by default and add no cost to the training loop when unused.
file_name : str = None,output_path : str = None: Filename prefix and output directory for saved artifacts.epoch_checkpoints : int = -1: Dump the whole fitted embedder via joblib every N epochs. Requiresfile_name.embedding_csv_epochs : int = -1: Write the current sample, variable and bin embeddings to CSV every N epochs.epoch_callback = None: Callable invoked asepoch_callback(epoch, autoencoder)after every epoch. It persists nothing itself - the caller decides what, if anything, to snapshot.
After initializing the Embedder object, the main functions for using POME are:
fit(X, y=None): Training POME on the given input dataframe, with the input format as specified above.get_embeddings(): Return computed embeddings in dataframe format. Output is a four-tuple of sample embeddings (position 0), variable embeddings (position 1), continuous-bin embeddings (position 2), and a mapping from sample name to its row in the attention matrix (position 3).transform(X): Embed new, unseen samples with the frozen trained encoder, without retraining.Xmust contain exactly the same variables as the training data plus the type column, and every sample must share at least one observed, known value with the training data — a sample whose values are all missing or unseen carries no signal and raises. Returns a(num_new_samples, embedding_dimension)dataframe.impute_all(na_value : float): Imputes all missing values specified byna_valuein the input dataset, and directly returns the imputed dataframe. Categorical values are imputed by scoring candidate values with the trained decoder, continuous ones by a regression head trained on the frozen embeddings. Requiresenable_imputation=True.
POME is released under the GPL-3.0 license, see LICENSE.