oscillode/blender_maxwell/node_trees/maxwell_sim_nodes/nodes/sources/point_dipole_source.py

134 lines
3.0 KiB
Python

import typing as typ
import tidy3d as td
import sympy as sp
import sympy.physics.units as spu
import bpy
from ... import contracts as ct
from ... import sockets
from .. import base
from ... import managed_objs
class PointDipoleSourceNode(base.MaxwellSimNode):
node_type = ct.NodeType.PointDipoleSource
bl_label = "Point Dipole Source"
####################
# - Sockets
####################
input_sockets = {
"Temporal Shape": sockets.MaxwellTemporalShapeSocketDef(),
"Center": sockets.PhysicalPoint3DSocketDef(),
"Interpolate": sockets.BoolSocketDef(
default_value=True,
),
}
output_sockets = {
"Source": sockets.MaxwellSourceSocketDef(),
}
managed_obj_defs = {
"sphere_empty": ct.schemas.ManagedObjDef(
mk=lambda name: managed_objs.ManagedBLObject(name),
name_prefix="point_dipole_",
)
}
####################
# - Properties
####################
pol_axis: bpy.props.EnumProperty(
name="Polarization Axis",
description="Polarization Axis",
items=[
("EX", "Ex", "Electric field in x-dir"),
("EY", "Ey", "Electric field in y-dir"),
("EZ", "Ez", "Electric field in z-dir"),
],
default="EX",
update=(lambda self, context: self.sync_prop("pol_axis")),
)
####################
# - UI
####################
def draw_props(self, context, layout):
layout.prop(self, "pol_axis", text="Pol Axis")
####################
# - Output Socket Computation
####################
@base.computes_output_socket(
"Source",
input_sockets={"Temporal Shape", "Center", "Interpolate"},
props={"pol_axis"},
)
def compute_source(self, input_sockets: dict[str, typ.Any], props: dict[str, typ.Any]) -> td.PointDipole:
pol_axis = {
"EX": "Ex",
"EY": "Ey",
"EZ": "Ez",
}[props["pol_axis"]]
temporal_shape = input_sockets["Temporal Shape"]
_center = input_sockets["Center"]
interpolate = input_sockets["Interpolate"]
center = tuple(spu.convert_to(_center, spu.um) / spu.um)
_res = td.PointDipole(
center=center,
source_time=temporal_shape,
interpolate=interpolate,
polarization=pol_axis,
)
return _res
####################
# - Preview
####################
@base.on_value_changed(
socket_name="Center",
input_sockets={"Center"},
managed_objs={"sphere_empty"},
)
def on_value_changed__center(
self,
input_sockets: dict,
managed_objs: dict[str, ct.schemas.ManagedObj],
):
_center = input_sockets["Center"]
center = tuple(spu.convert_to(_center, spu.um) / spu.um)
## TODO: Preview unit system?? Presume um for now
mobj = managed_objs["sphere_empty"]
bl_object = mobj.bl_object("EMPTY")
bl_object.location = center #tuple([float(el) for el in center])
@base.on_show_preview(
managed_objs={"sphere_empty"},
)
def on_show_preview(
self,
managed_objs: dict[str, ct.schemas.ManagedObj],
):
managed_objs["sphere_empty"].show_preview(
"EMPTY",
empty_display_type="SPHERE",
)
####################
# - Blender Registration
####################
BL_REGISTER = [
PointDipoleSourceNode,
]
BL_NODES = {
ct.NodeType.PointDipoleSource: (
ct.NodeCategory.MAXWELLSIM_SOURCES
)
}