Crea flussi di lavoro di piegamento delle proteine per accelerare la scoperta di farmaci su Amazon SageMaker

Crea flussi di lavoro per scoprire farmaci su Amazon SageMaker

Lo sviluppo di farmaci è un processo complesso e lungo che prevede la selezione di migliaia di candidati farmaci e l’utilizzo di metodi computazionali o sperimentali per valutare i principi attivi. Secondo McKinsey, un singolo farmaco può richiedere 10 anni e costare in media 2,6 miliardi di dollari per passare dall’identificazione del bersaglio di malattia, alla selezione del farmaco, alla validazione del bersaglio del farmaco e infine al lancio commerciale. La scoperta di farmaci è la componente di ricerca di questa pipeline che genera candidati farmaci con la più alta probabilità di essere efficaci con il minor danno possibile ai pazienti. I metodi di apprendimento automatico (ML) possono aiutare a identificare composti adatti in ogni fase del processo di scoperta dei farmaci, risultando in una prioritizzazione e un test dei farmaci più efficienti, risparmiando miliardi di dollari nei costi di sviluppo dei farmaci (per ulteriori informazioni, fare riferimento a AI nella ricerca biofarmaceutica: Un momento per concentrarsi e scalare).

I bersagli dei farmaci sono tipicamente entità biologiche chiamate proteine, i mattoni fondamentali della vita. La struttura tridimensionale di una proteina determina come interagisce con un composto farmaceutico; pertanto, la comprensione della struttura tridimensionale della proteina può apportare miglioramenti significativi al processo di sviluppo dei farmaci selezionando i composti farmaceutici che si adattano meglio alla struttura della proteina bersaglio. Un’altra area in cui la previsione della struttura delle proteine può essere utile è la comprensione della diversità delle proteine, in modo da selezionare solo i farmaci che colpiscono selettivamente specifiche proteine senza influenzare altre proteine nel corpo (per ulteriori informazioni, fare riferimento a Miglioramento della valutazione del bersaglio nella ricerca biomedica: le raccomandazioni GOT-IT). Strutture tridimensionali precise delle proteine bersaglio possono consentire la progettazione di farmaci con maggiore specificità e minore probabilità di interazioni incrociate con altre proteine.

Tuttavia, prevedere come le proteine si ripiegano nella loro struttura tridimensionale è un problema difficile e i metodi sperimentali tradizionali come la cristallografia a raggi X e la spettroscopia NMR possono richiedere tempo ed essere costosi. I recenti progressi nei metodi di apprendimento profondo per la ricerca sulle proteine hanno dimostrato di poter prevedere il ripiegamento proteico con una precisione notevole utilizzando reti neurali. Algoritmi di ripiegamento come AlphaFold2, ESMFold, OpenFold e RoseTTAFold possono essere utilizzati per costruire rapidamente modelli accurati delle strutture proteiche. Purtroppo, questi modelli richiedono molte risorse computazionali per essere eseguiti e i risultati possono essere difficili da confrontare su una scala di migliaia di strutture proteiche candidate. Una soluzione scalabile per l’utilizzo di questi vari strumenti consentirà ai ricercatori e ai team di R&S commerciali di incorporare rapidamente gli ultimi progressi nella previsione della struttura delle proteine, gestire i loro processi sperimentali e collaborare con partner di ricerca.

Amazon SageMaker è un servizio completamente gestito per preparare, creare, addestrare e distribuire modelli di apprendimento automatico di alta qualità in modo rapido, mettendo insieme un’ampia gamma di funzionalità appositamente progettate per l’apprendimento automatico. Offre un ambiente completamente gestito per l’apprendimento automatico, astrae l’infrastruttura, la gestione dei dati e i requisiti di scalabilità in modo da potersi concentrare sulla creazione, l’addestramento e il test dei modelli.

• Cos’è snowChat?
• Intelligenza Artificiale che aiuta nella lotta contro gli incendi boschivi
• Una nuova ricerca sull’IA propone il modello PanGu-Coder2 e il framework RRTF che potenziano in modo efficiente i modelli di lingua pre-addestrati per la generazione di codice

