dietnetwork 0.3.0__py3-none-any.whl

dietnet/__init__.py

@@ -0,0 +1,11 @@
+"""Diet Network for genetic ancestry prediction.
+
+The model, training loop and the statistics it is trained on are HussinLab's
+published code (Rochefort-Boulanger, Scicluna et al., AJHG 2026), shipped
+unmodified and called from here; see :mod:`dietnet.upstream`. This package adds
+what the published scripts leave to the user: building their HDF5 input from
+PLINK, allele-aware matching of new cohorts, portable model packages, ensembles,
+biobank-scale inference and a command line around all of it.
+"""
+
+__version__ = "0.3.0"

dietnet/_launch.py

@@ -0,0 +1,40 @@
+"""Run one of upstream's scripts, unmodified, under current torch.
+
+Usage (by :func:`dietnet.upstream.run`): ``python -m dietnet._launch SCRIPT [ARGS...]``.
+
+upstream's train.py saves its validation results (numpy float64 scalars) in its
+checkpoint, then reloads it with a bare ``torch.load``. Since torch 2.6 that
+call defaults to ``weights_only=True`` and refuses numpy scalars. Rather than
+switching the safe loader off, this allowlists exactly the numpy types such a
+checkpoint holds -- the remedy torch's own error message gives -- and then runs
+the script as ``python SCRIPT`` would, from its own directory.
+"""
+
+import runpy
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+
+
+def _allow_numpy_scalars() -> None:
+    try:
+        from numpy._core.multiarray import scalar  # numpy >= 2
+    except ImportError:  # pragma: no cover - numpy 1.x
+        from numpy.core.multiarray import scalar
+    torch.serialization.add_safe_globals(
+        [scalar, np.dtype, type(np.dtype(np.float64)), type(np.dtype(np.int64))]
+    )
+
+
+def main() -> None:
+    script = Path(sys.argv[1]).resolve()
+    _allow_numpy_scalars()
+    sys.argv = [str(script), *sys.argv[2:]]
+    sys.path.insert(0, str(script.parent))
+    runpy.run_path(str(script), run_name="__main__")
+
+
+if __name__ == "__main__":
+    main()

dietnet/_upstream/Dietnet_train_and_evaluate/compute_input_features_statistics.py

