FRIdata

Purpose

Generate sequences, coordinates, distograms, and embeddings from protein structures at scale. Supports PDB, AFDB, ESMatlas, and local/custom inputs.

Backend

Built on Dask for distributed, parallel processing on CPU and GPU clusters.

HPC Ready

Includes SLURM job scripts for the PLGrid infrastructure as examples.

How FRIdata works

Every dataset is defined by a database type (-d, --db) and a collection type (-c, --collection), plus optional --proteome and --version. Together they determine the dataset name (for example AFDB-subset--test) and how structures are resolved.

Use -t / --type to choose what to generate in a run: sequences, coordinates, distograms, embeddings, or all.

Database type (`-d`, `--db`)

Where structures come from.

Value	Meaning
`PDB`	RCSB PDB structures (download by ID)
`AFDB`	AlphaFold Database
`ESMatlas`	ESM Atlas
`other`	Local or custom files via `--input-path` or archives

Collection type (`-c`, `--collection`)

How much of that source to include.

Value	Meaning
`all`	Full database collection
`part`	AFDB proteome partition (requires `--proteome`; foldcomp-based)
`clust`	AFDB cluster partition (requires `--proteome`; foldcomp-based)
`subset`	User-defined ID list via `-i` / `--ids` (optionally `--input-path` for local structures)

Installation

Clone the repository

git clone https://github.com/Tomasz-Lab/FRIdata.git
cd FRIdata

Install Miniconda, then install mamba:

# prioritize the conda-forge channel
conda config --add channels conda-forge
conda update -n base --all
conda install -n base mamba

Create the environment (recommended). The setup script creates the conda environment, installs the pip dependencies, and installs a PyTorch build matched to your GPU driver:
```
# default install
./scripts/setup_env.sh

# custom environment name
./scripts/setup_env.sh -n fridata_gpu_verify

# CPU-only systems
./scripts/setup_env.sh --cpu
```

Activate the environment

# shell type is one of: bash, cmd.exe, dash, fish, nu, posix, powershell, tcsh, xonsh, zsh
eval "$(mamba shell hook --shell <shell type>)"
mamba activate fridata_env

(Optional) Run tests to verify the installation
```
pytest ./tests
```

Manual installation — if you prefer to run the steps yourself:

mamba env create -f fridata_env_conda.yml -n fridata_env
mamba run -n fridata_env python -m pip install -r requirements-fridata.txt
./scripts/install_pytorch.sh

Pip dependencies live in requirements-fridata.txt rather than fridata_env_conda.yml so environment creation works from read-only repository checkouts.

Troubleshooting — if you see Error opening for writing ".../mambaf..." while creating the environment, libmamba is trying to write a temporary pip requirements file next to the environment YAML. Use ./scripts/setup_env.sh instead, or ensure the directory containing fridata_env_conda.yml is writable.

Usage examples

Runnable walkthroughs for common setups. AFDB is one supported database; see Overview for the full -d / -c model.

AFDB subset (local)

Requires a directory of AFDB structures and a text file listing AFDB IDs (one per line, \n delimited). Assumes the installation steps above succeeded.

FRIDATA_PATH="<repository path>"
AFDB_PATH="<AFDB structures directory path>"
IDS_PATH="<AFDB IDs file path>"

cd ${FRIDATA_PATH}

EMBEDDER_TYPE=esm2_t33_650M_UR50D

# (macOS only) fix for OpenMP multiple runtime error
export KMP_DUPLICATE_LIB_OK=TRUE

PYTHONPATH='.' python3 -u ${FRIDATA_PATH}/fridata.py \
  generate_data \
  -t sequences,coordinates,distograms,embeddings \
  -d AFDB \
  -c subset \
  --version test \
  -i ${IDS_PATH} \
  --input-path ${AFDB_PATH} \
  -e ${EMBEDDER_TYPE}

For subset runs with --input-path, new datasets store canonical keys as {line_from_ids_file}_{chain} (for example A0A2K6V5L6_A), not the full AlphaFold CIF filename stem. The dataset's input_structures.idx maps each canonical key to the source structure filename. Older datasets created before this convention may still use long AF-style keys.

CLI installation

From the FRIdata directory, install it as a CLI tool:

python3 -m pip install -e .

# now available as
fridata <...>

Use ids_file tokens (e.g. plain UniProt) plus chain as the canonical dataset index keys.

HPC (PLGrid / SLURM)

Running FRIdata on HPC differs between CPU and GPU nodes. These instructions are valid for HPC hosted in PLGrid infrastructure; other infrastructures may require additional adjustments.

Prerequisites:

