Remove use of stdlib

RasmussenLab · Jul 22, 2024 · 922d968 · 922d968
1 parent 12d6dde
commit 922d968
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 4 deletions.
diff --git a/vamb/__main__.py b/vamb/__main__.py
@@ -18,7 +18,6 @@
 from torch.utils.data import DataLoader
 from functools import partial
 from loguru import logger
-from array import array
 
 _ncpu = os.cpu_count()
 DEFAULT_THREADS = 8 if _ncpu is None else min(_ncpu, 8)
@@ -1413,7 +1412,7 @@ def predict_taxonomy(
                 predicted_vector[i] > taxonomy_options.softmax_threshold
             )
             ranks = list(nodes_ar[threshold_mask][1:])
-            probs = array("f", predicted_vector[i][threshold_mask][1:])
+            probs = predicted_vector[i][threshold_mask][1:]
             tax = vamb.taxonomy.PredictedContigTaxonomy(
                 vamb.taxonomy.ContigTaxonomy(ranks), probs
             )

diff --git a/vamb/taxonomy.py b/vamb/taxonomy.py
@@ -1,7 +1,7 @@
 from typing import Optional, IO
 from pathlib import Path
 from vamb.parsecontigs import CompositionMetaData
-import array
+import numpy as np
 
 
 class ContigTaxonomy:
@@ -121,7 +121,7 @@ def parse_tax_file(
 class PredictedContigTaxonomy:
     slots = ["contig_taxonomy", "probs"]
 
-    def __init__(self, tax: ContigTaxonomy, probs: array.array[float]):
+    def __init__(self, tax: ContigTaxonomy, probs: np.ndarray):
         if len(probs) != len(tax.ranks):
             raise ValueError("The length of probs must equal that of ranks")
         self.tax = tax