data_loader.py

import datetime
import os
from tokenize import group
from typing import Literal, Tuple

import bottleneck as bn
import joblib
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
import torch as th
from einops import rearrange
from geopandas import gpd
from tqdm import tqdm

from imterp.densifier import Densifier
from utils import Scaler, adj_top_k, haversine, logger, mse, read_config, timed

ALWAYS_GEN = read_config('always_gen_dataset')
memory = joblib.Memory('./cache', verbose=0)


def check_generated(name: str) -> bool:
    """Check if a dataset is already generated by examing the existence of
        mandatory data files (X.z, sensors.csv, dates.z).

    Args:
        name (str): Dataset name

    Returns:
        bool: If generated.
    """
    expect_files = ['X.z', 'sensors.csv', 'dates.z']
    for f in expect_files:
        if not os.path.exists(f'./data/{name}/{f}'):
            logger.warning(f'File {f} not found in ./data/{name}/, will generate data.')
            return False
    return True


def validate_dataset(X: np.ndarray, coords: np.ndarray, dates: np.ndarray):
    assert X.shape[0] == coords.shape[0]
    assert X.shape[1] == dates.shape[0]


@memory.cache
def compute_A(coords: np.ndarray, norm: bool = False) -> np.ndarray:
    """Calculate the adjacency matrix from geo coordinates.
    This function is cached by joblib.memory.cache.

    Args:
        coords (np.ndarray): Coordinates of shape [n, 2] where n is the number of nodes.
        norm (bool): If normalize the adjacency matrix via exp func.

    Returns:
        np.ndarray: Adjacency matrix of shape [n, n].
    """
    A = np.zeros((coords.shape[0], coords.shape[0]))
    for i in range(coords.shape[0]):
        for j in range(coords.shape[0]):
            A[i, j] = haversine(coords[i, 0], coords[i, 1], coords[j, 0], coords[j, 1])
    if norm:
        A = np.exp(-A / A.max())
    return A


@timed
def gen_ushcn() -> None:
    logger.debug('Generating ushcn data...')
    flocation = open('./data/raw/ushcn/ushcn-v2.5-stations.txt')
    stations = []
    for index, line in enumerate(flocation):
        stations.append({
            'id': line[0:11].strip(),
            'lat': float(line[12:21]),
            'lng': float(line[21:30]),
            'state': line[38:40].strip(),
            'name': line[41:71].strip(),
        })
    df_location = pd.DataFrame(stations)
    df_location.to_csv('./data/ushcn/sensors.csv', index=True)
    year_begin = 1900
    year_end = 2022
    p = 12 * (year_end - year_begin + 1)  # Time steps
    n = len(stations)  # Number of sensors
    X = np.zeros((n, p))
    A = np.zeros((n, n))
    dates = []
    missing = np.zeros((n, p))
    for sindex, station in enumerate(stations):
        with open(f'./data/raw/ushcn/ushcn.v2.5.5.20231225/{station["id"]}.FLs.52j.prcp') as f:
            for line_idx, line in enumerate(f):
                year = int(line[12:16])
                if not (year >= year_begin and year <= year_end):
                    continue
                for month in range(12):
                    value = float(line[17 + month * 9:22 + month * 9])
                    if value == -9999:
                        missing[sindex, 12 * (year - year_begin) + month] = 1
                        X[sindex, 12 * (year - year_begin) + month] = 0
                    else:
                        X[sindex, 12 * (year - year_begin) + month] = value / 100  # the results precipitation in mm
    for year in range(year_begin, year_end + 1):
        for month in range(12):
            dates.append(datetime.datetime(year, month + 1, 1, 0, 0, 0))
    dates = np.array(dates)
    flocation.close()
    joblib.dump(X, './data/ushcn/X.z')  # z-lib compressoed
    # joblib.dump(A, './data/ushcn/A.z')
    joblib.dump(missing, './data/ushcn/missing.z')
    joblib.dump(dates, './data/ushcn/dates.z')


