force default_value to be float32

Author: jorenretel

bigwig_loader/download_example_data.py

@@ -19,20 +19,55 @@ def download_example_data() -> None:
 
 def get_reference_genome(reference_genome_path: Path = config.reference_genome) -> Path:
     compressed_file = reference_genome_path.with_suffix(".fasta.gz")
-    if reference_genome_path.exists():
-        return reference_genome_path
-    elif compressed_file.exists():
-        # subprocess.run(["bgzip", "-d", compressed_file])
-        unzip_gz_file(compressed_file, reference_genome_path)
-    else:
-        LOGGER.info("Need reference genome for tests. Downloading it from ENCODE.")
-        url = "https://www.encodeproject.org/files/GRCh38_no_alt_analysis_set_GCA_000001405.15/@@download/GRCh38_no_alt_analysis_set_GCA_000001405.15.fasta.gz"
-        urllib.request.urlretrieve(url, compressed_file)
+    if compressed_file.exists() and not reference_genome_path.exists():
         # subprocess.run(["bgzip", "-d", compressed_file])
         unzip_gz_file(compressed_file, reference_genome_path)
+
+    if (
+        reference_genome_path.exists()
+        and checksum_md5_for_path(reference_genome_path)
+        != config.reference_genome_checksum
+    ):
+        LOGGER.info(
+            f"Reference genome checksum mismatch, downloading again from {reference_genome_path}"
+        )
+        _download_genome(
+            url=config.reference_genome_url,
+            compressed_file_path=compressed_file,
+            uncompressed_file_path=reference_genome_path,
+            md5_checksum=config.reference_genome_checksum,
+        )
+
+    if not reference_genome_path.exists():
+        LOGGER.info(
+            f"Reference genome not found, downloading from {config.reference_genome_url}"
+        )
+        _download_genome(
+            url=config.reference_genome_url,
+            compressed_file_path=compressed_file,
+            uncompressed_file_path=reference_genome_path,
+            md5_checksum=config.reference_genome_checksum,
+        )
     return reference_genome_path
 
 
+def _download_genome(
+    url: str,
+    compressed_file_path: Path,
+    uncompressed_file_path: Path,
+    md5_checksum: str,
+) -> Path:
+    urllib.request.urlretrieve(url, compressed_file_path)
+    # subprocess.run(["bgzip", "-d", compressed_file])
+    unzip_gz_file(compressed_file_path, uncompressed_file_path)
+    this_checksum = checksum_md5_for_path(uncompressed_file_path)
+    if this_checksum != md5_checksum:
+        raise RuntimeError(
+            f"{uncompressed_file_path} has incorrect checksum: {this_checksum} vs. {md5_checksum}"
+        )
+    return uncompressed_file_path
+
+
 def unzip_gz_file(compressed_file_path: Path, output_file_path: Path) -> Path:
     with gzip.open(compressed_file_path, "rb") as gz_file:
         with open(output_file_path, "wb") as output_file:
@@ -52,6 +87,13 @@ def unzip_gz_file(compressed_file_path: Path, output_file_path: Path) -> Path:
 }
 
 
+def checksum_md5_for_path(path: Path, chunk_size: int = 10 * 1024 * 1024) -> str:
+    """return the md5sum"""
+    with path.open(mode="rb") as f:
+        checksum = checksum_md5(f, chunk_size=chunk_size)
+    return checksum
+
+
 def checksum_md5(f: BinaryIO, *, chunk_size: int = 10 * 1024 * 1024) -> str:
     """return the md5sum"""
     m = hashlib.md5(b"", usedforsecurity=False)
@@ -68,7 +110,7 @@ def get_example_bigwigs_files(bigwig_dir: Path = config.bigwig_dir) -> Path:
             file = bigwig_dir / fn
             if not file.exists():
                 urllib.request.urlretrieve(url, file)
-            with file.open(mode="rb") as f:
-                if checksum_md5(f) != md5:
-                    raise RuntimeError(f"{fn} has incorrect checksum!")
+            checksum = checksum_md5_for_path(file)
+            if checksum != md5:
+                raise RuntimeError(f"{fn} has incorrect checksum: {checksum} vs. {md5}")
     return bigwig_dir

bigwig_loader/intervals_to_values.py

@@ -87,11 +87,15 @@ def intervals_to_values(
         )
 
     if out is None:
+        logging.debug(f"Creating new out tensor with default value {default_value}")
+
         out = cp.full(
             (found_starts.shape[0], len(query_starts), sequence_length // window_size),
             default_value,
             dtype=cp.float32,
         )
+        logging.debug(out)
+
     else:
         logging.debug(f"Setting default value in output tensor to {default_value}")
         out.fill(default_value)
@@ -139,7 +143,7 @@ def intervals_to_values(
             sequence_length,
             max_number_intervals,
             window_size,
-            default_value,
+            cp.float32(default_value),
             default_value_isnan,
             out,
         ),
@@ -171,8 +175,10 @@ def kernel_in_python_with_window(
         int,
         int,
         int,
-        cp.ndarray,
         int,
+        float,
+        bool,
+        cp.ndarray,
     ],
 ) -> cp.ndarray:
     """Equivalent in python to cuda_kernel_with_window. Just for debugging."""
@@ -190,6 +196,8 @@ def kernel_in_python_with_window(
         sequence_length,
         max_number_intervals,
         window_size,
+        default_value,
+        default_value_isnan,
         out,
     ) = args
 
