Fix AF2 test with trucated MSA

Author: eachanjohnson

test/scripts/test-af2.sh

@@ -22,14 +22,10 @@ python -c '
 from yunta.structs.msa import MSA
 
 def make_stub(f):
-    msa = MSA.from_file(f)
+    msa = MSA.from_file(f).truncate(20)
     # Write a cropped version — first 60 columns
     with open(f.split(".")[0] + "_stub.a3m", "w") as f:
-        for i, line in enumerate(msa.lines):
-            line.sequence = line.sequence[:20]
-            print(str(line), file=f)
-            if i > 40:
-                break
+        msa.write(f)
 
 make_stub("test/inputs/DYR_YEAST.a3m")
 make_stub("test/inputs/CAPZA_YEAST.a3m")

yunta/structs/msa.py

@@ -250,6 +250,16 @@ def sequences(self) -> list[str]:
     def gap_fraction(self) -> list[float]:
         return [line.gap_fraction for line in self.lines]
 
+    def truncate(self, n: int) -> 'MSA':
+        return replace(self, lines=[
+            MSALine(
+                sequence=line.sequence[:n],
+                description=str(line.description),
+                name=str(line.name),
+            ) 
+            for line in self.lines
+        ])
+
     @classmethod
     def from_file(cls, file: str | TextIOWrapper) -> 'MSA':
         from bioino import FastaCollection
@@ -331,32 +341,31 @@ def write(self, file=sys.stdout) -> None:
         return None
 
 
+@dataclass
 class PairedMSA(MSA):
-
     """Paired MSA object which can be used for co-evolutionary analyses.
     """
+    chain_a_length: int
+    chain_b_length: int = field(init=False)
 
-    def __init__(self, 
-                 chain_a_length: int, 
-                 *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.chain_a_length = chain_a_length
+    def __post_init__(self):
+        super().__post_init__()
         self.chain_b_length = self.seq_length - self.chain_a_length
 
     def split(
         self
     ):
         msa1 = MSA([
             MSALine(
-                sequence=line.sequence[self.chain_a_length:],
+                sequence=line.sequence[:self.chain_a_length],
                 description=str(line.description[0]),
                 name=str(line.name[0]),
             ) 
             for line in self.lines
         ])
         msa2 = MSA([
             MSALine(
-                sequence=line.sequence[:self.chain_a_length],
+                sequence=line.sequence[self.chain_a_length:],
                 description=str(line.description[1]),
                 name=str(line.name[1]),
             )