fix(protspace): derive missing sequence length from fasta

apps/protspace/docs/annotations.md

@@ -20,10 +20,12 @@ Annotation sources have different requirements for protein identifiers:
 
 | Requirement | Sources | Works with `-f` FASTA? |
 | ----------- | ------- | ---------------------- |
-| **UniProt accession** | UniProt, Taxonomy, TED | No — accession needed |
+| **UniProt accession** | UniProt, Taxonomy, TED | No* — accession needed |
 | **Protein sequence** | InterPro, Biocentral, Pfam CLANS | Yes — provide `-f` |
 
-If your H5 keys are not valid UniProt accessions (e.g., `NCBI|...`, custom IDs), accession-dependent annotations will be empty. Sequence-dependent annotations can still work if you provide the original FASTA file with `-f`.
+\* `length` is the exception: when UniProt does not provide a sequence length, ProtSpace derives it from a matching FASTA sequence. A non-empty UniProt length remains authoritative, even when a FASTA sequence is available.
+
+If your H5 keys are not valid UniProt accessions, other accession-dependent annotations will be empty. Sequence-dependent annotations and the missing-length fallback can still work when the H5 keys match the identifiers parsed from the FASTA file supplied with `-f` (for example, the same bare accession or custom ID).
 
 ## Group Presets
 
@@ -75,7 +77,7 @@ With `--keep-tmp`, only API-fetched annotations are cached; the CSV is always re
 | `go_cc`                   | GO — Cellular Component              | `nucleus\|IDA;cytoplasm\|IEA`                                  |
 | `go_mf`                   | GO — Molecular Function              | `DNA binding\|IDA;protein binding\|IEA`                        |
 | `keyword`                 | UniProt keywords                     | `KW-0002 (3D-structure);KW-0025 (Alternative splicing)`        |
-| `length`                  | Sequence length (amino acids)        | `393`                                                          |
+| `length`                  | Sequence length (UniProt-preferred; FASTA fallback when missing; `*` and `-` are not counted as residues) | `393`                                                          |
 | `protein_existence`       | Evidence level for protein existence | `Evidence at protein level`                                    |
 | `protein_families`        | First protein family                 | `Protein kinase superfamily\|ISS`                              |
 | `reviewed`                | Swiss-Prot / TrEMBL                  | `true` / `false`                                               |

apps/protspace/src/protspace/cli/annotate.py

