Initial data generation tasks for SPA

impc_etl/jobs/ingest/data_ingestion.py

@@ -144,6 +144,32 @@ def copy_products_report():
 
     return tracking_products_mice_file_asset
 
+tracking_products_crispr_file_asset = create_input_asset(
+    "tracking/gentar-products_crispr-latest.json"
+)
+
+@asset.multi(
+    schedule=[tracking_directory_asset],
+    outlets=[tracking_products_crispr_file_asset],
+    dag_id=f"{dr_tag}_copy_products_crispr_report",
+)
+def copy_products_report():
+    """Gather tracking data from GenTar when tracking directory is available"""
+    source_file = f"{data_archive_path}/gentar-data-archive/product_reports/gentar-products_crispr-latest.json"
+    target_file = f"{input_data_path}/tracking/gentar-products_crispr.json"
+
+    task_logger.info(f"Copying tracking data from {source_file} to {target_file}")
+
+    # Ensure source file exists
+    if not os.path.exists(source_file):
+        raise FileNotFoundError(f"Source tracking file not found: {source_file}")
+
+    # Copy the file
+    shutil.copy(source_file, target_file)
+    task_logger.info(f"Successfully copied tracking data to {target_file}")
+
+    return tracking_products_mice_file_asset
+
 
 gene_interest_asset = create_input_asset("tracking/gene_interest.tsv")
 gene_interest_json_asset = create_input_asset("tracking/gene_interest.json")

impc_etl/jobs/load/impc_spa/impc_chromosome_mapper.py

@@ -0,0 +1,36 @@
+import logging
+import textwrap
+from airflow.sdk import Variable, asset
+from impc_etl.utils.airflow import create_input_asset, create_output_asset
+from impc_etl.utils.spark import with_spark_session
+
+task_logger = logging.getLogger("airflow.task")
+dr_tag = Variable.get("data_release_tag")
+
+gene_summary_service_json_asset = create_input_asset("output/impc_web_api/gene_summary_service_json")
+chromosome_map_json_asset = create_output_asset("impc_spa/chromosome-map.json")
+
+@asset.multi(
+    schedule=[gene_summary_service_json_asset],
+    outlets=[chromosome_map_json_asset],
+    dag_id=f"{dr_tag}_impc_spa_chromosome_mapper",
+    description=textwrap.dedent(
+        """IMPC SPA chromosome mapper DAG."""
+    ),
+    tags=["impc_spa", "chromosome map"],
+)
+@with_spark_session
+def impc_spa_chromosome_mapper():
+    import json
+    from pyspark.sql import SparkSession
+    from urllib.parse import unquote, urlparse
+
+    spark = SparkSession.builder.getOrCreate()
+    gene_summary_service_json_path = gene_summary_service_json_asset.uri
+    gene_summary_df = spark.read.json(gene_summary_service_json_path)
+    gene_summary_df = gene_summary_df.select("mgiGeneAccessionId", "chrName")
+    gene_list = map(lambda row: row.asDict(), gene_summary_df.collect())
+    chromosome_map_dict = {gene["mgiGeneAccessionId"]: gene["chrName"] for gene in gene_list}
+    output_path = unquote(urlparse(chromosome_map_json_asset.uri).path)
+    with open(output_path, "w") as output_file:
+        output_file.write(json.dumps(chromosome_map_dict))

impc_etl/jobs/load/impc_spa/impc_diseases_mapper.py

