reformated code

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()