Transforms

Design tasks

class rfproteina.transforms.design_task.ConditionedDesignTask(global_design_task_kwargs: dict | None = None, **kwargs)

Bases: DesignTaskABC

Base for tasks with shared conditioning modifiers.

Each modifier has a flat parameter dictionary. Constructor parameters take precedence over global parameters and shared defaults, in that order. Task-specific conditioning rates are configured in the design-task presets.

apply_design_task_modifiers(data: dict, *, cath_kwargs: dict | None = None) → dict

Apply configured modifiers with optional per-sample CATH parameters.

ensure_atomized_tokens_are_sequence_conditioned(data: dict) → dict

class rfproteina.transforms.design_task.NoMotifABC(global_design_task_kwargs: dict | None = None, **kwargs)

Bases: ConditionedDesignTask

can_apply(data: dict) → bool

Evaluate whether the design task can be applied to the data.

This method should check if the current data meets the requirements for applying this design task.

Parameters:

data – The input data dictionary containing the atom array and annotations.

Returns:

True if the task can be applied, False otherwise.

set_task_info(data: dict) → dict

class rfproteina.transforms.design_task.UnconditionalTask(global_design_task_kwargs: dict | None = None, **kwargs)

Bases: NoMotifABC

Sample a protein backbone & sequence unconditionally.

forward(data: dict) → dict

Apply a transformation to the input dictionary and return the transformed dictionary.

Parameters:

• data – The input dictionary to transform.

• *args – Additional positional arguments.

• **kwargs – Additional keyword arguments.

Returns:

The transformed dictionary.

class rfproteina.transforms.design_task.StructurePredictionTask(global_design_task_kwargs: dict | None = None, **kwargs)

Bases: NoMotifABC

Sample a protein backbone, given the sequence information (= structure prediction).

forward(data: dict) → dict

Apply a transformation to the input dictionary and return the transformed dictionary.

Parameters:

• data – The input dictionary to transform.

• *args – Additional positional arguments.

• **kwargs – Additional keyword arguments.

Returns:

The transformed dictionary.

class rfproteina.transforms.design_task.InverseFoldingTask(polymer_types: list[str] | None = None, **kwargs)

Bases: NoMotifABC

Design sequence and sidechains/bases from selected polymers’ backbone coordinates.

Other molecules stay fully conditioned; at least one designable atom is required.

Parameters:

polymer_types – Polymer types to design: “protein” (default), “rna”, or “dna”.

can_apply(data: dict) → bool

Evaluate whether the design task can be applied to the data.

This method should check if the current data meets the requirements for applying this design task.

Parameters:

data – The input data dictionary containing the atom array and annotations.

Returns:

True if the task can be applied, False otherwise.

forward(data: dict) → dict

Apply a transformation to the input dictionary and return the transformed dictionary.

Parameters:

• data – The input dictionary to transform.

• *args – Additional positional arguments.

• **kwargs – Additional keyword arguments.

Returns:

The transformed dictionary.

class rfproteina.transforms.design_task.ProteinFromDNAContextABC(global_design_task_kwargs: dict | None = None, **kwargs)

Bases: ConditionedDesignTask

Base for DNA-binder design tasks: the whole protein is generated from noise while a DNA chain is provided as context.

The DNA is supplied through the C_CRD / C_SEQ condition annotations rather than as a protein motif, so there is no motif region — hence the all-zeros motif mask stamped by set_task_info(). This mirrors the no-motif tasks, but subclassing ConditionedDesignTask directly keeps this a distinct setting. Subclasses decide whether the DNA structure is given as context or co-generated.

Design-task modifiers (RASA / H-bond / SSE / CATH conditioning) are left at the experiment’s global rates (datasets.global_design_task_kwargs); the DNA tasks intentionally participate in the same conditioning as the protein tasks.

can_apply(data: dict) → bool

Evaluate whether the design task can be applied to the data.

This method should check if the current data meets the requirements for applying this design task.

Parameters:

data – The input data dictionary containing the atom array and annotations.

Returns:

True if the task can be applied, False otherwise.

set_task_info(data: dict) → dict

Stamp task bookkeeping with an all-zeros motif mask (no protein motif).

class rfproteina.transforms.design_task.ProteinFromDNAStructureContextTask(global_design_task_kwargs: dict | None = None, **kwargs)

Bases: ProteinFromDNAContextABC

