ElisR/spherical_harmonics_visualisation

1import gradio as gr2import plotly.graph_objects as go3 4import numpy as np5from scipy.special import sph_harm6 7 8def simple_sph_plot(l, m):9    # Fix the input if invalid10    if abs(m) > l:11        m = 012 13    resolution = 20014 15    # Pick uniform angles16    theta = np.linspace(0, np.pi, resolution)17    phi = np.linspace(0, 2 * np.pi, resolution)18    theta, phi = np.meshgrid(theta, phi)19 20    # The spherical harmonic function to plot21    Y_lm = sph_harm(m, l, phi, theta).real  # Original22    #if m == 0:23    #    Y_lm = sph_harm(m, l, phi, theta).real24    #elif m < 0:25    #    Y_lm = np.sqrt(2) * (-1)**m * sph_harm(-m, l, phi, theta).imag26    #else:27    #    Y_lm = np.sqrt(2) * (-1)**m * sph_harm(m, l, phi, theta).real28 29    # Convert to Cartesian coordinates30    x = np.sin(theta) * np.cos(phi) * np.abs(Y_lm)31    y = np.sin(theta) * np.sin(phi) * np.abs(Y_lm)32    z = np.cos(theta) * np.abs(Y_lm)33 34    surf = go.Surface(35        x=x, y=y, z=z, surfacecolor=Y_lm, colorscale="RdBu", showscale=False36    )37 38    fig = go.Figure(data=[surf])39 40    # Make axis invisible41    invisible_axis = dict(42        showbackground=False,43        showline=False,44        zeroline=False,45        showgrid=False,46        showticklabels=False,47        title="",48    )49 50    camera = {51        "up": {"x": 0, "y": 0, "z": 1},52        "center": {"x": 0, "y": 0, "z": 0},53        "eye": {"x": 1.25, "y": 1.25, "z": 1.25},54    }55 56    fig.update_layout(57        **{58            f"scene": {59                "xaxis": invisible_axis,60                "yaxis": invisible_axis,61                "zaxis": invisible_axis,62            },63            "scene_camera": camera,64        }65    )66 67    return fig68 69 70outputs = gr.Plot()71inputs = [72    gr.Number(value=1, minimum=0, info="Angular Momentum"),73    gr.Number(value=1, info="z-Component"),74]75 76iface = gr.Interface(fn=simple_sph_plot, inputs=inputs, outputs=outputs)77 78iface.launch()79