### Setup and Run AtomWorks Test Suite

Source: https://github.com/rosettacommons/atomworks/blob/production/docs/installation.md

Commands to set up the AtomWorks test suite by downloading test data and configuring environment variables, followed by instructions to run the tests using pytest.

```bash
atomworks setup tests

# For running tests with the test pack:
PDB_MIRROR_PATH=tests/data/pdb
CCD_MIRROR_PATH=tests/data/ccd

cp .env.sample .env
# Then edit .env to set the paths above

# Run all tests (excluding very slow ones)
pytest tests -m "not very_slow"

# Run tests in parallel for faster execution
pytest tests -m "not very_slow" -n auto

# Run a specific test file
pytest tests/io/components/test_parser.py
```

--------------------------------

### Development Installation of AtomWorks

Source: https://github.com/rosettacommons/atomworks/blob/production/docs/installation.md

Clones the AtomWorks repository and installs it for development purposes using 'make install' or an editable pip install. Also includes instructions for setting up a fresh environment.

```bash
git clone https://github.com/RosettaCommons/atomworks.git
cd atomworks
make install  # or pip install -e ".[dev]"

git clone https://github.com/RosettaCommons/atomworks.git
cd atomworks
make env
```

--------------------------------

### Install AtomWorks via pip

Source: https://github.com/rosettacommons/atomworks/blob/production/docs/installation.md

Installs AtomWorks using pip, with options for base, ML, development, or combined dependencies. Also shows installation with Open Babel as an alternative to RDKit.

```bash
pip install atomworks
pip install "atomworks[ml]"
pip install "atomworks[dev]"
pip install "atomworks[ml,dev]"
pip install "atomworks[openbabel]"
pip install "atomworks[ml,openbabel,dev]"
```

--------------------------------

### Install AtomWorks using pip

Source: https://github.com/rosettacommons/atomworks/blob/production/docs/tutorial/index.md

This command installs the AtomWorks Python package using pip. Ensure you have Python and pip installed on your system. This is the quickest way to get started with AtomWorks.

```bash
pip install atomworks
```

--------------------------------

### Configure AF3-style Dataset with Hydra

Source: https://github.com/rosettacommons/atomworks/blob/production/docs/mirrors.md

An example Hydra configuration for setting up a dataset for training. It demonstrates filtering for specific chain types (D-polypeptide and L-polypeptide) and excluding certain ligands based on a regular expression.

```yaml
# NOTE: The below is a hydra config and the _target_ fields are the hydra syntax for instantiating a class.
#  You can use this without hyrda, but will then instead need to provide the corresponding arguments for the
#  _target_ objects directly.

 Chain type ids used below (from atomworks.enums.ChainType):
# 0=CyclicPseudoPeptide, 1=OtherPolymer, 2=PeptideNucleicAcid,
# 3=DNA, 4=DNA_RNA_HYBRID, 5=POLYPEPTIDE_D, 6=POLYPEPTIDE_L, 7=RNA,

```

--------------------------------

### AtomWorks Dataset Three-Step Pipeline Example

Source: https://github.com/rosettacommons/atomworks/blob/production/src/atomworks/ml/datasets/README.md

Shows the internal steps of processing a dataset example: raw data retrieval, loading into an AtomArray, and transformation into features. This clarifies the data flow within AtomWorks datasets.

```python
# What happens when you call dataset[0]:
raw_data = dataset._get_raw_data(0)           # Step 1: {"path": "/data/3ne2.cif", "label": 5}
loaded = dataset._apply_loader(raw_data)       # Step 2: {"atom_array": AtomArray, "label": 5}
transformed = dataset._apply_transform(loaded) # Step 3: {"features": Tensor, "label": Tensor}
```

--------------------------------

### Run AtomWorks Dataset Tests

Source: https://github.com/rosettacommons/atomworks/blob/production/docs/mirrors.md

Command to execute tests for the various datasets within the AtomWorks project. This command assumes that the test data has already been downloaded and set up using `atomworks setup tests`.

```bash
pytest tests/ml/conftest.py

```

--------------------------------

### Install AtomWorks Base and ML Dependencies

Source: https://github.com/rosettacommons/atomworks/blob/production/README.md

Installs the base atomworks package without PyTorch, or the full package with ML dependencies including PyTorch. Also shows installation with development and Open Babel dependencies.

```shell
pip install atomworks
pip install "atomworks[ml]"
pip install "atomworks[dev]"
pip install "atomworks[openbabel]"
pip install "atomworks[ml,openbabel,dev]"
```

```shell
uv pip install "atomworks[ml,openbabel,dev]"
```

--------------------------------

### Save Failed ML Example to Disk (Python)

Source: https://github.com/rosettacommons/atomworks/blob/production/docs/ml/utils.md

Saves a failed ML example to disk as a pickle file. It takes an example ID, a directory path, and optional data, RNG state, and error messages. This is useful for debugging and analyzing pipeline failures.

```python
from atomworks.ml.utils.debug import save_failed_example_to_disk

# Example usage:
example_id = "ex_123"
fail_dir = "/path/to/failures"
data_to_save = {"input": [1, 2, 3], "output": 4}
error_message = "Model prediction failed."

save_failed_example_to_disk(example_id, fail_dir, data=data_to_save, error_msg=error_message)
```

--------------------------------

### Download and Extract PDB Metadata with AtomWorks

Source: https://github.com/rosettacommons/atomworks/blob/production/docs/mirrors.md

Downloads pre-computed PDB metadata, including PN units and interfaces, and extracts it to a specified directory. This metadata is used for calculating sampling probabilities and filtering examples for model training.

```bash
atomworks setup metadata /path/to/metadata
```

--------------------------------

### Explore FileDataset Contents in Python

Source: https://github.com/rosettacommons/atomworks/blob/production/docs/auto_examples/dataset_exploration.md

Shows how to explore the contents of an initialized FileDataset. This includes counting the number of examples and printing the raw data for a few examples. This is useful for understanding the dataset's structure and content.