@@ -0,0 +1,89 @@
+import json
+import logging
+import textwrap
+from airflow.sdk import Variable, task, AssetAlias, chain, dag, Asset
+from impc_etl.utils.airflow import create_input_asset, get_data_release_work_dir
+from impc_etl.utils.impc_spa import generate_valid_json_from_file
+from impc_etl.utils.spark import with_spark_session
+
+task_logger = logging.getLogger("airflow.task")
+dr_tag = Variable.get("data_release_tag")
+
+gene_diseases_service_json_asset = create_input_asset("output/impc_web_api/gene_diseases_service_json")
+chromosome_map_json_asset = create_input_asset("output/impc_spa/chromosome-map.json")
+gene_diseases_asset = AssetAlias("impc_spa_diseases")
+
+@dag(
+    schedule=[gene_diseases_service_json_asset, chromosome_map_json_asset],
+    dag_id=f"{dr_tag}_impc_spa_diseases_mapper",
+    description=textwrap.dedent(
+        """IMPC SPA diseases mapper DAG."""
+    ),
+    tags=["impc_spa", "diseases"],
+)
+def impc_spa_gene_diseases_mapper():
+    @with_spark_session
+    @task
+    def process_gene_diseases():
+        from pyspark.sql import SparkSession
+        from pyspark.sql.functions import col, regexp_replace
+        from urllib.parse import unquote, urlparse
+
+        spark = SparkSession.builder.getOrCreate()
+
+        chromosome_map_json_path = unquote(urlparse(chromosome_map_json_asset.uri).path)
+        chromosome_map_json = json.loads(open(chromosome_map_json_path).read())
+        chromosome_map_df = spark.createDataFrame(chromosome_map_json.items(), ["mgiGeneAccessionId", "chromosome"])
+        gene_diseases_service_json_path = gene_diseases_service_json_asset.uri
+        gene_diseases_df = spark.read.json(gene_diseases_service_json_path)
+        result_df = chromosome_map_df.join(gene_diseases_df, "mgiGeneAccessionId", "left_outer")
+        result_df.show()
+        result_df = result_df.filter(col("associationCurated").isNotNull())
+        result_df = result_df.drop("chromosome")
+        result_df = result_df.withColumn("mgiGeneAccessionId", regexp_replace("mgiGeneAccessionId", ":", "_"))
+        (result_df
+         .repartition("mgiGeneAccessionId")
+         .write
+         .option("header", True)
+         .mode("overwrite")
+         .partitionBy("mgiGeneAccessionId", "associationCurated")
+         .json(f"{get_data_release_work_dir()}/output/impc_spa/gene_diseases_temp_json")
+         )
+        print("Finished")
+
+    @task(outlets=[gene_diseases_asset])
+    def process_temp_folder(*, outlet_events):
+        import os
+        import shutil
+        from glob import iglob
+        from urllib.parse import unquote, urlparse
+
+        chromosome_map_json_path = unquote(urlparse(chromosome_map_json_asset.uri).path)
+        chromosome_map_json = json.loads(open(chromosome_map_json_path).read())
+
+        input_path = f"{get_data_release_work_dir()}/output/impc_spa/gene_diseases_temp_json"
+        output_path = f"{get_data_release_work_dir()}/output/impc_spa"
+        for file_path in iglob(f"{input_path}/**/*.json", recursive=True):
+            filepath_parts = file_path.split("/")
+            filepath_parts.pop()
+            parent_dir = filepath_parts.pop()
+            association_status = parent_dir.split("=")[1]
+            parent_dir = filepath_parts.pop()
+            mgi_gene_accession_id = parent_dir.split("=")[1]
+            original_mgi_gene_accession_id = mgi_gene_accession_id.replace("_", ":")
+
+            chromosome = chromosome_map_json[original_mgi_gene_accession_id]
+            chromosome_folder = f"{output_path}/{chromosome}"
+            os.makedirs(chromosome_folder, exist_ok=True)
+
+            gene_dir_path = f"{chromosome_folder}/{mgi_gene_accession_id}"
+            os.makedirs(gene_dir_path, exist_ok=True)
+            file_name = 'associated-diseases' if association_status == "true" else 'predicted-diseases'
+            gene_external_links_path = f"{gene_dir_path}/{file_name}.json"
+            with open(gene_external_links_path, "w") as gene_file:
+                gene_file.write(generate_valid_json_from_file(file_path))
+        shutil.rmtree(input_path)
+        print("Finished")
+
+    chain(process_gene_diseases(), process_temp_folder())
+impc_spa_gene_diseases_mapper()

impc_etl/jobs/load/impc_spa/impc_gene_all_phenotype_data_mapper.py