@@ -218,7 +226,7 @@ def kernel_in_python_with_window(
     print("reduced_dim")
     print(reduced_dim)
 
-    out_vector = [0.0] * reduced_dim * batch_size * num_tracks
+    out_vector = [default_value] * reduced_dim * batch_size * num_tracks
 
     for thread in range(n_threads):
         batch_index = thread % batch_size
@@ -239,7 +247,8 @@ def kernel_in_python_with_window(
 
         cursor = found_start_index
         window_index = 0
-        summation = 0
+        summation = 0.0
+        valid_count = 0
 
         # cursor moves through the rows of the bigwig file
         # window_index moves through the sequence
@@ -265,19 +274,31 @@ def kernel_in_python_with_window(
             print("start index", start_index)
 
             if start_index >= window_end:
-                print("CONTINUE")
-                out_vector[i * reduced_dim + window_index] = summation / window_size
-                summation = 0
+                if default_value_isnan:
+                    if valid_count > 0:
+                        out_vector[i * reduced_dim + window_index] = (
+                            summation / valid_count
+                        )
+                    else:
+                        out_vector[i * reduced_dim + window_index] = default_value
+                else:
+                    summation = summation + (window_size - valid_count) * default_value
+                    out_vector[i * reduced_dim + window_index] = summation / window_size
+                summation = 0.0
+                valid_count = 0
                 window_index += 1
+                print("CONTINUE")
                 continue
 
             number = min(window_end, end_index) - max(window_start, start_index)
 
-            print(
-                f"Add {number} x {track_values[cursor]} = {number * track_values[cursor]} to summation"
-            )
-            summation += number * track_values[cursor]
-            print(f"Summation = {summation}")
+            if number > 0:
+                print(
+                    f"Add {number} x {track_values[cursor]} = {number * track_values[cursor]} to summation"
+                )
+                summation += number * track_values[cursor]
+                print(f"Summation = {summation}")
+                valid_count += number
 
             print("end_index", "window_end")
             print(end_index, window_end)
@@ -292,8 +313,19 @@ def kernel_in_python_with_window(
                     print(
                         "cursor + 1 >= found_end_index \t\t calculate average, reset summation and move to next window"
                     )
-                out_vector[i * reduced_dim + window_index] = summation / window_size
-                summation = 0
+                # out_vector[i * reduced_dim + window_index] = summation / window_size
+                if default_value_isnan:
+                    if valid_count > 0:
+                        out_vector[i * reduced_dim + window_index] = (
+                            summation / valid_count
+                        )
+                    else:
+                        out_vector[i * reduced_dim + window_index] = default_value
+                else:
+                    summation = summation + (window_size - valid_count) * default_value
+                    out_vector[i * reduced_dim + window_index] = summation / window_size
+                summation = 0.0
+                valid_count = 0
                 window_index += 1
             # move cursor
             if end_index < window_end:

bigwig_loader/settings.py

@@ -34,6 +34,7 @@ class Settings(BaseSettings):
     reference_genome: Path = (
         example_data_dir / "GRCh38_no_alt_analysis_set_GCA_000001405.15.fasta"
     )
+    reference_genome_checksum: str = "a6da8681616c05eb542f1d91606a7b2f"
     bigwig_dir: Path = example_data_dir / "bigwig"
 
     def __str__(self) -> str:

tests/test_intervals_to_values_window_function.py

@@ -37,32 +37,35 @@ def test_get_values_from_intervals_window(default_value) -> None:
     assert (values == expected).all()
 
 
-@pytest.mark.parametrize("default_value", [0.0, cp.nan])
+@pytest.mark.parametrize("default_value", [0.0, cp.nan, 5.6, 10.0, 7565])
 def test_get_values_from_intervals_edge_case_1(default_value) -> None:
     """Query start is somewhere in a "gap"."""
     track_starts = cp.asarray([1, 10, 12, 16], dtype=cp.int32)
     track_ends = cp.asarray([3, 12, 16, 20], dtype=cp.int32)
     track_values = cp.asarray([20.0, 30.0, 40.0, 50.0], dtype=cp.dtype("f4"))
     query_starts = cp.asarray([6], dtype=cp.int32)
     query_ends = cp.asarray([18], dtype=cp.int32)
-    reserved = cp.zeros((1, 4), dtype=cp.dtype("<f4"))
+    # reserved = cp.zeros((1, 4), dtype=cp.dtype("<f4"))
     values = intervals_to_values(
         track_starts,
         track_ends,
         track_values,
         query_starts,
         query_ends,
         default_value=default_value,
-        out=reserved,
         window_size=3,
     )
     x = default_value
     if isnan(default_value):
         expected = cp.asarray([[x, 30.0, 40.0, 46.666668]])
+    elif default_value != 0:
+        expected = cp.asarray([[x, (x + 30.0 + 30.0) / 3, 40.0, 46.666668]])
     else:
         expected = cp.asarray([[x, 20.0, 40.0, 46.666668]])
 
+    print("expected:")
     print(expected)
+    print("actual:")
     print(values)
 
     assert (
@@ -227,7 +230,9 @@ def test_get_values_from_intervals_batch_of_2(default_value) -> None:
         expected = cp.asarray(
             [[20.0, 26.666666, 26.666666, 0.0], [23.333334, 36.666668, 0.0, 33.333332]]
         )
+    print("expected:")
     print(expected)
+    print("actual:")
     print(values)
     assert cp.allclose(values, expected, equal_nan=True)
 
@@ -621,9 +626,11 @@ def test_combinations_window_size_batch_size_n_tracks_on_random_data(
         query_ends,
         sizes=sizes,
         window_size=1,
-        default_value=default_value,
+        default_value=cp.nan,
     )
 
+    cp.nan_to_num(values_with_window_size_1, copy=False, nan=default_value)
+
     reduced_dim = sequence_length // window_size
     full_matrix = values_with_window_size_1[:, :, : reduced_dim * window_size]
     full_matrix = full_matrix.reshape(