```Python
# Count the number of examples in the dataset
print(f"Dataset has {len(dataset)} examples.")

# Print the raw data of the first 5 examples
for i, example in enumerate(dataset):
    if i >= 5:
        break
    print(f"Example {i + 1}: {example}")
```

--------------------------------

### Load Example Structure with AtomWorks

Source: https://github.com/rosettacommons/atomworks/blob/production/docs/auto_examples/pocket_conditioning_transform.md

Loads a protein structure (myoglobin with heme ligand) using AtomWorks' parse function. It then prints the total number of atoms, unique non-polymer residue names, and the count of heme atoms. This setup is crucial for subsequent transformations.

```Python
import biotite.structure as struc
import numpy as np
from biotite.structure import AtomArray

# AtomWorks imports
from atomworks.io import parse
from atomworks.io.utils.testing import get_pdb_path_or_buffer
from atomworks.io.utils.visualize import view
from atomworks.ml.transforms._checks import check_atom_array_annotation
from atomworks.ml.transforms.base import Transform

# AtomWorks imports
from atomworks.io import parse
from atomworks.io.utils.testing import get_pdb_path_or_buffer
from atomworks.ml.transforms._checks import check_atom_array_annotation
from atomworks.ml.transforms.base import Transform

# Load example structure (myoglobin with heme ligand; our recurring test example)
example_pdb_id = "101m"
pdb_path = get_pdb_path_or_buffer(example_pdb_id)

parse_output = parse(pdb_path)
atom_array = parse_output["assemblies"]["1"][0]

print(f"Loaded structure: {len(atom_array)} atoms")
print(f"Non-polymer residues: {np.unique(atom_array.res_name[~atom_array.is_polymer])}")
print(f"Heme atoms: {np.sum(atom_array.res_name == 'HEM')}")
```

--------------------------------

### Create PandasDataset with Custom Loader and Transform in Python

Source: https://github.com/rosettacommons/atomworks/blob/production/docs/auto_examples/dataset_exploration.md

This code example shows how to instantiate a `PandasDataset` object in Python. It takes a pandas DataFrame (`interfaces_df`) as input data, assigns a name, and configures a custom loader using `create_loader_with_query_pn_units`. It also applies a transformation pipeline (`pipe`). The example concludes by printing the number of examples in the created dataset.

```Python
from atomworks.ml.datasets import PandasDataset
from atomworks.ml.datasets.loaders import create_loader_with_query_pn_units

dataset = PandasDataset(
    data=interfaces_df,
    name="interfaces_dataset",
    # We use a pre-built loader that takes in a list of column names and returns a loader function
    loader=create_loader_with_query_pn_units(pn_unit_iid_colnames=["pn_unit_1_iid", "pn_unit_2_iid"]),
    transform=pipe,
)

print(f"Created PandasDataset with {len(dataset)} examples")
```

--------------------------------

### Enable Failed Example Saving in PandasDataset (Python)

Source: https://github.com/rosettacommons/atomworks/blob/production/src/atomworks/ml/datasets/README.md

Configures a `PandasDataset` to save examples that fail during processing. This aids in debugging by storing the example ID, error message, RNG state, and timing information for reproducible analysis.

```python
dataset = PandasDataset(
    data="data.parquet",
    name="debug_dataset",
    save_failed_examples_to_dir="/tmp/failed_examples"
)
```

--------------------------------

### Create File-based Dataset (Python)

Source: https://context7.com/rosettacommons/atomworks/llms.txt

Shows how to create a `FileDataset` where each example corresponds to a file in a specified directory. This is useful for datasets where structures are pre-processed and stored as individual files. It supports custom file patterns and transforms.

```python
from atomworks.ml.datasets import FileDataset
from pathlib import Path

# Create dataset from directory of files
# dataset = FileDataset(
#     data_dir="/path/to/structures/",
#     name="structure_files",
#     pattern="*.cif.gz",  # File pattern to match
#     transform=my_transform, # Assuming my_transform is defined elsewhere
# )
```

--------------------------------

### Configure AtomWorks Environment Variables

Source: https://github.com/rosettacommons/atomworks/blob/production/docs/mirrors.md

Sets environment variables for PDB and CCD mirror paths. These variables should be configured in a `.env` file in the repository root to inform AtomWorks where to find the local data mirrors.

```bash
PDB_MIRROR_PATH=/path/to/pdb_mirror
CCD_MIRROR_PATH=/path/to/ccd_mirror
```

--------------------------------

### Compose ML Transforms into Pipelines (Python)

Source: https://context7.com/rosettacommons/atomworks/llms.txt

Demonstrates how to chain multiple machine learning transforms into a sequential pipeline using `Compose`. It includes examples of custom transforms like `AddNoise` and `ComputeFeatures`, as well as built-in routing transforms like `ConditionalRoute` and `RandomRoute`.

```python
from atomworks.ml.transforms.base import (
    Compose, Transform, Identity, RemoveKeys,
    AddData, ConditionalRoute, RandomRoute
)
from atomworks.io import parse
from atomworks.io.utils.testing import get_pdb_path_or_buffer
import numpy as np

# Define transforms
class AddNoise(Transform):
    def __init__(self, scale: float = 0.1):
        self.scale = scale

    def forward(self, data: dict) -> dict:
        atom_array = data["atom_array"].copy()
        noise = np.random.randn(*atom_array.coord.shape) * self.scale
        atom_array.coord = atom_array.coord + noise
        data["atom_array"] = atom_array
        return data

class ComputeFeatures(Transform):
    def forward(self, data: dict) -> dict:
        atom_array = data["atom_array"]
        data["features"] = {
            "n_atoms": len(atom_array),
            "ca_coords": atom_array[atom_array.atom_name == "CA"].coord,
        }
        return data

# Create pipeline using Compose
pipeline = Compose([
    AddNoise(scale=0.05),
    ComputeFeatures(),
    RemoveKeys(["extra_info"], require_keys_exist=False),
])

# Or use + operator
pipeline_alt = AddNoise(scale=0.05) + ComputeFeatures()

# Apply pipeline
result = parse(get_pdb_path_or_buffer("101m"))
data = {"atom_array": result["assemblies"]["1"][0]}
output = pipeline(data)

print(f"Pipeline length: {len(pipeline)}")
print(f"Output keys: {list(output.keys())}")
print(f"CA coordinates shape: {output['features']['ca_coords'].shape}")

# Conditional routing
conditional_pipeline = Compose([
    ConditionalRoute(
        condition_func=lambda d: d.get("mode", "inference"),
        transform_map={
            "train": AddNoise(scale=0.1),
            "inference": Identity(),
        }
    ),
    ComputeFeatures(),
])

# Random routing (probabilistic augmentation)
random_pipeline = RandomRoute(
    transforms=[AddNoise(scale=0.05), AddNoise(scale=0.1), Identity()],
    probs=[0.3, 0.3, 0.4],
)
```