@@ -0,0 +1,96 @@
+import logging
+import textwrap
+from airflow.sdk import Variable, task, AssetAlias, chain, dag, Asset
+from impc_etl.utils.airflow import create_input_asset, get_data_release_work_dir
+from impc_etl.utils.impc_spa import write_partitioned_data, generate_valid_json_from_file
+from impc_etl.utils.spark import with_spark_session
+
+task_logger = logging.getLogger("airflow.task")
+dr_tag = Variable.get("data_release_tag")
+
+datasets_metadata_service_json_asset = create_input_asset("output/impc_web_api/datasets_metadata_service_json")
+chromosome_map_json_asset = create_input_asset("output/impc_spa/chromosome-map.json")
+gene_all_phenotype_data_asset = AssetAlias("impc_spa_gene_all_phenotype_data")
+
+@dag(
+    schedule=[datasets_metadata_service_json_asset, chromosome_map_json_asset],
+    dag_id=f"{dr_tag}_impc_spa_all_phenotype_data_mapper",
+    description=textwrap.dedent(
+        """IMPC SPA gene all phenotype data mapper DAG."""
+    ),
+    tags=["impc_spa", "gene", "all phenotype data"],
+)
+def impc_spa_gene_all_phenotype_data_mapper():
+    @with_spark_session
+    @task
+    def process_parquet():
+        from pyspark.sql import SparkSession
+        from pyspark.sql.functions import col, regexp_replace
+        from urllib.parse import unquote, urlparse
+        import json
+
+        spark = SparkSession.builder.getOrCreate()
+
+        chromosome_map_json_path = unquote(urlparse(chromosome_map_json_asset.uri).path)
+        chromosome_map_json = json.loads(open(chromosome_map_json_path).read())
+        chromosome_map_df = spark.createDataFrame(chromosome_map_json.items(), ["mgiGeneAccessionId", "chromosome"])
+
+        dataset_df = spark.read.json(datasets_metadata_service_json_asset.uri)
+        result_df = chromosome_map_df.join(dataset_df, "mgiGeneAccessionId", "left_outer")
+        result_df = result_df.filter(col("pipelineStableId").isNotNull())
+        result_df = result_df.drop("chromosome")
+        result_df = result_df.withColumn("mgiGeneAccessionId", regexp_replace("mgiGeneAccessionId", ":", "_"))
+
+        (result_df
+         .repartition("mgiGeneAccessionId")
+         .write
+         .mode("overwrite")
+         .partitionBy("mgiGeneAccessionId", "pipelineStableId", "procedureStableId")
+         .json(f"{get_data_release_work_dir()}/output/impc_spa/all_ph_data_temp_json")
+         )
+
+        print("Finished")
+
+    @task(outlets=[gene_all_phenotype_data_asset])
+    def process_temp_folder(*, outlet_events):
+        import json
+        import os
+        import shutil
+        from glob import iglob
+        from urllib.parse import unquote, urlparse
+
+        chromosome_map_json_path = unquote(urlparse(chromosome_map_json_asset.uri).path)
+        chromosome_map_json = json.loads(open(chromosome_map_json_path).read())
+
+        input_path = f"{get_data_release_work_dir()}/output/impc_spa/all_ph_data_temp_json"
+        output_path = f"{get_data_release_work_dir()}/output/impc_spa"
+        for file_path in iglob(f"{input_path}/**/*.json", recursive=True):
+            filepath_parts = file_path.split("/")
+            filepath_parts.pop()
+            parent_dir = filepath_parts.pop()
+            procedure_stable_id = parent_dir.split("=")[1]
+            parent_dir = filepath_parts.pop()
+            pipeline_stable_id = parent_dir.split("=")[1]
+            parent_dir = filepath_parts.pop()
+            mgi_gene_accession_id = parent_dir.split("=")[1]
+            original_mgi_gene_accession_id = mgi_gene_accession_id.replace("_", ":")
+
+            chromosome = chromosome_map_json[original_mgi_gene_accession_id]
+            chromosome_folder = f"{output_path}/{chromosome}"
+            os.makedirs(chromosome_folder, exist_ok=True)
+
+            gene_dir_path = f"{chromosome_folder}/{mgi_gene_accession_id}"
+            os.makedirs(gene_dir_path, exist_ok=True)
+            general_gene_images_path = f"{gene_dir_path}/pipeline"
+            os.makedirs(general_gene_images_path, exist_ok=True)
+            pipeline_dir_path = f"{general_gene_images_path}/{pipeline_stable_id}"
+            os.makedirs(pipeline_dir_path, exist_ok=True)
+            # write data
+            file_to_be_generated_path = f"{pipeline_dir_path}/{procedure_stable_id}.json"
+            with open(file_to_be_generated_path, "w") as dataset_file:
+                dataset_file.write(generate_valid_json_from_file(file_path))
+        shutil.rmtree(input_path)
+        print("Finished")
+
+    chain(process_parquet(), process_temp_folder())
+impc_spa_gene_all_phenotype_data_mapper()

impc_etl/jobs/load/impc_spa/impc_gene_allele_crispr_mapper.py