In questo post, presentiamo una soluzione di apprendimento automatico completamente gestita con SageMaker che semplifica l’operazione dei flussi di lavoro di previsione della struttura del ripiegamento proteico. Prima discutiamo la soluzione a un livello elevato e la sua esperienza utente. Successivamente, vi guidiamo attraverso la configurazione semplice di flussi di lavoro ottimizzati per il calcolo di AlphaFold2 e OpenFold con SageMaker. Infine, dimostriamo come è possibile monitorare e confrontare le previsioni della struttura proteica come parte di un’analisi tipica. Il codice per questa soluzione è disponibile nel seguente repository GitHub.

Panoramica della soluzione

In questa soluzione, gli scienziati possono lanciare interattivamente esperimenti di ripiegamento proteico, analizzare la struttura tridimensionale, monitorare l’avanzamento del lavoro e tenere traccia degli esperimenti in Amazon SageMaker Studio.

La seguente schermata mostra una singola esecuzione di un flusso di lavoro di ripiegamento proteico con Amazon SageMaker Studio. Include la visualizzazione della struttura tridimensionale in un notebook, lo stato di esecuzione dei lavori di SageMaker nel flusso di lavoro e collegamenti ai parametri di input e ai dati di output e ai log.

Il diagramma seguente illustra l’architettura della soluzione a un livello elevato.

Per comprendere l’architettura, definiamo prima i componenti chiave di un esperimento di ripiegamento proteico come segue:

• File di sequenza target FASTA – Il formato FASTA è un formato basato su testo per rappresentare sequenze di nucleotidi o sequenze di amminoacidi (proteine), in cui i nucleotidi o gli amminoacidi sono rappresentati utilizzando codici a lettera singola.
• Banche dati genetiche – Una banca dati genetica è uno o più insiemi di dati genetici archiviati insieme a un software che consente agli utenti di recuperare dati genetici. Per eseguire gli algoritmi AlphaFold e OpenFold sono necessarie diverse banche dati genetiche, come BFD, MGnify, PDB70, PDB, PDB seqres, UniRef30 (precedentemente UniClust30), UniProt e UniRef90.
• Allineamento di sequenze multiple (MSA) – Un allineamento di sequenze è un modo di organizzare le sequenze primarie di una proteina per identificare regioni di somiglianza che possono essere una conseguenza di relazioni funzionali, strutturali o evolutive tra le sequenze. Le caratteristiche di input per le previsioni includono dati MSA.
• Previsione della struttura proteica – La struttura delle sequenze di bersaglio di input viene prevista con algoritmi di ripiegamento come AlphaFold2 e OpenFold che utilizzano un’architettura di trasformatore multitraccia addestrata su modelli di proteine noti.
• Visualizzazione e indicatori – Visualizza la struttura tridimensionale con la libreria py3Dmol come una visualizzazione 3D interattiva. È possibile utilizzare gli indicatori per valutare e confrontare le previsioni di struttura, in particolare la deviazione quadratico-media (RMSD) e il punteggio di modellazione del modello (TM-score).

Il flusso di lavoro contiene i seguenti passaggi:

1. Gli scienziati utilizzano l’IDE SageMaker ML basato sul web per esplorare la base di codice, creare flussi di lavoro di analisi delle sequenze proteiche nei notebook SageMaker Studio e eseguire pipeline di piegamento delle proteine tramite l’interfaccia utente grafica in SageMaker Studio o l’SDK di SageMaker.
2. I database genetici e strutturali richiesti da AlphaFold e OpenFold vengono scaricati prima della configurazione della pipeline utilizzando Amazon SageMaker Processing, una funzionalità di elaborazione dei dati di ML basata su calcolo effimero, in un bucket di Amazon Simple Storage Service (Amazon S3). Con SageMaker Processing, è possibile eseguire un lavoro a lungo termine con un calcolo adeguato senza dover configurare alcun cluster di calcolo e archiviazione e senza la necessità di spegnere il cluster. I dati vengono automaticamente salvati in una posizione specificata del bucket S3.
3. Viene configurato un sistema di file Amazon FSx for Lustre, con il repository dati che è la posizione del bucket S3 in cui vengono salvati i database. FSx for Lustre può scalare fino a centinaia di GB/s di throughput e milioni di IOPS con il recupero di file a bassa latenza. Quando si avvia un lavoro di stima, SageMaker monta il sistema di file FSx for Lustre nel file system dell’istanza, quindi avvia lo script.
4. Amazon SageMaker Pipelines viene utilizzato per orchestrare più esecuzioni di algoritmi di piegamento delle proteine. SageMaker Pipelines offre un’interfaccia visiva desiderata per la presentazione interattiva delle attività, la tracciabilità del progresso e la ripetibilità.
5. All’interno di una pipeline, vengono eseguiti due algoritmi di piegamento delle proteine computazionalmente intensivi – AlphaFold e OpenFold – con gli estimatori SageMaker. Questa configurazione supporta il montaggio di un sistema di file FSx for Lustre per la ricerca di database ad alto throughput negli algoritmi. Un singolo run di inferenza viene suddiviso in due fasi: una fase di costruzione MSA utilizzando un’istanza CPU ottimale e una fase di previsione della struttura utilizzando un’istanza GPU. Questi sottopassi, come SageMaker Processing nel Passaggio 2, sono effimeri, on-demand e completamente gestiti. I file di output del lavoro come file MSA, file di struttura pdb previsti e altri file di metadati vengono salvati in una posizione specificata di S3. Una pipeline può essere progettata per eseguire un singolo algoritmo di piegamento delle proteine o eseguire sia AlphaFold che OpenFold dopo una costruzione MSA comune.
6. Le esecuzioni della previsione del piegamento delle proteine vengono automaticamente tracciate da Amazon SageMaker Experiments per ulteriori analisi e confronti. I log dei lavori vengono conservati in Amazon CloudWatch per il monitoraggio.

Prerequisiti

Per seguire questo post ed eseguire questa soluzione, è necessario aver completato diversi prerequisiti. Consultare il repository GitHub per una spiegazione dettagliata di ciascun passaggio.

• Un dominio SageMaker e un profilo utente – Se non si dispone di un dominio SageMaker Studio, fare riferimento a Onboard to Amazon SageMaker Domain Using Quick Setup.
• Politiche IAM – L’utente deve avere la politica IAM AmazonSageMakerFullAccess di AWS Identity and Access Management (IAM) collegata, la capacità di creare immagini di container Docker in Amazon Elastic Container Registry (Amazon ECR) e file system FSx for Lustre creati. Vedere il file readme per ulteriori dettagli.
• Rete – Una VPC con un endpoint VPC di Amazon S3. Utilizziamo questa posizione VPC per fornire il file system FSx for Lustre e i lavori di SageMaker.
• Risorse Docker – Eseguire 00-prerequisite.ipynb dal repository per creare le immagini Docker, scaricare il database genetico su Amazon S3 e creare un file system FSx for Lustre con un’associazione di repository dati al bucket S3.

Esegui il piegamento delle proteine su SageMaker

Utilizziamo le capacità completamente gestite di SageMaker per eseguire lavori di piegamento delle proteine computazionalmente intensivi senza troppa overhead di infrastruttura. SageMaker utilizza immagini di container per eseguire script personalizzati per l’elaborazione generica dei dati, l’addestramento e l’hosting. È possibile avviare facilmente un lavoro effimero on-demand che esegue un programma con un’immagine di container con un paio di righe dell’SDK di SageMaker senza autogestire alcuna infrastruttura di calcolo. In particolare, il lavoro di stima di SageMaker offre flessibilità per quanto riguarda la scelta dell’immagine del contenitore, dello script di esecuzione e della configurazione dell’istanza, e supporta una vasta gamma di opzioni di archiviazione, inclusi i sistemi di file come FSx for Lustre. Il diagramma seguente illustra questa architettura.

Algoritmi di piegamento come AlphaFold e OpenFold utilizzano un’architettura di trasformatore multi-traccia addestrata su modelli proteici noti per prevedere la struttura di sequenze peptidiche sconosciute. Queste previsioni possono essere eseguite su istanze GPU per fornire la migliore throughput e la latenza più bassa. Tuttavia, le caratteristiche di input per queste previsioni includono dati MSA. Gli algoritmi MSA dipendono dalla CPU e possono richiedere diverse ore di tempo di elaborazione.