--------------------------------

### Create File-Based Dataset with Custom Loader

Source: https://github.com/rosettacommons/atomworks/blob/production/src/atomworks/ml/datasets/README.md

Demonstrates creating a FileDataset from a directory using a custom loader function. The loader parses a file and extracts a specific assembly. This is useful for datasets where each file represents a distinct example.

```python
from atomworks.ml.datasets import FileDataset
from atomworks.io import parse

# Define a simple loader function
def simple_loader(file_path):
    result = parse(file_path)
    return {"atom_array": result["assemblies"]["1"][0]}

# Create dataset from a directory
dataset = FileDataset.from_directory(
    directory="path/to/structures",
    name="my_structures",
    loader=simple_loader
)
```

--------------------------------

### Implement Fallback Dataset Wrapper (Python)

Source: https://github.com/rosettacommons/atomworks/blob/production/src/atomworks/ml/datasets/README.md

Wraps an existing dataset with `FallbackDatasetWrapper` to provide robustness during training. If an error occurs while loading an example, the wrapper attempts to load it from a specified fallback dataset, preventing training crashes.

```python
from atomworks.ml.datasets import FallbackDatasetWrapper

# Wrap dataset to enable fallback on errors
dataset_with_fallback = FallbackDatasetWrapper(
    dataset=my_dataset,
    fallback_dataset=my_dataset  # Can be same or different dataset
)
```

--------------------------------

### Create Loader with Query for Protein-Nucleic Acid Units

Source: https://github.com/rosettacommons/atomworks/blob/production/src/atomworks/ml/datasets/README.md

Shows how to use `create_loader_with_query_pn_units` to create a loader specifically designed for datasets involving protein-nucleic acid units. This factory function simplifies the setup for such common biological structures.

```python
from atomworks.ml.datasets.loaders import create_loader_with_query_pn_units

loader = create_loader_with_query_pn_units(
    pn_unit_iid_colnames=["pn_unit_1_iid", "pn_unit_2_iid"]
)
```

--------------------------------

### Create Base Structure Loading Function

Source: https://github.com/rosettacommons/atomworks/blob/production/docs/ml/datasets/datasets.md

Generates a standard loader function for loading structures from files. It requires column names for example ID, file path, and assembly ID, along with base path and file extension.

```python
from atomworks.ml.datasets.loaders import create_base_loader

loader = create_base_loader(
    example_id_colname="example_id",
    path_colname="path",
    assembly_id_colname="assembly_id",
    base_path="/data/structures",
    extension=".cif"
)
```

--------------------------------

### Residue Start Indices

Source: https://github.com/rosettacommons/atomworks/blob/production/docs/io/utils/selection.md

Gets the start (and optionally stop) indices of residues in an AtomArray, differentiating residues across different transformation IDs. This is a more robust version of biotite's residue start index function.

```APIDOC
## GET /atomworks/selection/get_residue_starts

### Description
Get the start (and optionally stop) indices of residues in an AtomArray.
This is a more robust version of `biotite.structure.residues.get_residue_starts()`
that additionally differentiates residues across different `transformation_id` values
when present. It is backwards compatible if the annotation is absent.

### Method
GET

### Endpoint
/atomworks/selection/get_residue_starts

### Parameters
#### Query Parameters
- **atom_array** (AtomArray | AtomArrayStack) - Required - Structure to analyze.
- **add_exclusive_stop** (bool) - Optional - Append an exclusive stop index at the end. Defaults to `False`.

### Response
#### Success Response (200)
- **residue_boundary_indices** (ndarray) - 1D array of residue boundary indices.

#### Response Example
```json
{
  "residue_boundary_indices": [0, 15, 30, 45]
}
```
```

--------------------------------

### Get Residue Start Indices

Source: https://github.com/rosettacommons/atomworks/blob/production/docs/io/utils/selection.md

Obtains the start and optionally stop indices of residues within an AtomArray. This function provides a more robust method than standard Biotite functions by differentiating residues across different transformation IDs when present, while maintaining backward compatibility.

```python
def get_residue_starts(atom_array: AtomArray | AtomArrayStack, add_exclusive_stop: bool = False) -> ndarray:
    """Get the start (and optionally stop) indices of residues in an AtomArray. This is a more robust version of `biotite.structure.residues.get_residue_starts()` that additionally differentiates residues across different `transformation_id` values when present. It is backwards compatible if the annotation is absent."""
    # ... implementation details ...
    pass
```

--------------------------------

### Initialize FileDataset by Scanning Directory

Source: https://github.com/rosettacommons/atomworks/blob/production/src/atomworks/ml/datasets/README.md

Shows how to create a FileDataset by scanning a directory recursively. It includes options for controlling search depth and filtering files based on a provided function.

```python
# From directory scan
dataset = FileDataset.from_directory(
    directory="path/to/structures",
    name="my_structures",
    max_depth=3,  # How deep to search subdirectories
    filter_fn=lambda path: path.suffix == ".cif"  # Optional filter
)
```

--------------------------------

### I/O Utilities

