3D visualization of the Cheetah accelerator segment

app.py

@@ -0,0 +1,44 @@
+import torch
+import cheetah
+import gradio as gr
+import os
+
+# Define the accelerator segment
+segment = cheetah.Segment(
+    elements=[
+        cheetah.Drift(length=torch.tensor(0.175)),
+        cheetah.Quadrupole(length=torch.tensor(0.122), name="AREAMQZM1"),
+        cheetah.Drift(length=torch.tensor(0.428)),
+        cheetah.Quadrupole(length=torch.tensor(0.122), name="AREAMQZM2"),
+        cheetah.Drift(length=torch.tensor(0.204)),
+        cheetah.VerticalCorrector(length=torch.tensor(0.02), name="AREAMCVM1"),
+        cheetah.Drift(length=torch.tensor(0.204)),
+        cheetah.Quadrupole(length=torch.tensor(0.122), name="AREAMQZM3"),
+        cheetah.Drift(length=torch.tensor(0.179)),
+        cheetah.HorizontalCorrector(length=torch.tensor(0.02), name="AREAMCHM1"),
+        cheetah.Drift(length=torch.tensor(0.45)),
+        cheetah.Screen(name="AREABSCR1"),
+    ]
+)
+
+# Generate the mesh
+mesh, _ = segment.to_mesh(
+    cuteness={cheetah.HorizontalCorrector: 2.0, cheetah.VerticalCorrector: 2.0}
+)
+
+# Export the mesh to a .glb file
+output_file = "ares_mesh.glb"
+mesh.export(file_obj=output_file, file_type="glb")
+
+# Define the Gradio interface
+def display_3d_model():
+    return output_file
+
+with gr.Blocks() as demo:
+    gr.Markdown("# 3D Visualization of Cheetah Accelerator Segment")
+    gr.Markdown("This app renders a 3D model of the accelerator segment defined using the Cheetah library.")
+    gr.File(file_types=[".glb"], label="3D Model", value=output_file)
+    gr.Model3D(value=output_file, label="3D Visualization")
+
+# Launch the Gradio app
+demo.launch()