Add support for center mask scoring over entire sequence

PiperOrigin-RevId: 796402809
Change-Id: Ic45642fd777f776a36a75a0d0ccc50542cf2edd1

Author: tomwardio

src/alphagenome/models/variant_scorers.py

@@ -138,8 +138,8 @@ class BaseVariantScorer(enum.Enum):
 })
 
 SUPPORTED_WIDTHS = immutabledict.immutabledict({
-    BaseVariantScorer.CENTER_MASK: [501, 2001, 10_001, 100_001, 200_001],
-    BaseVariantScorer.GENE_MASK_SPLICING: [101, 1_001, 10_001, None],
+    BaseVariantScorer.CENTER_MASK: [None, 501, 2001, 10_001, 100_001, 200_001],
+    BaseVariantScorer.GENE_MASK_SPLICING: [None, 101, 1_001, 10_001],
 })
 
 SUPPORTED_AGGREGATIONS = immutabledict.immutabledict({
@@ -166,7 +166,8 @@ class CenterMaskScorer:
 
   Attributes:
     requested_output: The requested output type (e.g., ATAC, DNASE, etc.)
-    width: The width of the mask around the variant.
+    width: The width of the mask around the variant. If None, the score is
+      computed over the entire sequence.
     aggregation_type: The aggregation type.
     base_variant_scorer: The base variant scorer.
     name: The name of the scorer (a composite of the above attributes that
@@ -183,7 +184,7 @@ class CenterMaskScorer:
   """
 
   requested_output: dna_output.OutputType
-  width: int
+  width: int | None
   aggregation_type: AggregationType
 
   @property

src/alphagenome/models/variant_scorers_test.py

@@ -64,6 +64,15 @@ def test_variant_scorers_to_proto(self):
                 requested_output=dna_model_pb2.OUTPUT_TYPE_PROCAP,
             )
         ),
+        dna_model_pb2.VariantScorer(
+            center_mask=dna_model_pb2.CenterMaskScorer(
+                width=None,
+                aggregation_type=(
+                    dna_model_pb2.AggregationType.AGGREGATION_TYPE_DIFF_LOG2_SUM
+                ),
+                requested_output=dna_model_pb2.OUTPUT_TYPE_ATAC,
+            )
+        ),
         dna_model_pb2.VariantScorer(
             gene_mask=dna_model_pb2.GeneMaskLFCScorer(
                 requested_output=dna_model_pb2.OUTPUT_TYPE_RNA_SEQ,
@@ -93,6 +102,11 @@ def test_variant_scorers_to_proto(self):
             aggregation_type=(variant_scorers.AggregationType.ACTIVE_MEAN),
             requested_output=dna_output.OutputType.PROCAP,
         ),
+        variant_scorers.CenterMaskScorer(
+            width=None,
+            aggregation_type=(variant_scorers.AggregationType.DIFF_LOG2_SUM),
+            requested_output=dna_output.OutputType.ATAC,
+        ),
         variant_scorers.GeneMaskLFCScorer(
             requested_output=dna_output.OutputType.RNA_SEQ,
         ),

src/alphagenome/protos/dna_model.proto

@@ -295,7 +295,7 @@ message IntervalScorer {
 // Variant scorer message for center mask scoring.
 message CenterMaskScorer {
   // The width of the mask around the variant.
-  int64 width = 1;
+  optional int64 width = 1;
 
   // The aggregation type.
   AggregationType aggregation_type = 2;