Generate the entire protein from noise; DNA structure and sequence are both provided as context.

forward(data: dict) → dict

Apply a transformation to the input dictionary and return the transformed dictionary.

Parameters:

• data – The input dictionary to transform.

• *args – Additional positional arguments.

• **kwargs – Additional keyword arguments.

Returns:

The transformed dictionary.

class rfproteina.transforms.design_task.ProteinFromDNASequenceContextTask(global_design_task_kwargs: dict | None = None, **kwargs)

Bases: ProteinFromDNAContextABC

Generate the entire protein from noise; only the DNA sequence is provided as context — the DNA structure is co-generated.

forward(data: dict) → dict

Apply a transformation to the input dictionary and return the transformed dictionary.

Parameters:

• data – The input dictionary to transform.

• *args – Additional positional arguments.

• **kwargs – Additional keyword arguments.

Returns:

The transformed dictionary.

class rfproteina.transforms.design_task.PartialLigandTask(*, min_coverage_fraction: float = 0.0, max_coverage_fraction: float = 1.0, p_fix_ligand_surroundings: float = 1.0, **kwargs)

Bases: ConditionedDesignTask

Intended only for OMOL examples where we have ligands and maybe some tip atoms. In this case, we always coordinate condition the non-ligands and sequence condition everything.

can_apply(data: dict) → bool

Evaluate whether the design task can be applied to the data.

This method should check if the current data meets the requirements for applying this design task.

Parameters:

data – The input data dictionary containing the atom array and annotations.

Returns:

True if the task can be applied, False otherwise.

sample_ligand_mask(atom_array: AtomArray, ligand_mask: ndarray) → ndarray

forward(data: dict) → dict

Apply a transformation to the input dictionary and return the transformed dictionary.

Parameters:

• data – The input dictionary to transform.

• *args – Additional positional arguments.

• **kwargs – Additional keyword arguments.

Returns:

The transformed dictionary.

class rfproteina.transforms.design_task.TipAtomPartialLigandTask(*, n_max_ligands: int = 3, n_hops_expected_ligand: int = 5, n_max_residues: int = 10, n_hops_expected_tip_atom: int = 1, ligand_scaling_factor: float = 5.0, tip_atom_scaling_factor: float = 1.5, p_seq_conditioned: float = 1.0, p_knockout_seq: float = 0.0, p_unindex_all: float = 0.9, p_merge_unindexed_groups: float = 0.1, p_split_unindexed_groups: float = 0.01, p_increase_group_id_gap: float = 0.1, p_internal_coords: float = 0.0, **kwargs)

Bases: ConditionedDesignTask

can_apply(data: dict) → bool

Evaluate whether the design task can be applied to the data.

This method should check if the current data meets the requirements for applying this design task.

Parameters:

data – The input data dictionary containing the atom array and annotations.

Returns:

True if the task can be applied, False otherwise.

sample_mask(atom_array: AtomArray) → tuple[ndarray, list[ndarray]]

forward(data: dict) → dict

Apply a transformation to the input dictionary and return the transformed dictionary.

Parameters:

• data – The input dictionary to transform.

• *args – Additional positional arguments.

• **kwargs – Additional keyword arguments.

Returns:

The transformed dictionary.

class rfproteina.transforms.design_task.TipAtomActiveSiteTask(tip_atom_rasa_threshold: float = 0.2, min_tip_atoms: int = 2, max_tip_atoms: int = 15, p_seq_conditioned: float = 1.0, min_protein_length: int = 50, p_unindex_all: float = 0.5, p_unindex_any: float = 0.05, p_merge_unindexed_groups: float = 0.1, p_split_unindexed_groups: float = 0.01, p_increase_group_id_gap: float = 0.1, p_internal_coords: float = 0.0, **kwargs)

Bases: ConditionedDesignTask

Designed to sample tip atoms around a protein pocket when no ligand is present. This is done by calculating the SASA of all atoms, and calculating the distance they lie from the convex hull of the protein chain. Atoms above a SASA threshold are sampled at rates proportional to their distance from the convex hull. In effect this samples tip atoms around protein pocket.

Parameters:

• tip_atom_rasa_threshold – rASA threshold to consider tip atoms for sampling

• min_tip_atoms – minimum number of tip atoms to sample

• max_tip_atoms – maximum number of tip atoms to sample

• p_seq_conditioned – probability of providing sequence conditioning to atoms with tip-atoms sampled NOTE: In proteina, providing tip atoms inherently gives away sequence information because of Atom37!