Source: https://github.com/rosettacommons/atomworks/blob/production/docs/io/utils.md

Utilities for input/output operations, including file handling, structure loading, and data conversion to CIF format.

```APIDOC
## apply_sharding_pattern()

### Description
Applies a sharding pattern to data.

### Method
(Not specified, likely a Python function call)

### Endpoint
N/A (Python function)

### Parameters
(Not specified in the provided text)

### Request Example
N/A

### Response
(Not specified in the provided text)
```

```APIDOC
## build_sharding_pattern()

### Description
Builds a sharding pattern.

### Method
(Not specified, likely a Python function call)

### Endpoint
N/A (Python function)

### Parameters
(Not specified in the provided text)

### Request Example
N/A

### Response
(Not specified in the provided text)
```

```APIDOC
## get_structure()

### Description
Retrieves a structure.

### Method
(Not specified, likely a Python function call)

### Endpoint
N/A (Python function)

### Parameters
(Not specified in the provided text)

### Request Example
N/A

### Response
(Not specified in the provided text)
```

```APIDOC
## load_any()

### Description
Loads data from any supported format.

### Method
(Not specified, likely a Python function call)

### Endpoint
N/A (Python function)

### Parameters
(Not specified in the provided text)

### Request Example
N/A

### Response
(Not specified in the provided text)
```

```APIDOC
## parse_sharding_pattern()

### Description
Parses a sharding pattern.

### Method
(Not specified, likely a Python function call)

### Endpoint
N/A (Python function)

### Parameters
(Not specified in the provided text)

### Request Example
N/A

### Response
(Not specified in the provided text)
```

```APIDOC
## read_any()

### Description
Reads data from any supported format.

### Method
(Not specified, likely a Python function call)

### Endpoint
N/A (Python function)

### Parameters
(Not specified in the provided text)

### Request Example
N/A

### Response
(Not specified in the provided text)
```

```APIDOC
## suppress_logging_messages()

### Description
Suppresses logging messages during I/O operations.

### Method
(Not specified, likely a Python function call)

### Endpoint
N/A (Python function)

### Parameters
(Not specified in the provided text)

### Request Example
N/A

### Response
(Not specified in the provided text)
```

```APIDOC
## to_cif_buffer()

### Description
Converts data to a CIF format buffer.

### Method
(Not specified, likely a Python function call)

### Endpoint
N/A (Python function)

### Parameters
(Not specified in the provided text)

### Request Example
N/A

### Response
(Not specified in the provided text)
```

```APIDOC
## to_cif_file()

### Description
Writes data to a file in CIF format.

### Method
(Not specified, likely a Python function call)

### Endpoint
N/A (Python function)

### Parameters
(Not specified in the provided text)

### Request Example
N/A

### Response
(Not specified in the provided text)
```

```APIDOC
## to_cif_string()

### Description
Converts data to a string in CIF format.

### Method
(Not specified, likely a Python function call)

### Endpoint
N/A (Python function)

### Parameters
(Not specified in the provided text)

### Request Example
N/A

### Response
(Not specified in the provided text)
```

--------------------------------

### Configure Binary Interfaces Dataset for AF3 Training

Source: https://github.com/rosettacommons/atomworks/blob/production/docs/mirrors.md

Sets up a StructuralDatasetWrapper for binary interfaces using the InterfacesDFParser. It employs the same AF3 transformation pipeline as the PN units dataset but with different spatial and contiguous crop probabilities. The dataset is loaded from a Parquet file with filters tailored for interfaces, including criteria for both interacting units and exclusion of specific ligands.

```yaml
- _target_: atomworks.ml.datasets.datasets.StructuralDatasetWrapper
      dataset_parser:
        _target_: atomworks.ml.datasets.parsers.InterfacesDFParser
      transform:
        _target_: atomworks.ml.pipelines.af3.build_af3_transform_pipeline
        is_inference: false
        n_recycles: 5
        crop_size: 256
        crop_spatial_probability: 1.0
        crop_contiguous_probability: 0.0
        diffusion_batch_size: 32
        template_lookup_path: ${paths.shared}/template_lookup.csv
        template_base_dir: ${paths.shared}/template
        # Optional MSAs (see Step 4)
        # protein_msa_dirs:
        #   - { dir: /path/to/msa, extension: .a3m.gz, directory_depth: 2 }
        # rna_msa_dirs:
        #   - { dir: /path/to/msa, extension: .afa, directory_depth: 0 }
      dataset:
        _target_: atomworks.ml.datasets.datasets.PandasDataset
        name: interfaces
        id_column: example_id
        data: /path/to/metadata/interfaces_df.parquet
        filters:
          - "deposition_date < '2022-01-01'"
          - "resolution < 5.0 and ~method.str.contains('NMR')"
          - "num_polymer_pn_units <= 20"
          - "cluster.notnull()"
          - "method in ['X-RAY_DIFFRACTION', 'ELECTRON_MICROSCOPY']"
          # Train only on D-polypeptide interfaces:
          - "pn_unit_1_type in [5, 6]"  # 5 = POLYPEPTIDE_D, 6 = POLYPEPTIDE_L
          - "pn_unit_2_type in [5, 6]"  # 5 = POLYPEPTIDE_D, 6 = POLYPEPTIDE_L
          - "~(pn_unit_1_non_polymer_res_names.notnull() and pn_unit_1_non_polymer_res_names.str.contains('${af3_excluded_ligands_regex}', regex=True))"
          - "~(pn_unit_2_non_polymer_res_names.notnull() and pn_unit_2_non_polymer_res_names.str.contains('${af3_excluded_ligands_regex}', regex=True))"
        columns_to_load: null
      cif_parser_args:
        cache_dir: null
      save_failed_examples_to_dir: null

```

--------------------------------

### get_ligand_of_interest_info

Source: https://github.com/rosettacommons/atomworks/blob/production/docs/io/transforms/categories.md

Extracts information about ligands of interest from a CIF block, referencing the PDB101 Small Molecule Ligands Guide.