An active grant valid on the HPC
Loaded modules: miniconda3 and gcc
Mandatory environment variables (ideally set in .bashrc):
- DEEPFRI_PATH — parent directory of this repo
- IDS_PATH — text file with AFDB indexes listed
- AFDB_PATH — path to AFDB structures (may be an empty directory — structures will be fetched there)
- DATA_PATH — parent directory of all generated output data
Optional environment variables (with defaults):
- COMMON_SLURM_PATH — defaults to $DEEPFRI_PATH/FRIdata/scripts/hpc/common_slurm.sh
- LAUNCH_WORKER_SLURM_PATH — defaults to $DEEPFRI_PATH/FRIdata/scripts/hpc/launch_workers_slurm.sh
- MEMORY_LIMIT — memory limit per Dask worker, defaults to 288GiB
- IP_INTERFACE — network interface where Dask workers connect, defaults to ens1f0
- CONDA_ENV_PATH — path to conda environment, defaults to $DEEPFRI_PATH/conda_dev

Preperation steps:

# 1. Clone the repo
git clone https://github.com/Tomasz-Lab/FRIdata.git
cd FRIdata

# 2. Update run permissions
chmod u+x -R scripts/hpc/cpu

# 3. Initialize SLURM (add --cpu when running on a CPU cluster)
./scripts/hpc/initialize_slurm.sh <path to .conda> [--cpu]

Then, schedule the sbatch script. Operations: sequences, coordinates, embeddings, distograms.

# CPU nodes
sbatch --cpus-per-task=<cpus> --time=<HH:MM:SS> --nodes=<nodes> \
  --account=<grant name> scripts/hpc/run_slurm.sh sequences,coordinates

# GPU nodes
sbatch --gres=gpu[:gpu-number] --time=<HH:MM:SS> --account=<grant name> \
  --nodes=1 --partition=<partition name> --cpus-per-task=<cpus> \
  scripts/hpc/run_slurm.sh embeddings

API reference

Primary documentation — start here for flags, subcommands, and accepted values.

FRIdata is organized into sub-commands. The most common is generate_data, which creates a dataset and generates the requested data types in one run.

Sub-commands

Command	Description
`generate_data`	Create a dataset and generate sequences, coordinates, distograms and embeddings
`create_dataset`	Create a protein dataset only
`generate_sequence`	Generate sequences from a dataset
`generate_distograms`	Generate distograms from a dataset
`generate_embeddings`	Generate embeddings from a dataset
`read_distograms`	Read distograms from a dataset
`read_pdbs`	Read PDBs (print to terminal or extract with `--to_directory`)
`verify_chains`	Verify chains in a dataset
`create_archive`	Build one merged PDB `.zip` from the dataset index
`create_dashboard`	Export a single-file HTML report for dataset indexes
`inspect_h5`	Inspect an `.h5` file (tree, keys, structures, numeric dumps)
`inspect_idx`	Pretty-print an `.idx` (JSON) file
`remove_dataset`	Remove all traces of a dataset

Dataset options

Flag	Description
`-p`, `--file-path`	Existing dataset directory (with `dataset.json`) or path to `dataset.json`
`-d`, `--db`	Database type (see choices below)
`-c`, `--collection`	Collection type (see choices below)
`-t`, `--type`	Comma-separated types to generate: `dataset`, `sequences`, `coordinates`, `distograms`, `embeddings`, or `all` (`generate_data` only)
`-e`, `--embedder`	Embedder type (see choices below)
`-i`, `--ids`	File with IDs used to create a subset
`-s`, `--seqres`	FASTA file to use as the sequence source
`--input-path`	Input directory or archive (zip/tar.gz) of protein files (pdb/cif)
`--archive`	Path to a tar.gz archive containing structure files
`--proteome`	Precise proteome of an AFDB `part` dataset
`--version`	String to differentiate datasets (default: current date)
`-b`, `--batch-size`	Batch size
`-o`, `--overwrite`	Overwrite existing files (default: false)
`--binary`	Download binary CIF in the PDB database
`--slurm`	Use the SLURM job scheduler
`-v`, `--verbose`	Enable verbose logging
`--log-file`	Path to a log file for file logging
`--config`	Path to a config JSON file (default: `./config.json`)

Choices

Option	Accepted values
`-d`, `--db`	`PDB`, `AFDB`, `ESMatlas`, `other`
`-c`, `--collection`	`all`, `clust`, `part`, `subset`
`-e`, `--embedder`	`esm2_t30_150M_UR50D`, `esm2_t33_650M_UR50D`, `esmc_300m`, `esmc_600m`, `gLM2_150M`, `gLM2_650M`

Command reference (--help)

Verbatim output of fridata <subcommand> --help for each subcommand. Run fridata <subcommand> --help locally to reproduce it.

generate_data