• min_protein_length – minimum length of protein chain to sample

• p_unindex_all – probability of unindexing all atoms in the protein chain

• p_unindex_any – probability of unindexing individual sampled atom. Will only do this if p_index_all was sampled to False

• p_merge_unindexed_groups – probability of merging unindexed groups

• p_split_unindexed_groups – probability of splitting unindexed groups

• p_increase_group_id_gap – probability of increasing the gap between unindexed group ids

annotate_for_task_selection(data: dict) → dict

Add any annotations to the data that are required for the task selection.

This method allows design tasks to add annotations to the data that they need to evaluate whether they can be applied. Usually the atom_array object of the data dictionary is annotated with any extra annotations.

Preferrably you should use the annotator module to add annotations to the data. Those can easily & lazily be generated & removed.

Parameters:

data – The input data dictionary containing the atom array and other fields.

Returns:

The data dictionary with additional annotations added.

can_apply(data: dict) → bool

Evaluate whether the design task can be applied to the data.

This method should check if the current data meets the requirements for applying this design task.

Parameters:

data – The input data dictionary containing the atom array and annotations.

Returns:

True if the task can be applied, False otherwise.

sample_protein_pocket(atom_array: AtomArray, eligible_tip_atom_mask: ndarray) → int

Calculates convex hull of the protein chain (self.sampled_chain) and samples a tip atom based on distance from the convex hull boundary and the rasa (stored in tip_atom_rasa annotation). Returns the index of the sampled tip atom.

sample_mask(atom_array: AtomArray) → ndarray

forward(data: dict) → dict

Apply a transformation to the input dictionary and return the transformed dictionary.

Parameters:

• data – The input dictionary to transform.

• *args – Additional positional arguments.

• **kwargs – Additional keyword arguments.

Returns:

The transformed dictionary.

class rfproteina.transforms.design_task.TipAtomTask(min_residues: int = 2, max_residues: int = 25, p_tip_atom: float = 0.8, n_hops_expected: int = 3, p_unindex_all: float = 0.9, p_merge_unindexed_groups: float = 0.1, p_split_unindexed_groups: float = 0.01, p_increase_group_id_gap: float = 0.1, p_internal_coords: float = 0.0, **kwargs)

Bases: ConditionedDesignTask

can_apply(data: dict) → bool

Evaluate whether the design task can be applied to the data.

This method should check if the current data meets the requirements for applying this design task.

Parameters:

data – The input data dictionary containing the atom array and annotations.

Returns:

True if the task can be applied, False otherwise.

sample_mask(atom_array: AtomArray) → tuple[ndarray, list[ndarray], int]

forward(data: dict) → dict

Apply a transformation to the input dictionary and return the transformed dictionary.

Parameters:

• data – The input dictionary to transform.

• *args – Additional positional arguments.

• **kwargs – Additional keyword arguments.

Returns:

The transformed dictionary.

class rfproteina.transforms.design_task.MotifScaffoldTask(min_tokens_per_motif: int = 1, max_tokens_per_motif: int = 50, max_unindexed_tokens: int = 50, min_coverage_fraction: float = 0.05, max_coverage_fraction: float = 1.0, p_unindex_any: float = 0.5, p_merge_unindexed_groups: float = 0.1, p_split_unindexed_groups: float = 0.01, p_increase_group_id_gap: float = 0.1, p_internal_coords: float = 0.0, p_backbone_only: float = 0.0, p_bb_only_seq_conditioned: float = 0.3, **kwargs)

Bases: ConditionedDesignTask

annotate_for_task_selection(data: dict) → dict

Add any annotations to the data that are required for the task selection.

This method allows design tasks to add annotations to the data that they need to evaluate whether they can be applied. Usually the atom_array object of the data dictionary is annotated with any extra annotations.

Preferrably you should use the annotator module to add annotations to the data. Those can easily & lazily be generated & removed.

Parameters:

data – The input data dictionary containing the atom array and other fields.

Returns:

The data dictionary with additional annotations added.

can_apply(data: dict) → bool

Evaluate whether the design task can be applied to the data.

This method should check if the current data meets the requirements for applying this design task.

Parameters:

data – The input data dictionary containing the atom array and annotations.

Returns:

True if the task can be applied, False otherwise.

sample_mask(atom_array: AtomArray) → tuple[ndarray, list[ndarray]]

