From 221231f2a52c96d92d9fd792459adf321fbbdfa9 Mon Sep 17 00:00:00 2001 From: pcbve1 Date: Thu, 23 Jul 2026 10:04:26 -0700 Subject: [PATCH 1/2] add chain_id parameter, fix atom name issue, specify torch-fourier-slice version --- pyproject.toml | 1 + src/mosaics/template_iterator.py | 60 ++++++++++++++------------------ 2 files changed, 27 insertions(+), 34 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index b708b99..1776074 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -49,6 +49,7 @@ dependencies = [ "ttsim3d", "teamtomo-basemodel", "Leopard-EM>=v1.0", + "torch-fourier-slice>=v0.4.0" ] [tool.hatch.metadata] diff --git a/src/mosaics/template_iterator.py b/src/mosaics/template_iterator.py index 7ca0a53..9845afa 100644 --- a/src/mosaics/template_iterator.py +++ b/src/mosaics/template_iterator.py @@ -358,7 +358,7 @@ def get_template_scattering_potential( atom_counts = self.structure_df.iloc[atom_idxs]["element"].value_counts() for atom, count in atom_counts.items(): atom = atom.upper() - potentials = get_a_param(atom) + potentials = get_a_param([atom]) potentials = torch.sum(potentials).item() total_scattering_potential += potentials * count @@ -433,20 +433,27 @@ def alternate_template_iter( class ChainTemplateIterator(BaseTemplateIterator): - """Iterates over each chain in the structure and removes specified atoms from it. + """Iterates over each chain or a selected chain in the structure and removes specified atoms from it. Attributes ---------- type : Literal["chain"] Discriminator field for differentiating between template iterator types. + chain_ids : list[str] + If populated, denotes the chains to be removed. When empty, sets all chains to be removed. Default is all. """ type: ClassVar[Literal["chain"]] = "chain" - + chain_ids: list[str] = ["all"] + @property def num_alternate_structures(self) -> int: """Get the number of alternate structures (i.e. number of chains).""" - return len(self.structure_df["chain"].unique()) + if self.chain_ids == ['all']: + return len(self.structure_df["chain"].unique()) + else: + return len(self.chain_ids) + def alternate_template_iter( self, inverted: bool = True @@ -459,10 +466,15 @@ def alternate_template_iter( If 'True', return the indexes of atoms to remove rather than keep. """ subset_df = self.subset_df_on_residues_and_atoms() - unique_chain_ids = subset_df["chain"].unique() + if self.chain_ids == ["all"]: + unique_chain_ids = subset_df["chain"].unique() + else: + unique_chain_ids = self.chain_ids for chain_id in unique_chain_ids: + print(f"Removing chain {chain_id}") residue_ids = subset_df[subset_df["chain"] == chain_id]["residue_id"] + residue_ids = residue_ids.unique().tolist() chain_ids = [chain_id] * len(residue_ids) @@ -497,11 +509,14 @@ class ResidueTemplateIterator(BaseTemplateIterator): Discriminator field for differentiating between template iterator types. randomize_chain_order : bool If 'True', randomize the order of chains in the structure. Default is 'False'. + chain_ids : list[str] + If populated, denotes the chains to be removed. When empty, sets all chains to be removed. Default is empty. """ type: ClassVar[Literal["residue"]] = "residue" num_residues_removed: Annotated[int, Field(gt=0)] residue_increment: Annotated[int, Field(gt=0)] + chain_ids: list[str] = ["all"] randomize_chain_order: bool = False _chain_order: list[str] @@ -509,38 +524,14 @@ class ResidueTemplateIterator(BaseTemplateIterator): def __init__(self, **data: Any): super().__init__(**data) - # The unique method should retain default order - self._chain_order = self.structure_df["chain"].unique() + # The unique method should retain default order, or an order specified by the user. + if self.chain_ids == ["all"]: + self._chain_order = self.structure_df["chain"].unique() + else: + self._chain_order = self.chain_ids if self.randomize_chain_order: np.random.shuffle(self._chain_order) - def set_chain_order(self, chain_order: list[str]) -> None: - """Set the order of chains to iterate over. - - Parameters - ---------- - chain_order : list[str] - List of chain identifiers, in desired order, to use when iterating - over the structure. - - Raises - ------ - ValueError - If the chain order does not contain all chains in the structure. - - Returns - ------- - None - """ - # Check that all the chains are present in the chain_order list - if set(chain_order) != set(self.structure_df["chain"].unique()): - warnings.warn( - "The provided chain order does not contain all chains within the " - "structure. If this was intentional (e.g. to tile using only a subset " - "of chains in the structure) then ignore this warning." - ) - - self._chain_order = chain_order @property def num_alternate_structures(self) -> int: @@ -608,6 +599,7 @@ def alternate_template_iter( for idx in window_iter: chain_ids = chains[idx] residue_ids = residues[idx] + print(f"Removing residues {residue_ids[0]} to {residue_ids[-1]}") # Merge the DataFrame to keep only positions where the chain and residue # pairs match the current window From 6184f829e465303845ea0ea81311174b3e21a0dd Mon Sep 17 00:00:00 2001 From: pcbve1 Date: Fri, 24 Jul 2026 10:19:04 -0700 Subject: [PATCH 2/2] remove print statements, add chain_id validator, and update dependencies --- pyproject.toml | 2 +- src/mosaics/template_iterator.py | 68 +++++++++++++++++++------------- 2 files changed, 42 insertions(+), 28 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 1776074..7a85661 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -49,7 +49,7 @@ dependencies = [ "ttsim3d", "teamtomo-basemodel", "Leopard-EM>=v1.0", - "torch-fourier-slice>=v0.4.0" + "torch-fourier-slice>=v0.5.2" ] [tool.hatch.metadata] diff --git a/src/mosaics/template_iterator.py b/src/mosaics/template_iterator.py index 9845afa..23c101a 100644 --- a/src/mosaics/template_iterator.py +++ b/src/mosaics/template_iterator.py @@ -2,8 +2,7 @@ from abc import abstractmethod from collections.abc import Iterator -from typing import Annotated, Any, ClassVar, Literal -import warnings +from typing import Annotated, Any, ClassVar, Literal, Union import numpy as np import pandas as pd @@ -440,21 +439,35 @@ class ChainTemplateIterator(BaseTemplateIterator): type : Literal["chain"] Discriminator field for differentiating between template iterator types. chain_ids : list[str] - If populated, denotes the chains to be removed. When empty, sets all chains to be removed. Default is all. + If populated, denotes the chains to be removed. When empty or None, sets all chains to be removed. Default is None. """ type: ClassVar[Literal["chain"]] = "chain" - chain_ids: list[str] = ["all"] + chain_ids: Union[list[str], None] = None + + def __init__(self, **data: Any): + super().__init__(**data) + + unique_chains = list(self.structure_df["chain"].unique()) + # Checks that all chains passed by the user are present in the template. + if self.chain_ids: + absent_chains = [ + chain_id + for chain_id in self.chain_ids + if chain_id not in unique_chains + ] + if absent_chains: + raise ValueError( + f"Chains {absent_chains} are not present in the template." + ) + else: + self.chain_ids = unique_chains @property def num_alternate_structures(self) -> int: """Get the number of alternate structures (i.e. number of chains).""" - if self.chain_ids == ['all']: - return len(self.structure_df["chain"].unique()) - else: - return len(self.chain_ids) + return len(self.chain_ids) - def alternate_template_iter( self, inverted: bool = True ) -> Iterator[tuple[list[str | None], list[int | None], torch.Tensor]]: @@ -466,13 +479,8 @@ def alternate_template_iter( If 'True', return the indexes of atoms to remove rather than keep. """ subset_df = self.subset_df_on_residues_and_atoms() - if self.chain_ids == ["all"]: - unique_chain_ids = subset_df["chain"].unique() - else: - unique_chain_ids = self.chain_ids - for chain_id in unique_chain_ids: - print(f"Removing chain {chain_id}") + for chain_id in self.chain_ids: residue_ids = subset_df[subset_df["chain"] == chain_id]["residue_id"] residue_ids = residue_ids.unique().tolist() @@ -510,27 +518,34 @@ class ResidueTemplateIterator(BaseTemplateIterator): randomize_chain_order : bool If 'True', randomize the order of chains in the structure. Default is 'False'. chain_ids : list[str] - If populated, denotes the chains to be removed. When empty, sets all chains to be removed. Default is empty. + If populated, denotes the chains to be removed. When empty or None, sets all chains to be removed. Default is None. """ type: ClassVar[Literal["residue"]] = "residue" num_residues_removed: Annotated[int, Field(gt=0)] residue_increment: Annotated[int, Field(gt=0)] - chain_ids: list[str] = ["all"] + chain_ids: Union[list[str] , None] = None randomize_chain_order: bool = False - _chain_order: list[str] - def __init__(self, **data: Any): super().__init__(**data) - - # The unique method should retain default order, or an order specified by the user. - if self.chain_ids == ["all"]: - self._chain_order = self.structure_df["chain"].unique() + # The unique method should retain default order. + unique_chains = list(self.structure_df["chain"].unique()) + # check if all chains are present in the dataframe + if self.chain_ids: + absent_chains = [ + chain_id + for chain_id in self.chain_ids + if chain_id not in unique_chains + ] + if absent_chains: + raise ValueError( + f"Chains {absent_chains} are not present in the template." + ) else: - self._chain_order = self.chain_ids + self.chain_ids = unique_chains if self.randomize_chain_order: - np.random.shuffle(self._chain_order) + np.random.shuffle(self.chain_ids) @property @@ -565,7 +580,7 @@ def chain_residue_pairs(self) -> list[tuple[str, int]]: # Chunk the df into groups based on chain and re-stich together in order # This will ensure that the chain order is respected df_list = [] - for chain in self._chain_order: + for chain in self.chain_ids: df_list.append(subset_df[subset_df["chain"] == chain]) ordered_df = pd.concat(df_list) @@ -599,7 +614,6 @@ def alternate_template_iter( for idx in window_iter: chain_ids = chains[idx] residue_ids = residues[idx] - print(f"Removing residues {residue_ids[0]} to {residue_ids[-1]}") # Merge the DataFrame to keep only positions where the chain and residue # pairs match the current window