```APIDOC
## get_ligand_of_interest_info

### Description
Extract ligand of interest information from a CIF block.

### Method
POST

### Endpoint
/rosettacommons/atomworks/io/transforms/categories/get_ligand_of_interest_info

### Parameters
#### Path Parameters
None

#### Query Parameters
* **cif_block** (CIFBlock) - Required - The CIF block to extract ligand information from.

#### Request Body
None

### Request Example
```json
{
  "cif_block": "<CIFBlock object>"
}
```

### Response
#### Success Response (200)
* **dict** (dict) - Dictionary containing ligand of interest information.

#### Response Example
```json
{
  "ligands": [
    {
      "chem_id": "LIG",
      "details": "Example ligand details"
    }
  ]
}
```
```

--------------------------------

### Download AtomWorks Test Data

Source: https://github.com/rosettacommons/atomworks/blob/production/docs/mirrors.md

Command to download and unpack the test data pack for AtomWorks. This pack includes a mini PDB, custom CCD, distillation data, and metadata necessary for testing the AF3 pipeline. It's recommended to run this command from the AtomWorks root directory to ensure correct placement of the downloaded files.

```bash
atomworks setup tests  # This will download the test pack to `tests/data` and unpack it there (~500 MB).

```

--------------------------------

### Compose Ligand Pocket Processing Pipeline with Python

Source: https://github.com/rosettacommons/atomworks/blob/production/docs/auto_examples/pocket_conditioning_transform.md

Demonstrates pipeline composition using `atomworks.ml.transforms.base.Compose` and the `+` operator. It shows how to create a complete ligand pocket processing pipeline by chaining `AnnotateLigandPockets` and `FeaturizePocketAtoms`. The example also illustrates applying the pipeline to fresh data and printing the results, including the number of pocket atoms found and the shape of the features. It further shows an alternative pipeline composition using the `+` operator.

```Python
from atomworks.ml.transforms.base import Compose

# Create a complete ligand pocket processing pipeline
ligand_pocket_pipeline = Compose(
    [
        AnnotateLigandPockets(pocket_distance=6.0, n_min_ligand_atoms=3),
        FeaturizePocketAtoms(feature_key="pocket_features"),
    ]
)

# Apply pipeline to fresh data
fresh_data = {"atom_array": atom_array}
pipeline_result = ligand_pocket_pipeline(fresh_data)

print("Pipeline Results:")
print(f"  Transforms applied: {[t.__class__.__name__ for t in ligand_pocket_pipeline.transforms]}")
print(f"  Pocket atoms found: {np.sum(pipeline_result['atom_array'].is_ligand_pocket)}")
print(f"  Features shape: {pipeline_result['pocket_features']['features'].shape}")

# Demonstrate the + operator
alternative_pipeline = AnnotateLigandPockets(n_min_ligand_atoms=8) + FeaturizePocketAtoms()
alt_result = alternative_pipeline({"atom_array": atom_array})
print(f"  Alternative (min 8 atoms): {np.sum(alt_result['atom_array'].is_ligand_pocket)} pocket atoms")
```

--------------------------------

### AtomSelection - Get mask

Source: https://github.com/rosettacommons/atomworks/blob/production/docs/io/utils/selection.md

Creates a boolean mask for an AtomArray based on the current AtomSelection.

```APIDOC
## CLASS AtomSelection

### Description
Represent a selection of atoms in a molecular structure.

### Method
`get_mask(atom_array: AtomArray)`

### Endpoint
N/A (Instance Method)

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
None

### Request Example
```python
# Assuming 'atom_array' is an instance of AtomArray
mask = atom_selection.get_mask(atom_array)
```

### Response
#### Success Response (200)
- **ndarray** (numpy.ndarray) - A boolean numpy array where True indicates a selected atom.

#### Response Example
```json
{
  "mask": [false, true, false, true, ...]
}
```
```

--------------------------------

### AtomSelection - Get indices

Source: https://github.com/rosettacommons/atomworks/blob/production/docs/io/utils/selection.md

Retrieves the indices of atoms within an AtomArray that match the current AtomSelection.

```APIDOC
## CLASS AtomSelection

### Description
Represent a selection of atoms in a molecular structure.

### Method
`get_idxs(atom_array: AtomArray)`

### Endpoint
N/A (Instance Method)

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
None

### Request Example
```python
# Assuming 'atom_array' is an instance of AtomArray
indices = atom_selection.get_idxs(atom_array)
```

### Response
#### Success Response (200)
- **ndarray** (numpy.ndarray) - A numpy array containing the indices of the selected atoms.

#### Response Example
```json
{
  "indices": [10, 15, 22]
}
```
```

--------------------------------

### Sync PDB Mirror with AtomWorks

Source: https://github.com/rosettacommons/atomworks/blob/production/docs/mirrors.md

Synchronizes the PDB mirror to a specified local directory. Supports downloading the entire PDB, specific PDB IDs, or a list of IDs from a file. The downloaded files follow the RCSB sharding pattern.

```bash
atomworks pdb sync /path/to/pdb_mirror

# Or download specific PDB IDs only
atomworks pdb sync /path/to/pdb_mirror --pdb-id 1A0I --pdb-id 7XYZ

# Or from a file of IDs (one per line)
atomworks pdb sync /path/to/pdb_mirror --pdb-ids-file /path/to/ids.txt
```

--------------------------------

### Initialize FileDataset from Explicit File List

Source: https://github.com/rosettacommons/atomworks/blob/production/src/atomworks/ml/datasets/README.md

Demonstrates initializing a FileDataset by providing an explicit list of file paths. This method is direct and suitable when the file locations are known and fixed.

```python
# From explicit file list
dataset = FileDataset(
    file_paths=["file1.cif", "file2.cif"],
    name="my_files",
    loader=my_loader_fn
)
```

--------------------------------

### Get Annotation

Source: https://github.com/rosettacommons/atomworks/blob/production/docs/io/utils/selection.md