sample_unindexed_mask(all_masks: list[ndarray], token_starts: ndarray) → ndarray

Samples a mask of tokens to be unindexed, ensuring only full coordinate-conditioned islands are unindexed.

NOTE: This function assumes that the coordinate-conditioned islands do not contain partial tokens, which is always true for the MotifScaffoldTask.

Parameters:

• all_masks – List of boolean masks for each coordinate-conditioned island (shape of each mask is (n_atoms,))

• token_starts – Boolean array indicating the start of each token (shape (n_atoms,))

Returns:

Boolean array of shape (n_atoms,) where True indicates the token is unindexed

Return type:

unindexed_mask

forward(data: dict) → dict

Apply a transformation to the input dictionary and return the transformed dictionary.

Parameters:

• data – The input dictionary to transform.

• *args – Additional positional arguments.

• **kwargs – Additional keyword arguments.

Returns:

The transformed dictionary.

class rfproteina.transforms.design_task.BinderHotspotWith3DTarget(hotspot_count_distribution: Literal['uniform', 'geometric'] = 'uniform', hotspot_geometric_p: float | None = None, max_hotspots_frac_to_provide: float = 0.1, hotspot_max_distance: float = 7.0, p_no_hotspots: float = 0.05, p_sample_ppi_crop: float = 0.75, ppi_crop_binder_fraction: float = 0.667, p_bb_only_coord_conditioning: float = 0.0, target_chain_types: Sequence[ChainType] = (ChainType.POLYPEPTIDE_L, ChainType.POLYPEPTIDE_D), **kwargs)

Bases: ConditionedDesignTask

annotate_for_task_selection(data: dict) → dict

Add any annotations to the data that are required for the task selection.

This method allows design tasks to add annotations to the data that they need to evaluate whether they can be applied. Usually the atom_array object of the data dictionary is annotated with any extra annotations.

Preferrably you should use the annotator module to add annotations to the data. Those can easily & lazily be generated & removed.

Parameters:

data – The input data dictionary containing the atom array and other fields.

Returns:

The data dictionary with additional annotations added.

can_apply(data: dict) → bool

Can be applied if query_pn_unit_iids has at least two distinct chains that form a valid binder/target interface: a protein binder and one or more target chains whose types are in target_chain_types (protein-only by default; add DNA/RNA for nucleic-acid targets). The target may span multiple chains, e.g. a double-stranded DNA target (two nucleic chains).

forward(data: dict) → dict

Apply a transformation to the input dictionary and return the transformed dictionary.

Parameters:

• data – The input dictionary to transform.

• *args – Additional positional arguments.

• **kwargs – Additional keyword arguments.

Returns:

The transformed dictionary.

class rfproteina.transforms.design_task.NucleicSecondaryStructureTask(min_nucleic_tokens: int = 2, p_show_partial_feats: float = 0.7, p_add_motif_scaffold_condition: float = 0.0, min_tokens_per_motif: int = 1, max_tokens_per_motif: int = 12, min_motif_coverage_fraction: float = 0.02, max_motif_coverage_fraction: float = 0.1, max_unindexed_tokens: int = 0, p_unindex_any: float = 0.0, p_merge_unindexed_groups: float = 0.0, p_split_unindexed_groups: float = 0.0, p_increase_group_id_gap: float = 0.0, p_internal_coords: float = 0.0, nucleic_ss_min_shown: float = 0.2, nucleic_ss_max_shown: float = 1.0, n_islands_min: int = 1, n_islands_max: int = 6, p_canonical_bp_filter: float = 0.0, na_only: bool = False, planar_only: bool = True, **kwargs)

Bases: ConditionedDesignTask

Apply NA secondary-structure conditioning, with optional small motif scaffolds.

When this task is selected, it computes bp_partners from the current structure and either reveals the full annotation or a partial island mask. It can also optionally add a motif-scaffold overlay with coordinate/sequence conditioning on a small fraction of polymer tokens.

can_apply(data: dict) → bool

Evaluate whether the design task can be applied to the data.

This method should check if the current data meets the requirements for applying this design task.

Parameters:

data – The input data dictionary containing the atom array and annotations.

Returns:

True if the task can be applied, False otherwise.

sample_motif_mask(atom_array: AtomArray) → tuple[ndarray, list[ndarray]]

Sample small polymer token islands to retain as motif scaffolds.

