154 lines
4.1 KiB
Python
154 lines
4.1 KiB
Python
import typing as typ
|
|
|
|
import tidy3d as td
|
|
import numpy as np
|
|
import sympy as sp
|
|
import sympy.physics.units as spu
|
|
|
|
import bpy
|
|
|
|
from ......utils import extra_sympy_units as spuex
|
|
from .... import contracts as ct
|
|
from .... import sockets
|
|
from .... import managed_objs
|
|
from ... import base
|
|
|
|
class GaussianPulseTemporalShapeNode(base.MaxwellSimNode):
|
|
node_type = ct.NodeType.GaussianPulseTemporalShape
|
|
|
|
bl_label = "Gaussian Pulse Temporal Shape"
|
|
#bl_icon = ...
|
|
|
|
####################
|
|
# - Sockets
|
|
####################
|
|
input_sockets = {
|
|
#"amplitude": sockets.RealNumberSocketDef(
|
|
# label="Temporal Shape",
|
|
#), ## Should have a unit of some kind...
|
|
"Freq Center": sockets.PhysicalFreqSocketDef(
|
|
default_value=500 * spuex.terahertz,
|
|
),
|
|
"Freq Std.": sockets.PhysicalFreqSocketDef(
|
|
default_value=200 * spuex.terahertz,
|
|
),
|
|
"Phase": sockets.PhysicalAngleSocketDef(),
|
|
"Delay rel. AngFreq": sockets.RealNumberSocketDef(
|
|
default_value=5.0,
|
|
),
|
|
"Remove DC": sockets.BoolSocketDef(
|
|
default_value=True,
|
|
),
|
|
}
|
|
output_sockets = {
|
|
"Temporal Shape": sockets.MaxwellTemporalShapeSocketDef(),
|
|
}
|
|
|
|
managed_obj_defs = {
|
|
"amp_time": ct.schemas.ManagedObjDef(
|
|
mk=lambda name: managed_objs.ManagedBLImage(name),
|
|
name_prefix="amp_time_",
|
|
)
|
|
}
|
|
|
|
####################
|
|
# - Properties
|
|
####################
|
|
plot_time_start: bpy.props.FloatProperty(
|
|
name="Plot Time Start (ps)",
|
|
description="The instance ID of a particular MaxwellSimNode instance, used to index caches",
|
|
default=0.0,
|
|
update=(lambda self, context: self.sync_prop("plot_time_start", context)),
|
|
)
|
|
plot_time_end: bpy.props.FloatProperty(
|
|
name="Plot Time End (ps)",
|
|
description="The instance ID of a particular MaxwellSimNode instance, used to index caches",
|
|
default=5,
|
|
update=(lambda self, context: self.sync_prop("plot_time_start", context)),
|
|
)
|
|
|
|
####################
|
|
# - UI
|
|
####################
|
|
def draw_props(self, context, layout):
|
|
layout.label(text="Plot Settings")
|
|
split = layout.split(factor=0.6)
|
|
|
|
col = split.column()
|
|
col.label(text="t-Range (ps)")
|
|
|
|
col = split.column()
|
|
col.prop(self, "plot_time_start", text="")
|
|
col.prop(self, "plot_time_end", text="")
|
|
|
|
####################
|
|
# - Output Socket Computation
|
|
####################
|
|
@base.computes_output_socket(
|
|
"Temporal Shape",
|
|
input_sockets={
|
|
"Freq Center", "Freq Std.", "Phase", "Delay rel. AngFreq",
|
|
"Remove DC",
|
|
}
|
|
)
|
|
def compute_source(self, input_sockets: dict) -> td.GaussianPulse:
|
|
if (
|
|
(_freq_center := input_sockets["Freq Center"]) is None
|
|
or (_freq_std := input_sockets["Freq Std."]) is None
|
|
or (_phase := input_sockets["Phase"]) is None
|
|
or (time_delay_rel_ang_freq := input_sockets["Delay rel. AngFreq"]) is None
|
|
or (remove_dc_component := input_sockets["Remove DC"]) is None
|
|
):
|
|
raise ValueError("Inputs not defined")
|
|
|
|
cheating_amplitude = 1.0
|
|
freq_center = spu.convert_to(_freq_center, spu.hertz) / spu.hertz
|
|
freq_std = spu.convert_to(_freq_std, spu.hertz) / spu.hertz
|
|
phase = spu.convert_to(_phase, spu.radian) / spu.radian
|
|
|
|
return td.GaussianPulse(
|
|
amplitude=cheating_amplitude,
|
|
phase=phase,
|
|
freq0=freq_center,
|
|
fwidth=freq_std,
|
|
offset=time_delay_rel_ang_freq,
|
|
remove_dc_component=remove_dc_component,
|
|
)
|
|
|
|
@base.on_show_plot(
|
|
managed_objs={"amp_time"},
|
|
props={"plot_time_start", "plot_time_end"},
|
|
output_sockets={"Temporal Shape"},
|
|
stop_propagation=True,
|
|
)
|
|
def on_show_plot(
|
|
self,
|
|
managed_objs: dict[str, ct.schemas.ManagedObj],
|
|
output_sockets: dict[str, typ.Any],
|
|
props: dict[str, typ.Any],
|
|
):
|
|
temporal_shape = output_sockets["Temporal Shape"]
|
|
plot_time_start = props["plot_time_start"] * 1e-15
|
|
plot_time_end = props["plot_time_end"] * 1e-15
|
|
|
|
times = np.linspace(plot_time_start, plot_time_end)
|
|
|
|
managed_objs["amp_time"].mpl_plot_to_image(
|
|
lambda ax: temporal_shape.plot_spectrum(times, ax=ax),
|
|
bl_select=True,
|
|
)
|
|
|
|
|
|
|
|
####################
|
|
# - Blender Registration
|
|
####################
|
|
BL_REGISTER = [
|
|
GaussianPulseTemporalShapeNode,
|
|
]
|
|
BL_NODES = {
|
|
ct.NodeType.GaussianPulseTemporalShape: (
|
|
ct.NodeCategory.MAXWELLSIM_SOURCES_TEMPORALSHAPES
|
|
)
|
|
}
|