Eseguire sia i passaggi di MSA che di previsione della struttura nello stesso ambiente di calcolo può essere inefficiente in termini di costo perché le costose risorse GPU rimangono inutilizzate durante l’esecuzione del passaggio MSA. Pertanto, ottimizziamo il flusso di lavoro in due fasi. Prima eseguiamo un lavoro di stimatore SageMaker su un’istanza CPU specificamente per calcolare l’allineamento MSA dato una particolare sequenza di input FASTA e database genetici di origine. Quindi eseguiamo un lavoro di stimatore SageMaker su un’istanza GPU per prevedere la struttura proteica con un dato allineamento MSA di input e un algoritmo di piegatura come AlphaFold o OpenFold.

Esegui la generazione di MSA

Per il calcolo di MSA, includiamo uno script personalizzato run_create_alignment.sh e lo script create_alignments.py che è adottato dall’esistente codice di previsione di AlphaFold run_alphafold.py. Si noti che questo script potrebbe dover essere aggiornato se il codice di AlphaFold di origine viene aggiornato. Lo script personalizzato viene fornito allo stimatore SageMaker tramite la modalità script. I componenti chiave dell’immagine del contenitore, l’implementazione della modalità script e la configurazione di un lavoro di stimatore SageMaker fanno parte anche del passaggio successivo dell’esecuzione degli algoritmi di piegatura e vengono descritti ulteriormente nella sezione successiva.

Esegui AlphaFold

Iniziamo eseguendo una previsione di struttura di AlphaFold con una singola sequenza proteica utilizzando SageMaker. Eseguire un lavoro di AlphaFold comporta tre semplici passaggi, come si può vedere in 01-run_stepbystep.ipynb. Prima, creiamo un’immagine del contenitore Docker basata sul Dockerfile di AlphaFold in modo da poter eseguire anche AlphaFold in SageMaker. Secondo, costruiamo lo script run_alphafold.sh che istruisce come AlphaFold deve essere eseguito. Terzo, costruiamo e avviamo uno stimatore SageMaker con lo script, il contenitore, il tipo di istanza, i dati e la configurazione per il lavoro.

Immagine del contenitore

I requisiti di runtime per un’immagine del contenitore per eseguire AlphaFold (OpenFold anche) in SageMaker possono essere notevolmente semplificati con il Dockerfile di AlphaFold. Dobbiamo solo aggiungere un pugno di semplici layer in cima per installare una libreria Python specifica di SageMaker in modo che un lavoro di SageMaker possa comunicare con l’immagine del contenitore. Vedere il codice seguente:

# In Dockerfile.alphafold
## Specifico di SageMaker
RUN pip3 install sagemaker-training --upgrade --no-cache-dir
ENV PATH="/opt/ml/code:${PATH}"
# questa variabile di ambiente è utilizzata dallo stimatore SageMaker per determinare la directory del nostro codice utente
ENV SAGEMAKER_SUBMIT_DIRECTORY /opt/ml/code

Script di input

<p.Successivamente forniamo lo script run_alphafold.sh che esegue run_alphafold.py dal repository di AlphaFold attualmente posizionato nel contenitore /app/alphafold/run_alphafold.py. Quando viene eseguito questo script, la posizione dei database genetici e la sequenza FASTA di input saranno popolate da SageMaker come variabili di ambiente (SM_CHANNEL_GENETIC e SM_CHANNEL_FASTA, rispettivamente). Per ulteriori informazioni, fare riferimento alla configurazione dei dati di input.

Lavoro di stimatore

<p.Successivamente creiamo un lavoro utilizzando uno stimatore SageMaker con i seguenti argomenti di input chiave, che istruiscono SageMaker ad eseguire uno script specifico utilizzando un contenitore specificato con il tipo o il conteggio di istanze, la tua opzione di networking preferita e altri parametri per il lavoro. vpc_subnet_ids e security_group_ids istruiscono il lavoro ad essere eseguito all’interno di una specifica VPC in cui si trova il sistema di file FSx for Lustre in modo da poter montare e accedere al file system nel lavoro SageMaker. Il percorso di output si riferisce a una posizione del bucket S3 in cui il prodotto finale di AlphaFold verrà caricato alla fine di un lavoro di successo da parte di SageMaker in modo automatico. Qui impostiamo anche un parametro DB_PRESET, ad esempio, da passare e accedere all’interno di run_alphafold.sh come variabile ambientale durante l’esecuzione. Vedere il codice seguente:

from sagemaker.estimator import Estimator
alphafold_image_uri=f'{account}.dkr.ecr.{region}.amazonaws.com/sagemaker-studio-alphafold:v2.3.0'
instance_type='ml.g5.2xlarge'
instance_count=1
vpc_subnet_ids=['subnet-xxxxxxxxx'] # è possibile utilizzare una VPC predefinita
security_group_ids=['sg-xxxxxxxxx']
env={'DB_PRESET': db_preset} # <full_dbs|reduced_dbs>
output_path='s3://%s/%s/job-output/'%(default_bucket, prefix)

estimator_alphafold = Estimator(
source_dir='src', # directory in cui si trovano run_alphafold.sh e altri file di runtime
entry_point='run_alphafold.sh', # il nostro script che esegue /app/alphafold/run_alphafold.py
image_uri=alphafold_image_uri, # immagine del contenitore da utilizzare
instance_count=instance_count, #
instance_type=instance_type,
subnets=vpc_subnet_ids,
security_group_ids=security_group_ids,
environment=env,
output_path=output_path,
...)

Infine, raccogliamo i dati e facciamo sapere al lavoro dove si trovano. Il canale dei dati fasta è definito come un input dati S3 che verrà scaricato da una posizione S3 nell’istanza di calcolo all’inizio del lavoro. Ciò consente una grande flessibilità per gestire e specificare la sequenza di input. D’altra parte, il canale dei dati genetic è definito come un FileSystemInput che verrà montato sull’istanza all’inizio del lavoro. L’uso di un file system FSx for Lustre come modo per portare vicino a 3 TB di dati evita di scaricare ripetutamente i dati da un bucket S3 a un’istanza di calcolo. Chiamiamo il metodo .fit per avviare un lavoro AlphaFold:

from sagemaker.inputs import FileSystemInput
file_system_id='fs-xxxxxxxxx'
fsx_mount_id='xxxxxxxx'
file_system_directory_path=f'/{fsx_mount_id}/{prefix}/alphafold-genetic-db' # dovrebbe essere il prefisso completo del repository dati S3

file_system_access_mode='ro' # Specifica la modalità di accesso (sola lettura)
file_system_type='FSxLustre' # Specifica il tipo di file system

genetic_db = FileSystemInput(
file_system_id=file_system_id,
file_system_type=file_system_type,
directory_path=file_system_directory_path,
file_system_access_mode=file_system_access_mode)

s3_fasta=sess.upload_data(path='sequence_input/T1030.fasta', # posizione FASTA localmente
key_prefix='alphafoldv2/sequence_input') # prefisso S3. Il bucket è il bucket predefinito di sagemaker
fasta = sagemaker.inputs.TrainingInput(s3_fasta,
distribution='FullyReplicated',
s3_data_type='S3Prefix',
input_mode='File')
data_channels_alphafold = {'genetic': genetic_db, 'fasta': fasta}

estimator_alphafold.fit(inputs=data_channels_alphafold,
wait=False) # wait=False restituisce la cella nel notebook; impostare su True per vedere i log durante l'evoluzione del lavoro

E’ tutto. Abbiamo appena inviato un lavoro a SageMaker per eseguire AlphaFold. I log e l’output, inclusi i file di predizione .pdb, verranno scritti su Amazon S3.

Esegui OpenFold

