Update run_deepsomatic

PiperOrigin-RevId: 574972035
yinshiyi · Oct 19, 2023 · d970d98 · d970d98
1 parent 7907512
commit d970d98
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 7 deletions.
diff --git a/scripts/run_deepsomatic.py b/scripts/run_deepsomatic.py
@@ -238,7 +238,10 @@
 )
 
 
-MODEL_TYPE_MAP = {}
+MODEL_TYPE_MAP = {
+    'WGS': '/opt/models/deepsomatic/wgs',
+    'PACBIO': '/opt/models/deepsomatic/pacbio',
+}
 
 # Current release version of DeepVariant.
 # Should be the same in dv_vcf_constants.py.
@@ -451,6 +454,7 @@ def postprocess_variants_command(
   command.extend(['--ref', '"{}"'.format(ref)])
   command.extend(['--infile', '"{}"'.format(infile)])
   command.extend(['--outfile', '"{}"'.format(outfile)])
+  command.extend(['--cpus 0'])
   if process_somatic:
     command.extend(['--process_somatic=true'])
   else:
@@ -532,10 +536,6 @@ def check_flags():
     if (
         not tf.io.gfile.exists(_CUSTOMIZED_MODEL.value + '.data-00000-of-00001')
         or not tf.io.gfile.exists(_CUSTOMIZED_MODEL.value + '.index')
-        or (
-            not _USE_SLIM_MODEL.value
-            and not tf.io.gfile.exists(_CUSTOMIZED_MODEL.value + '.meta')
-        )
     ):
       raise RuntimeError(
           'The model files {}* do not exist. Potentially '
@@ -560,8 +560,7 @@ def get_model_ckpt(model_type, customized_model):  # pylint: disable=unused-argu
   if customized_model is not None:
     return customized_model
   else:
-    raise ValueError('We have not released a DeepSomatic model yet.')
-    # return MODEL_TYPE_MAP[model_type]
+    return MODEL_TYPE_MAP[model_type]
 
 
 def create_all_commands_and_logfiles(

diff --git a/scripts/run_deepsomatic_test.py b/scripts/run_deepsomatic_test.py
@@ -100,6 +100,7 @@ def test_basic_commands(self, model_type, use_slim_model):
             '--infile '
             '"/tmp/deepsomatic_tmp_output/call_variants_output.tfrecord.gz" '
             '--outfile "your_vcf" '
+            '--cpus 0 '
             '--process_somatic=true '
             '--gvcf_outfile "your_gvcf" '
             '--nonvariant_site_tfrecord_path '
@@ -167,6 +168,7 @@ def test_process_somatic(self, model_type, use_slim_model, process_somatic):
               '--infile '
               '"/tmp/deepsomatic_tmp_output/call_variants_output.tfrecord.gz" '
               '--outfile "your_vcf" '
+              '--cpus 0 '
               '--process_somatic=true '
               '--gvcf_outfile "your_gvcf" '
               '--nonvariant_site_tfrecord_path '
@@ -182,6 +184,7 @@ def test_process_somatic(self, model_type, use_slim_model, process_somatic):
               '--infile '
               '"/tmp/deepsomatic_tmp_output/call_variants_output.tfrecord.gz" '
               '--outfile "your_vcf" '
+              '--cpus 0 '
               '--noprocess_somatic '
               '--gvcf_outfile "your_gvcf" '
               '--nonvariant_site_tfrecord_path '
@@ -248,6 +251,7 @@ def test_sample_name_command(self, sample_name_tumor, sample_name_normal):
         '--infile '
         '"/tmp/deepsomatic_tmp_output/call_variants_output.tfrecord.gz" '
         '--outfile "your_vcf" '
+        '--cpus 0 '
         '--process_somatic=true',
     )
 
@@ -448,6 +452,7 @@ def test_postprocess_variants_extra_args(
         '--infile '
         '"/tmp/deepsomatic_tmp_output/call_variants_output.tfrecord.gz" '
         '--outfile "your_vcf" '
+        '--cpus 0 '
         '--process_somatic=true '
         '--gvcf_outfile "your_gvcf" '
         '--nonvariant_site_tfrecord_path '