def load_data(
        start_train_year=1900, end_train_year=2010, end_year=2013,
        continental=False, DC=True, N_names=None, train_state=None,
        downsample=None):

    # load data, data array is (years x states x names)
    data, years, states, names = download_data(N_names, end_year)

    # get the latitude and longitude
    loc_file = os.path.join('data','names','state_latlon.csv')
    with open(loc_file) as infile:
        latlon = {state:(float(lat),float(lon)) for
                  state, lat, lon in
                  [line.strip().split(',') for line in infile]}

    # flatten out the first two dimensions
    flat_data, flat_years, flat_states, lat, lon = all_flat \
        = flatten(data, years, states, names, latlon)

    # split train and test
    train_inds, test_inds = get_train_test_ind(
        flat_years, flat_states, train_state, continental, DC)
    train_data, train_years, train_states, train_lat, train_lon = \
        map(lambda arr: arr[train_inds], all_flat)
    test_data, test_years, test_states, test_lat, test_lon = \
        map(lambda arr: arr[test_inds], all_flat)

    # downsample
    if downsample is not None:
        print("Downsampling data to ", downsample, " names per year/state")
        assert isinstance(downsample, int) and 0 > downsample
        from pgmult.utils import downsample_data
        downsample_train_data = downsample_data(train_data, downsample)
        downsample_test_data = downsample_data(test_data, downsample)
    else:
        downsample_train_data = downsample_test_data = None

    # put into a Dataset namedtuple
    Dataset = namedtuple("Dataset", ["K", "data", "years", "lat", "lon", "states", "names"])

    train = Dataset(N_names, train_data, train_years, train_lat, train_lon, train_states, names)
    test  = Dataset(N_names, test_data,  test_years,  test_lat,  test_lon,  test_states,  names)

    downsample_train = Dataset(
        N_names, downsample_train_data, train_years, train_lat, train_lon, train_states, names)
    downsample_test  = Dataset(
        N_names, downsample_test_data,  test_years,  test_lat, test_lon,  test_states, names)

    return downsample_train, downsample_test, train, test
dataset/ETHPy150Open HIPS/pgmult/experiments/census_gp.py/load_data/swapped operands of `>`