polymerlearn/utils/graph_prep.py

import os, glob, random
import torch
import numpy as np
import pandas as pd

from torch_geometric.data import Data, InMemoryDataset
from torch_geometric.loader import DataLoader
from torch_geometric.utils import to_networkx
from torch_geometric.transforms import NormalizeFeatures

from sklearn.model_selection import train_test_split, KFold
from sklearn.preprocessing import StandardScaler

from polymerlearn.utils.xyz2mol import int_atom, xyz2mol

def read_xyz_file_top_conformer(filename, look_for_charge=True):
    """
    Reads an xyz file and parses the first conformer at the top
    """

    atomic_symbols = []
    xyz_coordinates = []
    charge = 0

    with open(filename, "r") as file:

        for line_number, line in enumerate(file):
            #print(line)
            if line_number == 0:
                num_atoms = int(line)
            elif line_number == 1:
                if "charge=" in line:
                    charge = int(line.split("=")[1])
            elif line_number >= num_atoms + 2:
                break # Break after first conformation
            else:
                atomic_symbol, x, y, z = line.split()
                atomic_symbols.append(atomic_symbol)
                xyz_coordinates.append([float(x), float(y), float(z)])

    atoms = [int_atom(atom) for atom in atomic_symbols]

    return atoms, charge, xyz_coordinates

def convert_xyz_to_mol(filename):
    atoms, charge, xyz_coordinates = read_xyz_file_top_conformer(filename)

    mols = xyz2mol(
            atoms, 
            xyz_coordinates,
            charge = charge,
            use_graph=True,
            allow_charged_fragments=True,
            embed_chiral=False,
            use_huckel=False)

    return mols[0]

# Citation (C) = https://towardsdatascience.com/practical-graph-neural-networks-for-molecular-machine-learning-5e6dee7dc003
def get_atom_features(mol):
    '''
    Make atom features
        - Can be made more robust with background work
    '''
    # Cite: C
    features = []
    
    for atom in mol.GetAtoms():
        #atomic_number.append(atom.GetAtomicNum())
        charge = atom.GetFormalCharge()
        degree = atom.GetDegree()
        mass = atom.GetMass()
        is_aromatic = atom.GetIsAromatic()
        anum = atom.GetAtomicNum()
        explicit_hs = atom.GetNumExplicitHs()
        num_valence = atom.GetTotalValence()
        num_rad_electrons = atom.GetNumRadicalElectrons()

        #features.append([charge, degree, mass, is_aromatic, anum, explicit_hs, num_rad_electrons])
        features.append([charge, degree, mass, is_aromatic, explicit_hs, num_valence])
        #num_hs.append(atom.GetTotalNumHs(includeNeighbors=True))
        
    return torch.tensor(features).float()

def get_edge_index(mol, get_edge_attr = False):
    '''
    Gets edge index for a molecule
    '''
    # Cite: C
    row, col = [], []
    edge_attr = []

    for bond in mol.GetBonds():
        start, end = bond.GetBeginAtomIdx(), bond.GetEndAtomIdx()
        row += [start, end]
        col += [end, start]

        if get_edge_attr:
            btype = bond.GetBondTypeAsDouble()
            inring = int(bond.IsInRing())
            edge_attr.append([btype, inring])

    eidx = torch.tensor([row, col], dtype=torch.long)

    if get_edge_attr:
        edge_attr = torch.tensor(edge_attr).float()
        return eidx, edge_attr
        
    return eidx

