Update app.py

app.py

@@ -89,7 +89,7 @@ seafoam = Seafoam()
 custom_css = """
 <style>
     .file-upload-height {
-        height:320px !important;
+        height:330px !important;
         display: none;
     }
     .file-upload-height2 {
@@ -99,7 +99,7 @@ custom_css = """
         height: 380px !important;
     }
     #custom_plot {
-       height: 360px !important;
+       height: 280px !important;
    }
     #custom_plot2 {
        height: 480px !important;
@@ -119,20 +119,7 @@ def draw_mass_spectrum(peak_data_path):
     ms2 = spectrum_processing(ms2)
     Mz = np.array(ms2.mz)
     Intens = np.array(ms2.intensities)
-    plt.figure(figsize=(7,3))
-    for i in range(len(Mz)):
-            plt.axvline(x=Mz[i], ymin=0, ymax=Intens[i],c='red')
-    plt.xlabel("m/z")
-    plt.ylabel("Intensity")
-    plt.title("Mass Spectrum") 
-    return plt
-
-def draw_mass_spectrum2(peak_data_path):
-    ms2 = list(load_from_msp(peak_data_path.name))[0]
-    ms2 = spectrum_processing(ms2)
-    Mz = np.array(ms2.mz)
-    Intens = np.array(ms2.intensities)
-    plt.figure(figsize=(8,5))
+    plt.figure(figsize=(7,4))
     for i in range(len(Mz)):
             plt.axvline(x=Mz[i], ymin=0, ymax=Intens[i],c='red')
     plt.xlabel("m/z")
@@ -140,6 +127,12 @@ def draw_mass_spectrum2(peak_data_path):
     plt.title("Mass Spectrum") 
     return plt
     
+def search_structure_from_mass(structureDB,mass, ppm): 
+    structures=pd.DataFrame()
+    mmin = mass - mass*ppm/10**6 
+    mmax = mass + mass*ppm/10**6 
+    structures = structureDB[(structureDB['MonoisotopicMass'] >= mmin) & (structureDB['MonoisotopicMass'] <= mmax)]
+    return structures 
 
 dataset_repo = "Tingxie/CSU-MS2-DB"
 db_filename = "csu_ms2_db.db"