sample_unindexed_mask(all_masks: list[ndarray], token_starts: ndarray) → ndarray

Sample motif islands to unindex while preserving whole conditioned islands.

forward(data: dict) → dict

Apply a transformation to the input dictionary and return the transformed dictionary.

Parameters:

• data – The input dictionary to transform.

• *args – Additional positional arguments.

• **kwargs – Additional keyword arguments.

Returns:

The transformed dictionary.

class rfproteina.transforms.design_task.IgRegion(*values)

Bases: StrEnum

Allowed values of the per-atom ig_region CIF annotation.

NOT_IG = ''

FRAMEWORK = 'framework'

CDR1 = 'cdr1'

CDR2 = 'cdr2'

CDR3 = 'cdr3'

CDR4 = 'cdr4'

class rfproteina.transforms.design_task.IgFrameworkConditioningTask(cdr_expand_min: int = -2, cdr_expand_max: int = 5, p_fix_all_but_cdr3: float = 0.3, p_cdr_fixed: float = 0.2, p_design_cdr4: float = 0.1, p_framework_seq_cond: float = 0.8, p_framework_bb_only: float = 0.3, p_flexible_target: float = 0.1, p_target_intcoord: float = 0.5, p_target_bb_only: float = 0.3, p_separate_ig_coord_groups: float = 0.5, p_no_hotspots: float = 0.05, hotspot_max_distance: float = 7.0, max_hotspots_frac_to_provide: float = 0.1, hotspot_count_distribution: Literal['uniform', 'geometric'] = 'uniform', hotspot_geometric_p: float | None = None, **kwargs)

Bases: ConditionedDesignTask

Framework-conditioning task for antibodies and TCRs.

Conditions on the Ig framework and (optionally) the target/antigen while leaving CDR loops free for design. It:

• reads the native ig_region structure annotation,

• conditions framework geometry with internal coordinates,

• supports variable CDR boundary expansion/contraction,

• supports per-CDR fixation (internal coords + sequence conditioning), and

• supports separate rigid-body coordinate groups for multi-chain Ig molecules (e.g. scFv heavy + light).

Works for both antibodies and TCRs.

Parameters:

• cdr_expand_min – Minimum residue expansion per CDR side (negative = shrink).

• cdr_expand_max – Maximum residue expansion per CDR side.

• p_fix_all_but_cdr3 – Probability of fixing every CDR except CDR3, forcing CDR3 to be designed. When False the per-CDR probability p_cdr_fixed is used instead. CDR4 is excluded from this logic; see p_design_cdr4.

• p_cdr_fixed – Per-CDR probability of fixing (internal coords + seq conditioning). Applied independently to each CDR (except CDR4) when p_fix_all_but_cdr3 is not triggered.

• p_design_cdr4 – Probability of including CDR4 in the designed (free) region. Sampled once per example and applied to CDR4 on every Ig chain; otherwise CDR4 follows the framework conditioning policy.

• p_framework_seq_cond – Probability of providing sequence conditioning on the framework in addition to internal-coord conditioning.

• p_framework_bb_only – Probability of restricting framework internal-coord conditioning to backbone atoms only (CA/N/C/O) rather than all atoms.

• p_flexible_target – Probability of giving the target chain no coordinate conditioning at all (free antigen).

• p_target_intcoord – Probability of using internal-coord conditioning for the target (given it is conditioned). Otherwise absolute coords.

• p_target_bb_only – Probability of restricting target conditioning to backbone atoms only.

• p_separate_ig_coord_groups – When multiple Ig chains are present (e.g. heavy + light), probability of assigning each its own rigid-body coordinate group. Otherwise all Ig chains are in one group.

• p_no_hotspots – Probability of providing no hotspot conditioning.

• hotspot_max_distance – Å threshold for hotspot eligibility.

• max_hotspots_frac_to_provide – Hard cap on the fraction of interface atoms to label as hotspots.

• hotspot_count_distribution – "uniform" or "geometric".

• hotspot_geometric_p – Geometric distribution parameter (required when hotspot_count_distribution="geometric").

annotate_for_task_selection(data: dict) → dict

Add any annotations to the data that are required for the task selection.

This method allows design tasks to add annotations to the data that they need to evaluate whether they can be applied. Usually the atom_array object of the data dictionary is annotated with any extra annotations.

Preferrably you should use the annotator module to add annotations to the data. Those can easily & lazily be generated & removed.