def prepare_dataloader_graph_AG(
        A_mol_list, 
        G_mol_list, 
        Y = None, 
        add_A = None, 
        add_G = None, 
        get_edge_attr = False,
        device = None,
        atom_feat=None,
        normalize_features = False):
    '''
    Prepares a dataloader given a list of molecules

    Args:
        A_mol_list (list of lists): List of lists of RDKit Mol objects for each sample.
            Should look something like:
                [[A_11, A_12], [A_21, A_22, A_23, A_24], ..., [A_n1]]
        G_mol_list (list of lists): Same as A_mol_list but for Glycols
        Y (iterable, optional): Y (ground truth values) for each sample
        add_A (dict of lists, optional): Dictionary of lists of values to add for each acid. 
            Should be keyed on strings of names of variables with list values 
            corresponding to numerical values to add to the Data objects in the
            DataLoader.
        add_G (dict of lists, optional): Same as add_A but for glycols.
        atom_feat (Callable[[RdKit.Mol], torch.Tensor], optional): Function to output the 
            feature matrix for a given molecule.
        normalize_features (bool, optional): If true, normalize all node features in the Acid 
            and Glycol graphs
    '''

    if atom_feat == None:
        atom_feat = get_atom_features

    assert len(A_mol_list) == len(G_mol_list), 'A and G mol (RDKit) lists not same length'

    norm_feat = NormalizeFeatures(attrs = ['x'])

    # Cite: C
    data_list = []

    i = 0 # Counts total Amols, Gmols that we've added (i.e. whole samples)

    for Amols, Gmols in zip(A_mol_list, G_mol_list):

        acid_graphs = []
        j = 0 # Counts total number of acids for this sample
        for Amol in Amols:
            Ax = atom_feat(Amol)
            if get_edge_attr:
                Aedge_index, Aedge_attr = get_edge_index(Amol, get_edge_attr=True)
            else:
                Aedge_index = get_edge_index(Amol)

            # Support for additional arguments:
            add_args = {}
            if add_A is not None:
                for key, val in add_A.items():
                    add_args[key] = torch.tensor([val[i][j]]).float().to(device)

            if Y is not None:
                add_args['y'] = torch.tensor(Y[i]).float().to(device)
            if get_edge_attr:
                add_args['edge_attr'] = Aedge_attr.to(device)

            acid_data = Data(
                x=Ax.to(device), 
                edge_index = Aedge_index.to(device), 
                **add_args)
            if normalize_features:
                acid_data = norm_feat(acid_data)
            # All acid data should be in device
            acid_graphs.append(acid_data)
            j += 1


        glycol_graphs = []
        j = 0 # Counts total number of glycols for this sample
        for Gmol in Gmols:
            Gx = atom_feat(Gmol)

            if get_edge_attr:
                Gedge_index, Gedge_attr = get_edge_index(Gmol, get_edge_attr=True)
            else:
                Gedge_index = get_edge_index(Gmol)

            add_args = {}
            if add_A is not None:
                for key, val in add_G.items():
                    add_args[key] = torch.tensor([val[i][j]]).float()

            # Support for adding multiple
            if Y is not None:
                add_args['y'] = torch.tensor(Y[i]).float().to(device)
            if get_edge_attr:
                add_args['edge_attr'] = Gedge_attr.to(device)

            glycol_data = Data(
                x=Gx.to(device), 
                edge_index=Gedge_index.to(device), 
                **add_args)
            if normalize_features:
                glycol_data = norm_feat(glycol_data)
            # All glycol data should be in device
            glycol_graphs.append(glycol_data)
            j += 1

        data_list.append((acid_graphs, glycol_graphs))

        i += 1

    return data_list

def graph_dataloader_z_pos(
        A_charge_coords, 
        G_charge_coords,
        Y,
    ):
    '''
    Gets dataloader by nuclear charges (z) and positions (pos)
    - Useful for SchNet architecture
    '''

    data_list = []

    for A, G in zip(A_charge_coords, G_charge_coords):
        # List all atoms and their coordinates:
        acid_data_list = []
        for z, pos in A:
            adata = Data(
                z = torch.as_tensor(z).long(),
                pos = torch.as_tensor(pos),
            )
            acid_data_list.append(adata)

        glycol_data_list = []
        for z, pos in G:
            gdata = Data(
                z = torch.as_tensor(z).long(),
                pos = torch.as_tensor(pos),
            )
            glycol_data_list.append(gdata)
            
        data_list.append((acid_data_list, glycol_data_list))

    return data_list