@@ -350,7 +343,7 @@ def rank_lib(database_name,spectrum_path,instrument_type,adduct,parent_Mass,coll
     #top_10_results = pd.DataFrame({'SMILES':[x[0] for x in top_10_smiles],'Struture':images,'Score':[x[1] for x in top_10_smiles],'Rank':list(range(10))})
     return image_descrips
 
-def rank_user_lib(candidate_file,spectrum_path,instrument_type,adduct,collision_energy):
+def rank_user_lib(candidate_file,spectrum_path,instrument_type,adduct,parent_Mass,collision_energy):
     model_inference_low = ModelInference(config_path=config_path,
                                  pretrain_model_path=pretrain_model_path_low,
                                  device="cpu")
@@ -361,8 +354,21 @@ def rank_user_lib(candidate_file,spectrum_path,instrument_type,adduct,collision_
                                  pretrain_model_path=pretrain_model_path_high,
                                  device="cpu")
     collision_energy=float(collision_energy)
-    candidate = pd.read_csv(candidate_file.name)
-    candidate_smiles = list(candidate['SMILES'])
+    users_candidate = pd.read_csv(candidate_file.name)
+    user_candidate_smiles = list(users_candidate['SMILES'])
+    candidate = pd.DataFrame(columns=['SMILES','MonoisotopicMass'])
+    for smi in user_candidate_smiles:
+        mol = Chem.MolFromSmiles(smi)
+        MonoisotopicMass = ExactMolWt(mol)
+        candidate.loc[len(candidate.index)] = [smi,MonoisotopicMass]
+    parent_Mass=float(parent_Mass)
+    query_ms=parent_Mass-1.008
+    sub_reference_library = search_structure_from_mass(candidate,query_ms,10)
+    if len(sub_reference_library) == 0:
+        sub_reference_library = search_structure_from_mass(candidate,query_ms,30)
+    if len(sub_reference_library) == 0:
+        sub_reference_library = candidate.copy()
+    candidate_smiles = list(sub_reference_library['SMILES'])
     ms2 = list(load_from_msp(spectrum_path.name))[0]
     ms2 = spectrum_processing(ms2)
     ms2_embedding_low,ms2_embedding_median,ms2_embedding_high = MS2Embedding(ms2)
@@ -407,18 +413,21 @@ def rank_user_lib(candidate_file,spectrum_path,instrument_type,adduct,collision_
     return image_descrips
 
 
+
 with gr.Blocks(theme=seafoam) as demo:
     gr.HTML(custom_css)
-    gr.Markdown('<div style="font-size:100px; font-weight:bold;">🔍 CSU-MS2 web server </div>')
-
-    #gr.Markdown('<div style="font-size:20px;">You can select from a list of collected databases or upload your structural file by Clicking the button. The MS/MS spectrum to be identified can be uploaded in MSP format. You also need to set the necessary experimental parameters for the input spectrum, including accurate precursor mass and collision energy. Finally, by clicking the "Cross-Modal Retrieval" button, you can retrieve candidates from the selected database or your own structural file and access identification results that include compound structures, scores, and rankings.</div>')
+    gr.Markdown('<div style="font-size:50px; font-weight:bold;">🔍 CSU-MS2 web server </div>')
+    with gr.Row():
+        with gr.Column():
+            peak_data = gr.File(file_count="single", label="Upload MS/MS spectrum file in .msp format", elem_classes=".file-upload-height")
+        spectrum_output = gr.Plot(label="Mass Spectrum",elem_id="custom_plot")
+        peak_data.change(fn=draw_mass_spectrum, inputs=[peak_data], outputs=[spectrum_output])    
+    with gr.Row():
+        instru=gr.Dropdown(["HCD"], label="Instrument Type")
+        ionmode=gr.Dropdown(["[M+H]+"], label="Adduct Type")
+        par_ion_mass=gr.Textbox(label="Parent Ion Mass",placeholder="e.g., 180.00")
+        collision_e=gr.Textbox(label="collision energy", placeholder="e.g., 40")
     with gr.Tab(label="📶 Struture library", elem_id='custom_tab'):
-        with gr.Row():
-            with gr.Column():
-                peak_data = gr.File(file_count="single", label="Upload MS/MS spectrum file in .msp format", elem_classes=".file-upload-height")
-                draw_button = gr.Button("Draw Mass Spectrum")
-            with gr.Column():
-                spectrum_output = gr.Plot(label="Mass Spectrum",elem_id="custom_plot")
         with gr.Column():
             dataset = gr.Dropdown(["CSU_MS2_DB",
                          "BloodExp: blood exposome database", 
@@ -443,35 +452,19 @@ with gr.Blocks(theme=seafoam) as demo:
                           "STOFF-IDENT: database of water relevant substances", 
                           "T3DB: toxic exposome database", 
                           "TCMSP: traditional chinese medicine systems pharmacology database",
-                          "YMDB: yeast metabolome database"], label="Choose a structure library")
-            #gr.CheckboxGroup(choices=["HMDB", "Lipidmaps", "CHEMBL"], label="Choose a structure library"),
-        with gr.Row():
-            instru=gr.Dropdown(["HCD"], label="Instrument Type")
-            ionmode=gr.Dropdown(["[M+H]+"], label="Adduct Type")
-            par_ion_mass=gr.Textbox(label="Parent Ion Mass",placeholder="e.g., 180.00")
-            collision_e=gr.Textbox(label="collision energy", placeholder="e.g., 40")
+                          "YMDB: yeast metabolome database"], label="Choose a structure library")        
         with gr.Column():
             lib_button = gr.Button("Cross-Modal Retrieval")
             lib_output = gr.Gallery(height='auto',columns=4,elem_classes="gallery-height",label='Cross-modal retrieval results')
-            #lib_output = gr.Dataframe(type="pandas")
     with gr.Tab("📁 Upload structure file"):
         with gr.Row():
             with gr.Column():
-                use_dataset= gr.File(file_count="single", label="Upload the candidate structure file in csv format, columns=['SMIELS']",elem_classes="file-upload-height2")
-                user_peak_data=gr.File(file_count="single", label="Upload MS/MS spectrum file in .msp format", elem_classes="file-upload-height2")
-                user_draw_button = gr.Button("Draw Mass Spectrum")
-            with gr.Column():
-                user_spectrum_output = gr.Plot(label="Mass Spectrum",elem_id="custom_plot2")           
-        with gr.Row():
-            user_instru=gr.Dropdown(["HCD"], label="Instrument Type")
-            user_ionmode=gr.Dropdown(["[M+H]+"], label="Adduct Type")
-            user_collision_e=gr.Textbox(label="collision energy", placeholder="e.g., 40")
+                user_dataset= gr.File(file_count="single", label="Upload the candidate structure file in csv format, columns=['SMIELS']",elem_classes="file-upload-height2")       
         with gr.Column():
             user_button = gr.Button("Cross-Modal Retrieval")
             user_output = gr.Gallery(height='auto',columns=4,elem_classes="gallery-height",label='Cross-modal retrieval results')
-    draw_button.click(draw_mass_spectrum, inputs=[peak_data], outputs=[spectrum_output])
-    user_draw_button.click(draw_mass_spectrum2, inputs=[user_peak_data], outputs=[user_spectrum_output])
     lib_button.click(rank_lib, inputs=[dataset,peak_data,instru,ionmode,par_ion_mass,collision_e], outputs=lib_output)
-    user_button.click(rank_user_lib, inputs=[use_dataset,user_peak_data,user_instru,user_ionmode,user_collision_e], outputs=user_output)
+    user_button.click(rank_user_lib, inputs=[user_dataset,peak_data,instru,ionmode,par_ion_mass,collision_e], outputs=user_output)
 demo.launch(share=True)
 
+