Parameters:

data – The input data dictionary containing the atom array and other fields.

Returns:

The data dictionary with additional annotations added.

can_apply(data: dict) → bool

Evaluate whether the design task can be applied to the data.

This method should check if the current data meets the requirements for applying this design task.

Parameters:

data – The input data dictionary containing the atom array and annotations.

Returns:

True if the task can be applied, False otherwise.

forward(data: dict) → dict

Apply a transformation to the input dictionary and return the transformed dictionary.

Parameters:

• data – The input dictionary to transform.

• *args – Additional positional arguments.

• **kwargs – Additional keyword arguments.

Returns:

The transformed dictionary.

Design-task modifiers

rfproteina.transforms.design_task_modifiers.add_rasa_condition(data: dict, p_add_rasa_condition: float = 0.0, probe_radius: float = 1.4, atom_radii: str | ndarray = 'element', point_number: int = 100, p_include_non_ligand: float = 0.1, rasa_max: float = 0.1) → dict

Add binned RASA conditioning to an example.

Atoms with nonfinite RASA remain unconditioned.

Parameters:

• data – Example dictionary containing an atom array.

• p_add_rasa_condition – Probability of adding this condition at all. Defaults to 0.0.

• probe_radius – Van der Waals radius of the probe in Angstrom. Defaults to 1.4 for water.

• atom_radii – Atom radii used for the calculation. Defaults to "element", matching training and supporting both protein and ligand atoms.

• point_number – Number of points sampled by the Shrake-Rupley algorithm. Defaults to 100.

• p_include_non_ligand – Probability of conditioning non-ligand atoms. Defaults to 0.1.

• rasa_max – Inclusive upper boundary of the intermediate bin. Defaults to 0.1. The buried/intermediate boundary is 0.02.

Returns:

The example dictionary with RASA annotations added to its atom array.

rfproteina.transforms.design_task_modifiers.subsample_condition_enum(condition_values: ndarray, fraction: float, none_value: Any = 0, groups: Sequence[ndarray] = ()) → ndarray

Randomly clear a subset of the active condition values, keeping fraction of them.

Parameters:

• condition_values – Per-atom condition values (boolean mask or integer enum).

• fraction – Fraction of active (non-none_value) entries to keep. Must be in (0, 1].

• none_value – Value denoting “no condition”. Defaults to 0.

• groups – Index groups kept or cleared as a unit, e.g. the atoms of one aromatic ring. Groups may overlap, so an atom shared by a kept and a cleared ring stays; any active entry outside every group is its own group. fraction then counts groups rather than atoms.

Returns:

A copy of condition_values with 1 - fraction of its active entries reset to none_value. At least one active entry is always kept.

rfproteina.transforms.design_task_modifiers.add_hbond_condition(data: dict, p_add_hbond_condition: float = 0.0, p_hbond_subsample: float = 0.5, hbond_subsample_frac: float = 0.5) → dict

Subset the ground-truth hydrogen-bond conditioning already present on the AtomArray.

rfproteina.transforms.design_task_modifiers.add_pistack_condition(data: dict, p_add_pistack_condition: float = 0.0, p_pistack_subsample: float = 0.5, pistack_subsample_frac: float = 0.5) → dict

Subset the ground-truth pi-stacking conditioning already present on the AtomArray.

rfproteina.transforms.design_task_modifiers.add_pication_condition(data: dict, p_add_pication_condition: float = 0.0, p_pication_subsample: float = 0.5, pication_subsample_frac: float = 0.5) → dict

Subset the ground-truth pi-cation conditioning already present on the AtomArray.

rfproteina.transforms.design_task_modifiers.add_salt_condition(data: dict, p_add_salt_condition: float = 0.0, p_salt_subsample: float = 0.5, salt_subsample_frac: float = 0.5) → dict

Subset the ground-truth salt-bridge conditioning already present on the AtomArray.

rfproteina.transforms.design_task_modifiers.add_hydrophobic_condition(data: dict, p_add_hydrophobic_condition: float = 0.0, p_hydrophobic_subsample: float = 0.5, hydrophobic_subsample_frac: float = 0.5) → dict

Subset the ground-truth hydrophobic-contact conditioning already present on the AtomArray.

rfproteina.transforms.design_task_modifiers.add_metal_condition(data: dict, p_add_metal_condition: float = 0.0, p_metal_subsample: float = 0.5, metal_subsample_frac: float = 0.5) → dict