L’esecuzione di OpenFold in SageMaker segue un modello simile, come mostrato nella seconda metà di 01-run_stepbystep.ipynb. Aggiungiamo prima uno strato semplice per ottenere la libreria specifica di SageMaker per rendere l’immagine del contenitore compatibile con SageMaker sopra al Dockerfile di OpenFold. In secondo luogo, costruiamo uno script run_openfold.sh come punto di ingresso per il lavoro di SageMaker. In run_openfold.sh, eseguiamo il run_pretrained_openfold.py di OpenFold, che è disponibile nell’immagine del contenitore con gli stessi database genetici che abbiamo scaricato per AlphaFold e i pesi del modello di OpenFold (--openfold_checkpoint_path). Per quanto riguarda i luoghi dei dati di input, oltre al canale dei database genetici e al canale FASTA, introduciamo un terzo canale, SM_CHANNEL_PARAM, in modo da poter passare in modo flessibile i pesi del modello scelti dalla costruzione del job quando definiamo e inviamo un lavoro. Con l’estimatore di SageMaker, possiamo facilmente inviare lavori con diverse entry_point, image_uri, environment, inputs e altre configurazioni per OpenFold con la stessa firma. Per il canale dei dati, aggiungiamo un nuovo canale, param, come input Amazon S3 insieme all’uso degli stessi database genetici da FSx for Lustre e al file FASTA da Amazon S3. Questo, ancora una volta, ci consente di specificare facilmente il peso del modello da utilizzare dalla costruzione del job. Ecco il codice seguente:

s3_param=sess.upload_data(path='openfold_params/finetuning_ptm_2.pt',
key_prefix=f'{prefix}/openfold_params')
param = sagemaker.inputs.TrainingInput(s3_param,
distribution="FullyReplicated",
s3_data_type="S3Prefix",
input_mode='File')

data_channels_openfold = {"genetic": genetic_db, 'fasta': fasta, 'param': param}

estimator_openfold.fit(inputs=data_channels_openfold,
wait=False)

Per accedere all’output finale dopo il completamento del lavoro, eseguiamo i seguenti comandi:

!aws s3 cp {estimator_openfold.model_data} openfold_output/model.tar.gz
!tar zxfv openfold_output/model.tar.gz -C openfold_output/

Prestazioni in tempo di esecuzione

La seguente tabella mostra risparmi di costo del 57% e del 51% rispettivamente per AlphaFold e OpenFold, suddividendo l’allineamento MSA e gli algoritmi di piegatura in due lavori rispetto a un singolo calcolo. Ciò ci consente di adattare le risorse di calcolo per ogni lavoro: ml.m5.4xlarge per l’allineamento MSA e ml.g5.2xlarge per AlphaFold e OpenFold.

Dettagli del lavoro	Tipo di istanza	Sequenza FASTA di input	Tempo di esecuzione	Costo
Allineamento MSA + OpenFold	ml.g5.4xlarge	T1030	50 minuti	$1.69
Allineamento MSA + AlphaFold	ml.g5.4xlarge	T1030	65 minuti	$2.19
Allineamento MSA	ml.m5.4xlarge	T1030	46 minuti	$0.71
OpenFold	ml.g5.2xlarge	T1030	6 minuti	$0.15
AlphaFold	ml.g5.2xlarge	T1030	21 minuti	$0.53

Costruisci un flusso di lavoro ripetibile utilizzando SageMaker Pipelines

Con SageMaker Pipelines, possiamo creare un flusso di lavoro di machine learning che si occupa della gestione dei dati tra le fasi, dell’orchestrazione delle loro esecuzioni e della registrazione. SageMaker Pipelines ci fornisce anche un’interfaccia utente per visualizzare il nostro flusso di lavoro e eseguire facilmente il nostro flusso di lavoro di machine learning.

Un flusso di lavoro viene creato combinando un certo numero di fasi. In questo flusso di lavoro, combiniamo tre fasi di addestramento, che richiedono un estimatore di SageMaker. Gli estimatori definiti in questo notebook sono molto simili a quelli definiti in 01-run_stepbystep.ipynb, con l’eccezione che utilizziamo posizioni di Amazon S3 per indicare i nostri input e output. Le variabili dinamiche consentono a SageMaker Pipelines di eseguire le fasi una dopo l’altra e consentono anche all’utente di riprovare le fasi fallite. La seguente schermata mostra un Grafo Aciclico Diretto (DAG), che fornisce informazioni sui requisiti e sulle relazioni tra ogni fase del nostro flusso di lavoro.

Variabili dinamiche

