diff --git a/pyproject.toml b/pyproject.toml index b708b99..7a85661 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -49,6 +49,7 @@ dependencies = [ "ttsim3d", "teamtomo-basemodel", "Leopard-EM>=v1.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 7ca0a53..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 @@ -358,7 +357,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,21 +432,42 @@ 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 or None, sets all chains to be removed. Default is None. """ type: ClassVar[Literal["chain"]] = "chain" + 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).""" - return len(self.structure_df["chain"].unique()) - + return len(self.chain_ids) + def alternate_template_iter( self, inverted: bool = True ) -> Iterator[tuple[list[str | None], list[int | None], torch.Tensor]]: @@ -459,10 +479,10 @@ 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() - for chain_id in unique_chain_ids: + for chain_id in self.chain_ids: 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,50 +517,36 @@ 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 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: 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 - 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_ids = unique_chains 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." - ) + np.random.shuffle(self.chain_ids) - self._chain_order = chain_order @property def num_alternate_structures(self) -> int: @@ -574,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)