@timed
def gen_airdata(dataset_name: str, raw_csv: str = './data/raw/airdata/hourly_88101_2022.csv', states=None):
    logger.debug(f'Generating airdata {dataset_name} from {raw_csv}...')
    df_data = pd.read_csv(raw_csv)
    # Exploration code of geopandas
    # gdf_data = gpd.GeoDataFrame(df_data, geometry=gpd.points_from_xy(df_data['Longitude'], df_data['Latitude']))
    # gdf_data = gdf_data.set_crs('epsg:4326')

    # Filter stations outside the mainland US
    df_data = df_data.loc[(df_data['State Code'] != 15) & (df_data['State Code'] != 2)]

    # Filter (if specified states)
    if states is not None:
        df_data = df_data.loc[df_data['State Name'].isin(states)]

    df_data['station_hash'] = df_data['Latitude'].astype(str) + df_data['Longitude'].astype(str)
    df_data['datetime'] = pd.to_datetime(df_data['Date Local'] + ' ' + df_data['Time Local'])

    # Change the station whose minutes is 26 to 0
    df_data.loc[df_data['datetime'].dt.minute != 0,
                'datetime'] = df_data.loc[df_data['datetime'].dt.minute != 0,
                                          'datetime'].apply(lambda dt: dt.replace(minute=0))
    num_stations = df_data['station_hash'].nunique()
    num_timesteps = df_data['datetime'].nunique()
    df_data = df_data.sort_values(by=['station_hash', 'datetime'])
    X = np.zeros((num_stations, num_timesteps))
    time_idx = {t: i for i, t in enumerate(df_data['datetime'].unique())}
    stations = []
    dates = []
    for i, (station_hash, group) in tqdm(enumerate(df_data.groupby('station_hash'))):
        for j, (datetime, row) in enumerate(group.iterrows()):
            X[i, time_idx[row['datetime']]] = row['Sample Measurement']
        stations.append({
            'id': station_hash,
            'lat': group['Latitude'].values[0],
            'lng': group['Longitude'].values[0],
            'state': group['State Name'].values[0],
            'country': group['County Name'].values[0],
        })
    for i, (date, group) in enumerate(df_data.groupby('datetime')):
        dates.append(date.to_pydatetime())

    dates = np.array(dates)
    stations = pd.DataFrame(stations)
    os.makedirs(f'./data/{dataset_name}', exist_ok=True)
    stations.to_csv(f'./data/{dataset_name}/sensors.csv', index=False)
    joblib.dump(X, f'./data/{dataset_name}/X.z')
    joblib.dump(dates, f'./data/{dataset_name}/dates.z')


def load_common(dataset: str, time_step: int = 1, time_cutoff: int = 99999999):
    logger.debug(f'Loading {dataset} data...')
    X = joblib.load(f'./data/{dataset}/X.z')
    coords = pd.read_csv(f'./data/{dataset}/sensors.csv')[['lng', 'lat']].values
    dates = joblib.load(f'./data/{dataset}/dates.z')
    X, dates = X[:, :time_cutoff:time_step], dates[:time_cutoff:time_step]
    logger.info(
        f'Dimension: X: {X.shape}, X range: [{bn.nanmin(X):.4f}, {bn.nanmax(X):.4f}] coords: {coords.shape} dates: {dates.shape}'
    )
    return X, coords, dates


def load_ushcn() -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    if not check_generated('ushcn') or ALWAYS_GEN:
        gen_ushcn()
    logger.info('Loading uschn data...')
    X = joblib.load('./data/ushcn/X.z')
    # A = joblib.load('./data/ushcn/A.z')
    sensors = pd.read_csv('./data/ushcn/sensors.csv')
    coords = sensors[['lng', 'lat']].values
    dates = joblib.load('./data/ushcn/dates.z')
    # A = calculate_A(coords, norm=False)
    missing = joblib.load('./data/ushcn/missing.z')
    logger.info(
        f'Dimension: X: {X.shape}, X range: [{bn.nanmin(X):.4f}, {bn.nanmax(X):.4f}] coords: {coords.shape} dates: {dates.shape}'
    )
    return X, coords, dates


