CESPED：Cryo-EM Supervised Pose Estimation Dataset的实用工具

CESPED是一个专为Cryo-EM中的监督姿态估计设计的新数据集。您可以在https://arxiv.org/abs/2311.06194查看我们的手稿。

安装

cesped已在python 3.11上进行了测试。应使用pip自动安装

pip install cesped
#Or directy from the master branch
pip install git+https://github.com/rsanchezgarc/cesped

或克隆存储库

git clone https://github.com/rsanchezgarc/cesped
cd cesped
pip install .

基本用法

ParticlesDataset类

用于加载图像和姿态。

1. 获取可下载条目的列表

from cesped.particlesDataset import ParticlesDataset
listOfEntries = ParticlesDataset.getCESPEDEntries()

2. 加载给定条目

targetName, halfset = listOfEntries[0] #We will work with the first entry only

dataset = ParticlesDataset(targetName, halfset)

为快速测试，请使用targetName="TEST"和halfset=0。如果数据集尚未在benchmarkDir中（在defaultDataConfig.yaml中定义），它将自动下载。元数据（欧拉角、CTF等）使用Relion starfile格式存储，图像存储为.mrcs堆栈。

3. 用作常规数据集

dl = DataLoader(dataset, batch_size=32)
for batch in dl:
  iid, img, (rotMat, xyShiftAngs, confidence), metadata = batch
  
  #iid is the list of ids of the particles (string)
  #img is a batch of Bx1xNxN images
  #rotMat is a batch of rotation matrices Bx3x3
  #xyShiftAngs is a batch of image shifts in Angstroms Bx2
  #confidence is a batch of numbers, between 0 and 1, Bx1
  #metata is a dictionary of names:values for all the information about the particle
  
  #YOUR PYTORCH CODE HERE
  predRot = model(img)
  loss = loss_function(predRot, rotMat)
  loss.backward()
  optimizer.step()
  optimizer.zero_grad()
  

4. 一旦您的模型经过训练，您就可以更新ParticlesDataset的元数据并将其保存，以便在cryo-EM软件中使用

for iid, pred_rotmats, maxprob in predictions:
    #iid is the list of ids of the particles (string)
    #pred_rotmats is a batch of predicted rotation matrices Bx3x3
    #maxprob is a batch of numbers, between 0 and 1, Bx1, that indicates the confidence in the prediction (e.g. softmax values)

    particlesDataset.updateMd(ids=iid, angles=pred_rotmats,
                              shifts=torch.zeros(pred_rotmats.shape[0],2, device=pred_rotmats.device), #Or actual predictions if you have them
                              confidence=maxprob,
                              angles_format="rotmat")
particlesDataset.saveMd(outFname) #Save the metadata as an starfile, a common cryo-EM format

  

5. 最后，如果使用evaluateEntry脚本保存了halfset 0和halfset 1的预测，则可以计算评估结果。

python -m cesped.evaluateEntry  --predictionType SO3 --targetName 11120  \
--half0PredsFname particles_preds_0.star  --half1PredsFname particles_preds_1.star \
--n_cpus 12 --outdir evaluation/

evaluateEntry使用Relion进行重建，因此您需要安装它并编辑配置文件defaultRelionConfig.yaml或通过命令行参数提供Relion的安装位置

--mpirun /path/to/mpirun  --relionBinDir /path/to/relion/bin

或者，您可以使用我们提供的定义文件relionSingularity.def构建一个singularity镜像。

singularity build relionSingularity.sif relionSingularity.def

并编辑配置文件以指向singularity镜像文件的位置，或者使用命令行参数。

--singularityImgFile /path/to/relionSingularity.sif

跨平台使用。

其他深度学习框架的用户可以使用以下命令下载CESPED条目：

python -m cesped.particlesDataset download_entry -t 10166 --halfset 0

这将下载关联的starfile和mrcs文件到默认的基准目录（在defaultDataConfig.yaml中定义）。使用--benchmarkDir指定另一个目录。