Retrieves an annotation array from an AtomArray or AtomArrayStack. Allows specifying the annotation body dimension and a default value if the annotation is missing.

```APIDOC
## GET /atomworks/selection/get_annotation

### Description
Return an annotation array if present, otherwise `default`.
If `n_body` is `None`, the dimensionality is auto-detected by probing 1D then 2D annotation categories.

### Method
GET

### Endpoint
/atomworks/selection/get_annotation

### Parameters
#### Query Parameters
- **atom_array** (AtomArray | AtomArrayStack) - Required - Structure to query.
- **annot** (str) - Required - Annotation category name.
- **n_body** (int | None) - Optional - 1 for 1D annotations, 2 for 2D annotations; auto-detected if `None`.
- **default** (Any) - Optional - Value to return if the annotation is missing. Defaults to `None`.

### Response
#### Success Response (200)
- **annotation_array** (ndarray) - The requested annotation array or `default` if missing.

#### Response Example
```json
{
  "annotation_array": ["A", "B", "A", ...]
}
```
```

--------------------------------

### Test Featurization Pipeline with Python

Source: https://github.com/rosettacommons/atomworks/blob/production/docs/auto_examples/pocket_conditioning_transform.md

Tests the `FeaturizePocketAtoms` transform by creating a pipeline that first applies `AnnotateLigandPockets` and then `FeaturizePocketAtoms`. It demonstrates how to initialize the transforms, apply them sequentially to data, and print the shape, names, type, and counts of the generated features.

```Python
# Test featurization using a proper pipeline
# First apply the annotation transform, then the featurization
annotator = AnnotateLigandPockets(pocket_distance=6.0, n_min_ligand_atoms=5)
featurizer = FeaturizePocketAtoms()

# Apply both transforms in sequence
data = {"atom_array": atom_array}
annotated_data = annotator(data)
feature_result = featurizer(annotated_data)

features = feature_result["pocket_features"]
print(f"Generated features: {features['features'].shape}")
print(f"Feature names: {features['feature_names']}")
print(f"Feature type: {type(features['features'])}")
print(f"Pocket atoms (sum): {features['features'].sum():.0f}")
print(f"Non-pocket atoms: {len(features['features']) - features['features'].sum():.0f}")
```

--------------------------------

### Annotation Start/Stop Indices

Source: https://github.com/rosettacommons/atomworks/blob/production/docs/io/utils/selection.md

Computes the start and stop indices for segments in an AtomArray where specified annotations change. Allows for an optional exclusive stop index.

```APIDOC
## POST /atomworks/selection/annot_start_stop_idxs

### Description
Computes the start and stop indices for segments in an AtomArray where any of the specified annotation(s) change.

### Method
POST

### Endpoint
/atomworks/selection/annot_start_stop_idxs

### Parameters
#### Request Body
- **atom_array** (AtomArray | AtomArrayStack) - Required - The AtomArray to process.
- **annots** (str | list[str]) - Required - Annotation name or names to define segments.
- **add_exclusive_stop** (bool) - Optional - Append an exclusive stop index at the end. Defaults to `False`.

### Request Example
```json
{
  "atom_array": "<AtomArray object>",
  "annots": "chain_id",
  "add_exclusive_stop": true
}
```

### Response
#### Success Response (200)
- **start_stop_indices** (ndarray) - 1D array of start/stop indices that bound segments.

#### Response Example
```json
{
  "start_stop_indices": [0, 5, 10, 15]
}
```
```

--------------------------------

### Get ChEMBL Normalizer Instance

Source: https://github.com/rosettacommons/atomworks/blob/production/docs/io/tools/rdkit.md

Provides access to a cached ChEMBL normalizer instance. This is used for standardizing chemical structures according to ChEMBL conventions.

```python
from typing import Any

class Normalizer:
    pass

def get_chembl_normalizer() -> Normalizer:
    """Cached ChEMBL normalizer."""
    # Implementation details would go here
    pass
```

--------------------------------

### Assembly Utilities

Source: https://github.com/rosettacommons/atomworks/blob/production/docs/io/utils.md

Provides functions for building assemblies from asymmetric units.

```APIDOC
## build_assemblies_from_asym_unit()

### Description
Builds assemblies from an asymmetric unit.

### Method
(Not specified, likely a Python function call)

### Endpoint
N/A (Python function)

### Parameters
(Not specified in the provided text)

### Request Example
N/A

### Response
(Not specified in the provided text)
```

--------------------------------

### Create and Concatenate Pandas Datasets

Source: https://context7.com/rosettacommons/atomworks/llms.txt

Demonstrates creating multiple PandasDatasets from DataFrames and concatenating them into a single dataset. Supports accessing data by global index or example ID.

```python
import pandas as pd
from atomworks.data import PandasDataset, ConcatDatasetWithID

df1 = pd.DataFrame({"example_id": ["a", "b"], "value": [1, 2]})
df2 = pd.DataFrame({"example_id": ["c", "d"], "value": [3, 4]})

dataset1 = PandasDataset(data=df1, name="dataset1", id_column="example_id")
dataset2 = PandasDataset(data=df2, name="dataset2", id_column="example_id")

combined = ConcatDatasetWithID([dataset1, dataset2])
print(f"Combined size: {len(combined)}")

# Example access (commented out in original)
# example = combined[2]  # Gets "c" from dataset2
# idx = combined.id_to_idx("c")  # Returns global index
```

--------------------------------

### Get RNA Chemical Components

Source: https://github.com/rosettacommons/atomworks/blob/production/docs/io/utils/ccd.md

Returns a frozenset containing the CCD codes for all standard RNA chemical components. This function is similar to `na_chem_comps` but specifically for RNA.

```python
from atomworks.io.utils.ccd import rna_chem_comps

rna_components = rna_chem_comps()
print(rna_components)
```

--------------------------------

### Import Modules for Structure Visualization (Python)

Source: https://github.com/rosettacommons/atomworks/blob/production/docs/tutorial/visualisation.ipynb

