Read copolymer type labels (P1*, P2*, ...) as polymer

CHANGELOG.md

@@ -1,5 +1,11 @@
 # Changelog
 
+## Unreleased
+
+### Fixed
+- Copolymers: HiTPoly type labels P1*, P2*, ... (one per monomer block) are read as polymer, as in the
+  original analysis code. Before, any label other than PL*, AN* or CA* stopped `compute_shells` with a KeyError.
+
 ## 1.0.0 — 2026-10-25
 
 First public release: a rewrite of the research code (v0.5.3) used for the paper

src/onsagerdecomposer/io.py

@@ -5,8 +5,9 @@
   xyz_unwrapped_rdf_all.txt   one line per atom per frame: x,y,z  (unwrapped, Angstrom)
   box_len.txt                 cubic box length (Angstrom)
 
-The type label prefix sets the residue kind: PL* -> "solvent" (solvent molecule or
-polymer chain), AN* -> "anion", CA* -> "cation".
+The type label prefix sets the residue kind: AN* -> "anion", CA* -> "cation", and any other
+label starting with P -> "solvent" (solvent molecule or polymer chain: PL* for one monomer type,
+P1*, P2*, ... for the blocks of a copolymer).
 """
 import warnings
 from dataclasses import dataclass
@@ -15,7 +16,15 @@
 import numpy as np
 import pandas as pd
 
-KIND_PREFIXES = {"PL": "solvent", "AN": "anion", "CA": "cation"}
+def residue_kind(label):
+    """"anion" (AN*), "cation" (CA*) or "solvent" (any other P*: PL*, or P1*, P2*, ... in copolymers)."""
+    if label.startswith("AN"):
+        return "anion"
+    if label.startswith("CA"):
+        return "cation"
+    if label.startswith("P"):
+        return "solvent"
+    raise ValueError(f"unknown HiTPoly type label {label!r}: expected P* (solvent/polymer), AN* or CA*")
 
 # Atomic masses (pymatgen values) used for the centre-of-mass correction.
 MASSES = {"H": 1.00794, "Li": 6.941, "B": 10.811, "C": 12.0107, "N": 14.0067, "O": 15.9994, "F": 18.9984032,
@@ -35,7 +44,7 @@ class Topology:
     @property
     def kinds(self):
         """(n_atoms,) "solvent", "anion" or "cation" from the type prefix."""
-        return np.array([KIND_PREFIXES[t[:2]] for t in self.types])
+        return np.array([residue_kind(t) for t in self.types])
 
     @property
     def n_atoms(self):

tests/test_io.py

@@ -0,0 +1,26 @@
+import numpy as np
+import pytest
+
+from onsagerdecomposer import iter_hitpoly_frames, read_hitpoly_topology
+from onsagerdecomposer.io import residue_kind
+
+
+def test_residue_kinds_include_copolymer_blocks():
+    assert [residue_kind(t) for t in ["PL1", "P1A", "P2B", "AN1", "CA1"]] == \
+        ["solvent", "solvent", "solvent", "anion", "cation"]
+    with pytest.raises(ValueError, match="XX1"):
+        residue_kind("XX1")
+
+
+def test_copolymer_folder(tmp_path):
+    # two copolymer blocks (P1, P2), one anion atom, one cation; 2 frames
+    (tmp_path / "atom_names_rdf_all.txt").write_text("C,P1A,1\nO,P2B,1\nN,AN1,2\nLi,CA1,3\n")
+    (tmp_path / "box_len.txt").write_text("20.0")
+    frames = np.array([[[0, 0, 0], [2, 0, 0], [5, 5, 5], [8, 8, 8]],
+                       [[1, 0, 0], [3, 0, 0], [5, 5, 5], [8, 8, 8]]], dtype=float)
+    np.savetxt(tmp_path / "xyz_unwrapped_rdf_all.txt", frames.reshape(-1, 3), delimiter=",")
+    top = read_hitpoly_topology(tmp_path)
+    assert list(top.kinds) == ["solvent", "solvent", "anion", "cation"]
+    xyz = np.concatenate(list(iter_hitpoly_frames(tmp_path, top)))
+    com = (12.0107 * xyz[:, 0] + 15.9994 * xyz[:, 1]) / (12.0107 + 15.9994)     # both blocks are solvent
+    assert np.allclose(com, 0)