jpoptum/EM_Wave_Simulation

1import numpy as np2import plotly.graph_objs as go3from plotly.subplots import make_subplots4import streamlit as st5 6def calculate_wave(wavelength, amplitude, frequency, num_periods, time_step):7    k = 2*np.pi/wavelength8    omega = 2*np.pi*frequency9    period = 1/frequency10    time_array = np.arange(0, num_periods*period, time_step)11    wave = amplitude * np.sin(k * np.arange(0, 1, wavelength/1000)[:, None] - omega * time_array)12    return time_array, wave13 14def plot_3d_wave(wavelength, amplitude, frequency):15    fig = make_subplots(rows=1, cols=2,16                        specs=[[{'type': 'surface'}, {'type': 'surface'}]],17                        subplot_titles=('Electric field', 'Magnetic field'),18                        )19    time_step = 1/(100*frequency)20    num_periods = 521    time_array, wave = calculate_wave(wavelength, amplitude, frequency, num_periods, time_step)22    E_x = wave*np.cos(2*np.pi*time_array*frequency)23    E_y = np.zeros_like(E_x)24    E_z = np.zeros_like(E_x)25    B_x = np.zeros_like(E_x)26    B_y = -wave*np.sin(2*np.pi*time_array*frequency)27    B_z = np.zeros_like(E_x)28    29    fig.add_trace(go.Surface(x=wave, y=E_x, z=E_y, colorscale='Blues'), row=1, col=1)30    fig.add_trace(go.Surface(x=wave, y=B_x, z=B_y, colorscale='Reds'), row=1, col=2)31    fig.update_layout(scene_aspectratio=dict(x=1, y=1, z=1), width=800, height=400,32                      scene=dict(xaxis_title='Wave', yaxis_title='E_x', zaxis_title='E_y'),33                      scene2=dict(xaxis_title='Wave', yaxis_title='B_x', zaxis_title='B_y'),34                     )35    36    fig.update_xaxes(range=[0, wavelength], row=1, col=1)37    fig.update_xaxes(range=[0, wavelength], row=1, col=2)38    fig.update_yaxes(range=[-amplitude, amplitude], row=1, col=1)39    fig.update_yaxes(range=[-amplitude, amplitude], row=1, col=2)40    fig.update_zaxes(range=[-amplitude, amplitude], row=1, col=1)41    fig.update_zaxes(range=[-amplitude, amplitude], row=1, col=2)42    43    return fig44 45wavelength = st.slider('Wavelength', 0.1, 10.0, 1.0)46amplitude = st.slider('Amplitude', 0.1, 1.0, 0.5)47frequency = st.slider('Frequency', 0.1, 10.0, 1.0)48 49fig = plot_3d_wave(wavelength, amplitude, frequency)50 51st.plotly_chart(fig)52