SageMaker Pipelines è in grado di ricevere input dall’utente all’inizio di ogni esecuzione del flusso di lavoro. Definiamo le seguenti variabili dinamiche, che vorremmo modificare durante ogni esperimento:

• FastaInputS3URI – URI di Amazon S3 del file FASTA caricato tramite SDK, Boto3 o manualmente.
• FastFileName – Nome del file FASTA.
• db_preset – Selezione tra full_dbs o reduced_dbs.
• MaxTemplateDate – La fase MSA di AlphaFold cercherà i template disponibili prima della data specificata da questo parametro.
• ModelPreset – Seleziona tra i modelli AlphaFold, tra cui monomer, monomer_casp14, monomer_ptm e multimer.
• NumMultimerPredictionsPerModel – Numero di seed da eseguire per modello quando si utilizza il sistema multimerico.
• InferenceInstanceType – Tipo di istanza da utilizzare per le fasi di inferenza (sia AlphaFold che OpenFold). Il valore predefinito è ml.g5.2xlarge.
• MSAInstanceType – Tipo di istanza da utilizzare per la fase MSA. Il valore predefinito è ml.m5.4xlarge.

Vedi il codice seguente:

fasta_file = ParameterString(name="FastaFileName")
fasta_input = ParameterString(name="FastaInputS3URI")
pipeline_db_preset = ParameterString(name="db_preset",
default_value='full_dbs',
enum_values=['full_dbs', 'reduced_dbs'])
max_template_date = ParameterString(name="MaxTemplateDate")
model_preset = ParameterString(name="ModelPreset")
num_multimer_predictions_per_model = ParameterString(name="NumMultimerPredictionsPerModel")
msa_instance_type = ParameterString(name="MSAInstanceType", default_value='ml.m5.4xlarge')
instance_type = ParameterString(name="InferenceInstanceType", default_value='ml.g5.2xlarge')

Una pipeline SageMaker viene costruita definendo una serie di passaggi e collegandoli in un ordine specifico in cui l’output di un passaggio precedente diventa l’input per il passaggio successivo. I passaggi possono essere eseguiti in parallelo e definiti in modo da dipendere da un passaggio precedente. In questa pipeline, definiamo un passaggio MSA, che è la dipendenza per un passaggio di inferenza AlphaFold e un passaggio di inferenza OpenFold che vengono eseguiti in parallelo. Vedi il codice seguente:

step_msa = TrainingStep(
name="RunMSA",
step_args=pipeline_msa_args,
)

step_alphafold = TrainingStep(
name="RunAlphaFold",
step_args=pipeline_alphafold_default_args,
)
step_alphafold.add_depends_on([step_msa])