usage: fridata generate_data [-h] [--slurm] [-p FILE_PATH] [-d name] [-c name]
                             [-t [name]] [--proteome [name]]
                             [--version VERSION] [-i IDS] [-s SEQRES] [-o]
                             [-b BATCH_SIZE] [--binary]
                             [--input-path INPUT_PATH] [--archive ARCHIVE]
                             [-v] [-e name]

options:
  -h, --help            show this help message and exit
  --slurm               Use SLURM job scheduler
  -p FILE_PATH, --file-path FILE_PATH
                        Existing dataset directory (with dataset.json) or path
                        to dataset.json
  -d name, --db name    Database Types: PDB AFDB ESMatlas other
  -c name, --collection name
                        Collection Types: all clust part subset
  -t [name], --type [name]
                        Comma-separated types to generate: dataset, sequences,
                        coordinates, distograms, embeddings; or 'all'. For
                        generate_data only; ignored by create_dataset.
  --proteome [name]     Precise proteome of AFDB part dataset
  --version VERSION     String to differentiate datasets; default: current
                        date
  -i IDS, --ids IDS     File with ids to create subset
  -s SEQRES, --seqres SEQRES
                        fasta file to use as sequence source
  -o, --overwrite       Should overwrite existing files? Default - false
  -b BATCH_SIZE, --batch-size BATCH_SIZE
  --binary              Download binary CIF in PDB db
  --input-path INPUT_PATH
                        Path to input directory or archive (zip/tar.gz) with
                        protein files (pdb/cif)
  --archive ARCHIVE     Path to tar.gz archive containing structure files
  -v, --verbose         Enable verbose logging mode
  -e name, --embedder name
                        Embedder Types: esm2_t30_150M_UR50D
                        esm2_t33_650M_UR50D esmc_300m esmc_600m gLM2_150M
                        gLM2_650M

create_dataset

usage: fridata create_dataset [-h] [--slurm] -d name -c name [-t [name]]
                              [--proteome [name]] [--version VERSION] [-i IDS]
                              [-s SEQRES] [-o] [-b BATCH_SIZE] [--binary]
                              [--input-path INPUT_PATH] [--archive ARCHIVE]
                              [-v] [-e name]

options:
  -h, --help            show this help message and exit
  --slurm               Use SLURM job scheduler
  -d name, --db name    Database Types: PDB AFDB ESMatlas other
  -c name, --collection name
                        Collection Types: all clust part subset
  -t [name], --type [name]
                        Comma-separated types to generate: dataset, sequences,
                        coordinates, distograms, embeddings; or 'all'. For
                        generate_data only; ignored by create_dataset.
  --proteome [name]     Precise proteome of AFDB part dataset
  --version VERSION     String to differentiate datasets; default: current
                        date
  -i IDS, --ids IDS     File with ids to create subset
  -s SEQRES, --seqres SEQRES
                        fasta file to use as sequence source
  -o, --overwrite       Should overwrite existing files? Default - false
  -b BATCH_SIZE, --batch-size BATCH_SIZE
  --binary              Download binary CIF in PDB db
  --input-path INPUT_PATH
                        Path to input directory or archive (zip/tar.gz) with
                        protein files (pdb/cif)
  --archive ARCHIVE     Path to tar.gz archive containing structure files
  -v, --verbose         Enable verbose logging mode
  -e name, --embedder name
                        Embedder Types: esm2_t30_150M_UR50D
                        esm2_t33_650M_UR50D esmc_300m esmc_600m gLM2_150M
                        gLM2_650M

generate_sequence

usage: fridata generate_sequence [-h] [--slurm] -p FILE_PATH [-v]
                                 [--log-file LOG_FILE] [--ca_mask]
                                 [--no_substitution]

options:
  -h, --help            show this help message and exit
  --slurm               Use SLURM job scheduler
  -p FILE_PATH, --file-path FILE_PATH
  -v, --verbose         Enable verbose logging
  --log-file LOG_FILE   Path to log file for file logging
  --ca_mask             Require a carbon alpha atom to include an amino acid
                        in a sequence
  --no_substitution     Don't substitute non standard amino acids

generate_distograms

usage: fridata generate_distograms [-h] [--slurm] -p FILE_PATH [-v]
                                   [--log-file LOG_FILE]

options:
  -h, --help            show this help message and exit
  --slurm               Use SLURM job scheduler
  -p FILE_PATH, --file-path FILE_PATH
  -v, --verbose         Enable verbose logging
  --log-file LOG_FILE   Path to log file for file logging

generate_embeddings

usage: fridata generate_embeddings [-h] [--slurm] -p FILE_PATH [-v]
                                   [--log-file LOG_FILE] -e name

