Add test screens and 3d view

app.py

@@ -1,5 +1,11 @@
 import gradio as gr
-from zern_generator import generate_zern_wavefront
+from zern_generator import generate_zern_wavefront_fig, generate_batch
+
+css = """
+.container {
+    max-height: 40vh;
+}
+"""
 
 zerns = {}
 
@@ -9,13 +15,25 @@ def update_zerns(key, value):
     zerns[key] = value
 
 
+def get_zern_from_file(file):
+    with open(file) as f:
+        coeffs = [float(l) for l in f]
+    return tuple(coeffs)
+
+
 theme = gr.themes.Default(primary_hue=gr.themes.colors.red, secondary_hue=gr.themes.colors.pink)
-with gr.Blocks(theme=theme) as demo:
+with gr.Blocks(theme=theme, css=css) as demo:
     with gr.Row():
         zern_sliders = []
-        sliders_container = gr.Column()
+
         with gr.Column(scale=1):
+            # file_inp = gr.File(label="Загрузка коэффициентов из файла", height="2em")
+
+            resolution_num = gr.Number(value=256, label="Разрешение")
             num_zernikes = gr.Number(label="Количество коэффициентов Цернике", value=2, interactive=True)
+            sliders_container = gr.Column(elem_classes=["container"])
+            # file_inp.upload(get_zern_from_file, inputs=[file_inp], outputs=[*sliders_container.children])
+
 
             @gr.render(inputs=[num_zernikes], triggers=[num_zernikes.change])
             def render_count(count):
@@ -24,7 +42,7 @@ with gr.Blocks(theme=theme) as demo:
                 sliders_container.children.clear()
                 count = min(count, 100)
                 for i in range(count):
-                    sld = gr.Slider(key=i, label=f"Zernike {i}", minimum=-50, interactive=True,
+                    sld = gr.Slider(key=i, label=f"Zernike {i}", minimum=-100, interactive=True,
                                     step=0.01, value=0)
 
                     sld.change(update_zerns, inputs=[gr.Number(value=i, visible=False), sld])
@@ -32,13 +50,25 @@ with gr.Blocks(theme=theme) as demo:
 
                 return gr.update()
 
-        with gr.Column(scale=2):
-            plot = gr.Plot()
-            btn = gr.Button("Рассчитать")
+        with gr.Column(scale=5):
+            with gr.Tab(label="2D"):
+                plot_2d = gr.Plot()
+            with gr.Tab(label="3D"):
+                plot_3d = gr.Plot()
+            with gr.Tab(label="Batch"):
+
+                plot_batch = gr.Plot(label="Тестовые экраны (NOLL scheme)")
+
+            with gr.Row():
+                btn = gr.Button("Рассчитать")
+                btn_batch = gr.Button("Рассчитать тестовые экраны")
+
+
+            def on_button_click(resolution):
+                return generate_zern_wavefront_fig(resolution, *[zerns[key] for key in sorted(zerns)])
 
 
-            def on_button_click():
-                return generate_zern_wavefront(*[zerns[key] for key in sorted(zerns)])
-            btn.click(on_button_click, inputs=[], outputs=[plot])
+            btn.click(on_button_click, inputs=[resolution_num], outputs=[plot_2d, plot_3d])
+            btn_batch.click(generate_batch, inputs=[resolution_num], outputs=[plot_batch])
 
 demo.launch()

zern_generator.py

@@ -8,11 +8,7 @@ SAMPLE_SIZE = 256
 span = (-1, 1)
 
 
-def generate_zern_wavefront(*zerns):
-    ddx = np.linspace(*span, SAMPLE_SIZE)
-    ddy = np.linspace(*span, SAMPLE_SIZE)
-    xv, yv = np.meshgrid(ddx, ddy)
-
+def generate_zern_wavefront(xv, yv, zerns):
     num_zerns = len(zerns)
 
     degree = math.ceil((-3 + math.sqrt(8 * num_zerns + 1)) / 2)
@@ -20,10 +16,44 @@ def generate_zern_wavefront(*zerns):
 
     cart = RZern(degree)
     cart.make_cart_grid(xv, yv)
+    return np.array(cart.eval_grid(np.array(zerns), matrix=True))
+
+
+def generate_zern_wavefront_fig(resolution, *zerns):
+    ddx = np.linspace(*span, resolution)
+    ddy = np.linspace(*span, resolution)
+    xv, yv = np.meshgrid(ddx, ddy)
 
-    Phi = np.array(cart.eval_grid(np.array(zerns), matrix=True))
-    Phi[np.isnan(Phi)] = 0
-    fig = plt.figure(figsize=(20, 10))
-    plt.imshow(Phi)
+    Phi = generate_zern_wavefront(xv, yv, zerns)
+    # Phi[np.isnan(Phi)] = 0
+    fig_2d, ax_2d = plt.subplots()
+    fig_3d, ax_3d = plt.subplots(subplot_kw={"projection": "3d"})
+    a = ax_2d.imshow(Phi, extent=(*span, *span), cmap='gray')
+    ax_3d.plot_surface(xv, yv, Phi)
+    fig_2d.colorbar(a)
+    plt.show()
+    return fig_2d, fig_3d
+
+
+def generate_batch(resolution):
+    ddx = np.linspace(*span, resolution)
+    ddy = np.linspace(*span, resolution)
+    xv, yv = np.meshgrid(ddx, ddy)
+    fig = plt.figure()
+    fig.set_figwidth(30)
+    fig.set_figheight(20)
+    zern_index = 1
+    for i in range(5):
+        for j in range(6):
+            if j > i:
+                continue
+
+            zern = [0] * (zern_index - 1) + [1]
+            ax = plt.subplot2grid((5, 6), (i, j))
+            Phi = generate_zern_wavefront(xv, yv, zern)
+            ax.set_title(f"Zernike {zern_index-1}")
+            ax.imshow(Phi, extent=(*span, *span), cmap='gray')
+
+            zern_index += 1
     plt.show()
     return fig