feat: implement plane wave node w/viz

the visualization of Tidy3D's spherical coordinates was exceptionally harsh. See #63 for more information.

Closes #63.
main
Sofus Albert Høgsbro Rose 2024-05-06 22:10:45 +02:00
parent 0fbd3752b3
commit 9f8ff33e4f
Signed by: so-rose
GPG Key ID: AD901CB0F3701434
9 changed files with 196 additions and 115 deletions

Binary file not shown.

Binary file not shown.

View File

@ -56,7 +56,13 @@ from .flow_signals import FlowSignal
from .icons import Icon from .icons import Icon
from .mobj_types import ManagedObjType from .mobj_types import ManagedObjType
from .node_types import NodeType from .node_types import NodeType
from .sim_types import BoundCondType, NewSimCloudTask, SimSpaceAxis, manual_amp_time from .sim_types import (
BoundCondType,
NewSimCloudTask,
SimAxisDir,
SimSpaceAxis,
manual_amp_time,
)
from .socket_colors import SOCKET_COLORS from .socket_colors import SOCKET_COLORS
from .socket_types import SocketType from .socket_types import SocketType
from .tree_types import TreeType from .tree_types import TreeType
@ -96,6 +102,7 @@ __all__ = [
'NodeType', 'NodeType',
'BoundCondType', 'BoundCondType',
'NewSimCloudTask', 'NewSimCloudTask',
'SimAxisDir',
'SimSpaceAxis', 'SimSpaceAxis',
'manual_amp_time', 'manual_amp_time',
'NodeCategory', 'NodeCategory',

View File

@ -221,9 +221,9 @@ class LazyArrayRangeFlow:
""" """
if self.unit is not None: if self.unit is not None:
log.debug( log.debug(
'%s: Scaled to unit system: %s', '%s: Scaled to new unit system (new unit = %s)',
self, self,
str(unit_system), unit_system[spux.PhysicalType.from_unit(self.unit)],
) )
return LazyArrayRangeFlow( return LazyArrayRangeFlow(
start=spux.strip_unit_system( start=spux.strip_unit_system(

View File

@ -114,6 +114,61 @@ class SimSpaceAxis(enum.StrEnum):
return {SSA.X: 0, SSA.Y: 1, SSA.Z: 2}[self] return {SSA.X: 0, SSA.Y: 1, SSA.Z: 2}[self]
class SimAxisDir(enum.StrEnum):
"""Positive or negative direction along an injection axis."""
Plus = enum.auto()
Minus = enum.auto()
@staticmethod
def to_name(v: typ.Self) -> str:
"""Convert the enum value to a human-friendly name.
Notes:
Used to print names in `EnumProperty`s based on this enum.
Returns:
A human-friendly name corresponding to the enum value.
"""
SAD = SimAxisDir
return {
SAD.Plus: '+',
SAD.Minus: '-',
}[v]
@staticmethod
def to_icon(_: typ.Self) -> str:
"""Convert the enum value to a Blender icon.
Notes:
Used to print icons in `EnumProperty`s based on this enum.
Returns:
A human-friendly name corresponding to the enum value.
"""
return ''
@property
def plus_or_minus(self) -> int:
"""Get '+' or '-' literal corresponding to the direction.
Returns:
The appropriate literal.
"""
SAD = SimAxisDir
return {SAD.Plus: '+', SAD.Minus: '-'}[self]
@property
def true_or_false(self) -> bool:
"""Get 'True' or 'False' bool corresponding to the direction.
Returns:
The appropriate bool.
"""
SAD = SimAxisDir
return {SAD.Plus: True, SAD.Minus: False}[self]
#################### ####################
# - Boundary Condition Type # - Boundary Condition Type
#################### ####################

View File

@ -112,12 +112,20 @@ def write_modifier_geonodes(
iface_id = socket_infos[socket_name].bl_isocket_identifier iface_id = socket_infos[socket_name].bl_isocket_identifier
input_value = modifier_attrs['inputs'][socket_name] input_value = modifier_attrs['inputs'][socket_name]
if isinstance(input_value, spux.SympyType): if modifier_attrs['unit_system'] is not None and not isinstance(
bl_modifier[iface_id] = spux.scale_to_unit_system( input_value, bool
):
value_to_write = spux.scale_to_unit_system(
input_value, modifier_attrs['unit_system'] input_value, modifier_attrs['unit_system']
) )
else: else:
bl_modifier[iface_id] = input_value value_to_write = input_value
# Edge Case: int -> float
if isinstance(bl_modifier[iface_id], float) and isinstance(value_to_write, int):
bl_modifier[iface_id] = float(value_to_write)
else:
bl_modifier[iface_id] = value_to_write
modifier_altered = True modifier_altered = True
## TODO: More fine-grained alterations ## TODO: More fine-grained alterations

View File

@ -17,7 +17,7 @@
from . import ( from . import (
# astigmatic_gaussian_beam_source, # astigmatic_gaussian_beam_source,
# gaussian_beam_source, # gaussian_beam_source,
# plane_wave_source, plane_wave_source,
point_dipole_source, point_dipole_source,
temporal_shapes, temporal_shapes,
) )
@ -26,7 +26,7 @@ BL_REGISTER = [
*temporal_shapes.BL_REGISTER, *temporal_shapes.BL_REGISTER,
*point_dipole_source.BL_REGISTER, *point_dipole_source.BL_REGISTER,
# *uniform_current_source.BL_REGISTER, # *uniform_current_source.BL_REGISTER,
# *plane_wave_source.BL_REGISTER, *plane_wave_source.BL_REGISTER,
# *gaussian_beam_source.BL_REGISTER, # *gaussian_beam_source.BL_REGISTER,
# *astigmatic_gaussian_beam_source.BL_REGISTER, # *astigmatic_gaussian_beam_source.BL_REGISTER,
# *tfsf_source.BL_REGISTER, # *tfsf_source.BL_REGISTER,
@ -35,7 +35,7 @@ BL_NODES = {
**temporal_shapes.BL_NODES, **temporal_shapes.BL_NODES,
**point_dipole_source.BL_NODES, **point_dipole_source.BL_NODES,
# **uniform_current_source.BL_NODES, # **uniform_current_source.BL_NODES,
# **plane_wave_source.BL_NODES, **plane_wave_source.BL_NODES,
# **gaussian_beam_source.BL_NODES, # **gaussian_beam_source.BL_NODES,
# **astigmatic_gaussian_beam_source.BL_NODES, # **astigmatic_gaussian_beam_source.BL_NODES,
# **tfsf_source.BL_NODES, # **tfsf_source.BL_NODES,

View File

@ -14,152 +14,156 @@
# You should have received a copy of the GNU Affero General Public License # You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>. # along with this program. If not, see <http://www.gnu.org/licenses/>.
import math
import typing as typ import typing as typ
import bpy
import sympy as sp import sympy as sp
import tidy3d as td import tidy3d as td
from blender_maxwell.assets.geonodes import GeoNodes, import_geonodes
from blender_maxwell.utils import bl_cache, logger
from blender_maxwell.utils import extra_sympy_units as spux
from ... import contracts as ct from ... import contracts as ct
from ... import managed_objs, sockets from ... import managed_objs, sockets
from .. import base, events from .. import base, events
def convert_vector_to_spherical(
v: sp.MatrixBase,
) -> tuple[str, str, sp.Expr, sp.Expr]:
"""Converts a vector (maybe normalized) to spherical coordinates from an arbitrary choice of injection axis.
Injection axis is chosen to minimize `theta`
"""
x, y, z = v
injection_axis = max(
('x', abs(x)), ('y', abs(y)), ('z', abs(z)), key=lambda item: item[1]
)[0]
## Select injection axis that minimizes 'theta'
if injection_axis == 'x':
direction = '+' if x >= 0 else '-'
theta = sp.acos(x / sp.sqrt(x**2 + y**2 + z**2))
phi = sp.atan2(z, y)
elif injection_axis == 'y':
direction = '+' if y >= 0 else '-'
theta = sp.acos(y / sp.sqrt(x**2 + y**2 + z**2))
phi = sp.atan2(x, z)
else:
direction = '+' if z >= 0 else '-'
theta = sp.acos(z / sp.sqrt(x**2 + y**2 + z**2))
phi = sp.atan2(y, x)
return injection_axis, direction, theta, phi
class PlaneWaveSourceNode(base.MaxwellSimNode): class PlaneWaveSourceNode(base.MaxwellSimNode):
node_type = ct.NodeType.PlaneWaveSource node_type = ct.NodeType.PlaneWaveSource
bl_label = 'Plane Wave Source' bl_label = 'Plane Wave Source'
use_sim_node_name = True
#################### ####################
# - Sockets # - Sockets
#################### ####################
input_sockets: typ.ClassVar = { input_sockets: typ.ClassVar = {
'Temporal Shape': sockets.MaxwellTemporalShapeSocketDef(), 'Temporal Shape': sockets.MaxwellTemporalShapeSocketDef(),
'Center': sockets.PhysicalPoint3DSocketDef(), 'Center': sockets.ExprSocketDef(
'Direction': sockets.Real3DVectorSocketDef(default_value=sp.Matrix([0, 0, -1])), shape=(3,),
'Pol Angle': sockets.PhysicalAngleSocketDef(), mathtype=spux.MathType.Real,
physical_type=spux.PhysicalType.Length,
default_value=sp.Matrix([0, 0, 0]),
),
'Spherical': sockets.ExprSocketDef(
shape=(2,),
mathtype=spux.MathType.Real,
physical_type=spux.PhysicalType.Angle,
default_value=sp.Matrix([0, 0]),
),
'Pol ∡': sockets.ExprSocketDef(
physical_type=spux.PhysicalType.Angle,
default_value=0,
),
} }
output_sockets: typ.ClassVar = { output_sockets: typ.ClassVar = {
'Source': sockets.MaxwellSourceSocketDef(), 'Angled Source': sockets.MaxwellSourceSocketDef(),
} }
managed_obj_types: typ.ClassVar = { managed_obj_types: typ.ClassVar = {
'plane_wave_source': managed_objs.ManagedBLMesh, 'mesh': managed_objs.ManagedBLMesh,
'modifier': managed_objs.ManagedBLModifier,
} }
####################
# - Properties
####################
injection_axis: ct.SimSpaceAxis = bl_cache.BLField(ct.SimSpaceAxis.X, prop_ui=True)
injection_direction: ct.SimAxisDir = bl_cache.BLField(
ct.SimAxisDir.Plus, prop_ui=True
)
####################
# - UI
####################
def draw_props(self, _: bpy.types.Context, layout: bpy.types.UILayout):
layout.prop(self, self.blfields['injection_axis'], expand=True)
layout.prop(self, self.blfields['injection_direction'], expand=True)
#################### ####################
# - Output Socket Computation # - Output Socket Computation
#################### ####################
@events.computes_output_socket( @events.computes_output_socket(
'Source', 'Angled Source',
input_sockets={'Temporal Shape', 'Center', 'Direction', 'Pol Angle'}, props={'sim_node_name', 'injection_axis', 'injection_direction'},
input_sockets={'Temporal Shape', 'Center', 'Spherical', 'Pol ∡'},
unit_systems={'Tidy3DUnits': ct.UNITS_TIDY3D}, unit_systems={'Tidy3DUnits': ct.UNITS_TIDY3D},
scale_input_sockets={ scale_input_sockets={
'Center': 'Tidy3DUnits', 'Center': 'Tidy3DUnits',
'Spherical': 'Tidy3DUnits',
'Pol ∡': 'Tidy3DUnits',
}, },
) )
def compute_source(self, input_sockets: dict): def compute_source(self, props, input_sockets, unit_systems):
direction = input_sockets['Direction']
pol_angle = input_sockets['Pol Angle']
injection_axis, dir_sgn, theta, phi = convert_vector_to_spherical(
input_sockets['Direction']
)
size = { size = {
'x': (0, math.inf, math.inf), ct.SimSpaceAxis.X: (0, td.inf, td.inf),
'y': (math.inf, 0, math.inf), ct.SimSpaceAxis.Y: (td.inf, 0, td.inf),
'z': (math.inf, math.inf, 0), ct.SimSpaceAxis.Z: (td.inf, td.inf, 0),
}[injection_axis] }[props['injection_axis']]
# Display the results # Display the results
return td.PlaneWave( return td.PlaneWave(
name=props['sim_node_name'],
center=input_sockets['Center'], center=input_sockets['Center'],
source_time=input_sockets['Temporal Shape'],
size=size, size=size,
direction=dir_sgn, source_time=input_sockets['Temporal Shape'],
angle_theta=theta, direction=props['injection_direction'].plus_or_minus,
angle_phi=phi, angle_theta=input_sockets['Spherical'][0],
pol_angle=pol_angle, angle_phi=input_sockets['Spherical'][1],
pol_angle=input_sockets['Pol ∡'],
) )
##################### ####################
## - Preview # - Preview - Changes to Input Sockets
##################### ####################
# @events.on_value_changed( @events.on_value_changed(
# socket_name={'Center', 'Direction'}, # Trigger
# input_sockets={'Center', 'Direction'}, prop_name='preview_active',
# managed_objs={'plane_wave_source'}, # Loaded
# ) managed_objs={'mesh'},
# def on_value_changed__center_direction( props={'preview_active'},
# self, )
# input_sockets: dict, def on_preview_changed(self, managed_objs, props):
# managed_objs: dict[str, ct.schemas.ManagedObj], """Enables/disables previewing of the GeoNodes-driven mesh, regardless of whether a particular GeoNodes tree is chosen."""
# ): mesh = managed_objs['mesh']
# _center = input_sockets['Center']
# center = tuple([float(el) for el in spu.convert_to(_center, spu.um) / spu.um])
# _direction = input_sockets['Direction'] # Push Preview State to Managed Mesh
# direction = tuple([float(el) for el in _direction]) if props['preview_active']:
# ## TODO: Preview unit system?? Presume um for now mesh.show_preview()
else:
mesh.hide_preview()
# # Retrieve Hard-Coded GeoNodes and Analyze Input @events.on_value_changed(
# geo_nodes = bpy.data.node_groups[GEONODES_PLANE_WAVE] # Trigger
# geonodes_interface = analyze_geonodes.interface(geo_nodes, direc='INPUT') socket_name={'Center', 'Spherical', 'Pol ∡'},
prop_name={'injection_axis', 'injection_direction'},
# # Sync Modifier Inputs run_on_init=True,
# managed_objs['plane_wave_source'].sync_geonodes_modifier( # Loaded
# geonodes_node_group=geo_nodes, managed_objs={'mesh', 'modifier'},
# geonodes_identifier_to_value={ props={'injection_axis', 'injection_direction'},
# geonodes_interface['Direction'].identifier: direction, input_sockets={'Temporal Shape', 'Center', 'Spherical', 'Pol ∡'},
# ## TODO: Use 'bl_socket_map.value_to_bl`! unit_systems={'BlenderUnits': ct.UNITS_BLENDER},
# ## - This accounts for auto-conversion, unit systems, etc. . scale_input_sockets={
# ## - We could keep it in the node base class... 'Center': 'BlenderUnits',
# ## - ...But it needs aligning with Blender, too. Hmm. },
# }, )
# ) def on_inputs_changed(self, managed_objs, props, input_sockets, unit_systems):
# Push Input Values to GeoNodes Modifier
# # Sync Object Position managed_objs['modifier'].bl_modifier(
# managed_objs['plane_wave_source'].bl_object('MESH').location = center managed_objs['mesh'].bl_object(location=input_sockets['Center']),
'NODES',
# @events.on_show_preview( {
# managed_objs={'plane_wave_source'}, 'node_group': import_geonodes(GeoNodes.SourcePlaneWave),
# ) 'unit_system': unit_systems['BlenderUnits'],
# def on_show_preview( 'inputs': {
# self, 'Inj Axis': props['injection_axis'].axis,
# managed_objs: dict[str, ct.schemas.ManagedObj], 'Direction': props['injection_direction'].true_or_false,
# ): 'theta': input_sockets['Spherical'][0],
# managed_objs['plane_wave_source'].show_preview('MESH') 'phi': input_sockets['Spherical'][1],
# self.on_value_changed__center_direction() 'Pol Angle': input_sockets['Pol ∡'],
},
},
)
#################### ####################

View File

@ -779,7 +779,14 @@ def scaling_factor(unit_from: spu.Quantity, unit_to: spu.Quantity) -> Number:
@functools.cache @functools.cache
def unit_str_to_unit(unit_str: str) -> Unit | None: def unit_str_to_unit(unit_str: str) -> Unit | None:
expr = sp.sympify(unit_str).subs(UNIT_BY_SYMBOL) # Edge Case: Manually Parse Degrees
## -> sp.sympify('degree') actually produces the sp.degree() function.
## -> Therefore, we must special case this particular unit.
if unit_str == 'degree':
expr = spu.degree
else:
expr = sp.sympify(unit_str).subs(UNIT_BY_SYMBOL)
if expr.has(spu.Quantity): if expr.has(spu.Quantity):
return expr return expr