def load_uspm25():
    if not check_generated('uspm25') or ALWAYS_GEN:
        gen_airdata('uspm25', './data/raw/airdata/hourly_88101_2022.csv')
    X, coords, dates = load_common('uspm25', time_step=4)
    return X, coords, dates


def load_ustmp():
    if not check_generated('ustmp') or ALWAYS_GEN:
        gen_airdata('ustmp', './data/raw/airdata/hourly_TEMP_2022.csv')
    return load_common('ustmp', time_cutoff=99999)

def load_catmp():
    if not check_generated('catmp') or ALWAYS_GEN:
        gen_airdata('catmp', './data/raw/airdata/hourly_TEMP_2022.csv', states=['California'])
    X, coords, dates = load_common('catmp', time_cutoff=99999, time_step=4)
    valid_sensors = np.count_nonzero(X, axis=1) > (0.1 * X.shape[1])
    X = X[valid_sensors, :]
    coords = coords[valid_sensors, :]
    return X, coords, dates


def load_cnpm25():
    if not check_generated('cnpm25') or ALWAYS_GEN:
        raise NotImplementedError('Data generation for chnpm25 is not implemented.')
    X, coords, dates = load_common('cnpm25', time_step=6)  # turn hourly data input 6 hour data
    # Remove invalid coordinates
    index_to_remove = []
    for i, coord in enumerate(coords):
        try:
            lng, lat = float(coord[0]), float(coord[1])
            if np.isnan(lng) or np.isnan(lat):
                index_to_remove.append(i)
        except:
            index_to_remove.append(i)
    X = np.delete(X, index_to_remove, axis=0)
    coords = np.delete(coords, index_to_remove, axis=0)
    coords = coords.astype(float)
    logger.info(f'After removing invalid coordinates, X: {X.shape}, coords: {coords.shape}')
    return X, coords, dates


def load_dtaqi():
    if not check_generated('dtaqi') or ALWAYS_GEN:
        raise NotImplementedError('Data generation for dtaqi is not implemented.')
    X, coords, dates = load_common('dtaqi', time_step=1, time_cutoff=5000)
    valid_sensors = np.count_nonzero(X, axis=1) > (0.1 * X.shape[1])
    X = X[valid_sensors, :]
    coords = coords[valid_sensors, :]
    return X, coords, dates


def load_dtpm25():
    if not check_generated('dtpm25') or ALWAYS_GEN:
        raise NotImplementedError('Data generation for dtaqi is not implemented.')
    X, coords, dates = load_common('dtaqi', time_step=1, time_cutoff=5000)
    return X, coords, dates


load_funcs = {
    'ushcn': load_ushcn,
    'uspm25': load_uspm25,
    'ustmp': load_ustmp,
    'cnpm25': load_cnpm25,
    'dtaqi': load_dtaqi,
    'dtpm25': load_dtpm25,
    'catmp': load_catmp,
}


class AuxInfo:

    def __init__(self, coords: np.ndarray, A: np.ndarray):

        # Spatial boundary info used for Spatial Encoding
        self.all_coords = coords
        self.min_lng = np.min(coords[:, 0])
        self.max_lng = np.max(coords[:, 0])
        self.min_lat = np.min(coords[:, 1])
        self.max_lat = np.max(coords[:, 1])

        # Degree-based scaler for PNA Conv
        self.deg = np.mean(np.log(np.sum(A, axis=0) + 1))


