uci_density_estimation.py

# Copyright © 2021 Dar Gilboa, Ari Pakman and Thibault Vatter
# This file is part of the mdma library and licensed under the terms of the MIT license.
# For a copy, see the LICENSE file in the root directory.

from mdma import fit
import torch as t
import numpy as np

from experiments.UCI.gas import GAS
from experiments.UCI.hepmass import HEPMASS
from experiments.UCI.miniboone import MINIBOONE
from experiments.UCI.power import POWER


def fit_uci():
  h = fit.get_default_h()
  data = load_dataset(h)
  outs = fit.fit_mdma(h, data)


def load_dataset(h):
  if h.dataset == 'gas':
    dataset = GAS(h.data_dir + '/gas/ethylene_CO.pickle')
  elif h.dataset == 'hepmass':
    dataset = HEPMASS(h.data_dir + '/hepmass')
  elif h.dataset == 'miniboone':
    dataset = MINIBOONE(h.data_dir + '/miniboone/data.npy')
  elif h.dataset == 'power':
    dataset = POWER(h.data_dir + '/power/data.npy')
  else:
    raise RuntimeError()

  if h.missing_data_pct > 0:
    # create missing data masks
    mask = np.random.rand(*dataset.trn.x.shape) > h.missing_data_pct
    data_and_mask = np.array([dataset.trn.x, mask]).swapaxes(0, 1)
    dataset_train = t.utils.data.TensorDataset(t.tensor(data_and_mask).float())
    val_mask = np.random.rand(*dataset.val.x.shape) > h.missing_data_pct
    val_data_and_mask = np.array([dataset.val.x, val_mask]).swapaxes(0, 1)
    dataset_valid = t.utils.data.TensorDataset(
        t.tensor(val_data_and_mask).float())
  else:
    dataset_train = t.utils.data.TensorDataset(
        t.tensor(np.expand_dims(dataset.trn.x, 1)).float())
    dataset_valid = t.utils.data.TensorDataset(
        t.tensor(np.expand_dims(dataset.val.x, 1)).float())

  data_loader_train = t.utils.data.DataLoader(dataset_train,
                                              batch_size=h.batch_size,
                                              shuffle=True)

  data_loader_valid = t.utils.data.DataLoader(dataset_valid,
                                              batch_size=h.batch_size,
                                              shuffle=False)

  dataset_test = t.utils.data.TensorDataset(t.tensor(dataset.tst.x).float())
  data_loader_test = t.utils.data.DataLoader(dataset_test,
                                             batch_size=h.batch_size,
                                             shuffle=False)

  h.d = dataset.n_dims
  h.M = len(dataset_train)
  h.M_val = len(dataset_valid)

  return [data_loader_train, data_loader_valid, data_loader_test]


if __name__ == '__main__':
  fit_uci()