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 .mobj_types import ManagedObjType
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_types import SocketType
from .tree_types import TreeType
@ -96,6 +102,7 @@ __all__ = [
'NodeType',
'BoundCondType',
'NewSimCloudTask',
'SimAxisDir',
'SimSpaceAxis',
'manual_amp_time',
'NodeCategory',

View File

@ -221,9 +221,9 @@ class LazyArrayRangeFlow:
"""
if self.unit is not None:
log.debug(
'%s: Scaled to unit system: %s',
'%s: Scaled to new unit system (new unit = %s)',
self,
str(unit_system),
unit_system[spux.PhysicalType.from_unit(self.unit)],
)
return LazyArrayRangeFlow(
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]
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
####################

View File

@ -112,12 +112,20 @@ def write_modifier_geonodes(
iface_id = socket_infos[socket_name].bl_isocket_identifier
input_value = modifier_attrs['inputs'][socket_name]
if isinstance(input_value, spux.SympyType):
bl_modifier[iface_id] = spux.scale_to_unit_system(
if modifier_attrs['unit_system'] is not None and not isinstance(
input_value, bool
):
value_to_write = spux.scale_to_unit_system(
input_value, modifier_attrs['unit_system']
)
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
## TODO: More fine-grained alterations

View File

@ -17,7 +17,7 @@
from . import (
# astigmatic_gaussian_beam_source,
# gaussian_beam_source,
# plane_wave_source,
plane_wave_source,
point_dipole_source,
temporal_shapes,
)
@ -26,7 +26,7 @@ BL_REGISTER = [
*temporal_shapes.BL_REGISTER,
*point_dipole_source.BL_REGISTER,
# *uniform_current_source.BL_REGISTER,
# *plane_wave_source.BL_REGISTER,
*plane_wave_source.BL_REGISTER,
# *gaussian_beam_source.BL_REGISTER,
# *astigmatic_gaussian_beam_source.BL_REGISTER,
# *tfsf_source.BL_REGISTER,
@ -35,7 +35,7 @@ BL_NODES = {
**temporal_shapes.BL_NODES,
**point_dipole_source.BL_NODES,
# **uniform_current_source.BL_NODES,
# **plane_wave_source.BL_NODES,
**plane_wave_source.BL_NODES,
# **gaussian_beam_source.BL_NODES,
# **astigmatic_gaussian_beam_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
# along with this program. If not, see <http://www.gnu.org/licenses/>.
import math
import typing as typ
import bpy
import sympy as sp
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 managed_objs, sockets
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):
node_type = ct.NodeType.PlaneWaveSource
bl_label = 'Plane Wave Source'
use_sim_node_name = True
####################
# - Sockets
####################
input_sockets: typ.ClassVar = {
'Temporal Shape': sockets.MaxwellTemporalShapeSocketDef(),
'Center': sockets.PhysicalPoint3DSocketDef(),
'Direction': sockets.Real3DVectorSocketDef(default_value=sp.Matrix([0, 0, -1])),
'Pol Angle': sockets.PhysicalAngleSocketDef(),
'Center': sockets.ExprSocketDef(
shape=(3,),
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 = {
'Source': sockets.MaxwellSourceSocketDef(),
'Angled Source': sockets.MaxwellSourceSocketDef(),
}
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
####################
@events.computes_output_socket(
'Source',
input_sockets={'Temporal Shape', 'Center', 'Direction', 'Pol Angle'},
'Angled Source',
props={'sim_node_name', 'injection_axis', 'injection_direction'},
input_sockets={'Temporal Shape', 'Center', 'Spherical', 'Pol ∡'},
unit_systems={'Tidy3DUnits': ct.UNITS_TIDY3D},
scale_input_sockets={
'Center': 'Tidy3DUnits',
'Spherical': 'Tidy3DUnits',
'Pol ∡': 'Tidy3DUnits',
},
)
def compute_source(self, input_sockets: dict):
direction = input_sockets['Direction']
pol_angle = input_sockets['Pol Angle']
injection_axis, dir_sgn, theta, phi = convert_vector_to_spherical(
input_sockets['Direction']
)
def compute_source(self, props, input_sockets, unit_systems):
size = {
'x': (0, math.inf, math.inf),
'y': (math.inf, 0, math.inf),
'z': (math.inf, math.inf, 0),
}[injection_axis]
ct.SimSpaceAxis.X: (0, td.inf, td.inf),
ct.SimSpaceAxis.Y: (td.inf, 0, td.inf),
ct.SimSpaceAxis.Z: (td.inf, td.inf, 0),
}[props['injection_axis']]
# Display the results
return td.PlaneWave(
name=props['sim_node_name'],
center=input_sockets['Center'],
source_time=input_sockets['Temporal Shape'],
size=size,
direction=dir_sgn,
angle_theta=theta,
angle_phi=phi,
pol_angle=pol_angle,
source_time=input_sockets['Temporal Shape'],
direction=props['injection_direction'].plus_or_minus,
angle_theta=input_sockets['Spherical'][0],
angle_phi=input_sockets['Spherical'][1],
pol_angle=input_sockets['Pol ∡'],
)
#####################
## - Preview
#####################
# @events.on_value_changed(
# socket_name={'Center', 'Direction'},
# input_sockets={'Center', 'Direction'},
# managed_objs={'plane_wave_source'},
# )
# def on_value_changed__center_direction(
# self,
# input_sockets: dict,
# managed_objs: dict[str, ct.schemas.ManagedObj],
# ):
# _center = input_sockets['Center']
# center = tuple([float(el) for el in spu.convert_to(_center, spu.um) / spu.um])
####################
# - Preview - Changes to Input Sockets
####################
@events.on_value_changed(
# Trigger
prop_name='preview_active',
# Loaded
managed_objs={'mesh'},
props={'preview_active'},
)
def on_preview_changed(self, managed_objs, props):
"""Enables/disables previewing of the GeoNodes-driven mesh, regardless of whether a particular GeoNodes tree is chosen."""
mesh = managed_objs['mesh']
# _direction = input_sockets['Direction']
# direction = tuple([float(el) for el in _direction])
# ## TODO: Preview unit system?? Presume um for now
# Push Preview State to Managed Mesh
if props['preview_active']:
mesh.show_preview()
else:
mesh.hide_preview()
# # Retrieve Hard-Coded GeoNodes and Analyze Input
# geo_nodes = bpy.data.node_groups[GEONODES_PLANE_WAVE]
# geonodes_interface = analyze_geonodes.interface(geo_nodes, direc='INPUT')
# # Sync Modifier Inputs
# managed_objs['plane_wave_source'].sync_geonodes_modifier(
# geonodes_node_group=geo_nodes,
# geonodes_identifier_to_value={
# geonodes_interface['Direction'].identifier: direction,
# ## TODO: Use 'bl_socket_map.value_to_bl`!
# ## - This accounts for auto-conversion, unit systems, etc. .
# ## - We could keep it in the node base class...
# ## - ...But it needs aligning with Blender, too. Hmm.
# },
# )
# # Sync Object Position
# managed_objs['plane_wave_source'].bl_object('MESH').location = center
# @events.on_show_preview(
# managed_objs={'plane_wave_source'},
# )
# def on_show_preview(
# self,
# managed_objs: dict[str, ct.schemas.ManagedObj],
# ):
# managed_objs['plane_wave_source'].show_preview('MESH')
# self.on_value_changed__center_direction()
@events.on_value_changed(
# Trigger
socket_name={'Center', 'Spherical', 'Pol ∡'},
prop_name={'injection_axis', 'injection_direction'},
run_on_init=True,
# Loaded
managed_objs={'mesh', 'modifier'},
props={'injection_axis', 'injection_direction'},
input_sockets={'Temporal Shape', 'Center', 'Spherical', 'Pol ∡'},
unit_systems={'BlenderUnits': ct.UNITS_BLENDER},
scale_input_sockets={
'Center': 'BlenderUnits',
},
)
def on_inputs_changed(self, managed_objs, props, input_sockets, unit_systems):
# Push Input Values to GeoNodes Modifier
managed_objs['modifier'].bl_modifier(
managed_objs['mesh'].bl_object(location=input_sockets['Center']),
'NODES',
{
'node_group': import_geonodes(GeoNodes.SourcePlaneWave),
'unit_system': unit_systems['BlenderUnits'],
'inputs': {
'Inj Axis': props['injection_axis'].axis,
'Direction': props['injection_direction'].true_or_false,
'theta': input_sockets['Spherical'][0],
'phi': input_sockets['Spherical'][1],
'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
def unit_str_to_unit(unit_str: str) -> Unit | None:
# 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):
return expr