@@ -57,14 +57,19 @@ def annotate(
 
     from protspace.data.annotations.encoding import stamp_format_version
     from protspace.data.annotations.manager import ProteinAnnotationManager
-    from protspace.data.io.fasta import is_fasta_file
-    from protspace.data.loaders.h5 import EMBEDDING_EXTENSIONS
+    from protspace.data.io.fasta import is_fasta_file, parse_fasta
+    from protspace.data.loaders.h5 import EMBEDDING_EXTENSIONS, parse_identifier
 
     # Extract identifiers from input
+    sequences = None
     if is_fasta_file(input):
         from protspace.data.loaders.query import extract_identifiers_from_fasta
 
         headers = extract_identifiers_from_fasta(input)
+        sequences = {
+            parse_identifier(header): sequence
+            for header, sequence in parse_fasta(input).items()
+        }
     elif input.suffix.lower() in EMBEDDING_EXTENSIONS:
         from protspace.data.loaders.h5 import _collect_datasets
 
@@ -100,6 +105,7 @@ def annotate(
         headers=headers,
         annotations=annotations_list,
         output_path=None,
+        sequences=sequences,
     ).to_pd()
 
     if not scores:

apps/protspace/src/protspace/cli/prepare.py

@@ -492,7 +492,10 @@ def prepare(
                     if not h5s:
                         logger.warning(f"No embedding files in: {path}")
                         continue
-                    embedding_sets.append(load_h5(h5s, name_override=name_override))
+                    emb_set = load_h5(h5s, name_override=name_override)
+                    if fasta_for_similarity:
+                        emb_set.fasta_path = fasta_for_similarity
+                    embedding_sets.append(emb_set)
                 elif path.suffix.lower() in EMBEDDING_EXTENSIONS:
                     emb_set = load_h5([path], name_override=name_override)
                     # Attach FASTA path from -f flag if provided (for sequence reuse)

apps/protspace/src/protspace/data/annotations/manager.py

@@ -39,6 +39,18 @@
 logger = logging.getLogger(__name__)
 
 
+def _resolve_fasta_sequence_length(
+    identifier: str,
+    length: object,
+    sequences: dict[str, str] | None,
+) -> object:
+    """Return a FASTA-derived length only when the existing value is empty."""
+    sequence = sequences.get(identifier, "") if sequences else ""
+    if length or not sequence:
+        return length
+    return str(sum(character not in "*-" for character in sequence))
+
+
 class ProteinAnnotationManager:
     """Orchestrator for protein annotation extraction workflow."""
 
@@ -132,6 +144,7 @@ def to_pd(self) -> pd.DataFrame:
             if self.sources_to_fetch["uniprot"]
             else cached_uniprot
         )
+        uniprot_annotations = self._fill_missing_fasta_lengths(uniprot_annotations)
         taxonomy_annotations = (
             self._fetch_taxonomy(uniprot_annotations, failed_sources)
             if self.sources_to_fetch["taxonomy"]
@@ -199,6 +212,36 @@ def to_pd(self) -> pd.DataFrame:
 
         return df
 
+    def _fill_missing_fasta_lengths(
+        self, proteins: list[ProteinAnnotations]
+    ) -> list[ProteinAnnotations]:
+        """Fill empty sequence lengths from matching local FASTA sequences."""
+        if not self.sequences:
+            return proteins
+
+        result = []
+        for protein in proteins:
+            length = protein.annotations.get("length")
+            resolved_length = _resolve_fasta_sequence_length(
+                protein.identifier,
+                length,
+                self.sequences,
+            )
+            if resolved_length == length:
+                result.append(protein)
+                continue
+
+            result.append(
+                ProteinAnnotations(
+                    identifier=protein.identifier,
+                    annotations={
+                        **protein.annotations,
+                        "length": resolved_length,
+                    },
+                )
+            )
+        return result
+
     def _fetch_uniprot(self, failed_sources: list) -> list[ProteinAnnotations]:
         """Fetch UniProt annotations."""
         try:
@@ -210,9 +253,13 @@ def _fetch_uniprot(self, failed_sources: list) -> list[ProteinAnnotations]:
         except Exception as e:
             failed_sources.append(f"UniProt ({str(e)})")
             logger.warning(f"Failed to retrieve UniProt annotations: {e}")
-            # Create minimal annotation set with just identifiers
+            # Preserve the same row schema as normal UniProt responses so later
+            # formatting cannot drop columns based on which protein comes first.
             return [
-                ProteinAnnotations(identifier=header, annotations={"organism_id": ""})
+                ProteinAnnotations(
+                    identifier=header,
+                    annotations=dict.fromkeys(UNIPROT_ANNOTATIONS, ""),
+                )
                 for header in self.headers
             ]
 

apps/protspace/src/protspace/data/processors/pipeline.py

@@ -350,7 +350,10 @@ def _fetch_annotations(
         self, headers: list[str], embedding_sets: list[EmbeddingSet] = None
     ) -> pd.DataFrame:
         """Fetch annotations from APIs with incremental caching support."""
-        from protspace.data.annotations.manager import ProteinAnnotationManager
+        from protspace.data.annotations.manager import (
+            ProteinAnnotationManager,
+            _resolve_fasta_sequence_length,
+        )
 
         # Extract sequences from FASTA files (if available) to avoid re-fetching
         sequences = self._extract_sequences(embedding_sets) if embedding_sets else {}
@@ -419,6 +422,22 @@ def _fetch_annotations(
                     else:
                         api_df = cached_df
 
+                    if "length" in api_df.columns and sequences:
+                        identifier_col = api_df.columns[0]
+                        api_df = api_df.copy()
+                        api_df["length"] = [
+                            _resolve_fasta_sequence_length(
+                                identifier,
+                                length,
+                                sequences,
+                            )
+                            for identifier, length in zip(
+                                api_df[identifier_col],
+                                api_df["length"],
+                                strict=True,
+                            )
+                        ]
+
                     # Warn if cached annotations are all empty
                     data_cols = [c for c in api_df.columns if c != "identifier"]
                     if data_cols:

apps/protspace/tests/test_annotate_cli.py

@@ -0,0 +1,44 @@
+import pandas as pd
+from typer.testing import CliRunner
+
+from protspace.cli.app import app
+from protspace.data.annotations.retrievers.uniprot_retriever import (
+    ProteinAnnotations,
+    UniProtRetriever,
+)
+
+
+def test_annotate_fasta_derives_missing_length_from_normalized_sequence(
+    tmp_path, monkeypatch
+):
+    """The FASTA-backed CLI path must supply sequences to the annotation manager."""
+    fasta = tmp_path / "input.fasta"
+    fasta.write_text(">custom|custom_protein|description\nMPEPTIDE\n")
+    output = tmp_path / "annotations.parquet"
+
+    monkeypatch.setattr(
+        UniProtRetriever,
+        "fetch_annotations",
+        lambda self: [
+            ProteinAnnotations(
+                identifier="custom_protein",
+                annotations={"length": ""},
+            )
+        ],
+    )
+
+    result = CliRunner().invoke(
+        app,
+        [
+            "annotate",
+            "-i",
+            str(fasta),
+            "-a",
+            "length",
+            "-o",
+            str(output),
+        ],
+    )
+
+    assert result.exit_code == 0, result.output
+    assert pd.read_parquet(output).loc[0, "length"] == "8"

apps/protspace/tests/test_annotation_manager.py

@@ -439,6 +439,83 @@ def test_transform_protein_families_with_semicolon(self):
 class TestIntegration:
     """Integration tests for complete workflows."""
 
+    @patch("src.protspace.data.annotations.manager.UniProtRetriever")
+    def test_uses_fasta_sequence_length_when_uniprot_length_is_missing(
+        self, mock_uniprot_retriever
+    ):
+        mock_uniprot_retriever.return_value.fetch_annotations.return_value = [
+            ProteinAnnotations(
+                identifier="custom_protein",
+                annotations={"length": ""},
+            )
+        ]
+        extractor = ProteinAnnotationExtractor(
+            headers=["custom_protein"],
+            annotations=["length"],
+            sequences={"custom_protein": "MPEPTIDE"},
+        )
+
+        result = extractor.to_pd()
+
+        assert result.loc[0, "length"] == "8"
+
+    @patch("src.protspace.data.annotations.manager.UniProtRetriever")
+    def test_preserves_uniprot_length_when_fasta_sequence_is_available(
+        self, mock_uniprot_retriever
+    ):
+        mock_uniprot_retriever.return_value.fetch_annotations.return_value = [
+            ProteinAnnotations(
+                identifier="P01308",
+                annotations={"length": "110"},
+            )
+        ]
+        extractor = ProteinAnnotationExtractor(
+            headers=["P01308"],
+            annotations=["length"],
+            sequences={"P01308": "MPEPTIDE"},
+        )
+
+        result = extractor.to_pd()
+
+        assert result.loc[0, "length"] == "110"
+
+    @patch("src.protspace.data.annotations.manager.UniProtRetriever")
+    def test_counts_residues_not_fasta_gap_or_terminator_markers(
+        self, mock_uniprot_retriever
+    ):
+        mock_uniprot_retriever.return_value.fetch_annotations.return_value = [
+            ProteinAnnotations(
+                identifier="custom_protein",
+                annotations={"length": ""},
+            )
+        ]
+        extractor = ProteinAnnotationExtractor(
+            headers=["custom_protein"],
+            annotations=["length"],
+            sequences={"custom_protein": "M-PEP*"},
+        )
+
+        result = extractor.to_pd()
+
+        assert result.loc[0, "length"] == "4"
+
+    @patch("src.protspace.data.annotations.manager.UniProtRetriever")
+    def test_retains_length_column_when_uniprot_request_fails(
+        self, mock_uniprot_retriever
+    ):
+        mock_uniprot_retriever.return_value.fetch_annotations.side_effect = (
+            RuntimeError("offline")
+        )
+        extractor = ProteinAnnotationExtractor(
+            headers=["no_sequence", "custom_protein"],
+            annotations=["length"],
+            sequences={"custom_protein": "MPEPTIDE"},
+        )
+
+        result = extractor.to_pd()
+
+        assert result["length"].tolist() == ["", "8"]
+
     @patch("src.protspace.data.annotations.manager.TaxonomyRetriever")
     @patch("src.protspace.data.annotations.manager.UniProtRetriever")
     def test_to_pd_complete_workflow(

apps/protspace/tests/test_backend_switch.py

@@ -206,3 +206,40 @@ def test_embed_cli_rejects_unknown_backend(tmp_path):
     )
 
     assert result.exit_code != 0
+
+
+def test_prepare_directory_h5_attaches_fasta_to_embedding_set(tmp_path, monkeypatch):
+    h5_dir = tmp_path / "embeddings"
+    h5_dir.mkdir()
+    with h5py.File(h5_dir / "model.h5", "w") as h5_file:
+        h5_file.attrs["model_name"] = "prot_t5"
+        h5_file.create_dataset("custom_protein", data=np.ones(4, dtype=np.float32))
+
+    fasta = tmp_path / "input.fasta"
+    fasta.write_text(">custom_protein\nMPEPTIDE\n")
+    captured = {}
+
+    def capture_run(self, embedding_sets):
+        captured["fasta_path"] = embedding_sets[0].fasta_path
+        return self.config.output_path
+
+    monkeypatch.setattr(
+        "protspace.data.processors.pipeline.ReductionPipeline.run", capture_run
+    )
+
+    result = CliRunner().invoke(
+        app,
+        [
+            "prepare",
+            "-i",
+            str(h5_dir),
+            "-f",
+            str(fasta),
+            "-o",
+            str(tmp_path / "out"),
+            "--no-log",
+        ],
+    )
+
+    assert result.exit_code == 0, result.output
+    assert captured["fasta_path"] == fasta