Imports necessary modules from atomworks.io, biotite, and io for loading, parsing, and visualizing molecular structures. It includes functions to fetch PDB files locally or from RCSB and utilities for PyMOL visualization.

```python
import io

from biotite.database import rcsb

from atomworks.io.utils.testing import get_pdb_path
from atomworks.io.utils.visualize import view


def get_example_path_or_buffer(pdb_id: str) -> str | io.StringIO:
    try:
        # ... if file is locally available
        return get_pdb_path(pdb_id)
    except FileNotFoundError:
        # ... otherwise, fetch the file from RCSB
        return rcsb.fetch(pdb_id, format="cif")


from atomworks.io.utils.io_utils import load_any
from atomworks.io.utils.visualize import view_pymol
```

--------------------------------

### Get DNA Chemical Components

Source: https://github.com/rosettacommons/atomworks/blob/production/docs/io/utils/ccd.md

Retrieves a frozenset containing the three-letter codes for all standard DNA chemical components. This is useful for filtering or validating DNA-related entries.

```python
from atomworks.io.utils.ccd import dna_chem_comps

dna_components = dna_chem_comps()
print(dna_components)
```

--------------------------------

### Initialize PandasDataset with Filtering and Column Selection

Source: https://github.com/rosettacommons/atomworks/blob/production/src/atomworks/ml/datasets/README.md

Explains how to initialize a PandasDataset from a file or DataFrame, specifying an ID column, applying multiple filters, and selecting specific columns for loading. This optimizes data loading and access.

```python
dataset = PandasDataset(
    data="path/to/data.parquet",  # Or a DataFrame
    name="my_dataset",
    id_column="example_id",  # Column to use for ID-based access
    filters=[                # Optional pandas query filters
        "resolution < 2.5",
        "method == 'X-RAY_DIFFRACTION'"
    ],
    columns_to_load=["path", "label"],  # Load subset of columns (efficient for Parquet)
    loader=my_loader_fn,
    transform=my_transform_pipeline
)
```

--------------------------------

### Get RDKit Valence Checker Instance

Source: https://github.com/rosettacommons/atomworks/blob/production/docs/io/tools/rdkit.md

Provides access to a cached RDKit valence checker instance. This checker is used to validate the valences of atoms within an RDKit molecule.

```python
from typing import Any

class RDKitValidation:
    pass

def get_valence_checker() -> RDKitValidation:
    """Cached RDKit valence checker."""
    # Implementation details would go here
    pass
```

--------------------------------

### Get Unknown CCD Code for Chemical Component Type

Source: https://github.com/rosettacommons/atomworks/blob/production/docs/io/utils/ccd.md

Retrieves the CCD code for an unknown chemical component type. This function can be used when a specific component type is not directly recognized.

```python
from atomworks.io.utils.ccd import get_unknown_ccd_code_for_chem_comp_type

unknown_code = get_unknown_ccd_code_for_chem_comp_type("SOME_UNKNOWN_TYPE")
print(unknown_code)
```

--------------------------------

### Create Loader with Query PN Units (Python)

Source: https://github.com/rosettacommons/atomworks/blob/production/src/atomworks/ml/datasets/README.md

Initializes a data loader for PDB structures, specifically designed for cropping operations using query protein units. It takes column names for interface units, a base path for structure files, and the file extension.

```python
loader = create_loader_with_query_pn_units(
    pn_unit_iid_colnames=["pn_unit_1_iid", "pn_unit_2_iid"],  # For interfaces
    base_path="/data/pdb",
    extension=".cif.gz"
)
```

--------------------------------

### Get Chain Type from Chemical Component Type

Source: https://github.com/rosettacommons/atomworks/blob/production/docs/io/utils/ccd.md

Retrieves the ChainType enum corresponding to a given chemical component type string. This is useful for categorizing components within a molecular structure.

```python
from atomworks.io.utils.ccd import get_chain_type_from_chem_comp_type

chain_type = get_chain_type_from_chem_comp_type("L-PEPTIDE LINKING")
print(chain_type)
```

--------------------------------

### Configure PN Units Dataset for AF3 Training

Source: https://github.com/rosettacommons/atomworks/blob/production/docs/mirrors.md

Sets up a StructuralDatasetWrapper for protein-nucleic acid (PN) units using the PNUnitsDFParser. It applies a transformation pipeline built with `build_af3_transform_pipeline` for training (is_inference=false), specifying parameters like number of recycles, crop size, and probabilities. The dataset is loaded from a Parquet file with various filters applied to select specific data points, including deposition date, resolution, number of PN units, and unit types.

```yaml
_target_: atomworks.ml.datasets.datasets.ConcatDatasetWithID
  datasets:
    # Single PN units
    - _target_: atomworks.ml.datasets.datasets.StructuralDatasetWrapper
      dataset_parser:
        _target_: atomworks.ml.datasets.parsers.PNUnitsDFParser
      transform:
        _target_: atomworks.ml.pipelines.af3.build_af3_transform_pipeline
        is_inference: false
        n_recycles: 5  # This means that we will subsample 5 random sets from the MSA for each example.
        crop_size: 256
        crop_contiguous_probability: 0.3333333333333333
        crop_spatial_probability: 0.6666666666666666
        diffusion_batch_size: 32
        # Optional templates (if available)
        template_lookup_path: ${paths.shared}/template_lookup.csv
        template_base_dir: ${paths.shared}/template
        # Optional MSAs (see Step 4)
        # protein_msa_dirs:
        #   - { dir: /path/to/msa, extension: .a3m.gz, directory_depth: 2 }
        # rna_msa_dirs:
        #   - { dir: /path/to/msa, extension: .afa, directory_depth: 0 }
      dataset:
        _target_: atomworks.ml.datasets.datasets.PandasDataset
        name: pn_units
        id_column: example_id
        data: /path/to/metadata/pn_units_df.parquet
        filters:
          - "deposition_date < '2022-01-01'"
          - "resolution < 5.0 and ~method.str.contains('NMR')"
          - "num_polymer_pn_units <= 20"
          - "cluster.notnull()"
          - "method in ['X-RAY_DIFFRACTION', 'ELECTRON_MICROSCOPY']"
          # Train only on D-polypeptides:
          - "q_pn_unit_type in [5, 6]"  # 5 = POLYPEPTIDE_D, 6 = POLYPEPTIDE_L
          # Exclude ligands from AF3 excluded set:
          - "~(q_pn_unit_non_polymer_res_names.notnull() and q_pn_unit_non_polymer_res_names.str.contains('${af3_excluded_ligands_regex}', regex=True))"
        columns_to_load: null
      save_failed_examples_to_dir: null

```