options:
  -h, --help            show this help message and exit
  --slurm               Use SLURM job scheduler
  -p FILE_PATH, --file-path FILE_PATH
  -v, --verbose         Enable verbose logging
  --log-file LOG_FILE   Path to log file for file logging
  -e name, --embedder name
                        Embedder Types: esm2_t30_150M_UR50D
                        esm2_t33_650M_UR50D esmc_300m esmc_600m gLM2_150M
                        gLM2_650M

read_distograms

usage: fridata read_distograms [-h] [--slurm] -p FILE_PATH [-v]
                               [--log-file LOG_FILE]

options:
  -h, --help            show this help message and exit
  --slurm               Use SLURM job scheduler
  -p FILE_PATH, --file-path FILE_PATH
  -v, --verbose         Enable verbose logging
  --log-file LOG_FILE   Path to log file for file logging

read_pdbs

usage: fridata read_pdbs [-h] [--print] [--to_directory TO_DIRECTORY] [-i IDS]
                         [--slurm] -p FILE_PATH [-v] [--log-file LOG_FILE]

options:
  -h, --help            show this help message and exit
  --print               Print PDB files to the terminal
  --to_directory TO_DIRECTORY
                        Extract PDB files to the provided directory
  -i IDS, --ids IDS     File with ids to extract
  --slurm               Use SLURM job scheduler
  -p FILE_PATH, --file-path FILE_PATH
  -v, --verbose         Enable verbose logging
  --log-file LOG_FILE   Path to log file for file logging

verify_chains

usage: fridata verify_chains [-h] [--slurm] -p FILE_PATH [-v]
                             [--log-file LOG_FILE]

options:
  -h, --help            show this help message and exit
  --slurm               Use SLURM job scheduler
  -p FILE_PATH, --file-path FILE_PATH
  -v, --verbose         Enable verbose logging
  --log-file LOG_FILE   Path to log file for file logging

create_archive

usage: fridata create_archive [-h] [--keep-shard-zips] [--slurm] -p FILE_PATH
                              [-v] [--log-file LOG_FILE]

options:
  -h, --help            show this help message and exit
  --keep-shard-zips     Keep per-shard zip files under _staging/ after merge
                        (default: remove staging)
  --slurm               Use SLURM job scheduler
  -p FILE_PATH, --file-path FILE_PATH
  -v, --verbose         Enable verbose logging
  --log-file LOG_FILE   Path to log file for file logging

Working directory: a `mate` helper (if you use it) often cd's into deepFRI2-toolbox-dev; after that, pass an absolute -p/--file-path to the dataset directory (the folder containing dataset.json), or cd to <data_path>/datasets first if you want to use a short folder name.
Invocation: run `python fridata.py create_archive ...` from the repo unless fridata.py is installed on PATH.

create_dashboard

usage: fridata create_dashboard [-h] [--root ROOT] [--index-types INDEX_TYPES]
                                [--output-dir OUTPUT_DIR]
                                dataset

positional arguments:
  dataset               Dataset: filesystem path (directory or dataset.json),
                        slug (folder suffix after --), or full folder name
                        under --root when not an existing path

options:
  -h, --help            show this help message and exit
  --root ROOT           Override datasets root (default: <data_path>/datasets)
  --index-types INDEX_TYPES
                        Comma-separated index types
                        (dataset,sequences,coordinates,embeddings,distograms)
                        or 'all'
  --output-dir OUTPUT_DIR
                        Output directory for reports (default:
                        <repo_root>/reports)

inspect_h5

usage: fridata inspect_h5 [-h]
                          [--mode {preview,structures,keys,distograms,coordinates,embeddings}]
                          [--name KEYS] [--save [FILEPATH]]
                          file

positional arguments:
  file                  Path to h5 file

options:
  -h, --help            show this help message and exit
  --mode {preview,structures,keys,distograms,coordinates,embeddings}
                        preview: H5 tree; keys: id list; structures: full PDB
                        text (files group); distograms|coordinates|embeddings:
                        numeric preview (or full dump with --name)
  --name KEYS           Comma-separated protein keys. Only with --mode
                        structures|distograms|coordinates|embeddings. For
                        numeric modes, emits a full array dump. Use inspect_h5
                        --mode keys to list keys.
  --save [FILEPATH]     Write output to a file instead of vi. Omit FILEPATH to
                        write a persistent temp file and print its path.

inspect_idx

usage: fridata inspect_idx [-h] file

positional arguments:
  file        Path to idx file

options:
  -h, --help  show this help message and exit

remove_dataset

usage: fridata remove_dataset [-h] name

positional arguments:
  name        Dataset name (e.g. AFDB-subset--20250609_1333)

options:
  -h, --help  show this help message and exit