def get_AG_info(data, ac = (20,33), gc = (34,46)):
    '''
    Gets acid/glycol info from a dataframe containing input in the Eastman fashion
    '''

    ac_tuple = False
    gc_tuple = False
    if type(ac) == tuple:
        ac_tuple = True
    
    if type(gc) == tuple:
        gc_tuple = True

    # Decompose the data into included names
    if ac_tuple:
        acid_names = pd.Series([c[1:] for c in data.columns[ac[0]:ac[1]].tolist()])
    else:
        acid_names = pd.Series([c[1:] for c in data[ac].columns.tolist()])
    if gc_tuple:     
        glycol_names = pd.Series([c[1:] for c in data.columns[gc[0]:gc[1]].tolist()])
    else:
        glycol_names = pd.Series([c[1:] for c in data[gc].columns.tolist()])

    # Holds all names of acids and glycols
    acid_included = []
    glycol_included = []

    # Keep track of percents in each acid, glycol
    acid_pcts = []
    glycol_pcts = []

    # Get relevant names and percentages of acid/glycols
    for i in range(data.shape[0]):

        if ac_tuple:
            acid_hit = (data.iloc[i,ac[0]:ac[1]].to_numpy() > 0)
        else:
            acid_hit = (data[ac].iloc[i].to_numpy() > 0)
        if gc_tuple:
            glycol_hit = (data.iloc[i,gc[0]:gc[1]].to_numpy() > 0)
        else:
            glycol_hit = (data[gc].iloc[i].to_numpy() > 0)

        # Add to percentage lists:
        if ac_tuple:
            acid_pcts.append(data.iloc[i,ac[0]:ac[1]][acid_hit].tolist())
        else:
            acid_pcts.append(data[ac].iloc[i][acid_hit].tolist())
        if gc_tuple:
            glycol_pcts.append(data.iloc[i,gc[0]:gc[1]][glycol_hit].tolist())
        else:
            glycol_pcts.append(data[gc].iloc[i][glycol_hit].tolist())

        acid_pos = acid_names[np.argwhere(acid_hit).flatten()].tolist()
        glycol_pos = glycol_names[np.argwhere(glycol_hit).flatten()].tolist()

        acid_included.append(acid_pos)
        glycol_included.append(glycol_pos)

    return acid_included, glycol_included, acid_pcts, glycol_pcts

def list_mask(L, mask):
    '''
    Transform a list with a boolean mask
    Args:
        L (list): List to be masked
        mask (iterable): Mask to apply on list L
    '''
    return [L[int(i)] for i in range(len(mask)) if mask[i]]

def split_validation(train_mask, val_num):

    train_mask, val_mask = train_test_split(train_mask, test_size=val_num, random_state=14)

    return train_mask, val_mask



# base_structure_dir = os.path.join('/home/sai/Eastman_Project',
#     'ReadyToEnsemble',
#     'Structures',
#     'AG',
#     'xyz')
base_structure_dir = os.path.join('..',
    'Structures',
    'AG',
    'xyz')