要列出可下载的条目以及已下载的条目，可以使用：

python -m cesped.particlesDataset list_entries

可以使用以下方法执行数据集条目的预处理：

python -m cesped.particlesDataset preprocess_entry --t 10166 --halfset 0 --o /tmp/dumpedData/ --ctf_correction "phase_flip"

其中--t是目标名称。使用-h显示可用的预处理操作列表。

可以使用Python包starstack轻松访问原始数据，它依赖于mrcfile和starfile包。预测应写入包含新预测的欧拉角的star文件。

一旦保存了0和1半集合的预测，就可以计算评估。

python -m cesped.evaluateEntry  --predictionType SO3 --targetName 11120  \
--half0PredsFname particles_preds_0.star  --half1PredsFname particles_preds_1.star \
--n_cpus 12 --outdir evaluation/

Image2Sphere实验

这些实验使用lightning和lightningCLI实现。您可以在以下位置找到配置文件：

YOUR_DIR/cesped/configs/

您也可以在这里找到它：

import cesped
cesped.default_configs_dir

训练

为了在一个目标上训练模型，您运行：

python -m cesped.trainEntry --data.halfset <HALFSET> --data.targetName <TARGETNAME> --trainer.default_root_dir <OUTDIR>

其中<HALFSET>为0或1，<TARGETNAME>为可以使用ParticlesDataset.getCESPEDEntries()找到的列表之一。

包含的目标有

* Nyquist频率为1.5 Å/pixel；分辨率估计在通常的阈值0.143。
报告的FSCR_0.143值直接来自relion_refine日志，而掩蔽的FSCR_0.143值来自relion_postprocess日志。

此外，条目TEST是EMPIAR-11120的小子集。

不要忘记更改配置文件或通过命令行或环境变量提供不同的值。此外，[--config CONFIG_NAME.yaml]还允许使用（一个或多个）自定义yaml文件覆盖默认值。使用-h查看可配置参数列表。其中一些最重要的参数包括。

• trainer.default_root_dir. 存储检查点和日志的目录，从defaultTrainerConfig.yaml
• optimizer.lr. 学习率，从defaultOptimizerConfig.yaml
• data.benchmarkDir. 存储基准条目的目录，从defaultDataConfig.yaml。建议在配置文件中更改此值。
• data.num_data_workers. 数据加载的工作者数量，从defaultDataConfig.yaml
• data.batch_size. 从defaultDataConfig.yaml

推理

默认情况下，使用 python -m cesped.trainEntry 时，在训练完成后，在单个 GPU 上对补集进行推理，预测的 star 文件可找到于 <OUTDIR>/lightning_logs/version_<\d>/predictions_[0,1].star。为了手动运行姿态预测代码（并使用所有 GPU），您可以运行

python -m cesped.inferEntry --data.halfset <HALFSET> --data.targetName <TARGETNAME> --ckpt_path <PATH_TO_CHECKPOINT> \
--outFname /path/to/output/starfile.star

评估

5. 与之前相同，如果使用 evaluateEntry 脚本保存了 0 组和 1 组的预测，则可以计算评估。

python -m cesped.evaluateEntry  --predictionType SO3 --targetName 11120  \
--half0PredsFname particles_preds_0.star  --half1PredsFname particles_preds_1.star \
--n_cpus 12 --outdir evaluation/

API

有关 API 文档，请参阅 文档文件夹

Relion Singularity

可以使用以下命令构建具有 MPI 支持的 relion_reconstruct singularity 容器。

singularity build relionSingulary.sif relionSingulary.def 

然后，可以使用以下命令计算 Relion 重建

singularity exec relionSingulary.sif mpirun -np 4 relion_reconstruct_mpi --ctf --pad 2 --i input_particles.star --o output_map.mrc
#Or the following command
./relionSingulary.sif  4 --ctf --pad 2 --i input_particles.star --o output_map.mrc #This uses 4 mpis

然而，典型用户通常不需要手动执行容器。所有操作都在 evaluateEntry.py 脚本中透明进行