@@ -0,0 +1,102 @@
+import logging
+import textwrap
+from airflow.sdk import Variable, task, dag, chain
+
+from impc_etl.utils.airflow import create_input_asset, get_data_release_work_dir
+from impc_etl.utils.impc_spa import create_gentar_crispr_report_df
+from impc_etl.utils.spark import with_spark_session
+
+task_logger = logging.getLogger("airflow.task")
+dr_tag = Variable.get("data_release_tag")
+
+gentar_products_crispr_latest_json_output_asset = create_input_asset("tracking/gentar-products_crispr.json")
+chromosome_map_json_asset = create_input_asset("output/impc_spa/chromosome-map.json")
+
+@dag(
+    schedule=[gentar_products_crispr_latest_json_output_asset, chromosome_map_json_asset],
+    dag_id=f"{dr_tag}_impc_spa_gene_allele_crispr_mapper",
+    description=textwrap.dedent(
+        """IMPC SPA gene allele crispr mapper DAG."""
+    ),
+    tags=["impc_spa", "allele", "crispr"],
+)
+def impc_spa_gene_allele_crispr_mapper():
+    @with_spark_session
+    @task
+    def process_allele_crispr_data():
+        import json
+        from pyspark.sql import SparkSession
+        from pyspark.sql.functions import col, regexp_replace
+        from urllib.parse import unquote, urlparse
+
+        spark = SparkSession.builder.getOrCreate()
+
+        chromosome_map_json_path = unquote(urlparse(chromosome_map_json_asset.uri).path)
+        chromosome_map_json = json.loads(open(chromosome_map_json_path).read())
+        chromosome_map_df = spark.createDataFrame(chromosome_map_json.items(), ["mgiGeneAccessionId", "chromosome"])
+
+        product_crispr_df = create_gentar_crispr_report_df(spark, gentar_products_crispr_latest_json_output_asset.uri)
+
+        result_df = chromosome_map_df.join(product_crispr_df, "mgiGeneAccessionId", "left_outer")
+        result_df = result_df.filter(col("alleleSuperscript").isNotNull())
+        result_df = result_df.drop("chromosome")
+        result_df = result_df.withColumn("mgiGeneAccessionId", regexp_replace("mgiGeneAccessionId", ":", "_"))
+        (result_df
+         .repartition("mgiGeneAccessionId")
+         .write
+         .mode("overwrite")
+         .partitionBy("mgiGeneAccessionId")
+         .json(f"{get_data_release_work_dir()}/output/impc_spa/allele_crispr_temp_json")
+         )
+
+        print("Finished")
+    @task
+    def process_temp_folder():
+        import json
+        import os
+        import shutil
+        from glob import iglob
+        from urllib.parse import unquote, urlparse
+
+        chromosome_map_json_path = unquote(urlparse(chromosome_map_json_asset.uri).path)
+        chromosome_map_json = json.loads(open(chromosome_map_json_path).read())
+
+        input_path = f"{get_data_release_work_dir()}/output/impc_spa/allele_crispr_temp_json"
+        output_path = f"{get_data_release_work_dir()}/output/impc_spa"
+        for file_path in iglob(f"{input_path}/**/*.json"):
+            filepath_parts = file_path.split("/")
+            filepath_parts.pop()
+            parent_dir = filepath_parts.pop()
+            mgi_gene_accession_id = parent_dir.split("=")[1]
+            original_mgi_gene_accession_id = mgi_gene_accession_id.replace("_", ":")
+            chromosome = chromosome_map_json[original_mgi_gene_accession_id]
+            chromosome_folder = f"{output_path}/{chromosome}"
+            os.makedirs(chromosome_folder, exist_ok=True)
+            gene_dir_path = f"{chromosome_folder}/{mgi_gene_accession_id}"
+            os.makedirs(gene_dir_path, exist_ok=True)
+
+            allele_dict = {}
+            # JSON file can have multiple objects for distinct alleles
+            allele_data = open(file_path, 'r')
+            for json_obj_str in allele_data.readlines():
+                allele_obj = json.loads(json_obj_str)
+                if allele_obj["alleleSuperscript"] in allele_dict:
+                    allele_dict[allele_obj["alleleSuperscript"]].append(json_obj_str)
+                else:
+                    allele_dict[allele_obj["alleleSuperscript"]] = [json_obj_str]
+
+
+            general_alleles_dir_path = f"{gene_dir_path}/alleles"
+            os.makedirs(general_alleles_dir_path, exist_ok=True)
+            for original_allele_name, allele_json_list in allele_dict.items():
+                allele_name = original_allele_name.replace("/", "_")
+                allele_dir_path = f"{general_alleles_dir_path}/{allele_name}"
+                os.makedirs(allele_dir_path, exist_ok=True)
+                file_to_be_generated_path = f"{allele_dir_path}/crispr.json"
+                with open(file_to_be_generated_path, "w") as allele_file:
+                    allele_file.write(f"[{','.join(allele_json_list)}]")
+        shutil.rmtree(input_path)
+        print("Finished")
+
+    chain(process_allele_crispr_data(), process_temp_folder())
+impc_spa_gene_allele_crispr_mapper()