# def load_data(dataset: str, train_split=0.8, num_unknown=300):
#     """Data loading function for D_GCN. DEPRECATED!!!
#     """
#     if dataset == 'ushcn':
#         X = load_ushcn()
#     A = compute_A(X, norm=False)
#     rand = np.random.RandomState(2023)
#     split = int(X.shape[1] * train_split)  # Split on temporal dimension
#     X_train = X[:, :split]  # split temporal
#     X_test = X[:, split:]
#     A = np.exp(-A / 1e5)
#     # Unknown locations to be removed from the training set
#     unknown_set = set(rand.choice(list(range(X.shape[0])), num_unknown, replace=False))
#     full_set = set(list(range(X.shape[0])))
#     known_set = full_set - unknown_set
#     X_train_s = X_train[list(known_set), :]  # mask locations
#     A_s = A[:, list(known_set)][list(known_set), :]  # masked adjacency matrix
#     return A, X, X_train, X_test, X_train_s, A_s, unknown_set, full_set, known_set


class GKDataLoader(object):
    """
       Temporal ────▶                        n x t
Spatial┌───────────────────────────────┬──────────┐
   │   │                        masked │          │
   │   ├ ─ ─ ─ Train ─ ─ ─ ─ ─ ─ ─ ─ ─ ┤          │
   ▼   │     Masked ->                 │   known  │ known_rate
       │  Masked + Visible     visible │          │
       │                               │          │
       ├───────────────────────────────┼──────────┤
       │             Test              │   eval   │
       │           test loss           │   loss   │ unknown_rate
       └───────────────────────────────┴──────────┘
                  train_rate               Eval
                                 known -> known + unknown
    """

    def __init__(
            self,
            dataset: str,
            batch_size: int = 2,
            p: int = 7,  # Temporal length per training sample
            max_nodes: int = 1500,
            unknown_rate: float = 0.3,
            masked_rate: float = 0.3,
            train_rate: float = 0.7,
            adj_k: float = 3,
            temporal_sr: int = 1,  # Temporal super resolution rate. Must be integer
            ignore0: bool = True,
            rand_seed: int = 2024,
            normalized: bool = True,
            outdir: str = None,  # Base out dir for producing data
    ) -> None:
        self.dataset = dataset
        self.batch_size = batch_size
        self.unknown_rate = unknown_rate
        self.masked_rate = masked_rate
        self.train_rate = train_rate
        self.adj_k = adj_k
        self.p = p
        self.rand = np.random.RandomState(rand_seed)
        self.outdir = outdir

        X, coords, dates = load_funcs[self.dataset]()

        A = compute_A(coords, norm=True)

        # A = np.where(A != 0, np.exp(-A / A.max()), 0)
        self.scaler = Scaler(X)

        if normalized:
            X = self.scaler.norm(X)  # Scale to 0 - 1 (sort of)

        if X.shape[0] > max_nodes:
            logger.warning(f'Number of nodes {X.shape[0]} exceeds max_node {max_nodes}, will sample {max_nodes} nodes.')
            samples = self.rand.choice(list(range(X.shape[0])), max_nodes, replace=False)
            X = X[samples, :]
            A = A[samples, :][:, samples]
            coords = coords[samples, :]

        self.coords = coords
        self.X = X
        self.A = A

        train_eval_split = int(self.X.shape[1] * self.train_rate)
        self.train_eval_split = train_eval_split
        self.X_train_all = self.X[:, :train_eval_split]
        self.X_eval_all = self.X[:, train_eval_split:]

        self.num_nodes = self.X_train_all.shape[0]
        self.train_total_t = self.X_train_all.shape[1]

        self.dates = dates
        self.dates_train = dates[:train_eval_split]
        self.dates_eval = dates[train_eval_split:]

        num_nodes = self.num_nodes
        self.temporal_sr = temporal_sr

        self.unknown_set = self.rand.choice(
            list(range(X.shape[0])),
            int(self.unknown_rate * X.shape[0]),
            replace=False,
        )
        self.unknown_set = set(self.unknown_set)
        self.known_set = set(range(num_nodes)) - set(self.unknown_set)

        self.info = AuxInfo(coords, adj_top_k(A, adj_k))

        gdf_dir = f'./data/{self.dataset}/terrain.shp'
        if not os.path.exists(gdf_dir):
            logger.error(f'Terrain file {gdf_dir} not found.')
            raise FileNotFoundError()
        self.terrain_gdf: gpd.GeoDataFrame = gpd.read_file(f'./data/{self.dataset}/terrain.shp')

    def sample(self):
        """Sample batch data for training.
        Here num_nodes equals len(known_set).

        Returns:
            X_batch: [batch_size, 1, num_nodes, p]
            Y_batch: [batch_size, 1, num_nodes, p * temporal_sr]
            A_first / A_sub: [num_nodes, num_nodes]
            masked_set: Set of masked nodes
            coords: Coordinates of sampled nodes [num_nodes, 2]

        """
        t_sr = self.temporal_sr
        num_nodes = self.X.shape[0]
        known_set = self.known_set
        masked_set = self.rand.choice(list(range(len(known_set))),
                                      int(self.masked_rate * len(known_set)),
                                      replace=False)
        visible_set = known_set - set(masked_set)
        self.masked_set = set(masked_set)

        A_known = self.A[list(known_set), :][:, list(known_set)]
        # A for the first layer, masked
        A_first = A_known.copy()
        A_first[list(masked_set), :][:, list(masked_set)] = 0
        A_first = adj_top_k(A_first, self.adj_k)  # Also k-nn A_first

        # A for the subsequent layers, allow k-nearest neightbor of the sensor to pass messages
        A_sub = adj_top_k(A_known, self.adj_k)

        A_known_k = adj_top_k(A_known, self.adj_k)
        X_train = self.X_train_all[list(known_set), :]
        X_test = self.X_train_all[list(self.known_set), :]

        X_batch = np.zeros((self.batch_size, 1, X_train.shape[0], self.p))
        Y_batch = np.zeros((self.batch_size, 1, X_train.shape[0], self.p * t_sr))

        for b in range(self.batch_size):
            # end = self.rand.randint(self.p * self.temporal_sr, X_train.shape[1])
            start = self.rand.randint(0, X_train.shape[1] - self.p * t_sr)
            X_batch[b, 0, :, :] = X_train[:, start:start + self.p * t_sr:t_sr]
            Y_batch[b, 0, :, :] = X_test[:, start:start + self.p * t_sr]
            # X_batch[b, 0, :, :] = X_train[:, end - self.p * self.temporal_sr:end:self.temporal_sr]
            # X_test_batch[b, 0, :, :] = X_test[:, end - self.p:end]
            # X_batch[b, 0, :, :] = X_train[:, start:start + self.p]
            # X_test_batch[b, 0, :, :] = X_test[:, start:start + self.p]

        # Y_batch = X_test_batch.copy()
        X_batch[:, :, list(masked_set), :] = 0

        coords = self.coords[list(known_set), :]

        return X_batch, Y_batch, A_first, A_sub, masked_set, coords

    def sample_eval(self):
        """
        Sample batch data for evaluation of data imputation.
            The procedure is already done during the Trainer class.

        Here num_nodes = all nodes in the training sets (up to max_nodes)
        If temporal super resolution is applied, the total length will be
            total_timesteps * (1 - train_rate) / temporal_sr

        Returns:
            X_batch_groups: [group, batch, 1, num_nodes, p]
            Y_batch_groups: [group, batch, 1, num_nodes, p * temporal_sr]
            A_first / A_sub: [num_nodes, num_nodes]
        """
        t_sr = self.temporal_sr
        num_nodes = self.X_eval_all.shape[0]
        total_t = self.X_eval_all.shape[1] // t_sr
        batches = total_t // self.p
        groups = batches // self.batch_size

        total_t_output = groups * self.batch_size * self.p
        X_unknowned = self.X_eval_all.copy()
        X_unknowned[list(self.unknown_set), :] = 0

        X_batch_groups = np.zeros((groups, self.batch_size, 1, num_nodes, self.p))
        Y_batch_groups = np.zeros((groups, self.batch_size, 1, num_nodes, self.p * t_sr))

        for g in range(groups):
            for b in range(self.batch_size):
                start = g * self.batch_size * self.p + b * self.p
                X_batch_groups[g, b, 0, :, :] = X_unknowned[:, start:start + self.p * t_sr:t_sr]
                Y_batch_groups[g, b, 0, :, :] = self.X_eval_all[:, start:start + self.p * t_sr]

        A_first = self.A.copy()
        A_first[list(self.unknown_set), :][:, list(self.unknown_set)] = 0
        A_first = adj_top_k(A_first, self.adj_k)

        A_sub = adj_top_k(self.A, self.adj_k)

        coords = self.coords

        return X_batch_groups, Y_batch_groups, A_first, A_sub, coords

    def _make_group_data(self, data: np.ndarray):
        """Make group data for model forward pass.

        Args:
            data (np.ndarray): Input data of shape [num_nodes, num_timesteps]

        Returns:
            np.ndarray: Grouped data of shape [num_groups, batch_size, 1, num_nodes, num_timesteps]
        """
        total_t = data.shape[1]
        batches = total_t // self.p
        groups = batches // self.batch_size

        grouped_data = np.zeros((groups, self.batch_size, 1, data.shape[0], self.p))
        for g in range(groups):
            for b in range(self.batch_size):
                start = g * self.batch_size * self.p + b * self.p
                grouped_data[g, b, 0, :, :] = data[:, start:start + self.p]

        return grouped_data

    def _get_nodesource(self, node_source: str):
        if isinstance(node_source, str):
            if node_source == 'known':
                source_ids = list(self.known_set)
            elif node_source == 'all':
                source_ids = list(range(self.X.shape[0]))
            else:
                logger.error(f'Unknown source {source_ids}. Could be "known", "all" or a list of index.')
                raise ValueError()
        elif isinstance(source_ids, list) or isinstance(node_source, set):
            source_ids = list(source_ids)

        return source_ids

    def _get_timesource(self, time_source):
        if isinstance(time_source, str):
            if time_source == 'train':
                source_ts = list(range(self.train_eval_split))
            elif time_source == 'eval':
                source_ts = list(range(self.train_eval_split, self.X.shape[1]))
            else:
                logger.error(f'Unknown source {source_ts}. Could be "train" or "eval".')
                raise ValueError()
        else:
            source_ts = list(source_ts)
        return source_ts

    def sample_densify(
        self,
        ratio: float,
        iter: int = 1,
        plot: bool = False,
        node_source='known',
        time_source='eval',
        densify_unknown=False,
        densify_uniform=False,
    ):
        """
        Sample data with densification. No ground truth could be provided.
            Used in the runner. The data will be densified, and then data imputation, then apply interpolation.

        Args:
            ratio (float): Densify ratio
            iter (int, optional): Number of iterations of CVT. Defaults to 1.
            plot (bool, optional): If output plot. Defaults to False.
            node_source (str, optional): The sensor source for densification. Can be 'known' or 'all' or
                a list of index. Defaults to 'known'.
            time_source (str, optional): The source for time. Can be 'train' or 'eval'
                Defaults to 'eval'.
        """

        # see type of source
        source_ids = self._get_nodesource(node_source)
        source_ts = self._get_timesource(time_source)

        densifier = Densifier(
            self.coords[source_ids, :],
            terrain_gdf=self.terrain_gdf,
            grid_size=1000,
            outdir=self.outdir,
        )

        if densify_unknown:
            logger.warning('Will use unknown nodes as densified.')
            densified_set = self.unknown_set
            self.densified_set = densified_set
            densified_coords = self.coords
            X = self.X[:, source_ts]
            Y = X.copy()

            if self.temporal_sr > 1:
                # X = X[:, ::self.temporal_sr]
                X = X[:, source_ts]

            A_first = self.A.copy()
            X[list(densified_set), :] = 0
            A_first[list(densified_set), :][:, list(densified_set)] = 0
            A_first = adj_top_k(A_first, self.adj_k)
            A_sub = adj_top_k(self.A, self.adj_k)
            X_batch_groups = self._make_group_data(X)
            Y_batch_groups = self._make_group_data(Y)
            X_dates = self.dates[source_ts]

            if self.temporal_sr > 1:
                X_dates = X_dates[::self.temporal_sr]
            return X_batch_groups, Y_batch_groups, A_first, A_sub, densified_coords, X_dates

        # densified_coords: including original coords and densified coords
        densified_coords, densified_set = densifier.densify(ratio, iter, plot, uniform=densify_uniform)

        A = compute_A(densified_coords, norm=True)
        A_first = A.copy()
        A_first[list(densified_set), :][:, list(densified_set)] = 0
        A_first = adj_top_k(A_first, self.adj_k)

        A_sub = adj_top_k(A, self.adj_k)

        X = self.X[source_ids, :][:, source_ts]
        Y = X.copy()

        if self.temporal_sr > 1:
            X = X[:, ::self.temporal_sr]


        # Put zero values for the densified node.
        X_densified = np.concatenate([X, np.zeros((len(densified_set), X.shape[1]))], axis=0)

        # Y offers the values in between
        Y_densified = np.concatenate([Y, np.zeros((len(densified_set), Y.shape[1]))], axis=0)

        X_batch_groups = self._make_group_data(X_densified)
        Y_batch_groups = self._make_group_data(Y_densified)

        X_dates = self.dates[source_ts]

        if self.temporal_sr > 1:
            X_dates = X_dates[::self.temporal_sr]

        self.densified_set = densified_set
        return X_batch_groups, Y_batch_groups, A_first, A_sub, densified_coords, X_dates

    def sample_uncertainty(self, node_source='known', time_source='eval'):
        """Sample data for uncertainty estimation. The data will be split into two half, and get masked
            seperately.

        Args:
            node_source (str, optional): The sensor source for densification. Can be 'known' or 'all' or
                a list of index. Defaults to 'known'.
            time_source (str, optional): The source for time. Can be 'train' or 'eval'
                Defaults to 'eval'.

        Returns:
            _type_: _description_
        """
        source_ids = self._get_nodesource(node_source)
        source_ts = self._get_timesource(time_source)
        X_all = self.X[source_ids, :][:, source_ts]
        A_all = self.A[source_ids, :][:, source_ids]

        if self.temporal_sr > 1:
            X_all = X_all[:, ::self.temporal_sr]

        X_index_all = list(range(X_all.shape[0]))
        masked_half_0 = self.rand.choice(X_index_all, int(X_all.shape[0] / 2), replace=False).tolist()
        masked_half_1 = list(set(X_index_all) - set(masked_half_0))

        coords_0 = self.coords[source_ids, :]
        coords_1 = coords_0.copy()

        X_0 = X_all.copy()
        X_0[masked_half_1, :] = 0
        X_group_0 = self._make_group_data(X_0)

        X_1 = X_all.copy()
        X_1[masked_half_0, :] = 0
        X_group_1 = self._make_group_data(X_1)

        A_first_0 = A_all.copy()
        A_first_0[masked_half_1, :][:, masked_half_1] = 0
        A_first_0 = adj_top_k(A_first_0, self.adj_k)

        A_first_1 = A_all.copy()
        A_first_1[masked_half_0, :][:, masked_half_0] = 0
        A_first_1 = adj_top_k(A_first_1, self.adj_k)

        A_sub_0 = adj_top_k(A_all, self.adj_k)
        A_sub_1 = A_sub_0.copy()

        X_dates = self.dates[source_ts][:X_group_0.shape[0] * X_group_0.shape[1] * X_group_0.shape[4]][::self.temporal_sr]

        return X_group_0, A_first_0, A_sub_0, coords_0, masked_half_0, X_group_1, A_first_1, A_sub_1, coords_1, masked_half_1, X_dates


if __name__ == '__main__':
    load_catmp()
    # for k, v in load_funcs.items():
    #     logger.info(f'Testing loader for {k}...')
    #     v()