Epi json filter annotation

src/jobs/PlinkFilterSNPAnnotation.cpp

@@ -0,0 +1,76 @@
+//
+// Created by juli on 09.01.24.
+//
+
+#include "PlinkFilterSNPAnnotation.hpp"
+#include "../util/TimeLogger.hpp"
+#include "../util/helper_functions.hpp"
+
+namespace epi {
+    PlinkFilterSNPAnnotation::PlinkFilterSNPAnnotation(std::vector<std::string> annotations, bool exclude_annotations,
+                                                       std::string input_path, std::string output_path,
+                                                       std::string ext_path, int num_threads) {
+
+        this->annotations = annotations;
+        this->exclude_annotations = exclude_annotations;
+
+        this->input_path = input_path;
+        this->output_path = output_path;
+        this->ext_path = ext_path;
+        this->num_threads = num_threads;
+    }
+
+    void PlinkFilterSNPAnnotation::run(std::shared_ptr<DataModel> data) {
+        TimeLogger logger ("filter for annotations");
+
+        Logger::logLine("find entries that have at least one of the selected annotations");
+
+        std::unordered_set<std::string> with_annotation_names;
+        { // delete with_annotation after this block
+            std::unordered_set<SNP_t, SNP_t::SNPHash> with_annotation;
+            for (const auto &anno: annotations) {
+                auto snps = data->snpStorage->by_annotation(anno);
+                with_annotation.insert(snps.begin(), snps.end());
+            }
+            with_annotation_names.reserve(with_annotation.size());
+            for (const auto &snp: with_annotation) {
+                with_annotation_names.insert(data->snpStorage->snp_get_name(snp));
+            }
+        }
+
+        CSVParser geno_parser;
+        geno_parser.parse(input_path + ".bim", '\t');
+
+        std::ofstream exclude_file (output_path + ".withanno");
+        // find all entries that do not start with rs...
+        for (size_t i = 0; i < geno_parser.num_rows(); ++i) {
+            const auto& name = geno_parser.cell(i, 1);
+            if (with_annotation_names.contains(name)) {
+                exclude_file << name << '\n';
+            }
+        }
+        exclude_file.close();
+
+        if (exclude_annotations) {
+            Logger::logLine("Exclude those variants from the data");
+        } else {
+            Logger::logLine("Keep those variants from the data");
+        }
+        run_plink_command(ext_path, {
+                "--threads",
+                std::to_string(num_threads),
+                "--bfile",
+                input_path,
+                exclude_annotations ? "--exclude" : "--extract",
+                output_path + ".withanno",
+                "--noweb",
+                "--make-bed",
+                "--out",
+                output_path
+        });
+
+        remove_plink_temp_files(output_path, { ".log", ".hh", ".withanno", ".nosex" });
+
+        logger.stop();
+    }
+} // epi

src/jobs/PlinkFilterSNPAnnotation.hpp

@@ -0,0 +1,33 @@
+//
+// Created by juli on 09.01.24.
+//
+
+#ifndef NEEDL_PLINKFILTERSNPANNOTATION_HPP
+#define NEEDL_PLINKFILTERSNPANNOTATION_HPP
+
+#include "Job.hpp"
+
+namespace epi {
+
+    class PlinkFilterSNPAnnotation : public Job {
+    public:
+        PlinkFilterSNPAnnotation(std::vector<std::string> annotations, bool exclude_annotations, std::string input_path, std::string output_path, std::string ext_path, int num_threads);
+        void run(std::shared_ptr<DataModel> data) override;
+
+
+    private:
+        std::vector<std::string> annotations;
+        bool exclude_annotations;
+        std::string input_path;
+        std::string output_path;
+        std::string ext_path;
+        int num_threads;
+    };
+
+} // epi
+
+#ifdef HEADER_ONLY
+#include "PlinkFilterSNPAnnotation.cpp"
+#endif
+
+#endif //NEEDL_PLINKFILTERSNPANNOTATION_HPP

test/model/src/epiJSON.cpp

@@ -37,6 +37,7 @@
 #include "../../../src/jobs/PlinkConvertJsonToPlink.hpp"
 #include "../../../src/jobs/PlinkConvertPlinkBinToLINDENInput.hpp"
 #include "../../../src/jobs/PlinkShufflePhenotype.hpp"
+#include "../../../src/jobs/PlinkFilterSNPAnnotation.hpp"
 
 using namespace epi;
 
@@ -143,6 +144,12 @@ int main(int argc, char **argv) {
     bool all_filters = false;
     app.add_flag("--all-filters", all_filters, "Applies all available filters.");
 
+    std::vector<std::string> annotations_keep;
+    app.add_option("--keep-annotation", annotations_keep, "Can be set multiple times. Only keeps SNPs that have at least one of the given annotations in dbSNP.");
+
+    std::vector<std::string> annotations_exclude;
+    app.add_option("--exclude-annotation", annotations_exclude, "Can be set multiple times. Excludes all SNPs that have at least one of the given annotations in dbSNP.");
+
 
     // Parse the options.
     CLI11_PARSE(app, argc, argv);
@@ -279,10 +286,30 @@ int main(int argc, char **argv) {
         current_input_file = outfile;
     }
 
-    if (filter_snp_network) {
+    if (!annotations_exclude.empty() && !annotations_keep.empty()) {
+        throw epi::Error("--keep-annotation and --exclude-annotation cannot be used together.");
+    }
+    bool do_anno_filtering = !annotations_exclude.empty() || !annotations_keep.empty();
+
+    if (filter_snp_network || do_anno_filtering) {
         // create snp network
         seq.add(std::make_shared<epi::InternalBimLoader>(current_input_file + ".bim"));
         seq.add(std::make_shared<epi::DbSNPAnnotator>(data_directory));
+    }
+
+    if (do_anno_filtering) {
+        std::vector<std::string> annotations_for_filtering;
+        annotations_for_filtering.insert(annotations_for_filtering.end(), annotations_keep.begin(), annotations_keep.end());
+        annotations_for_filtering.insert(annotations_for_filtering.end(), annotations_exclude.begin(), annotations_exclude.end());
+
+        std::string outfile = output_directory + "filtered_annotations";
+        seq.add(std::make_shared<epi::PlinkFilterSNPAnnotation>(annotations_for_filtering, !annotations_exclude.empty(), current_input_file, outfile, ext_directory, num_threads));
+        seq.add(std::make_shared<epi::PlinkCollectDatasetStats>(outfile, phenotype, "filter_annotations"));
+        seq.add(std::make_shared<epi::PlinkRemoveTempFiles>(current_input_file, std::vector<std::string>{".bim", ".bed", ".fam" }));
+        current_input_file = outfile;
+    }
+
+    if (filter_snp_network) {
         seq.add(std::make_shared<epi::SameAnnotationConnector>());
         seq.add(std::make_shared<epi::BioGridConnector>(data_directory));
         seq.add(std::make_shared<epi::NetworkStatsPrinter>(false));