@@ -0,0 +1,221 @@
+import argparse
+import os
+import multiprocessing as mp
+import time
+from datetime import datetime
+
+import numpy as np
+
+import h5py
+
+
+def main():
+    start_time = time.time()
+    args = parse_args()
+
+    # Flag to compute only mean (True or False)
+    mean_only = args.mean_only
+
+    # Load indices of every fold
+    partition_data = np.load(args.partition, allow_pickle=True)
+    indices_byfold = partition_data['folds_indexes']
+
+    # Dataset
+    dataset_file = h5py.File(args.dataset, 'r')
+
+    # Compute inp feat stats by fold
+    means_by_fold = []
+    sds_by_fold = []
+    for fold in range(len(indices_byfold)):
+        print('\n---\nComputing input features statistic of fold', str(fold))
+        fold_start_time = time.time()
+
+        # Indices of the fold
+        fold_indices = indices_byfold[fold]
+
+        # Genotypes of training samples
+        train_indices = np.sort(fold_indices[0]) # sort is a hdf5 requirement
+        x_train = dataset_file['inputs'][train_indices]
+
+        # Compute statistics of every features in parallel
+        ncpus = args.ncpus
+        print('Computing input features statistics of {} SNPs for {} '
+              'samples using {} cpu(s)'.format(
+              x_train.shape[1], x_train.shape[0], ncpus))
+
+        pool = mp.Pool(processes=ncpus) #this starts ncpus nb of processes
+
+        # Results that will be returned by processes
+        results = []
+
+        # The statistics of every snp feature is computed by a different process
+        nb_snps = x_train.shape[1]
+        for i in range(nb_snps):
+            # Compute snps mean
+            if mean_only :
+                results.append(
+                        pool.apply_async(
+                            compute_feature_mean, args=(i, x_train[:,i])))
+
+            # Compute mean and standard deviation (sd)
+            else:
+                results.append(
+                        pool.apply_async(
+                            compute_feature_mean_and_sd, args=(i, x_train[:,i])))
+
+        # Get order (snp nb) and mean of features from multiprocess results
+        print('Computed statistics of all input features, getting the results')
+        results_values = np.array([np.array(p.get()) for p in results])
+
+        # Order stats based on snp nb
+        print('Ordering {} means returned by multiprocessing'.format(
+            len(results_values[:,1])))
+        idx = results_values[:,0].astype(int)
+        ordered_idx_pos = idx.argsort()
+
+        # Ordered means
+        means = results_values[:,1][ordered_idx_pos]
+
+        # Ordered sd
+        if not mean_only:
+            print('Ordering {} standard deviations returned by multiprocessing'.format(
+                len(results_values[:,2])))
+            sds = results_values[:,2][ordered_idx_pos]
+
+        # Append fold statistics
+        means_by_fold.append(means)
+        sds_by_fold.append(sds)
+
+        print('Computed fold input features stats in {} seconds\n---\n'.format(
+            time.time()-fold_start_time))
+
+    # Close access to h5py dataset file
+    dataset_file.close()
+
+    # Write input features stats to file
+    if args.out is not None:
+        stats_filename = args.out
+    else:
+        stats_filename = 'input_features_stats_' \
+                + args.dataset.split('/')[-1][0:-5] \
+                + '_' \
+                + datetime.now().strftime("%Y_%m_%d_%Hh%Mmin%Ssec")
+
+    stats_fullpath = os.path.join(args.exp_path, stats_filename)
+
+    if mean_only:
+        np.savez(stats_fullpath,
+                 means_by_fold=means_by_fold)
+
+    else:
+        np.savez(stats_fullpath,
+                 means_by_fold=means_by_fold,
+                 sd_by_fold=sds_by_fold)
+
+    print('\n---')
+    print('Input features stats were saved to {}'.format(stats_fullpath))
+    print('Execution time {} seconds\n---\n'.format(
+          time.time() - start_time))
+
+
+def compute_feature_mean(i, x):
+    """
+    i is the snp nb (from 0 to nb_snps-1)
+    x is the vector of genotypes for the snp (dim : nb_samples x 1)
+    """
+    # Set to false missing value (missing values are encoded with -1)
+    mask = (x >= 0)
+
+    # Compute mean :
+    # in numerator, missing values are 0
+    # in the denominator we divise by the nb of non missing values
+    mean = np.sum(x*mask) / np.sum(mask)
+
+    return (i,mean)
+
+
+def compute_feature_mean_and_sd(i, x):
+    """
+    i is the snp nb (from 0 to nb_snps-1)
+    x is the vector of genotypes for the snp (dim : nb_samples x 1)
+    """
+    # Compute mean
+    _,mean = compute_feature_mean(i,x)
+
+    # --- Compute sd ---
+    mask = (x >= 0) #set missing values (-1) to False
+
+    numerator = np.sum((x*mask - mean*mask)**2)
+
+    denominator = np.sum(mask) - 1
+
+    sd = np.sqrt(numerator/denominator)
+
+    sd += 1e-6
+
+    return (i,mean,sd)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+            description=('Compute features means and standard deviations '
+                         'for missing values filing and input normalization '
+                         'at training time')
+            )
+
+    # Path
+    parser.add_argument(
+            '--exp-path',
+            type=str,
+            required=True,
+            help=('Path to directory where input features statistics will '
+                  'be written')
+            )
+
+    # Files
+    parser.add_argument(
+            '--dataset',
+            type=str,
+            required=True,
+            help=('Hdf5 dataset created with create_dataset.py. '
+                  'Provide full path')
+            )
+
+    parser.add_argument(
+            '--partition',
+            type=str,
+            required=True,
+            help=('Npz dataset partition returned by partition_data.py '
+                  'Provide full path')
+            )
+
+    # Which statistic to compute (mean or mean+sd)
+    parser.add_argument(
+            '--mean-only',
+            action='store_true',
+            help=('Use this flag to compute only input features means '
+                  'and not the standard deviations')
+            )
+
+    parser.add_argument(
+            '--ncpus',
+            type=int,
+            default=1,
+            help=('Number of cpus for parallel computation of means '
+                  'and of standard deviations. Default: %(default)i')
+            )
+
+    # Output
+    parser.add_argument(
+            '--out',
+            type=str,
+            help=('Optional filename for input features statistics file. '
+                  'If not provided the file will be named '
+                  'input_features_stats_datasetFilename_date')
+            )
+
+    return parser.parse_args()
+
+
+if __name__ == '__main__':
+    main()

dietnet/_upstream/Dietnet_train_and_evaluate/generate_embedding.py

