From cd3cd4ba9ecf45a6c1b71aada31882432d8a2b37 Mon Sep 17 00:00:00 2001 From: Mohamed Date: Fri, 11 Oct 2024 14:55:58 +0200 Subject: [PATCH] reformated code --- .../nanover-omni/src/nanover/omni/UI.py | 140 ++++++++++++++---- 1 file changed, 112 insertions(+), 28 deletions(-) diff --git a/python-libraries/nanover-omni/src/nanover/omni/UI.py b/python-libraries/nanover-omni/src/nanover/omni/UI.py index ecdf8368..b6d2779e 100644 --- a/python-libraries/nanover-omni/src/nanover/omni/UI.py +++ b/python-libraries/nanover-omni/src/nanover/omni/UI.py @@ -1,10 +1,32 @@ import gradio as gr + imd_runner = None -def run_simulation(simulation_type, input_files, trajectory_files, state_file, verbosity, show_progression, server_name, - port, simulation_fps, frame_interval, force_interval, start_paused, include_velocities, - include_forces, record_stats, stats_file, stats_fps, record_trajectory, - trajectory_output_file, record_shared_state, shared_state_file): - ''' + + +def run_simulation( + simulation_type, + input_files, + trajectory_files, + state_file, + verbosity, + show_progression, + server_name, + port, + simulation_fps, + frame_interval, + force_interval, + start_paused, + include_velocities, + include_forces, + record_stats, + stats_file, + stats_fps, + record_trajectory, + trajectory_output_file, + record_shared_state, + shared_state_file, +): + """ Runs a simulation with the given parameters. :param simulation_type: choose between "From xml" and "From recording" @@ -30,17 +52,18 @@ def run_simulation(simulation_type, input_files, trajectory_files, state_file, v :param shared_state_file: shared state file :return: a string with the simulation type and settings - ''' + """ from omni import OmniRunner from playback import PlaybackSimulation from openmm import OpenMMSimulation from record import record_from_server + global imd_runner # Initialize simulation files list simulation_files = list() for i in range(len(input_files)): - if input_files[i].endswith('.xml'): + if input_files[i].endswith(".xml"): # Create OpenMMSimulation from XML file simulation = OpenMMSimulation.from_xml_path(input_files[i], name=str(i)) simulation.frame_interval = frame_interval @@ -55,20 +78,31 @@ def run_simulation(simulation_type, input_files, trajectory_files, state_file, v trajectory_files = [] if trajectory_files is None else trajectory_files for i in range(len(trajectory_files)): # Create PlaybackSimulation from trajectory and state files - recording_playbacks.append(PlaybackSimulation(name='recording-playback_'+str(i), - traj=trajectory_files[i], - state=state_file[i])) + recording_playbacks.append( + PlaybackSimulation( + name="recording-playback_" + str(i), + traj=trajectory_files[i], + state=state_file[i], + ) + ) # Create OmniRunner with basic server - imd_runner = OmniRunner.with_basic_server(*tuple(simulation_files+recording_playbacks), name=server_name, port=port) + imd_runner = OmniRunner.with_basic_server( + *tuple(simulation_files + recording_playbacks), name=server_name, port=port + ) imd_runner.next() if start_paused: imd_runner.pause() if record_trajectory or record_shared_state: # Record from server if required - record_from_server(f"localhost:{imd_runner.app_server.port}",trajectory_output_file, shared_state_file) + record_from_server( + f"localhost:{imd_runner.app_server.port}", + trajectory_output_file, + shared_state_file, + ) return f"Simulation started with type: {simulation_type}, settings: {locals()}" + def stop_simulation(): global imd_runner try: @@ -78,32 +112,48 @@ def stop_simulation(): except NameError as e: return e return "Simulation stopped!" + + def create_ui(): with gr.Blocks(theme=gr.themes.Soft(), fill_height=True) as demo: gr.Markdown("# Nanover IMD Python Server GUI") # Radio button to select simulation type - simulation_type = gr.Radio(["From xml", "From recording"], label="Simulation Type", value="From xml") + simulation_type = gr.Radio( + ["From xml", "From recording"], label="Simulation Type", value="From xml" + ) with gr.Row(): with gr.Column(visible=True) as realtime_col: # File input for live simulation - input_files = gr.File(label="Input Files (for From xml)", file_count="multiple") + input_files = gr.File( + label="Input Files (for From xml)", file_count="multiple" + ) with gr.Column(visible=False) as playback_col: # File inputs for playback simulation - trajectory_files = gr.File(label="Trajectory Files (for playback)", file_count="multiple") - state_file = gr.File(label="State File (for playback)", file_count="multiple") + trajectory_files = gr.File( + label="Trajectory Files (for playback)", file_count="multiple" + ) + state_file = gr.File( + label="State File (for playback)", file_count="multiple" + ) with gr.Row(): with gr.Column(): gr.Markdown("## Verbosity") # Radio button for verbosity level - verbosity = gr.Radio(["Normal", "Verbose", "Super verbose"], label="Verbosity Level", value="Normal") + verbosity = gr.Radio( + ["Normal", "Verbose", "Super verbose"], + label="Verbosity Level", + value="Normal", + ) show_progression = gr.Checkbox(label="Show simulation progression") gr.Markdown("## Network") # Textbox for server name - server_name = gr.Textbox(label="Server name", value="NanoVer-RS iMD Server") + server_name = gr.Textbox( + label="Server name", value="NanoVer-RS iMD Server" + ) # Number input for port port = gr.Number(label="Port", value=38801) @@ -115,7 +165,9 @@ def create_ui(): force_interval = gr.Number(label="Force interval", value=5) # Checkboxes for simulation options start_paused = gr.Checkbox(label="Start simulation paused") - include_velocities = gr.Checkbox(label="Include the velocities in the frames") + include_velocities = gr.Checkbox( + label="Include the velocities in the frames" + ) include_forces = gr.Checkbox(label="Include the forces in the frames") with gr.Column(): @@ -127,10 +179,15 @@ def create_ui(): stats_fps = gr.Number(label="Statistics FPS", value=4) with gr.Group(): record_trajectory = gr.Checkbox(label="Record trajectory") - trajectory_output_file = gr.Textbox(label="Trajectory output file", value=f"{server_name.value}.traj") + trajectory_output_file = gr.Textbox( + label="Trajectory output file", + value=f"{server_name.value}.traj", + ) with gr.Group(): record_shared_state = gr.Checkbox(label="Record shared state") - shared_state_file = gr.Textbox(label="Shared state file", value=f"{server_name.value}.state") + shared_state_file = gr.Textbox( + label="Shared state file", value=f"{server_name.value}.state" + ) # Buttons to run and stop the simulation run_button = gr.Button("Run the selected file!") @@ -139,22 +196,49 @@ def create_ui(): def toggle_visibility(choice): # Toggle visibility of columns based on simulation type - return gr.update(visible=choice == "From xml"), gr.update(visible=choice == "From recording") + return gr.update(visible=choice == "From xml"), gr.update( + visible=choice == "From recording" + ) # Change event for simulation type radio button - simulation_type.change(toggle_visibility, inputs=[simulation_type], outputs=[realtime_col, playback_col]) + simulation_type.change( + toggle_visibility, + inputs=[simulation_type], + outputs=[realtime_col, playback_col], + ) # Click events for run and stop buttons run_button.click( run_simulation, - inputs=[simulation_type, input_files, trajectory_files, state_file, verbosity, show_progression, server_name, - port, simulation_fps, frame_interval, force_interval, start_paused, include_velocities, - include_forces, record_stats, stats_file, stats_fps, record_trajectory, - trajectory_output_file, record_shared_state, shared_state_file], - outputs=output + inputs=[ + simulation_type, + input_files, + trajectory_files, + state_file, + verbosity, + show_progression, + server_name, + port, + simulation_fps, + frame_interval, + force_interval, + start_paused, + include_velocities, + include_forces, + record_stats, + stats_file, + stats_fps, + record_trajectory, + trajectory_output_file, + record_shared_state, + shared_state_file, + ], + outputs=output, ) stop_button.click(stop_simulation, outputs=output) return demo + + # Launch the Gradio interface if __name__ == "__main__": ui = create_ui()