Subset the ground-truth metal-coordination conditioning already present on the AtomArray.

rfproteina.transforms.design_task_modifiers.add_secondary_structure_condition(data: dict[str, Any], p_add_secondary_structure_condition: float = 0.0, max_secondary_structure_frac_to_provide: float = 0.4, min_ss_island_len: int = 1, max_ss_island_len: int = 10, n_islands_max: int = 3) → dict[str, Any]

Add secondary structure conditioning.

Uses P-SEA to annotate secondary structure, samples islands from residues with valid SSE annotations, and sets C_SEC condition. Sampling is done at residue granularity (not token granularity) so that residues spanning multiple tokens — e.g., atomized PTM residues — receive a uniform value, satisfying the Level.RESIDUE invariant on the secondary-structure condition.

Parameters:

• data – Data dict containing atom_array.

• p_add_secondary_structure_condition – Probability of adding this condition at all. Defaults to 0.0.

• max_secondary_structure_frac_to_provide – Max fraction of valid SS residues to condition. Defaults to 0.4.

• min_ss_island_len – Minimum island length (in residues). Defaults to 1.

• max_ss_island_len – Maximum island length (in residues). Defaults to 10.

• n_islands_max – Maximum number of islands. Defaults to 3.

Returns:

Updated data dict with C_SEC conditioning.

rfproteina.transforms.design_task_modifiers.diversify_sequence_conditioning(data: dict[str, Any], p_exact: float = 0.5, p_zero_only: float = 0.16666666666666666, p_positive_only: float = 0.16666666666666666, p_zero_and_positive: float = 0.16666666666666666, dirichlet_concentration: float = 3.0, exclusion_count_alpha: float = 2.0, p_full_probability_vector: float = 0.75, p_augment_unconditioned: float = 0.0, min_augmentation_fraction: float = 0.05, max_augmentation_fraction: float = 0.2, p_zero_only_whole_chain: float = 0.0) → dict[str, Any]

Sample calibrated sequence conditions on residues whose identity is not otherwise revealed.

Existing exact conditions are diversified, while an independently sampled subset of unconditioned residues can receive additional conditioning. Atomized residues and residues with conditioned non-backbone coordinates are never modified. A zero-only draw may cover a whole chain by broadcasting identities absent from that chain.

Parameters:

• data – Example dictionary containing an atom array.

• p_exact – Relative weight for retaining ordinary exact sequence conditioning.

• p_zero_only – Relative weight for exposing only exact-zero exclusions.

• p_positive_only – Relative weight for exposing positive probabilities without exact zeros.

• p_zero_and_positive – Relative weight for exposing both exclusions and positive probabilities.

• dirichlet_concentration – Concentration of the restricted prior before truth size-biasing.

• exclusion_count_alpha – Exponent of the truncated power-law distributions for exclusion and exposed-positive counts. Counts range from one through the maximum valid for the mode.

• p_full_probability_vector – Probability of exposing every positive probability; otherwise expose a nonempty proper subset.

• p_augment_unconditioned – Per-example probability of adding conditions to otherwise unconditioned eligible residues.

• min_augmentation_fraction – Minimum fraction of eligible unconditioned residues selected.

• max_augmentation_fraction – Maximum fraction of eligible unconditioned residues selected.

• p_zero_only_whole_chain – Conditional probability that a zero-only draw uses one shared set of exclusions across a whole chain.

rfproteina.transforms.design_task_modifiers.add_cath_condition(data: dict[str, Any], p_add_cath_c_condition: float = 0.0, p_add_cath_ca_condition: float = 0.0, p_add_cath_cat_condition: float = 0.0, min_retained_chain_fraction: float = 0.75) → dict[str, Any]

Independently sample each chain’s training CATH conditioning from stored labels.

Unselected labels are cleared, including on skipped or ineligible examples. Stored labels retain their domain boundaries. Eligible chains can receive different levels in the same example. Inference does not use this modifier.

Parameters:

• data – Data dict containing atom_array with optional per-residue CATH annotations.

• p_add_cath_c_condition – Per-chain probability of revealing the class only.

• p_add_cath_ca_condition – Per-chain probability of revealing class and architecture.

• p_add_cath_cat_condition – Per-chain probability of revealing class, architecture and topology.

• min_retained_chain_fraction – Minimum fraction of the labelled chain retained by cropping.