@@ -0,0 +1,216 @@
+import argparse
+import os
+import time
+from datetime import datetime
+import multiprocessing as mp
+
+import numpy as np
+
+import h5py
+
+import helpers.dataset_utils as du
+
+
+# Additive encoding : genotypes 0, 1 and 2
+NB_POSSIBLE_GENOTYPES = 3
+
+# Class label key in dataset file
+CLASS_LABEL = 'class_labels'
+
+
+def main():
+    start_time = time.time()
+    args = parse_args()
+
+    # Hdf5 dataset
+    dataset_file = h5py.File(args.dataset, 'r')
+
+    # Indices by fold
+    partition_data = np.load(args.partition, allow_pickle=True)
+    indices_byfold = partition_data['folds_indexes']
+
+    print('\n---')
+    print('Computing embedding by fold using {} cpu(s)'.format(args.ncpus))
+    if args.ncpus > 1:
+        pool = mp.Pool(processes=args.ncpus) #this starts ncpus nb of processes
+
+    # Results that will be returned by processes (embedding for a given fold)
+    results = []
+
+    for fold in range(len(indices_byfold)):
+        # Fold indices
+        fold_indices = indices_byfold[fold]
+        # Get fold data
+        (_,_,_,
+         x_train, y_train, _,
+         x_valid, y_valid, _,
+         x_test, y_test,_) = du.get_fold_data(fold, indices_byfold, dataset_file, label=CLASS_LABEL)
+
+        # Embedding on train+valid sets
+        if args.include_valid:
+            print('Including valid samples in emb computation')
+            x = np.concatenate((x_train, x_valid))
+            y = np.concatenate((y_train, y_valid))
+        # Embedding on valid set
+        elif args.only_valid:
+            x = x_valid
+            y = y_valid
+        # Embedding on test set
+        elif args.only_test:
+            x = x_test
+            y = y_test
+        # Embedding on training set
+        else:
+            x = x_train
+            y = y_train
+
+        # Compute embedding
+        #emb = compute_fold_embedding(x, y)
+
+        print('Computing embedding using {} samples'.format(len(y)))
+        # One cpu
+        if args.ncpus == 1:
+            results.append(compute_fold_embedding(x, y, fold))
+        else:
+            results.append(pool.apply_async(
+                compute_fold_embedding, args=(x,y,fold)))
+
+    # Close access to dataset file
+    dataset_file.close()
+
+    # Results values: [(fold0, embedding0), (fold1, embedding1) ...]
+    print('Computed embedding of every fold, getting the results')
+    if args.ncpus > 1:
+        results_values = [p.get() for p in results]
+    else:
+        results_values = results
+
+    # Order results by fold
+    print('Putting embeddings by fold order')
+    folds_order = np.array([i[0] for i in results_values])
+
+    # Order counts (nb samples) per class
+    classes_counts = np.array([i[1] for i in results_values])
+    classes_counts_ordered = classes_counts[folds_order]
+
+    # Order embeddings
+    embeddings = np.array([i[2] for i in results_values])
+    embeddings_ordered = embeddings[folds_order]
+
+    print('Computed embedding of {} folds with following '
+            'number of samples per class for each fold: {}'.format(
+              len(folds_order), classes_counts_ordered))
+    print('---')
+
+    # Write embedding to file
+    if args.out is not None:
+        emb_filename = args.out
+    else:
+        emb_filename = 'genotype_class_freq_embedding_' \
+                + args.dataset.split('/')[-1][0:-5] \
+                + '_' \
+                + datetime.now().strftime("%Y_%m_%d_%Hh%Mmin%Ssec")
+
+    emb_fullpath = os.path.join(args.exp_path, emb_filename)
+
+    print('\n---\nEmbeddings generated in {} seconds'.format(
+        time.time()-start_time))
+
+    print('Saving results to {}'.format(emb_fullpath))
+    np.savez(emb_fullpath, emb=embeddings_ordered)
+
+    print('---\n')
+
+
+def compute_fold_embedding(xs, ys, fold):
+    # Total number of classes
+    nb_class = ys.max() + 1 #class 0
+    class_to_count = {}
+
+    # Compute sum of genotypes (0-1-2) per class
+    xs = xs.transpose() # rows are snps, col are inds
+    embedding = np.zeros((xs.shape[0],nb_class*NB_POSSIBLE_GENOTYPES))
+    for c in range(nb_class):
+        # Select genotypes for samples of same class
+        class_genotypes = xs[:,ys==c]
+        nb = class_genotypes.shape[1] #nb of samples in that class
+        #print('Class:', c, 'NB:', nb)
+        class_to_count[c] = nb
+        for genotype in range(NB_POSSIBLE_GENOTYPES):
+            col = NB_POSSIBLE_GENOTYPES*c+genotype
+            embedding[:,col] = (class_genotypes == genotype).sum(axis=1)/nb
+
+    return fold, class_to_count, embedding
+
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+            description='Generate embedding'
+            )
+
+    parser.add_argument(
+            '--exp-path',
+            type=str,
+            required=True,
+            help='Path to directory where embedding will be saved. '
+            )
+
+    parser.add_argument(
+            '--dataset',
+            type=str,
+            required=True,
+            help=('Hdf5 dataset created with create_dataset.py '
+                  'Provide full path')
+            )
+
+    parser.add_argument(
+            '--partition',
+            type=str,
+            required=True,
+            help=('Npz dataset partition returned by partition_data.py '
+                  'Provide full path')
+            )
+
+    parser.add_argument(
+            '--ncpus',
+            type=int,
+            default=1,
+            help=('Number of cpus available to compute folds embedding in '
+                  'parallel. Default:%(default)i')
+            )
+
+    parser.add_argument(
+            '--include-valid',
+            action='store_true',
+            help=('Use this flag to include samples from the validation set '
+                  'in the embedding computation.')
+            )
+
+    parser.add_argument(
+            '--only-valid',
+            action='store_true',
+            help=('Use this flag to compute embedding only on samples '
+                  ' in the validation set')
+            )
+
+    parser.add_argument(
+            '--only-test',
+            action='store_true',
+            help=('Use this flag to compute embedding only on samples '
+                  ' in the test set')
+            )
+
+    parser.add_argument(
+            '--out',
+            type=str,
+            help=('Optional filename for embedding file. If not provided '
+                  'the file will be named '
+                  'genotype_class_freq_embedding_datasetFilename_date')
+            )
+
+    return parser.parse_args()
+
+
+if __name__ == '__main__':
+    main()