step_openfold = TrainingStep(
name="RunOpenFold",
step_args=pipeline_openfold_args,
)
step_openfold.add_depends_on([step_msa]

Per mettere insieme tutti i passaggi, chiamiamo la classe Pipeline e forniamo un nome di pipeline, variabili di input della pipeline e i passaggi individuali:

pipeline_name = f"ProteinFoldWorkflow"
pipeline = Pipeline(
name=pipeline_name,
parameters=[
fasta_input,
instance_type,
msa_instance_type,
pipeline_db_preset
],
steps=[step_msa, step_alphafold, step_openfold],
)

pipeline.upsert(role_arn=role, # esegui questo se è la prima volta che imposti la pipeline
description='Protein_Workflow_MSA')

Esegui la pipeline

Nell’ultima cella del notebook 02-define_pipeline.ipynb, mostriamo come eseguire una pipeline utilizzando il SDK SageMaker. Le variabili dinamiche descritte in precedenza sono fornite nel seguente modo:

!mkdir ./sequence_input/
!curl 'https://www.predictioncenter.org/casp14/target.cgi?target=T1030&view=sequence' > ./sequence_input/T1030.fasta
fasta_file_name = 'T1030.fasta'

pathName = f'./sequence_input/{fasta_file_name}'
s3_fasta=sess.upload_data(path=pathName,
key_prefix='alphafoldv2/sequence_input')

PipelineParameters={
'FastaInputS3URI':s3_fasta,
'db_preset': 'full_dbs',
'FastaFileName': fasta_file_name,
'MaxTemplateDate': '2020-05-14',
'ModelPreset': 'monomer',
'NumMultimerPredictionsPerModel': '5',
'InferenceInstanceType':'ml.g5.2xlarge',
'MSAInstanceType':'ml.m5.4xlarge'
}
execution = pipeline.start(execution_display_name='SDK-Executetd',
execution_description='This pipeline was executed via SageMaker SDK',
parameters=PipelineParameters
)

Tieni traccia degli esperimenti e confronta le strutture proteiche

Per il nostro esperimento, utilizziamo una sequenza proteica di esempio dalla competizione CASP14, che fornisce un meccanismo indipendente per la valutazione dei metodi di modellazione della struttura proteica. Il target T1030 è derivato dalla proteina PDB 6P00 e ha 237 amminoacidi nella sequenza primaria. Eseguiamo la pipeline SageMaker per prevedere la struttura proteica di questa sequenza di input con gli algoritmi OpenFold e AlphaFold.

Quando la pipeline è completa, scarichiamo i file .pdb predetti da ciascun lavoro di folding e visualizziamo la struttura nel notebook utilizzando py3Dmol, come nel notebook 04-compare_alphafold_openfold.ipynb.

La seguente schermata mostra la previsione del lavoro di predizione AlphaFold.

La struttura prevista viene confrontata con la sua struttura di riferimento di base nota con codice PDB 6poo archiviato in RCSB. Analizziamo le prestazioni di previsione rispetto al codice PDB di base 6poo con tre metriche: RMSD, RMSD con sovrapposizione e punteggio di modellazione del template, come descritto in Confronto delle strutture.

.	Sequenza di input	Confronto con	RMSD	RMSD con sovrapposizione	Punteggio di modellazione del template
AlphaFold	T1030	6poo	247,26	3,87	0,3515

Gli algoritmi di folding vengono ora confrontati tra loro per più sequenze FASTA: T1030, T1090 e T1076. Le nuove sequenze target potrebbero non avere la struttura pdb di base nei database di riferimento e quindi è utile confrontare la variabilità tra gli algoritmi di folding.

.	Sequenza di input	Confronto con	RMSD	RMSD con sovrapposizione	Punteggio di modellazione del template
AlphaFold	T1030	OpenFold	73,21	24,8	0,0018
AlphaFold	T1076	OpenFold	38,71	28,87	0,0047
AlphaFold	T1090	OpenFold	30,03	20,45	0,005

La seguente schermata mostra le esecuzioni di ProteinFoldWorkflow per le tre sequenze di input FASTA con SageMaker Pipeline:

Registriamo anche le metriche con SageMaker Experiments come nuove esecuzioni dello stesso esperimento create dal pipeline:

from sagemaker.experiments.run import Run, load_run
metric_type='compare:'
experiment_name = 'proteinfoldworkflow'
with Run(experiment_name=experiment_name, run_name=input_name_1, sagemaker_session=sess) as run:
run.log_metric(name=metric_type + "rmsd_cur", value=rmsd_cur_one, step=1)
run.log_metric(name=metric_type + "rmds_fit", value=rmsd_fit_one, step=1)
run.log_metric(name=metric_type + "tm_score", value=tmscore_one, step=1)

Quindi analizziamo e visualizziamo queste esecuzioni sulla pagina Esperimenti in SageMaker Studio.

Il grafico seguente rappresenta il valore RMSD tra AlphaFold e OpenFold per le tre sequenze: T1030, T1076 e T1090.

Conclusioni

In questo post, descriviamo come è possibile utilizzare SageMaker Pipelines per impostare ed eseguire flussi di lavoro di piegatura delle proteine con due popolari algoritmi di previsione della struttura: AlphaFold2 e OpenFold. Abbiamo dimostrato un’architettura di soluzione con prestazioni ottimali di diversi job che separa i requisiti di calcolo per la generazione di MSA dalla previsione della struttura. Abbiamo anche evidenziato come sia possibile visualizzare, valutare e confrontare le strutture 3D previste delle proteine in SageMaker Studio.

Per iniziare con i flussi di lavoro di piegatura delle proteine su SageMaker, fare riferimento al codice di esempio nel repository GitHub.