diff --git a/src/blender_maxwell/node_trees/maxwell_sim_nodes/contracts/flow_kinds/lazy_array_range.py b/src/blender_maxwell/node_trees/maxwell_sim_nodes/contracts/flow_kinds/lazy_array_range.py index 867fb1d..cc3da7d 100644 --- a/src/blender_maxwell/node_trees/maxwell_sim_nodes/contracts/flow_kinds/lazy_array_range.py +++ b/src/blender_maxwell/node_trees/maxwell_sim_nodes/contracts/flow_kinds/lazy_array_range.py @@ -388,6 +388,12 @@ class LazyArrayRangeFlow: self.stop.subs({sym: symbol_values[sym.name] for sym in self.symbols}) ) + def realize_step_size( + self, + symbol_values: dict[spux.Symbol, typ.Any] = MappingProxyType({}), + ) -> ArrayFlow | LazyValueFuncFlow: + return (self.realize_stop() - self.realize_start()) / self.steps + def realize( self, symbol_values: dict[spux.Symbol, typ.Any] = MappingProxyType({}), diff --git a/src/blender_maxwell/node_trees/maxwell_sim_nodes/contracts/node_types.py b/src/blender_maxwell/node_trees/maxwell_sim_nodes/contracts/node_types.py index 4ffefc1..8e6826b 100644 --- a/src/blender_maxwell/node_trees/maxwell_sim_nodes/contracts/node_types.py +++ b/src/blender_maxwell/node_trees/maxwell_sim_nodes/contracts/node_types.py @@ -55,10 +55,7 @@ class NodeType(blender_type_enum.BlenderTypeEnum): JSONFileExporter = enum.auto() # Sources - ## Sources / Temporal Shapes - PulseTemporalShape = enum.auto() - WaveTemporalShape = enum.auto() - ExprTemporalShape = enum.auto() + TemporalShape = enum.auto() ## Sources / PointDipoleSource = enum.auto() PlaneWaveSource = enum.auto() diff --git a/src/blender_maxwell/node_trees/maxwell_sim_nodes/nodes/analysis/math/map_math.py b/src/blender_maxwell/node_trees/maxwell_sim_nodes/nodes/analysis/math/map_math.py index 9f52c9e..b6a6cf0 100644 --- a/src/blender_maxwell/node_trees/maxwell_sim_nodes/nodes/analysis/math/map_math.py +++ b/src/blender_maxwell/node_trees/maxwell_sim_nodes/nodes/analysis/math/map_math.py @@ -374,7 +374,7 @@ class MapMathNode(base.MaxwellSimNode): # - Properties #################### operation: MapOperation = bl_cache.BLField( - prop_ui=True, enum_cb=lambda self, _: self.search_operations() + enum_cb=lambda self, _: self.search_operations() ) @property diff --git a/src/blender_maxwell/node_trees/maxwell_sim_nodes/nodes/analysis/viz.py b/src/blender_maxwell/node_trees/maxwell_sim_nodes/nodes/analysis/viz.py index 050bd80..ed605a2 100644 --- a/src/blender_maxwell/node_trees/maxwell_sim_nodes/nodes/analysis/viz.py +++ b/src/blender_maxwell/node_trees/maxwell_sim_nodes/nodes/analysis/viz.py @@ -22,7 +22,7 @@ import jaxtyping as jtyp import matplotlib.axis as mpl_ax import sympy as sp -from blender_maxwell.utils import bl_cache, image_ops, logger +from blender_maxwell.utils import bl_cache, image_ops, logger, sim_symbols from blender_maxwell.utils import extra_sympy_units as spux from ... import contracts as ct @@ -210,8 +210,8 @@ class VizNode(base.MaxwellSimNode): input_sockets: typ.ClassVar = { 'Expr': sockets.ExprSocketDef( active_kind=ct.FlowKind.LazyValueFunc, - symbols={_x := sp.Symbol('x', real=True)}, - default_value=2 * _x, + default_symbols=[sim_symbols.x], + default_value=2 * sim_symbols.x.sp_symbol, ), } output_sockets: typ.ClassVar = { @@ -295,8 +295,8 @@ class VizNode(base.MaxwellSimNode): #################### @events.on_value_changed( socket_name='Expr', - input_sockets={'Expr'}, run_on_init=True, + input_sockets={'Expr'}, input_socket_kinds={'Expr': {ct.FlowKind.Info, ct.FlowKind.Params}}, input_sockets_optional={'Expr': True}, ) diff --git a/src/blender_maxwell/node_trees/maxwell_sim_nodes/nodes/sources/__init__.py b/src/blender_maxwell/node_trees/maxwell_sim_nodes/nodes/sources/__init__.py index 6a83d80..cb007d4 100644 --- a/src/blender_maxwell/node_trees/maxwell_sim_nodes/nodes/sources/__init__.py +++ b/src/blender_maxwell/node_trees/maxwell_sim_nodes/nodes/sources/__init__.py @@ -19,11 +19,11 @@ from . import ( gaussian_beam_source, plane_wave_source, point_dipole_source, - temporal_shapes, + temporal_shape, ) BL_REGISTER = [ - *temporal_shapes.BL_REGISTER, + *temporal_shape.BL_REGISTER, *plane_wave_source.BL_REGISTER, *point_dipole_source.BL_REGISTER, # *uniform_current_source.BL_REGISTER, @@ -32,7 +32,7 @@ BL_REGISTER = [ # *tfsf_source.BL_REGISTER, ] BL_NODES = { - **temporal_shapes.BL_NODES, + **temporal_shape.BL_NODES, **plane_wave_source.BL_NODES, **point_dipole_source.BL_NODES, # **uniform_current_source.BL_NODES, diff --git a/src/blender_maxwell/node_trees/maxwell_sim_nodes/nodes/sources/plane_wave_source.py b/src/blender_maxwell/node_trees/maxwell_sim_nodes/nodes/sources/plane_wave_source.py index 994aca3..496881c 100644 --- a/src/blender_maxwell/node_trees/maxwell_sim_nodes/nodes/sources/plane_wave_source.py +++ b/src/blender_maxwell/node_trees/maxwell_sim_nodes/nodes/sources/plane_wave_source.py @@ -28,6 +28,8 @@ from ... import contracts as ct from ... import managed_objs, sockets from .. import base, events +log = logger.get(__name__) + class PlaneWaveSourceNode(base.MaxwellSimNode): """An infinite-extent angled source simulating an plane wave with linear polarization. @@ -73,7 +75,6 @@ class PlaneWaveSourceNode(base.MaxwellSimNode): } managed_obj_types: typ.ClassVar = { - 'mesh': managed_objs.ManagedBLMesh, 'modifier': managed_objs.ManagedBLModifier, } @@ -132,18 +133,14 @@ class PlaneWaveSourceNode(base.MaxwellSimNode): # Trigger prop_name='preview_active', # Loaded - managed_objs={'mesh'}, + managed_objs={'modifier'}, 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'] - - # Push Preview State to Managed Mesh if props['preview_active']: - mesh.show_preview() + managed_objs['modifier'].show_preview() else: - mesh.hide_preview() + managed_objs['modifier'].hide_preview() @events.on_value_changed( # Trigger @@ -151,7 +148,7 @@ class PlaneWaveSourceNode(base.MaxwellSimNode): prop_name={'injection_axis', 'injection_direction'}, run_on_init=True, # Loaded - managed_objs={'mesh', 'modifier'}, + managed_objs={'modifier'}, props={'injection_axis', 'injection_direction'}, input_sockets={'Temporal Shape', 'Center', 'Spherical', 'Pol ∡'}, unit_systems={'BlenderUnits': ct.UNITS_BLENDER}, @@ -162,7 +159,6 @@ class PlaneWaveSourceNode(base.MaxwellSimNode): 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), @@ -175,6 +171,7 @@ class PlaneWaveSourceNode(base.MaxwellSimNode): 'Pol Angle': input_sockets['Pol ∡'], }, }, + location=input_sockets['Center'], ) diff --git a/src/blender_maxwell/node_trees/maxwell_sim_nodes/nodes/sources/temporal_shape.py b/src/blender_maxwell/node_trees/maxwell_sim_nodes/nodes/sources/temporal_shape.py new file mode 100644 index 0000000..242aa0b --- /dev/null +++ b/src/blender_maxwell/node_trees/maxwell_sim_nodes/nodes/sources/temporal_shape.py @@ -0,0 +1,197 @@ +# blender_maxwell +# Copyright (C) 2024 blender_maxwell Project Contributors +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . + +"""Implements the `TemporalShapeNode`.""" + +import typing as typ + +import bpy +import sympy as sp +import sympy.physics.units as spu +import tidy3d as td + +from blender_maxwell.utils import extra_sympy_units as spux +from blender_maxwell.utils import logger, sim_symbols + +from ... import contracts as ct +from ... import managed_objs, sockets +from .. import base, events + +log = logger.get(__name__) + + +_max_e_socket_def = sockets.ExprSocketDef( + mathtype=spux.MathType.Complex, + physical_type=spux.PhysicalType.EField, + default_value=1 + 0j, +) +_offset_socket_def = sockets.ExprSocketDef(default_value=5, abs_min=2.5) + + +class TemporalShapeNode(base.MaxwellSimNode): + """Declare a source-time dependence for use in simulation source nodes.""" + + node_type = ct.NodeType.TemporalShape + bl_label = 'Temporal Shape' + + #################### + # - Sockets + #################### + input_sockets: typ.ClassVar = { + 'μ Freq': sockets.ExprSocketDef( + physical_type=spux.PhysicalType.Freq, + default_unit=spux.THz, + default_value=500, + ), + 'σ Freq': sockets.ExprSocketDef( + physical_type=spux.PhysicalType.Freq, + default_unit=spux.THz, + default_value=200, + ), + } + input_socket_sets: typ.ClassVar = { + 'Pulse': { + 'max E': _max_e_socket_def, + 'Offset Time': _offset_socket_def, + 'Remove DC': sockets.BoolSocketDef(default_value=True), + }, + 'Constant': { + 'max E': _max_e_socket_def, + 'Offset Time': _offset_socket_def, + }, + 'Symbolic': { + 't Range': sockets.ExprSocketDef( + active_kind=ct.FlowKind.LazyArrayRange, + physical_type=spux.PhysicalType.Time, + default_unit=spu.picosecond, + default_min=0, + default_max=10, + default_steps=100, + ), + 'Envelope': sockets.ExprSocketDef( + default_symbols=[sim_symbols.t_ps], + default_value=10 * sim_symbols.t_ps.sp_symbol, + ), + }, + } + + output_sockets: typ.ClassVar = { + 'Temporal Shape': sockets.MaxwellTemporalShapeSocketDef(), + } + + managed_obj_types: typ.ClassVar = { + 'plot': managed_objs.ManagedBLImage, + } + + def draw_info(self, _: bpy.types.Context, layout: bpy.types.UILayout) -> None: + if self.active_socket_set != 'Symbolic': + box = layout.box() + row = box.row() + row.alignment = 'CENTER' + row.label(text='Parameter Scale') + + # Split + split = box.split(factor=0.3, align=False) + + ## LHS: Parameter Names + col = split.column() + col.alignment = 'RIGHT' + col.label(text='Off t:') + + ## RHS: Parameter Units + col = split.column() + col.label(text='1 / 2π·σ(𝑓)') + + #################### + # - FlowKind: Value + #################### + @events.computes_output_socket( + 'Temporal Shape', + # Loaded + props={'active_socket_set'}, + input_sockets={ + 'max E', + 'μ Freq', + 'σ Freq', + 'Offset Time', + 'Remove DC', + 't Range', + 'Envelope', + }, + input_socket_kinds={ + 't Range': ct.FlowKind.LazyArrayRange, + 'Envelope': ct.FlowKind.LazyValueFunc, + }, + input_sockets_optional={ + 'max E': True, + 'Offset Time': True, + 'Remove DC': True, + 't Range': True, + 'Envelope': True, + }, + unit_systems={'Tidy3DUnits': ct.UNITS_TIDY3D}, + scale_input_sockets={ + 'max E': 'Tidy3DUnits', + 'μ Freq': 'Tidy3DUnits', + 'σ Freq': 'Tidy3DUnits', + 't Range': 'Tidy3DUnits', + 'Offset Time': 'Tidy3DUnits', + }, + ) + def compute_temporal_shape( + self, props, input_sockets, unit_systems + ) -> td.GaussianPulse: + match props['active_socket_set']: + case 'Pulse': + return td.GaussianPulse( + amplitude=sp.re(input_sockets['max E']), + phase=sp.im(input_sockets['max E']), + freq0=input_sockets['μ Freq'], + fwidth=input_sockets['σ Freq'], + offset=input_sockets['Offset Time'], + remove_dc_component=input_sockets['Remove DC'], + ) + + case 'Constant': + return td.ContinuousWave( + amplitude=sp.re(input_sockets['max E']), + phase=sp.im(input_sockets['max E']), + freq0=input_sockets['μ Freq'], + fwidth=input_sockets['σ Freq'], + offset=input_sockets['Offset Time'], + ) + + case 'Symbolic': + lzrange = input_sockets['t Range'] + envelope_ps = input_sockets['Envelope'].func_jax + + return td.CustomSourceTime.from_values( + freq0=input_sockets['μ Freq'], + fwidth=input_sockets['σ Freq'], + values=envelope_ps( + lzrange.rescale_to_unit(spu.ps).realize_array.values + ), + dt=input_sockets['t Range'].realize_step_size(), + ) + + +#################### +# - Blender Registration +#################### +BL_REGISTER = [ + TemporalShapeNode, +] +BL_NODES = {ct.NodeType.TemporalShape: (ct.NodeCategory.MAXWELLSIM_SOURCES)} diff --git a/src/blender_maxwell/node_trees/maxwell_sim_nodes/nodes/sources/temporal_shapes/__init__.py b/src/blender_maxwell/node_trees/maxwell_sim_nodes/nodes/sources/temporal_shapes/__init__.py deleted file mode 100644 index 673c2d3..0000000 --- a/src/blender_maxwell/node_trees/maxwell_sim_nodes/nodes/sources/temporal_shapes/__init__.py +++ /dev/null @@ -1,29 +0,0 @@ -# blender_maxwell -# Copyright (C) 2024 blender_maxwell Project Contributors -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . - -# from . import expr_temporal_shape, pulse_temporal_shape, wave_temporal_shape -from . import pulse_temporal_shape, wave_temporal_shape - -BL_REGISTER = [ - *pulse_temporal_shape.BL_REGISTER, - *wave_temporal_shape.BL_REGISTER, - # *expr_temporal_shape.BL_REGISTER, -] -BL_NODES = { - **pulse_temporal_shape.BL_NODES, - **wave_temporal_shape.BL_NODES, - # **expr_temporal_shape.BL_NODES, -} diff --git a/src/blender_maxwell/node_trees/maxwell_sim_nodes/nodes/sources/temporal_shapes/expr_temporal_shape.py b/src/blender_maxwell/node_trees/maxwell_sim_nodes/nodes/sources/temporal_shapes/expr_temporal_shape.py deleted file mode 100644 index 5b7d824..0000000 --- a/src/blender_maxwell/node_trees/maxwell_sim_nodes/nodes/sources/temporal_shapes/expr_temporal_shape.py +++ /dev/null @@ -1,21 +0,0 @@ -# blender_maxwell -# Copyright (C) 2024 blender_maxwell Project Contributors -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . - -#################### -# - Blender Registration -#################### -BL_REGISTER = [] -BL_NODES = {} diff --git a/src/blender_maxwell/node_trees/maxwell_sim_nodes/nodes/sources/temporal_shapes/pulse_temporal_shape.py b/src/blender_maxwell/node_trees/maxwell_sim_nodes/nodes/sources/temporal_shapes/pulse_temporal_shape.py deleted file mode 100644 index b55e39c..0000000 --- a/src/blender_maxwell/node_trees/maxwell_sim_nodes/nodes/sources/temporal_shapes/pulse_temporal_shape.py +++ /dev/null @@ -1,177 +0,0 @@ -# blender_maxwell -# Copyright (C) 2024 blender_maxwell Project Contributors -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . - -"""Implements the `PulseTemporalShapeNode`.""" - -import functools -import typing as typ - -import bpy -import sympy as sp -import sympy.physics.units as spu -import tidy3d as td - -from blender_maxwell.utils import extra_sympy_units as spux - -from .... import contracts as ct -from .... import managed_objs, sockets -from ... import base, events - - -class PulseTemporalShapeNode(base.MaxwellSimNode): - node_type = ct.NodeType.PulseTemporalShape - bl_label = 'Gaussian Pulse Temporal Shape' - - #################### - # - Sockets - #################### - input_sockets: typ.ClassVar = { - 'max E': sockets.ExprSocketDef( - mathtype=spux.MathType.Complex, - physical_type=spux.PhysicalType.EField, - default_value=1 + 0j, - ), - 'μ Freq': sockets.ExprSocketDef( - physical_type=spux.PhysicalType.Freq, - default_unit=spux.THz, - default_value=500, - ), - 'σ Freq': sockets.ExprSocketDef( - physical_type=spux.PhysicalType.Freq, - default_unit=spux.THz, - default_value=200, - ), - 'Offset Time': sockets.ExprSocketDef(default_value=5, abs_min=2.5), - 'Remove DC': sockets.BoolSocketDef( - default_value=True, - ), - } - output_sockets: typ.ClassVar = { - 'Temporal Shape': sockets.MaxwellTemporalShapeSocketDef(), - 'E(t)': sockets.ExprSocketDef(active_kind=ct.FlowKind.Array), - } - - managed_obj_types: typ.ClassVar = { - 'plot': managed_objs.ManagedBLImage, - } - - #################### - # - UI - #################### - def draw_info(self, _: bpy.types.Context, layout: bpy.types.UILayout) -> None: - box = layout.box() - row = box.row() - row.alignment = 'CENTER' - row.label(text='Parameter Scale') - - # Split - split = box.split(factor=0.3, align=False) - - ## LHS: Parameter Names - col = split.column() - col.alignment = 'RIGHT' - col.label(text='Off t:') - - ## RHS: Parameter Units - col = split.column() - col.label(text='1 / 2π·σ(𝑓)') - - #################### - # - FlowKind: Value - #################### - @events.computes_output_socket( - 'Temporal Shape', - input_sockets={ - 'max E', - 'μ Freq', - 'σ Freq', - 'Offset Time', - 'Remove DC', - }, - unit_systems={'Tidy3DUnits': ct.UNITS_TIDY3D}, - scale_input_sockets={ - 'max E': 'Tidy3DUnits', - 'μ Freq': 'Tidy3DUnits', - 'σ Freq': 'Tidy3DUnits', - }, - ) - def compute_temporal_shape(self, input_sockets, unit_systems) -> td.GaussianPulse: - return td.GaussianPulse( - amplitude=sp.re(input_sockets['max E']), - phase=sp.im(input_sockets['max E']), - freq0=input_sockets['μ Freq'], - fwidth=input_sockets['σ Freq'], - offset=input_sockets['Offset Time'], - remove_dc_component=input_sockets['Remove DC'], - ) - - #################### - # - FlowKind: LazyValueFunc / Info / Params - #################### - @events.computes_output_socket( - 'E(t)', - kind=ct.FlowKind.LazyValueFunc, - output_sockets={'Temporal Shape'}, - ) - def compute_time_to_efield_lazy(self, output_sockets) -> td.GaussianPulse: - temporal_shape = output_sockets['Temporal Shape'] - jax_amp_time = functools.partial(ct.manual_amp_time, temporal_shape) - - ## TODO: Don't just partial() it up, do it property in the ParamsFlow! - ## -> Right now it's recompiled every time. - - return ct.LazyValueFuncFlow( - func=jax_amp_time, - func_args=[spux.PhysicalType.Time], - supports_jax=True, - ) - - @events.computes_output_socket( - 'E(t)', - kind=ct.FlowKind.Info, - ) - def compute_time_to_efield_info(self) -> td.GaussianPulse: - return ct.InfoFlow( - dim_names=['t'], - dim_idx={ - 't': ct.LazyArrayRangeFlow( - start=sp.S(0), stop=sp.oo, steps=0, unit=spu.second - ) - }, - output_name='E', - output_shape=None, - output_mathtype=spux.MathType.Complex, - output_unit=spu.volt / spu.um, - ) - - @events.computes_output_socket( - 'E(t)', - kind=ct.FlowKind.Params, - ) - def compute_time_to_efield_params(self) -> td.GaussianPulse: - sym_time = sp.Symbol('t', real=True, nonnegative=True) - return ct.ParamsFlow(func_args=[sym_time], symbols={sym_time}) - - -#################### -# - Blender Registration -#################### -BL_REGISTER = [ - PulseTemporalShapeNode, -] -BL_NODES = { - ct.NodeType.PulseTemporalShape: (ct.NodeCategory.MAXWELLSIM_SOURCES_TEMPORALSHAPES) -} diff --git a/src/blender_maxwell/node_trees/maxwell_sim_nodes/nodes/sources/temporal_shapes/wave_temporal_shape.py b/src/blender_maxwell/node_trees/maxwell_sim_nodes/nodes/sources/temporal_shapes/wave_temporal_shape.py deleted file mode 100644 index 10eb77f..0000000 --- a/src/blender_maxwell/node_trees/maxwell_sim_nodes/nodes/sources/temporal_shapes/wave_temporal_shape.py +++ /dev/null @@ -1,169 +0,0 @@ -# blender_maxwell -# Copyright (C) 2024 blender_maxwell Project Contributors -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . - -"""Implements the `WaveTemporalShapeNode`.""" - -import functools -import typing as typ - -import bpy -import sympy as sp -import sympy.physics.units as spu -import tidy3d as td - -from blender_maxwell.utils import extra_sympy_units as spux - -from .... import contracts as ct -from .... import managed_objs, sockets -from ... import base, events - - -class WaveTemporalShapeNode(base.MaxwellSimNode): - node_type = ct.NodeType.WaveTemporalShape - bl_label = 'Continuous Wave Temporal Shape' - - #################### - # - Sockets - #################### - input_sockets: typ.ClassVar = { - 'max E': sockets.ExprSocketDef( - mathtype=spux.MathType.Complex, - physical_type=spux.PhysicalType.EField, - default_value=1 + 0j, - ), - 'μ Freq': sockets.ExprSocketDef( - physical_type=spux.PhysicalType.Freq, - default_unit=spux.THz, - default_value=500, - ), - 'σ Freq': sockets.ExprSocketDef( - physical_type=spux.PhysicalType.Freq, - default_unit=spux.THz, - default_value=200, - ), - 'Offset Time': sockets.ExprSocketDef(default_value=5, abs_min=2.5), - } - output_sockets: typ.ClassVar = { - 'Temporal Shape': sockets.MaxwellTemporalShapeSocketDef(), - 'E(t)': sockets.ExprSocketDef(active_kind=ct.FlowKind.Array), - } - - managed_obj_types: typ.ClassVar = { - 'plot': managed_objs.ManagedBLImage, - } - - #################### - # - UI - #################### - def draw_info(self, _: bpy.types.Context, layout: bpy.types.UILayout) -> None: - box = layout.box() - row = box.row() - row.alignment = 'CENTER' - row.label(text='Parameter Scale') - - # Split - split = box.split(factor=0.3, align=False) - - ## LHS: Parameter Names - col = split.column() - col.alignment = 'RIGHT' - col.label(text='Off t:') - - ## RHS: Parameter Units - col = split.column() - col.label(text='1 / 2π·σ(𝑓)') - - #################### - # - FlowKind: Value - #################### - @events.computes_output_socket( - 'Temporal Shape', - input_sockets={ - 'max E', - 'μ Freq', - 'σ Freq', - 'Offset Time', - }, - unit_systems={'Tidy3DUnits': ct.UNITS_TIDY3D}, - scale_input_sockets={ - 'max E': 'Tidy3DUnits', - 'μ Freq': 'Tidy3DUnits', - 'σ Freq': 'Tidy3DUnits', - }, - ) - def compute_temporal_shape(self, input_sockets, unit_systems) -> td.GaussianPulse: - return td.ContinuousWave( - amplitude=sp.re(input_sockets['max E']), - phase=sp.im(input_sockets['max E']), - freq0=input_sockets['μ Freq'], - fwidth=input_sockets['σ Freq'], - offset=input_sockets['Offset Time'], - ) - - #################### - # - FlowKind: LazyValueFunc / Info / Params - #################### - @events.computes_output_socket( - 'E(t)', - kind=ct.FlowKind.LazyValueFunc, - output_sockets={'Temporal Shape'}, - ) - def compute_time_to_efield_lazy(self, output_sockets) -> td.GaussianPulse: - temporal_shape = output_sockets['Temporal Shape'] - jax_amp_time = functools.partial(ct.manual_amp_time, temporal_shape) - - return ct.LazyValueFuncFlow( - func=jax_amp_time, - func_args=[spux.PhysicalType.Time], - supports_jax=True, - ) - - @events.computes_output_socket( - 'E(t)', - kind=ct.FlowKind.Info, - ) - def compute_time_to_efield_info(self) -> td.GaussianPulse: - return ct.InfoFlow( - dim_names=['t'], - dim_idx={ - 't': ct.LazyArrayRangeFlow( - start=sp.S(0), stop=sp.oo, steps=0, unit=spu.second - ) - }, - output_name='E', - output_shape=None, - output_mathtype=spux.MathType.Complex, - output_unit=spu.volt / spu.um, - ) - - @events.computes_output_socket( - 'E(t)', - kind=ct.FlowKind.Params, - ) - def compute_time_to_efield_params(self) -> td.GaussianPulse: - sym_time = sp.Symbol('t', real=True, nonnegative=True) - return ct.ParamsFlow(func_args=[sym_time], symbols={sym_time}) - - -#################### -# - Blender Registration -#################### -BL_REGISTER = [ - WaveTemporalShapeNode, -] -BL_NODES = { - ct.NodeType.WaveTemporalShape: (ct.NodeCategory.MAXWELLSIM_SOURCES_TEMPORALSHAPES) -} diff --git a/src/blender_maxwell/node_trees/maxwell_sim_nodes/sockets/expr.py b/src/blender_maxwell/node_trees/maxwell_sim_nodes/sockets/expr.py index 0d5f673..4a3925a 100644 --- a/src/blender_maxwell/node_trees/maxwell_sim_nodes/sockets/expr.py +++ b/src/blender_maxwell/node_trees/maxwell_sim_nodes/sockets/expr.py @@ -23,7 +23,7 @@ import bpy import pydantic as pyd import sympy as sp -from blender_maxwell.utils import bl_cache, logger +from blender_maxwell.utils import bl_cache, logger, sim_symbols from blender_maxwell.utils import extra_sympy_units as spux from .. import contracts as ct @@ -936,8 +936,11 @@ class ExprSocketDef(base.SocketDef): physical_type: spux.PhysicalType = spux.PhysicalType.NonPhysical default_unit: spux.Unit | None = None - # symbols: list[sim_symbols.SimSymbol] = frozenset() - symbols: frozenset[spux.SympyExpr] = frozenset() + default_symbols: list[sim_symbols.SimSymbol] = pyd.Field(default_factory=list) + + @property + def symbols(self) -> set[sp.Symbol]: + return {sim_symbol.sp_symbol for sim_symbol in self.default_symbols} # FlowKind: Value default_value: spux.SympyExpr = 0 diff --git a/src/blender_maxwell/utils/extra_sympy_units.py b/src/blender_maxwell/utils/extra_sympy_units.py index 13106e1..4d6eb38 100644 --- a/src/blender_maxwell/utils/extra_sympy_units.py +++ b/src/blender_maxwell/utils/extra_sympy_units.py @@ -173,6 +173,17 @@ class MathType(enum.StrEnum): MT.Complex: complex, }[self] + @property + def symbolic_set(self) -> type: + MT = MathType + return { + MT.Bool: sp.Set([sp.S(False), sp.S(True)]), + MT.Integer: sp.Integers, + MT.Rational: sp.Rationals, + MT.Real: sp.Reals, + MT.Complex: sp.Complexes, + }[self] + @staticmethod def to_str(value: typ.Self) -> type: return { diff --git a/src/blender_maxwell/utils/sim_symbols.py b/src/blender_maxwell/utils/sim_symbols.py index a090c6f..ab3b26e 100644 --- a/src/blender_maxwell/utils/sim_symbols.py +++ b/src/blender_maxwell/utils/sim_symbols.py @@ -16,6 +16,7 @@ import dataclasses import enum +import sys import typing as typ import sympy as sp @@ -23,9 +24,19 @@ import sympy as sp from . import extra_sympy_units as spux -class SimSymbolNames(enum.StrEnum): +#################### +# - Simulation Symbols +#################### +class SimSymbolName(enum.StrEnum): LowerA = enum.auto() - LowerLambda = enum.auto() + LowerT = enum.auto() + LowerX = enum.auto() + LowerY = enum.auto() + LowerZ = enum.auto() + + # Physics + Wavelength = enum.auto() + Frequency = enum.auto() @staticmethod def to_name(v: typ.Self) -> str: @@ -37,11 +48,28 @@ class SimSymbolNames(enum.StrEnum): Returns: A human-friendly name corresponding to the enum value. """ - SSN = SimSymbolNames + return SimSymbolName(v).name + + @property + def name(self) -> str: + SSN = SimSymbolName return { SSN.LowerA: 'a', - SSN.LowerLambda: 'λ', - }[v] + SSN.LowerT: 't', + SSN.LowerX: 'x', + SSN.LowerY: 'y', + SSN.LowerZ: 'z', + SSN.Wavelength: 'wl', + SSN.Frequency: 'freq', + }[self] + + @property + def name_pretty(self) -> str: + SSN = SimSymbolName + return { + SSN.Wavelength: 'λ', + SSN.Frequency: '𝑓', + }.get(self, self.name) @staticmethod def to_icon(_: typ.Self) -> str: @@ -58,19 +86,214 @@ class SimSymbolNames(enum.StrEnum): @dataclasses.dataclass(kw_only=True, frozen=True) class SimSymbol: - name: SimSymbolNames = SimSymbolNames.LowerLambda + """A declarative representation of a symbolic variable. + + `sympy`'s symbols aren't quite flexible enough for our needs: The symbols that we're transporting often need exact domain information, an associated unit dimension, and a great deal of determinism in checks thereof. + + This dataclass is UI-friendly, as it only uses field type annotations/defaults supported by `bl_cache.BLProp`. + It's easy to persist, easy to transport, and has many helpful properties which greatly simplify working with symbols. + """ + + sim_node_name: SimSymbolName = SimSymbolName.LowerX mathtype: spux.MathType = spux.MathType.Real - ## TODO: - ## -> Physical Type: Track unit dimension information on the side. - ## -> Domain: Ability to constrain mathtype ex. (-pi,pi] - ## -> Shape: For using sp.MatrixSymbol w/predefined rows/cols. + physical_type: spux.PhysicalType = spux.PhysicalType.NonPhysical + + ## TODO: Shape/size support? Incl. MatrixSymbol. + + # Domain + interval_finite: tuple[float, float] = (0, 1) + interval_inf: tuple[bool, bool] = (True, True) + interval_closed: tuple[bool, bool] = (False, False) + + #################### + # - Properties + #################### + @property + def domain(self) -> sp.Interval | sp.Set: + """Return the domain of valid values for the symbol. + + For integer/rational/real symbols, the domain is an interval defined using the `interval_*` properties. + This interval **must** have the property`start <= stop`. + + Otherwise, the domain is the symbolic set corresponding to `self.mathtype`. + """ + if self.mathtype in [ + spux.MathType.Integer, + spux.MathType.Rational, + spux.MathType.Real, + ]: + return sp.Interval( + start=self.interval_finite[0] if not self.interval_inf[0] else -sp.oo, + end=self.interval_finite[1] if not self.interval_inf[1] else sp.oo, + left_open=( + True if self.interval_inf[0] else not self.interval_closed[0] + ), + right_open=( + True if self.interval_inf[1] else not self.interval_closed[1] + ), + ) + + return self.mathtype.symbolic_set + + #################### + # - Properties + #################### + @property + def sp_symbol(self) -> sp.Symbol: + """Return a symbolic variable corresponding to this `SimSymbol`. + + As much as possible, appropriate `assumptions` are set in the constructor of `sp.Symbol`, insofar as they can be determined. + + However, the assumptions system alone is rather limited, and implementations should therefore also strongly consider transporting `SimSymbols` directly, instead of `sp.Symbol`. + This allows making use of other properties like `self.domain`, when appropriate. + """ + # MathType Domain Constraint + ## -> We must feed the assumptions system. + mathtype_kwargs = {} + match self.mathtype: + case spux.MathType.Integer: + mathtype_kwargs |= {'integer': True} + case spux.MathType.Rational: + mathtype_kwargs |= {'rational': True} + case spux.MathType.Real: + mathtype_kwargs |= {'real': True} + case spux.MathType.Complex: + mathtype_kwargs |= {'complex': True} + + # Interval Constraints + if isinstance(self.domain, sp.Interval): + # Assumption: Non-Zero + if ( + ( + self.domain.left == 0 + and self.domain.left_open + or self.domain.right == 0 + and self.domain.right_open + ) + or self.domain.left > 0 + or self.domain.right < 0 + ): + mathtype_kwargs |= {'nonzero': True} + + # Assumption: Positive/Negative + if self.domain.left >= 0: + mathtype_kwargs |= {'positive': True} + elif self.domain.right <= 0: + mathtype_kwargs |= {'negative': True} + + # Construct the Symbol + return sp.Symbol(self.sim_node_name.name, **mathtype_kwargs) + + +#################### +# - Common Sim Symbols +#################### +class CommonSimSymbol(enum.StrEnum): + """A set of pre-defined symbols that might commonly be used in the context of physical simulation. + + Each entry maps directly to a particular `SimSymbol`. + + The enum is compatible with `BLField`, making it easy to declare a UI-driven dropdown of symbols that behave as expected. + + Attributes: + Wavelength: A symbol representing a real-valued wavelength. + Implicitly, this symbol often represents "vacuum wavelength" in particular. + Wavelength: A symbol representing a real-valued frequency. + Generally, this is the non-angular frequency. + """ + + X = enum.auto() + TimePS = enum.auto() + WavelengthNM = enum.auto() + FrequencyTHZ = 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. + """ + return CommonSimSymbol(v).sim_symbol_name.name + + @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 '' + + #################### + # - Properties + #################### + @property + def sim_symbol_name(self) -> str: + SSN = SimSymbolName + CSS = CommonSimSymbol + return { + CSS.X: SSN.LowerX, + CSS.TimePS: SSN.LowerT, + CSS.WavelengthNM: SSN.Wavelength, + CSS.FrequencyTHZ: SSN.Frequency, + }[self] @property - def sp_symbol(self): - mathtype_kwarg = {} - match self.mathtype: - case spux.MathType.Real: - mathtype_kwarg = {} + def sim_symbol(self) -> SimSymbol: + """Retrieve the `SimSymbol` associated with the `CommonSimSymbol`.""" + CSS = CommonSimSymbol + return { + CSS.X: SimSymbol( + sim_node_name=self.sim_symbol_name, + mathtype=spux.MathType.Real, + physical_type=spux.PhysicalType.NonPhysical, + ## TODO: Unit of Picosecond + interval_finite=(sys.float_info.min, sys.float_info.max), + interval_inf=(True, True), + interval_closed=(False, False), + ), + CSS.TimePS: SimSymbol( + sim_node_name=self.sim_symbol_name, + mathtype=spux.MathType.Real, + physical_type=spux.PhysicalType.Time, + ## TODO: Unit of Picosecond + interval_finite=(0, sys.float_info.max), + interval_inf=(False, True), + interval_closed=(True, False), + ), + CSS.WavelengthNM: SimSymbol( + sim_node_name=self.sim_symbol_name, + mathtype=spux.MathType.Real, + physical_type=spux.PhysicalType.Length, + ## TODO: Unit of Picosecond + interval_finite=(0, sys.float_info.max), + interval_inf=(False, True), + interval_closed=(False, False), + ), + CSS.FrequencyTHZ: SimSymbol( + sim_node_name=self.sim_symbol_name, + mathtype=spux.MathType.Real, + physical_type=spux.PhysicalType.Freq, + ## TODO: Unit of THz + interval_finite=(0, sys.float_info.max), + interval_inf=(False, True), + interval_closed=(False, False), + ), + }[self] - return sp.Symbol(self.name, **mathtype_kwarg) + +#################### +# - Selected Direct Access +#################### +x = CommonSimSymbol.X.sim_symbol +t_ps = CommonSimSymbol.TimePS.sim_symbol +wl_nm = CommonSimSymbol.WavelengthNM.sim_symbol +freq_thz = CommonSimSymbol.FrequencyTHZ.sim_symbol