class GraphDataset:
    '''
    Generates a graph dataset based on input from Eastman data

    Args:
        data (pd.DataFrame): DataFrame containing Eastman data format.
        Y_target (pd.DataFrame): Target values to predict with the dataset.
        structure_dir (str, optional): Location of base structures, i.e. xyz files.
            Assumes that all xyz files are named exactly like the molecules in the
            input data given from Eastman.
        add_features (np.array or torch.Tensor): Additional features to be added to
            each sample's final embedding before processing. These are GLOBAL features
            and are not per-atom features (TODO: make per-atom features).
        ac (tuple, len 2): Bottom and top bounds for all acids on the table given as
            data. Must be column indices.
        ac (list, str): List of columns for all acids
        gc (tuple, len 2): Bottom and top bounds for all glycols on the table given as
            data. Must be column indices.
        gc (list, str): List of columns for all glycols
        test_size (float): Proportion of data to be used as testing data (if using 
            train/test split).
        val_size (float): Proportion of the data to be used as validation data. If None,
            does not make a validation split.
        get_edge_attr (bool): Option to use predefined edge features on the graphs.
        bound_filter (list, length 2): Filters the dataset by some bound on Y value, 
            i.e. controls for outliers
            TODO: implementation for multiple Y values
        exclude_inds (list of ints): List of indices to exclude in the dataframe.
        device (str): Device name at which to run torch calculations on. Supports GPU.
        standard_scale (bool, optional): Whether to perform standard scaling for the 
            add_features at split time. Cannot be done as a preprocessing step (i.e. 
            incorporated to add_features) because the scales should depend only on trainig
            data. Therefore, scaling parameters must be recomputed at each split. 
            Default False.
        ss_mask (list/ndarray of bools): Mask over the variables that need to be standard
            scaled. May be used if you want to scale some variables (like AN, OHN) but not
            others (like Mw).
        normalize_features (bool, optional): If True, normalize the X values using 
            `torch_geometric.transforms.NormalizeFeatures` for each graph. (:default: :obj:`True`)
    '''

    def __init__(self,
            data,
            Y_target,
            structure_dir = base_structure_dir,
            add_features = None,
            ac = (20,33),
            gc = (34,46),
            test_size = 0.25,
            val_size = None,
            get_edge_attr  = False,
            bound_filter = None,
            exclude_inds = None,
            device = None,
            standard_scale = False,
            ss_mask = None,
            z_pos_loaders = False,
            kelvin=False,
            normalize_features = False,
        ):

        self.add_features = add_features
        self.get_edge_attr = get_edge_attr
        self.val_size = val_size
        self.test_size = test_size
        self.device = device
        self.standard_scale = standard_scale
        self.ss_mask = ss_mask
        self.z_pos_loaders = z_pos_loaders
        self.normalize_features = normalize_features
        if self.add_features is not None:
            if self.add_features.ndim == 1:
                self.add_features = self.add_features[:, np.newaxis] # Turn to column vector

        if type(ac) == tuple:
            self.ac_tuple = True
        else:
            self.ac_tuple = False
        
        if type(gc) == tuple:
            self.gc_tuple = True
        else:
            self.gc_tuple = False
        
        # Decompose the data into included names
        if self.ac_tuple:
            self.acid_names = pd.Series([c[1:] for c in data.columns[ac[0]:ac[1]].tolist()])
        else:
            self.acid_names = pd.Series([c[1:] for c in data[ac].columns.tolist()])
        if self.gc_tuple:     
            self.glycol_names = pd.Series([c[1:] for c in data.columns[gc[0]:gc[1]].tolist()])
        else:
            self.glycol_names = pd.Series([c[1:] for c in data[gc].columns.tolist()])

        # Holds all names of acids and glycols
        acid_included = []
        glycol_included = []

        # Keep track of percents in each acid, glycol
        acid_pcts = []
        glycol_pcts = []

        # Get relevant names and percentages of acid/glycols
        for i in range(data.shape[0]):

            if self.ac_tuple:
                acid_hit = (data.iloc[i,ac[0]:ac[1]].to_numpy() > 0)
            else:
                acid_hit = (data[ac].iloc[i].to_numpy() > 0)
            if self.gc_tuple:
                glycol_hit = (data.iloc[i,gc[0]:gc[1]].to_numpy() > 0)
            else:
                glycol_hit = (data[gc].iloc[i].to_numpy() > 0)

            # Add to percentage lists:
            if self.ac_tuple:
                acid_pcts.append(data.iloc[i,ac[0]:ac[1]][acid_hit].tolist())
            else:
                acid_pcts.append(data[ac].iloc[i][acid_hit].tolist())
            if self.gc_tuple:
                glycol_pcts.append(data.iloc[i,gc[0]:gc[1]][glycol_hit].tolist())
            else:
                glycol_pcts.append(data[gc].iloc[i][glycol_hit].tolist())

            acid_pos = self.acid_names[np.argwhere(acid_hit).flatten()].tolist()
            glycol_pos = self.glycol_names[np.argwhere(glycol_hit).flatten()].tolist()

            acid_included.append(acid_pos)
            glycol_included.append(glycol_pos)

        # Read all xyz files into generators, get lowest energy conformation (index 0)

        self.acid_mols = []
        self.glycol_mols = []

        if self.z_pos_loaders: # Load the Z-pos structure

            for i in range(len(acid_included)):
                A_sub = []
                for j in range(len(acid_included[i])):
                    Acharge, _, Acoords = read_xyz_file_top_conformer(os.path.join(structure_dir, acid_included[i][j] + '.xyz'))
                    A_sub.append((Acharge, Acoords))

                self.acid_mols.append(A_sub)

                G_sub = []
                for j in range(len(glycol_included[i])):
                    Gcharge, _, Gcoords = read_xyz_file_top_conformer(os.path.join(structure_dir, glycol_included[i][j] + '.xyz'))
                    G_sub.append((Acharge, Acoords))
                
                self.glycol_mols.append(G_sub)

        else:

            for i in range(len(acid_included)):
                self.acid_mols.append(
                    [convert_xyz_to_mol(os.path.join(structure_dir, acid_included[i][j] + '.xyz')) for j in range(len(acid_included[i]))]
                )

                self.glycol_mols.append(
                    [convert_xyz_to_mol(os.path.join(structure_dir, glycol_included[i][j] + '.xyz')) for j in range(len(glycol_included[i]))]
                )

        # Set Y (target)
        Y = data.loc[:,Y_target]
        if kelvin:
            if 'Tg' in Y_target:
                Y['Tg'] = Y['Tg'] + 273.15


        # Mask data for empty entries
        non_nan_mask = Y.notna()

        self.exclude_inds = exclude_inds

        if self.exclude_inds is not None:
            # Set up exclude-by-index mask:
            inds_lookup = set(self.exclude_inds)
            exclude_by_index = [not (i in inds_lookup) for i in range(Y.shape[0])]
        else:
            exclude_by_index = [True] * Y.shape[0]

        self.bound_filter = bound_filter
        
        if self.bound_filter is not None:
            non_nan_mask = non_nan_mask & (Y > bound_filter[0]) & (Y < bound_filter[1]) & exclude_by_index
        
        non_nan_mask['res_bool'] = False
        non_nan_mask.loc[non_nan_mask[Y_target].all(1), 'res_bool'] = True

        non_nan_mask = non_nan_mask['res_bool'].values
        self.total_samples = sum(non_nan_mask)

        # Mask acid, glycols:
        self.acid_mols = list_mask(self.acid_mols, non_nan_mask)
        self.glycol_mols = list_mask(self.glycol_mols, non_nan_mask)

        # Mask Y:
        self.Y = Y[non_nan_mask].values

        # Mask data:
        self.data = data.loc[non_nan_mask,:]

        # Mask percentages of acids and glycols:
        self.acid_pcts = list_mask(acid_pcts, non_nan_mask)
        self.glycol_pcts = list_mask(glycol_pcts, non_nan_mask)

        # Mask additional features:
        if self.add_features is not None: # Mask additional features, if needed
            self.add_features = list_mask(self.add_features, non_nan_mask)

        # Make dataloader
        rangeL = list(range(len(self.acid_mols)))
        train_mask, test_mask = train_test_split(rangeL, test_size = test_size, random_state=14)

        # Support validation splitting:
        if self.val_size is not None:
            adj_valsize = (self.val_size) / (self.val_size + (1 - self.test_size))
            train_mask, val_mask = train_test_split(train_mask, test_size = adj_valsize, random_state=14)
        else:
            val_mask = None

        self.split_by_indices(train_mask=train_mask, test_mask=test_mask, val_mask=val_mask)

        # if self.standard_scale:
        #     add_features = np.array(self.add_features)

        print(Y)

    def get_train_batch(self, size: int):
        '''
        Perform manual batching of graph dataset

        Args:
            size (int): Size of the batch to be retrieved
        '''

        # Randomly sample the training data
        sample_inds = random.sample(list(np.arange(len(self.Ytrain))), k = size)

        if self.add_features is None:
            return [self.train_data[i] for i in sample_inds], torch.tensor([self.Ytrain[i] for i in sample_inds]).float(), None
        else:
            train_masked = [self.train_data[i] for i in sample_inds]
            Y_masked = torch.tensor(np.array([self.Ytrain[i] for i in sample_inds])).float()
            add_masked = [self.add_train[i] for i in sample_inds]
            return train_masked, Y_masked, add_masked

    def get_test(self, test_inds = None):
        '''
        Get test data nbased on the current internal split

        Args:
            test_inds (any, optional): If provided, is returned with all other
                values.
        '''
        if test_inds is not None:
            return self.test_data, \
                torch.tensor(np.array(self.Ytest)).float().to(self.device), \
                self.add_test, \
                test_inds
        else:
            return self.test_data, \
                torch.tensor(np.array(self.Ytest)).float().to(self.device), \
                torch.tensor(self.add_test).float().to(self.device)

    def get_validation(self):
        '''
        Get the validation data based on current internal splits
        '''
        if self.val_data is None: # No validation split present
            return None
        else:
            # val_data should already be on device
            return self.val_data, \
                torch.tensor(self.Yval).float().to(self.device), \
                torch.tensor(self.add_val).float().to(self.device)

    def Kfold_CV(self, folds, val = False, val_size = None):
        '''
        Generator that wraps SKLearn's K-fold cross validation

        Note that the yield of this function is the testing data, you must perform batching
            of the dataset object (get_train_batch) to get the training data. Rationale 
            behind this is to allow you to train multiple epochs while repeatedly batching
            the training data under one iteration of the Kfold_CV function. 
        
        Args:
            folds (int): Number of folds for the cross validation.
            val (bool): Should be set to True if using validation split.
            val_size (float, 0<=x<=1): Proportion of whole dataset that is used for
                validation split on each fold.

        Yield:
            tuple(tuple(train_data, Ytrain, add_train), tuple(test_data, Ytest, add_test))
        '''

        inds = np.arange(self.total_samples)
        kfold = KFold(n_splits=folds, shuffle = True)

        for train_inds, test_inds in kfold.split(inds):
            if val: # Split the validation:
                val_size = val_size if val_size is not None else self.val_size
                val_adj = val_size / (val_size + (len(train_inds) / len(test_inds)))
                train_inds, val_inds = train_test_split(train_inds, test_size = val_adj, random_state=14)
                self.val_inds = val_inds # No set if val is not true
            else:
                val_inds = None

            self.split_by_indices(train_mask = train_inds, test_mask = test_inds, 
                val_mask = val_inds)
            self.train_inds = train_inds
            self.test_inds = test_inds
            yield self.get_test(test_inds)

    def make_dataloader_by_mask(self, mask):
        '''
        Makes an internal dataloader based on some given list of indices
            - Not technically a mask
            - Makes no internal updates
        '''

        # Perform all train masking: -------------------------------
        Ymask = [self.Y[int(i)] for i in mask]
        mask_Amols = [self.acid_mols[int(i)] for i in mask]
        mask_Gmols = [self.glycol_mols[int(i)] for i in mask]


        if self.z_pos_loaders:
            data = graph_dataloader_z_pos(mask_Amols, mask_Gmols, Ymask)
        else:
            add_A = {'pct': [self.acid_pcts[i] for i in mask]}
            add_G = {'pct': [self.glycol_pcts[i] for i in mask]}

            data = prepare_dataloader_graph_AG(mask_Amols, mask_Gmols, Ymask,
                            add_A = add_A, add_G = add_G, device = self.device,
                            normalize_features = self.normalize_features)

        return data

    def get_additional_by_mask(self, mask):
        '''
        Get additional elements based on given list of indices
        '''
        return [self.add_features[int(i)] for i in mask]

    def get_Y_by_mask(self, mask):
        return [self.Y[int(i)] for i in mask]

    def split_by_indices(self, train_mask, test_mask, val_mask = None):
        '''
        Resets train_data, test_data, Ytrain, and Ytest for internal use
        Splits the data given train_mask and test_mask and stores dataloaders in
            self.train_data and self.test_data
        '''

        self.train_mask = train_mask
        self.test_mask = test_mask
        self.val_mask = val_mask

        self.Ytrain = [self.Y[int(i)] for i in train_mask]
        self.Ytest = [self.Y[int(i)] for i in test_mask]

        self.train_data = self.make_dataloader_by_mask(train_mask)
        self.test_data = self.make_dataloader_by_mask(test_mask)

        if self.val_mask is not None:
            self.Yval = [self.Y[int(i)] for i in val_mask]
            self.val_data = self.make_dataloader_by_mask(val_mask)
        else:
            self.val_data = None

        # Perform all test masking:  --------------------------------
        # self.Ytest = [self.Y[int(i)] for i in test_mask]
        # self.test_Amols = [self.acid_mols[int(i)] for i in test_mask]
        # self.test_Gmols = [self.glycol_mols[int(i)] for i in test_mask]

        # add_A = {'pct': [self.acid_pcts[i] for i in test_mask]}
        # add_G = {'pct': [self.glycol_pcts[i] for i in test_mask]}
        # self.test_data = prepare_dataloader_graph_AG(self.test_Amols, self.test_Gmols, self.Ytest,
        #                 add_A = add_A, add_G = add_G)

        if self.add_features is not None:
            self.add_train = [self.add_features[int(i)] for i in train_mask]
            self.add_test = [self.add_features[int(i)] for i in test_mask]
            if self.val_mask is not None:
                self.add_val = [self.add_features[int(i)] for i in val_mask]

            if self.standard_scale:
                if self.ss_mask is not None:
                    self.add_train = np.array(self.add_train)
                    self.add_test = np.array(self.add_test)

                    # Scale only the variables masked in by the ss_mask:
                    ss = StandardScaler().fit(self.add_train[:,self.ss_mask])
                    self.add_train[:,self.ss_mask] = ss.transform(self.add_train[:,self.ss_mask])
                    self.add_test[:,self.ss_mask] = ss.transform(self.add_test[:,self.ss_mask])

                    if self.val_mask is not None:
                        self.add_val[:,self.ss_mask] = ss.transform(self.add_val[:,self.ss_mask]) 
                        self.add_val = list(self.add_val)

                    self.add_train = list(self.add_train)
                    self.add_test = list(self.add_test)

                else:
                    # Standard scale based on new split
                    ss = StandardScaler().fit(np.array(self.add_train))
                    #   Fit to only training data, as is customary
                    self.add_train = list(ss.transform(self.add_train))
                    self.add_test = list(ss.transform(self.add_test))
                    if self.val_mask is not None:
                        self.add_val = list(ss.transform(self.add_val))

        else:
            self.add_train = None
            self.add_test = None
            if self.val_mask is not None:
                self.add_val = None

# Misc. testing functions:
def test_xyz2mol():
    print(read_xyz_file_top_conformer(os.path.join(base_structure_dir, 'IPA.xyz')))

    mol = convert_xyz_to_mol(os.path.join(base_structure_dir, 'IPA.xyz'))

    print([a.GetAtomicNum() for a in mol.GetAtoms()])

def test_dataset():

    data = pd.read_csv('../../dataset/pub_data.csv')

    print(data)

    base_structure_dir = os.path.join('..', '..',
        'Structures',
        'AG',
        'xyz'
    )

    dataset = GraphDataset(data = data, Y_target=['IV'], z_pos_loaders = True,
        structure_dir = base_structure_dir, kelvin=True)

    #print(dataset.Y)

if __name__ == '__main__':

    test_dataset()