dietnet/_upstream/Dietnet_train_and_evaluate/helpers/dataset_utils.py

@@ -0,0 +1,200 @@
+import math
+
+import numpy as np
+
+import h5py
+
+import torch
+
+class FoldDataset(torch.utils.data.Dataset):
+    # These variables are set in train.py
+    dataset_file = None #path to h5py dataset file
+    f = None # Dataset file handler (h5py.File in reading mode)
+    task_handler = None # Classification or regression
+    data_x = None
+    data_y = None
+    data_samples = None
+
+    def __init__(self, set_indexes):
+        self.set_indexes = set_indexes
+
+    def __len__(self):
+        return len(self.set_indexes)
+
+    def __getitem__(self, index):
+        # dataloader index goes from 0 to length of set
+        # data_x, data_y, data_samples groups data from all sets (train + valid + test)
+        # to get the item corresponding to dataloader index, we get the index of the sample
+        # in the whole dataset
+        file_index = self.set_indexes[index]
+
+        # Input features
+        #x = np.array(self.f['inputs'][file_index], dtype=np.int8)
+        x = self.data_x[file_index]
+
+
+        # Label
+        #y = self.task_handler.get_label(self.f, file_index)
+        y = self.data_y[file_index]
+        """
+        if self.task == 'classification':
+            #y = (self.f['class_labels'][file_index]).astype(np.int64)
+            y = np.array(self.f['class_labels'][file_index],
+                         dtype=np.int64)
+        elif self.task == 'regression':
+            #y = (self.f['regression_labels'][file_index]).astype(np.float32)
+            y = np.array(self.f['regression_labels'][file_index],
+                         dtype=np.float32)
+        """
+        # Sample id
+        #sample = (self.f['samples'][file_index]).astype(np.str_)
+        # SAMPLES ARE INDEX in dataset file (hdf5)
+        sample = file_index
+
+
+        return x, y, sample
+
+    def get_samples(self):
+        indexes = np.sort(self.set_indexes)
+        samples = (self.f['samples'][indexes]).astype(np.str_)
+
+        return samples
+
+class FoldDataset_forBigdataset(torch.utils.data.Dataset):
+    # These variables are set in train.py
+    dataset_file = None #path to h5py dataset file
+    f = None # Dataset file handler (h5py.File in reading mode)
+    task_handler = None # Classification or regression
+
+    def __init__(self, set_indexes):
+        self.set_indexes = set_indexes
+
+    def __len__(self):
+        return len(self.set_indexes)
+
+    def __getitem__(self, index):
+        # Data of all sets (train, valid, test) is in one file
+        # so we convert the index to match file index
+        file_index = self.set_indexes[index]
+
+        # Input features
+        x = np.array(self.f['inputs'][file_index], dtype=np.int8)
+
+        # Label
+        y = self.task_handler.get_label(self.f, file_index)
+        """
+        if self.task == 'classification':
+            #y = (self.f['class_labels'][file_index]).astype(np.int64)
+            y = np.array(self.f['class_labels'][file_index],
+                         dtype=np.int64)
+        elif self.task == 'regression':
+            #y = (self.f['regression_labels'][file_index]).astype(np.float32)
+            y = np.array(self.f['regression_labels'][file_index],
+                         dtype=np.float32)
+        """
+        # Sample id
+        #sample = (self.f['samples'][file_index]).astype(np.str_)
+        # SAMPLES ARE INDEX
+        sample = np.array(index)
+
+
+        return x, y, sample
+
+    def get_samples(self):
+        indexes = np.sort(self.set_indexes)
+        samples = (self.f['samples'][indexes]).astype(np.str_)
+
+        return samples
+
+
+class IndepTestDataset(torch.utils.data.Dataset):
+    
+    # These variables are set in test_independent_dataset.py
+    dataset_file = None # path to h5py dataset file
+    f = None # Dataset file handler (h5py.File in reading mode)
+    task_handler = None # Classification or regression
+    data_x = None
+
+    def __init__(self, samples):
+        self.set_indexes = np.arange(len(samples))
+
+    def __len__(self):
+        return len(self.set_indexes)
+
+    def __getitem__(self, index):
+        # Input features
+        #x = np.array(self.f['inputs'][file_index], dtype=np.int8)
+        x = self.data_x[index]
+
+        return index,x
+
+
+
+def get_fold_data(which_fold, folds_indexes, data, label='labels'):
+    # Indices of each set for the fold (0:train, 1:valid, 2:test)
+    fold_indexes = folds_indexes[which_fold]
+    train_indexes = np.sort(fold_indexes[0]) # sort is a hdf5 requirement
+    valid_indexes = np.sort(fold_indexes[1])
+    test_indexes = np.sort(fold_indexes[2])
+
+    # Get data (x,y,samples) of each set (train, valid, test)
+    x_train = data['inputs'][train_indexes]
+    y_train = data[label][train_indexes]
+    samples_train = data['samples'][train_indexes]
+
+    x_valid = data['inputs'][valid_indexes]
+    y_valid = data[label][valid_indexes]
+    samples_valid = data['samples'][valid_indexes]
+
+    x_test = data['inputs'][test_indexes]
+    y_test = data[label][test_indexes]
+    samples_test = data['samples'][test_indexes]
+
+    return train_indexes, valid_indexes, test_indexes,\
+           x_train, y_train, samples_train,\
+           x_valid, y_valid, samples_valid,\
+           x_test, y_test, samples_test
+
+
+def compute_norm_values(x):
+    """
+    x is a tensor
+    """
+    # Non missing values
+    mask = (x >= 0)
+
+    # Compute mean of every column (feature)
+    with torch.no_grad():
+        per_feature_mean = torch.sum(x*mask, dim=0) / torch.sum(mask, dim=0)
+
+        print('Computed per feature mean')
+
+        # S.d. of every column (feature)
+        per_feature_sd = torch.sqrt(
+                torch.sum((x*mask-mask*per_feature_mean)**2, dim=0) / \
+                        (torch.sum(mask, dim=0) - 1)
+                        )
+        per_feature_sd += 1e-6
+
+        print('Computed per feature sd')
+
+    return per_feature_mean, per_feature_sd
+
+
+def replace_missing_values(x, per_feature_mean):
+    """
+    x and per_feature_mean are tensors
+    """
+    mask = (x >= 0)
+
+    for i in range(x.shape[0]):
+        x[i] =  mask[i]*x[i] + (~mask[i])*per_feature_mean
+
+
+def normalize(x, per_feature_mean, per_feature_sd):
+    """
+    x, per_feature_mean and per_feature_sd are tensors
+    """
+    x_norm = (x - per_feature_mean) / per_feature_sd
+
+    return x_norm