--------------------------------

### Get CCD Component from Biotite

Source: https://github.com/rosettacommons/atomworks/blob/production/docs/io/utils/ccd.md

Retrieves a specific chemical component from Biotite's built-in Chemical Component Dictionary using its CCD code. It returns the component as an AtomArray object.

```python
from atomworks.io.utils.ccd import get_ccd_component_from_biotite

# Retrieve the ALA component from Biotite's CCD
atom_array = get_ccd_component_from_biotite("ALA")
print(atom_array)
```

--------------------------------

### Get Annotation Array

Source: https://github.com/rosettacommons/atomworks/blob/production/docs/io/utils/selection.md

Retrieves an annotation array from an AtomArray or AtomArrayStack. If the annotation is not found, it returns a specified default value. It can auto-detect the dimensionality (1D or 2D) of the annotation if not explicitly provided.

```python
def get_annotation(atom_array: AtomArray | AtomArrayStack, annot: str, n_body: int | None = None, default: Any = None) -> ndarray:
    """Return an annotation array if present, otherwise `default`. If `n_body` is `None`, the dimensionality is auto-detected by probing 1D then 2D annotation categories."""
    # ... implementation details ...
    pass
```

--------------------------------

### Create Base Loader Factory

Source: https://github.com/rosettacommons/atomworks/blob/production/src/atomworks/ml/datasets/README.md

Demonstrates using `create_base_loader` to generate a loader function for simple datasets. It configures parameters like column names, base paths, file extensions, and sharding patterns for efficient data loading.

```python
from atomworks.ml.datasets.loaders import create_base_loader, create_loader_with_query_pn_units

# Basic loader for simple datasets
loader = create_base_loader(
    example_id_colname="example_id",
    path_colname="structure_path",
    assembly_id_colname="assembly_id",
    base_path="/data/pdb",           # Prepended to paths
    extension=".cif.gz",              # Added if missing
    sharding_pattern="/1:2/",         # e.g., "3ne2" → "3n/3ne2.cif.gz"
    parser_args={"remove_waters": True}
)
```

--------------------------------

### Testing AnnotateLigandPockets Function in Python

Source: https://github.com/rosettacommons/atomworks/blob/production/docs/auto_examples/pocket_conditioning_transform.md

This Python code snippet demonstrates how to test the standalone `annotate_ligand_pockets` function. It shows the instantiation of an `AtomArray` (assumed to be defined elsewhere) and the subsequent call to the function with specified parameters. It also highlights the use of AtomWorks' 'query' syntax for data manipulation.

```Python
# Test the functional version
result_array = annotate_ligand_pockets(
    atom_array, pocket_distance=6.0, n_min_ligand_atoms=5, annotation_name="is_ligand_pocket"
)

# Here, we are using AtomWork's "query" syntax for convenience, which operates similar to Pandas DataFrame queries

```

--------------------------------

### Get Amino Acid Chemical Components

Source: https://github.com/rosettacommons/atomworks/blob/production/docs/io/utils/ccd.md

Retrieves a frozenset containing the three-letter codes for all standard amino acid chemical components. This is useful for filtering or validating amino acid entries.

```python
from atomworks.io.utils.ccd import aa_chem_comps

amino_acids = aa_chem_comps()
print(amino_acids)
```

--------------------------------

### Concatenate Datasets with ID Preservation (Python)

Source: https://context7.com/rosettacommons/atomworks/llms.txt

Demonstrates the use of `ConcatDatasetWithID` to combine multiple datasets while ensuring that example IDs are preserved and unique. This allows for seamless merging of datasets with consistent ID lookups.

```python
from atomworks.ml.datasets import ConcatDatasetWithID, PandasDataset
import pandas as pd

# Assuming PandasDataset and other necessary components are defined and instantiated
```

--------------------------------

### Selection Utilities

Source: https://github.com/rosettacommons/atomworks/blob/production/docs/io/utils.md

Utilities for creating and manipulating atom selections, including selection stacks and annotations.

```APIDOC
## AtomSelection

### Description
Represents a selection of atoms.

### Method
(Not specified, likely a Python class)

### Endpoint
N/A (Python class)

### Parameters
(Not specified in the provided text)

### Request Example
N/A

### Response
(Not specified in the provided text)
```

```APIDOC
## AtomSelectionStack

### Description
Manages a stack of atom selections.

### Method
(Not specified, likely a Python class)

### Endpoint
N/A (Python class)

### Parameters
(Not specified in the provided text)

### Request Example
N/A

### Response
(Not specified in the provided text)
```

```APIDOC
## annot_start_stop_idxs()

### Description
Annotates start and stop indices.

### Method
(Not specified, likely a Python function call)

### Endpoint
N/A (Python function)

### Parameters
(Not specified in the provided text)

### Request Example
N/A

### Response
(Not specified in the provided text)
```

```APIDOC
## get_annotation()

### Description
Retrieves annotation information.

### Method
(Not specified, likely a Python function call)

### Endpoint
N/A (Python function)

### Parameters
(Not specified in the provided text)

### Request Example
N/A

### Response
(Not specified in the provided text)
```

```APIDOC
## get_residue_starts()

### Description
Gets the starting indices of residues.

### Method
(Not specified, likely a Python function call)

### Endpoint
N/A (Python function)

### Parameters
(Not specified in the provided text)

### Request Example
N/A

### Response
(Not specified in the provided text)
```