models/dhs_ridge_resnet.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "toc": true
   },
   "source": [
    "<h1>Table of Contents<span class=\"tocSkip\"></span></h1>\n",
    "<div class=\"toc\"><ul class=\"toc-item\"><li><span><a href=\"#Imports-and-Constants\" data-toc-modified-id=\"Imports-and-Constants-1\"><span class=\"toc-item-num\">1&nbsp;&nbsp;</span>Imports and Constants</a></span></li><li><span><a href=\"#Load-Saved-Data\" data-toc-modified-id=\"Load-Saved-Data-2\"><span class=\"toc-item-num\">2&nbsp;&nbsp;</span>Load Saved Data</a></span><ul class=\"toc-item\"><li><span><a href=\"#Load-labels\" data-toc-modified-id=\"Load-labels-2.1\"><span class=\"toc-item-num\">2.1&nbsp;&nbsp;</span>Load labels</a></span></li><li><span><a href=\"#Load-loc_dict\" data-toc-modified-id=\"Load-loc_dict-2.2\"><span class=\"toc-item-num\">2.2&nbsp;&nbsp;</span>Load <code>loc_dict</code></a></span></li><li><span><a href=\"#country_indices-and-country_labels\" data-toc-modified-id=\"country_indices-and-country_labels-2.3\"><span class=\"toc-item-num\">2.3&nbsp;&nbsp;</span><code>country_indices</code> and <code>country_labels</code></a></span></li><li><span><a href=\"#Get-folds\" data-toc-modified-id=\"Get-folds-2.4\"><span class=\"toc-item-num\">2.4&nbsp;&nbsp;</span>Get folds</a></span></li></ul></li><li><span><a href=\"#General-Training-Code\" data-toc-modified-id=\"General-Training-Code-3\"><span class=\"toc-item-num\">3&nbsp;&nbsp;</span>General Training Code</a></span></li><li><span><a href=\"#OOC\" data-toc-modified-id=\"OOC-4\"><span class=\"toc-item-num\">4&nbsp;&nbsp;</span>OOC</a></span><ul class=\"toc-item\"><li><span><a href=\"#Concatenated-MS,-NL-Features\" data-toc-modified-id=\"Concatenated-MS,-NL-Features-4.1\"><span class=\"toc-item-num\">4.1&nbsp;&nbsp;</span>Concatenated MS, NL Features</a></span></li><li><span><a href=\"#Transfer-features\" data-toc-modified-id=\"Transfer-features-4.2\"><span class=\"toc-item-num\">4.2&nbsp;&nbsp;</span>Transfer features</a></span></li></ul></li><li><span><a href=\"#OOC-keep_frac\" data-toc-modified-id=\"OOC-keep_frac-5\"><span class=\"toc-item-num\">5&nbsp;&nbsp;</span>OOC keep_frac</a></span><ul class=\"toc-item\"><li><span><a href=\"#Concatenated-MS,-NL-Features\" data-toc-modified-id=\"Concatenated-MS,-NL-Features-5.1\"><span class=\"toc-item-num\">5.1&nbsp;&nbsp;</span>Concatenated MS, NL Features</a></span></li></ul></li><li><span><a href=\"#Incountry\" data-toc-modified-id=\"Incountry-6\"><span class=\"toc-item-num\">6&nbsp;&nbsp;</span>Incountry</a></span><ul class=\"toc-item\"><li><span><a href=\"#Concatenated-MS,-NL-Features\" data-toc-modified-id=\"Concatenated-MS,-NL-Features-6.1\"><span class=\"toc-item-num\">6.1&nbsp;&nbsp;</span>Concatenated MS, NL Features</a></span></li><li><span><a href=\"#Transfer-Features\" data-toc-modified-id=\"Transfer-Features-6.2\"><span class=\"toc-item-num\">6.2&nbsp;&nbsp;</span>Transfer Features</a></span></li></ul></li></ul></div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Prerequisites:\n",
    "1. Run the first couple of sections from `models/baselines_dhs.ipynb` to create `data/dhs_image_hists.npz`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Imports and Constants"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import defaultdict\n",
    "from glob import glob\n",
    "import pickle\n",
    "import os\n",
    "import re\n",
    "import sys\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "sys.path.append('../')\n",
    "from batchers import batcher, dataset_constants\n",
    "from models.linear_model import ridge_cv, train_linear_logo\n",
    "from utils.general import load_npz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "FOLDS = ['A', 'B', 'C', 'D', 'E']\n",
    "SPLITS = ['train', 'val', 'test']\n",
    "DATASET_NAME = '2009-17'\n",
    "LOGS_ROOT_DIR = '../logs/'\n",
    "COUNTRIES = dataset_constants.DHS_COUNTRIES\n",
    "\n",
    "KEEPS = [0.05, 0.1, 0.25, 0.5]\n",
    "SEEDS = [123, 456, 789]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MODEL_DIRS = {\n",
    "    'resnet_ms_A': '2009-17A_18preact_ms_samescaled_b64_fc01_conv01_lr0001',\n",
    "    'resnet_ms_B': '2009-17B_18preact_ms_samescaled_b64_fc001_conv001_lr0001',\n",
    "    'resnet_ms_C': '2009-17C_18preact_ms_samescaled_b64_fc001_conv001_lr001',\n",
    "    'resnet_ms_D': '2009-17D_18preact_ms_samescaled_b64_fc001_conv001_lr01',\n",
    "    'resnet_ms_E': '2009-17E_18preact_ms_samescaled_b64_fc01_conv01_lr001',\n",
    "    'resnet_nl_A': '2009-17A_18preact_nl_random_b64_fc1.0_conv1.0_lr0001',\n",
    "    'resnet_nl_B': '2009-17B_18preact_nl_random_b64_fc1.0_conv1.0_lr0001',\n",
    "    'resnet_nl_C': '2009-17C_18preact_nl_random_b64_fc1.0_conv1.0_lr0001',\n",
    "    'resnet_nl_D': '2009-17D_18preact_nl_random_b64_fc1.0_conv1.0_lr01',\n",
    "    'resnet_nl_E': '2009-17E_18preact_nl_random_b64_fc1.0_conv1.0_lr0001',\n",
    "    'resnet_rgb_A': '2009-17A_18preact_rgb_same_b64_fc001_conv001_lr01',\n",
    "    'resnet_rgb_B': '2009-17B_18preact_rgb_same_b64_fc001_conv001_lr0001',\n",
    "    'resnet_rgb_C': '2009-17C_18preact_rgb_same_b64_fc001_conv001_lr0001',\n",
    "    'resnet_rgb_D': '2009-17D_18preact_rgb_same_b64_fc1.0_conv1.0_lr01',\n",
    "    'resnet_rgb_E': '2009-17E_18preact_rgb_same_b64_fc001_conv001_lr0001',\n",
    "\n",
    "    'resnet_rgb_transfer': 'transfer_2009-17nl_nlcenter_18preact_rgb_b64_fc001_conv001_lr0001',\n",
    "    'resnet_ms_transfer': 'transfer_2009-17nl_nlcenter_18preact_ms_b64_fc001_conv001_lr0001',\n",
    "\n",
    "    'incountry_resnet_ms_A': 'DHSIncountry/incountryA_18preact_ms_samescaled_b64_fc01_conv01_lr001',\n",
    "    'incountry_resnet_ms_B': 'DHSIncountry/incountryB_18preact_ms_samescaled_b64_fc1_conv1_lr001',\n",
    "    'incountry_resnet_ms_C': 'DHSIncountry/incountryC_18preact_ms_samescaled_b64_fc1.0_conv1.0_lr0001',\n",
    "    'incountry_resnet_ms_D': 'DHSIncountry/incountryD_18preact_ms_samescaled_b64_fc001_conv001_lr0001',\n",
    "    'incountry_resnet_ms_E': 'DHSIncountry/incountryE_18preact_ms_samescaled_b64_fc001_conv001_lr0001',\n",
    "    'incountry_resnet_nl_A': 'DHSIncountry/incountryA_18preact_nl_random_b64_fc1.0_conv1.0_lr0001',\n",
    "    'incountry_resnet_nl_B': 'DHSIncountry/incountryB_18preact_nl_random_b64_fc1.0_conv1.0_lr0001',\n",
    "    'incountry_resnet_nl_C': 'DHSIncountry/incountryC_18preact_nl_random_b64_fc1.0_conv1.0_lr0001',\n",
    "    'incountry_resnet_nl_D': 'DHSIncountry/incountryD_18preact_nl_random_b64_fc1.0_conv1.0_lr0001',\n",
    "    'incountry_resnet_nl_E': 'DHSIncountry/incountryE_18preact_nl_random_b64_fc01_conv01_lr001',\n",
    "}\n",
    "\n",
    "# add in keep_frac/seed dirs\n",
    "KEEP_MODEL_DIRS = sorted(glob(os.path.join(LOGS_ROOT_DIR, '*keep*seed*')))\n",
    "\n",
    "for model_dir in KEEP_MODEL_DIRS:\n",
    "    model_dir = os.path.basename(model_dir)\n",
    "    regex = r'2009-17(\\w)_18preact_(\\w+)_keep(.+)_seed(\\w+)_b64.+'\n",
    "    m = re.match(regex, model_dir)\n",
    "    fold, bands_name, keep, seed = m.groups()\n",
    "    model_name = 'resnet_{b}_{f}_keep{k}_seed{s}'.format(\n",
    "        b=bands_name,\n",
    "        f=fold,\n",
    "        k=keep,\n",
    "        s=seed)\n",
    "    MODEL_DIRS[model_name] = model_dir"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load Saved Data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "file_path = '../data/dhs_image_hists.npz'\n",
    "npz = load_npz(file_path)\n",
    "\n",
    "labels = npz['labels']\n",
    "locs = npz['locs']\n",
    "years = npz['years']\n",
    "\n",
    "num_examples = len(labels)\n",
    "assert np.all(np.asarray([len(labels), len(locs), len(years)]) == num_examples)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load `loc_dict`\n",
    "\n",
    "`loc_dict` has the format:\n",
    "```python\n",
    "{\n",
    "    (lat, lon): {\n",
    "        'cluster': 1,\n",
    "        'country': 'malawi',\n",
    "        'country_year': 'malawi_2012',  # surveyID\n",
    "        'households': 25,\n",
    "        'urban': False,\n",
    "        'wealth': -0.513607621192932,\n",
    "        'wealthpooled': -0.732255101203918,\n",
    "        'year': 2012\n",
    "    }, ...\n",
    "}\n",
    "```\n",
    "\n",
    "NOTE: `year` and `country_year` might differ in the year. `country_year` is the survey ID, which says which year the survey started. However, sometimes the DHS surveys cross the year-boundary, in which case `country_year` will remain the same but `year` will be the next year."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loc_dict_path = '../data/dhs_loc_dict.pkl'\n",
    "with open(loc_dict_path, 'rb') as f:\n",
    "    loc_dict = pickle.load(f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## `country_indices` and `country_labels`\n",
    "\n",
    "`country_indices` is a dictionary that maps a country name to a sorted `np.array` of its indices\n",
    "```python\n",
    "{ 'malawi': np.array([ 8530,  8531,  8532, ..., 10484, 10485, 10486]), ... }\n",
    "```\n",
    "\n",
    "`country_labels` is a `np.array` that shows which country each example belongs to\n",
    "```python\n",
    "np.array([0, 0, 0, 0, ..., 22, 22, 22])\n",
    "```\n",
    "where countries are indexed by their position in `dataset_constants.DHS_COUNTRIES`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "country_indices = defaultdict(list)  # country => np.array of indices\n",
    "country_labels = np.zeros(num_examples, dtype=np.int32)  # np.array of country labels\n",
    "\n",
    "for i, loc in enumerate(locs):\n",
    "    country = loc_dict[tuple(loc)]['country']\n",
    "    country_indices[country].append(i)\n",
    "\n",
    "for i, country in enumerate(COUNTRIES):\n",
    "    country_indices[country] = np.asarray(country_indices[country])\n",
    "    indices = country_indices[country]\n",
    "    country_labels[indices] = i"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get folds\n",
    "\n",
    "```python\n",
    "*_folds = {\n",
    "    'A': { # (ooc) list of country names, (incountry) np.array of indices\n",
    "           'train': [...],\n",
    "           'val': [...],\n",
    "           'test': [...],\n",
    "    }, ...\n",
    "}\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ooc_folds = {\n",
    "    f: {split: dataset_constants.SURVEY_NAMES[f'2009-17{f}'][split] for split in SPLITS}\n",
    "    for f in FOLDS\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('../data/dhs_incountry_folds.pkl', 'rb') as f:\n",
    "    incountry_folds = pickle.load(f)\n",
    "\n",
    "incountry_group_labels = np.zeros(num_examples, dtype=np.int32)\n",
    "for i, fold in enumerate(FOLDS):\n",
    "    test_indices = incountry_folds[fold]['test']\n",
    "    incountry_group_labels[test_indices] = i"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# General Training Code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_indices_for_countries(countries):\n",
    "    indices =  np.sort(np.concatenate([\n",
    "        country_indices[country] for country in countries\n",
    "    ]))\n",
    "    return indices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def countries_to_nums(countries):\n",
    "    '''\n",
    "    Args\n",
    "    - countries: list or set of str, names of countries\n",
    "\n",
    "    Returns: nums, list of int\n",
    "    '''\n",
    "    nums = []\n",
    "    for c in countries:\n",
    "        num = COUNTRIES.index(c)\n",
    "        nums.append(num)\n",
    "    return nums"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# OOC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ridgecv_ooc_wrapper(model_name, savedir):\n",
    "    features_dict = {}\n",
    "    for f in FOLDS:\n",
    "        model_fold_name = f'{model_name}_{f}'\n",
    "        model_dir = MODEL_DIRS[model_fold_name]\n",
    "        npz = load_npz(os.path.join(LOGS_ROOT_DIR, 'DHS_OOC', model_dir, 'features.npz'),\n",
    "                       check={'labels': labels})\n",
    "        features = npz['features']\n",
    "        for country in dataset_constants.SURVEY_NAMES[f'2009-17{f}']['test']:\n",
    "            features_dict[country] = features\n",
    "\n",
    "    ridge_cv(\n",
    "        features=features_dict,\n",
    "        labels=labels,\n",
    "        group_labels=country_labels,\n",
    "        group_names=COUNTRIES,\n",
    "        do_plot=True,\n",
    "        savedir=savedir,\n",
    "        save_weights=True,\n",
    "        save_dict=dict(locs=locs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = 'resnet_ms'\n",
    "savedir = '../logs/dhs_resnet/ms/'\n",
    "ridgecv_ooc_wrapper(model_name, savedir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = 'resnet_nl'\n",
    "savedir = '../logs/dhs_resnet/nl/'\n",
    "ridgecv_ooc_wrapper(model_name, savedir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Concatenated MS, NL Features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ridgecv_ooc_concat_wrapper(model_names, savedir):\n",
    "    features_dict = {}\n",
    "    for f in FOLDS:\n",
    "        concat_features = []  # list of np.array, each shape [N, D_i]\n",
    "        for model_name in model_names:\n",
    "            model_dir = MODEL_DIRS[f'{model_name}_{f}']\n",
    "            npz = load_npz(os.path.join(LOGS_ROOT_DIR, 'DHS_OOC', model_dir, 'features.npz'),\n",
    "                           check={'labels': labels})\n",
    "            concat_features.append(npz['features'])\n",
    "        concat_features = np.concatenate(concat_features, axis=1) # shape [N, D_1 + ... + D_m]\n",
    "        for country in dataset_constants.SURVEY_NAMES[f'2009-17{f}']['test']:\n",
    "            features_dict[country] = concat_features\n",
    "\n",
    "    ridge_cv(\n",
    "        features=features_dict,\n",
    "        labels=labels,\n",
    "        group_labels=country_labels,\n",
    "        group_names=COUNTRIES,\n",
    "        do_plot=True,\n",
    "        savedir=savedir,\n",
    "        save_weights=True,\n",
    "        save_dict=dict(locs=locs),\n",
    "        verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "model_names = ['resnet_ms', 'resnet_nl']\n",
    "savedir = os.path.join(LOGS_ROOT_DIR, 'dhs_resnet', 'msnl_concat')\n",
    "ridgecv_ooc_concat_wrapper(model_names, savedir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Transfer features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ridgecv_ooc_transfer_wrapper(model_name, savedir):\n",
    "    model_dir = MODEL_DIRS[model_name]\n",
    "    features = load_npz(os.path.join(LOGS_ROOT_DIR, model_dir, 'features.npz'),\n",
    "                        check={'labels': labels})['features']\n",
    "    ridge_cv(\n",
    "        features=features,\n",
    "        labels=labels,\n",
    "        group_labels=country_labels,\n",
    "        group_names=COUNTRIES,\n",
    "        savedir=savedir,\n",
    "        save_weights=False,\n",
    "        save_dict=dict(locs=locs),\n",
    "        verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = 'resnet_rgb_transfer'\n",
    "savedir = '../logs/dhs_resnet/rgb_transfer/'\n",
    "ridgecv_ooc_transfer_wrapper(model_name, savedir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = 'resnet_ms_transfer'\n",
    "savedir = '../logs/dhs_resnet/ms_transfer/'\n",
    "ridgecv_ooc_transfer_wrapper(model_name, savedir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# OOC keep_frac"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_split_idxs(dataset: str):\n",
    "    train_tfrecord_paths = batcher.get_tfrecord_paths(dataset, 'train')\n",
    "    val_tfrecord_paths = batcher.get_tfrecord_paths(dataset, 'val')\n",
    "    test_tfrecord_paths = batcher.get_tfrecord_paths(dataset, 'test')\n",
    "    all_tfrecord_paths = batcher.get_tfrecord_paths(dataset, 'all')\n",
    "\n",
    "    path_to_idx = {path: idx for idx, path in enumerate(all_tfrecord_paths)}\n",
    "    train_idxs = [path_to_idx[path] for path in train_tfrecord_paths]\n",
    "    val_idxs = [path_to_idx[path] for path in val_tfrecord_paths]\n",
    "    test_idxs = [path_to_idx[path] for path in test_tfrecord_paths]\n",
    "    return train_idxs, val_idxs, test_idxs\n",
    "\n",
    "def get_keep_indices(keep_frac: float, seed: int, dataset: str, test_country: str):\n",
    "    '''\n",
    "    Args\n",
    "    - keep_frac: float, fraction of non-test-country data to use for training\n",
    "    - seed: int\n",
    "    - dataset: str, one of the keys of dataset_constants.SIZES[dataset]\n",
    "    - test_country: str\n",
    "    '''\n",
    "    train_idxs, val_idxs, test_idxs = get_split_idxs(dataset)\n",
    "    test_country_idxs = get_indices_for_countries([test_country]).tolist()\n",
    "    test_other_idxs = sorted(set(test_idxs) - set(test_country_idxs))  # sort for determinism\n",
    "    \n",
    "    num_train = int(dataset_constants.SIZES[dataset]['train'] * keep_frac)\n",
    "    num_val = int(dataset_constants.SIZES[dataset]['val'] * keep_frac)\n",
    "    num_test = int(len(test_other_idxs) * keep_frac)\n",
    "\n",
    "    np.random.seed(seed)\n",
    "    train_idxs = np.random.choice(train_idxs, size=num_train, replace=False)\n",
    "    val_idxs = np.random.choice(val_idxs, size=num_val, replace=False)\n",
    "    test_other_idxs = np.random.choice(test_other_idxs, size=num_test, replace=False)\n",
    "\n",
    "    return np.sort(np.concatenate([train_idxs, val_idxs, test_other_idxs, test_country_idxs]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_ridgecv_keep(model_name, labels, locs, country_labels, folds, keep_frac, seed, savedir):\n",
    "    '''\n",
    "    For every country C (the test country):\n",
    "      1. uses leave-one-country-out CV on all other countries\n",
    "         to tune ridge model alpha parameter\n",
    "      2. using best alpha, trains ridge model on all countries except C\n",
    "      3. runs trained ridge model on C\n",
    "    Saves predictions for each country on test.\n",
    "\n",
    "    Args\n",
    "    - model_name: str, format 'resnet_{bands}', e.g. 'resnet_ms'\n",
    "    - labels: np.array, shape [num_examples]\n",
    "    - locs: np.array, shape [num_examples, 2]\n",
    "    - country_labels: np.array, shape [num_examples]\n",
    "    - folds: dict, fold (str) => dict\n",
    "    - keep_frac: float, fraction of non-test-country data to use for training\n",
    "    - seed: int\n",
    "    - savedir: str\n",
    "    '''\n",
    "    test_preds = np.zeros(num_examples, dtype=np.float32)\n",
    "    all_countries_set = set(COUNTRIES)\n",
    "\n",
    "    for f in FOLDS:\n",
    "        print('Fold:', f)\n",
    "        model_fold_name = f'{model_name}_{f}'\n",
    "        model_dir = MODEL_DIRS[model_fold_name]\n",
    "        model_dir = model_dir.replace('b64', f'keep{keep_frac}_seed{seed}_b64')\n",
    "\n",
    "        npz = load_npz(\n",
    "            os.path.join(LOGS_ROOT_DIR, 'DHS_OOC', model_dir, 'features.npz'),\n",
    "            check=dict(labels=labels, locs=locs))\n",
    "        features = npz['features']\n",
    "        dataset = '2009-17' + f\n",
    "\n",
    "        for test_country in folds[f]['test']:\n",
    "            print('test country:', test_country)\n",
    "            keep_subset_indices = get_keep_indices(\n",
    "                keep_frac=keep_frac,\n",
    "                seed=seed,\n",
    "                dataset=dataset,\n",
    "                test_country=test_country)\n",
    "\n",
    "            test_country_set = {test_country}\n",
    "            cv_countries_set = all_countries_set - test_country_set\n",
    "            test_indices = get_indices_for_countries(test_country_set)\n",
    "            test_preds[test_indices] = train_linear_logo(\n",
    "                features=features[keep_subset_indices],\n",
    "                labels=labels[keep_subset_indices],\n",
    "                group_labels=country_labels[keep_subset_indices],\n",
    "                cv_groups=countries_to_nums(cv_countries_set),\n",
    "                test_groups=countries_to_nums(test_country_set),\n",
    "                plot=False,\n",
    "                group_names=COUNTRIES)\n",
    "\n",
    "    # save preds on the test set\n",
    "    os.makedirs(savedir, exist_ok=True)\n",
    "    npz_path = os.path.join(savedir, 'test_preds_keep{k}_seed{s}.npz'.format(k=keep_frac, s=seed))\n",
    "    print('Saving preds to:', savedir)\n",
    "    np.savez_compressed(npz_path, test_preds=test_preds, labels=labels, locs=locs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = 'resnet_ms'\n",
    "savedir = os.path.join(LOGS_ROOT_DIR, 'dhs_resnet', 'ms')\n",
    "\n",
    "for keep_frac in KEEPS:\n",
    "    for seed in SEEDS:\n",
    "        print('-------------------------------------------')\n",
    "        print(f'---------- keep: {keep_frac:0.02g}, seed: {seed} ----------')\n",
    "        print('-------------------------------------------')\n",
    "        run_ridgecv_keep(\n",
    "            model_name=model_name,\n",
    "            labels=labels,\n",
    "            locs=locs,\n",
    "            country_labels=country_labels,\n",
    "            folds=ooc_folds,\n",
    "            keep_frac=keep_frac,\n",
    "            seed=seed,\n",
    "            savedir=savedir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = 'resnet_nl'\n",
    "savedir = os.path.join(LOGS_ROOT_DIR, 'dhs_resnet', 'nl')\n",
    "\n",
    "for keep_frac in KEEPS:\n",
    "    for seed in SEEDS:\n",
    "        print('-------------------------------------------')\n",
    "        print(f'---------- keep: {keep_frac:0.02g}, seed: {seed} ----------')\n",
    "        print('-------------------------------------------')\n",
    "        run_ridgecv_keep(\n",
    "            model_name=model_name,\n",
    "            labels=labels,\n",
    "            locs=locs,\n",
    "            country_labels=country_labels,\n",
    "            folds=ooc_folds,\n",
    "            keep_frac=keep_frac,\n",
    "            seed=seed,\n",
    "            savedir=savedir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Concatenated MS, NL Features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_ridgecv_keep_concat(model_names, concat_model_name, labels, locs,\n",
    "                            country_labels, folds, keep_frac, seed):\n",
    "    '''\n",
    "    For every country C (the test country):\n",
    "      1. uses leave-one-country-out CV on all other countries\n",
    "         to tune ridge model alpha parameter\n",
    "      2. using best alpha, trains ridge model on all countries except C\n",
    "      3. runs trained ridge model on C\n",
    "    Saves predictions for each country on test.\n",
    "\n",
    "    Args\n",
    "    - model_names: list of str\n",
    "    - concat_model_name: str\n",
    "    - labels: np.array, shape [num_examples]\n",
    "    - locs: np.array, shape [num_examples, 2]\n",
    "    - country_labels: np.array, shape [num_examples]\n",
    "    - folds: dict, fold (str) => dict\n",
    "    - keep_frac: float, fraction of non-test-country data to use for training\n",
    "    - seed: int\n",
    "    '''\n",
    "    test_preds = np.zeros(num_examples, dtype=np.float32)\n",
    "    all_countries_set = set(COUNTRIES)\n",
    "\n",
    "    for f in FOLDS:\n",
    "        print('Fold:', f)\n",
    "        concat_features = []  # list of np.array, each shape [N, D_i]\n",
    "        for model_name in model_names:\n",
    "            model_fold_name = f'{model_name}_{f}'\n",
    "            model_dir = MODEL_DIRS[model_fold_name].replace('b64', f'keep{keep_frac}_seed{seed}_b64')\n",
    "            npz = load_npz(\n",
    "                os.path.join(LOGS_ROOT_DIR, 'DHS_OOC', model_dir, 'features.npz'),\n",
    "                check=dict(labels=labels, locs=locs))\n",
    "            concat_features.append(npz['features'])\n",
    "        concat_features = np.concatenate(concat_features, axis=1) # shape [N, D_1 + ... + D_m]\n",
    "        dataset = '2009-17' + f\n",
    "\n",
    "        do_plot = True\n",
    "        for test_country in folds[f]['test']:\n",
    "            print('test country:', test_country)\n",
    "            keep_subset_indices = get_keep_indices(\n",
    "                keep_frac=keep_frac,\n",
    "                seed=seed,\n",
    "                dataset=dataset,\n",
    "                test_country=test_country)\n",
    "\n",
    "            test_country_set = {test_country}\n",
    "            cv_countries_set = all_countries_set - test_country_set\n",
    "            test_indices = get_indices_for_countries(test_country_set)\n",
    "            test_preds[test_indices] = train_linear_logo(\n",
    "                features=concat_features[keep_subset_indices],\n",
    "                labels=labels[keep_subset_indices],\n",
    "                group_labels=country_labels[keep_subset_indices],\n",
    "                cv_groups=countries_to_nums(cv_countries_set),\n",
    "                test_groups=countries_to_nums(test_country_set),\n",
    "                plot=do_plot,\n",
    "                group_names=COUNTRIES)\n",
    "            do_plot = False\n",
    "\n",
    "    # save preds on the test set\n",
    "    log_dir = os.path.join(LOGS_ROOT_DIR, 'dhs_resnet', concat_model_name)\n",
    "    os.makedirs(log_dir, exist_ok=True)\n",
    "    npz_path = os.path.join(log_dir, 'test_preds_keep{k}_seed{s}.npz'.format(k=keep_frac, s=seed))\n",
    "    np.savez_compressed(npz_path, test_preds=test_preds, locs=locs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for keep_frac in KEEPS:\n",
    "    for seed in SEEDS:\n",
    "        print('-------------------------------------------')\n",
    "        print(f'---------- keep: {keep_frac:0.02g}, seed: {seed} ----------')\n",
    "        print('-------------------------------------------')\n",
    "        run_ridgecv_keep_concat(\n",
    "            model_names=['resnet_ms', 'resnet_nl'],\n",
    "            concat_model_name='msnl_concat',\n",
    "            labels=labels,\n",
    "            locs=locs,\n",
    "            country_labels=country_labels,\n",
    "            folds=ooc_folds,\n",
    "            keep_frac=keep_frac,\n",
    "            seed=seed)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Incountry"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ridgecv_incountry_wrapper(model_name, savedir):\n",
    "    features_dict = {}\n",
    "    for f in FOLDS:\n",
    "        model_fold_name = f'{model_name}_{f}'\n",
    "        model_dir = MODEL_DIRS[model_fold_name]\n",
    "        npz = load_npz(os.path.join(LOGS_ROOT_DIR, model_dir, 'features.npz'),\n",
    "                       check={'labels': labels})\n",
    "        features_dict[f] = npz['features']\n",
    "\n",
    "    ridge_cv(\n",
    "        features=features_dict,\n",
    "        labels=labels,\n",
    "        group_labels=incountry_group_labels,\n",
    "        group_names=FOLDS,\n",
    "        savedir=savedir,\n",
    "        save_weights=True,\n",
    "        verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = 'incountry_resnet_ms'\n",
    "savedir = '../logs/dhs_resnet/incountry_ms/'\n",
    "ridgecv_incountry_wrapper(model_name, savedir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = 'incountry_resnet_nl'\n",
    "savedir = '../logs/dhs_resnet/incountry_nl/'\n",
    "ridgecv_incountry_wrapper(model_name, savedir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Concatenated MS, NL Features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ridgecv_incountry_concat_wrapper(model_names, savedir):\n",
    "    features_dict = {}\n",
    "    for i, f in enumerate(FOLDS):\n",
    "        concat_features = []  # list of np.array, each shape [N, D_i]\n",
    "        for model_name in model_names:\n",
    "            model_dir = MODEL_DIRS[f'{model_name}_{f}']\n",
    "            npz = load_npz(os.path.join(LOGS_ROOT_DIR, model_dir, 'features.npz'),\n",
    "                           check={'labels': labels})\n",
    "            concat_features.append(npz['features'])\n",
    "        concat_features = np.concatenate(concat_features, axis=1) # shape [N, D_1 + ... + D_m]\n",
    "        features_dict[f] = concat_features\n",
    "\n",
    "    ridge_cv(\n",
    "        features=features_dict,\n",
    "        labels=labels,\n",
    "        group_labels=incountry_group_labels,\n",
    "        group_names=FOLDS,\n",
    "        savedir=savedir,\n",
    "        save_weights=True,\n",
    "        verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_names = ['incountry_resnet_ms', 'incountry_resnet_nl']\n",
    "savedir = os.path.join(LOGS_ROOT_DIR, 'dhs_resnet', 'incountry_msnl_concat')\n",
    "ridgecv_incountry_concat_wrapper(model_names, savedir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Transfer Features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ridgecv_incountry_transfer_wrapper(model_name, savedir):\n",
    "    model_dir = MODEL_DIRS[model_name]\n",
    "    features = load_npz(os.path.join(LOGS_ROOT_DIR, model_dir, 'features.npz'),\n",
    "                        check={'labels': labels})['features']\n",
    "    ridge_cv(\n",
    "        features=features,\n",
    "        labels=labels,\n",
    "        folds=incountry_folds,\n",
    "        savedir=savedir,\n",
    "        save_weights=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = 'resnet_rgb_transfer'\n",
    "savedir = '../logs/dhs_resnet/incountry_rgb_transfer/'\n",
    "ridgecv_incountry_transfer_wrapper(model_name, savedir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "model_name = 'resnet_ms_transfer'\n",
    "savedir = '../logs/dhs_resnet/incountry_ms_transfer/'\n",
    "ridgecv_incountry_transfer_wrapper(model_name, savedir)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.3"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": true,
   "toc_position": {
    "height": "calc(100% - 180px)",
    "left": "10px",
    "top": "150px",
    "width": "284px"
   },
   "toc_section_display": true,
   "toc_window_display": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}