From 624914f8cb08dfae6d579e8194652689b69bc8a4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sofus=20Albert=20H=C3=B8gsbro=20Rose?= Date: Mon, 27 May 2024 19:46:11 +0200 Subject: [PATCH] feat: lazy pml BC and better FuncFlow __or__ --- .../contracts/flow_kinds/lazy_func.py | 18 +- .../bound_cond_nodes/absorbing_bound_cond.py | 8 +- .../bounds/bound_cond_nodes/pml_bound_cond.py | 267 ++++++++++++++++-- .../structures/primitives/box_structure.py | 12 +- 4 files changed, 265 insertions(+), 40 deletions(-) diff --git a/src/blender_maxwell/node_trees/maxwell_sim_nodes/contracts/flow_kinds/lazy_func.py b/src/blender_maxwell/node_trees/maxwell_sim_nodes/contracts/flow_kinds/lazy_func.py index f5c6163..ccf798f 100644 --- a/src/blender_maxwell/node_trees/maxwell_sim_nodes/contracts/flow_kinds/lazy_func.py +++ b/src/blender_maxwell/node_trees/maxwell_sim_nodes/contracts/flow_kinds/lazy_func.py @@ -254,6 +254,8 @@ class FuncFlow: ) supports_jax: bool = False + concatenated: bool = False + #################### # - Functions #################### @@ -464,12 +466,19 @@ class FuncFlow: Returns: A lazy function that takes all arguments of both inputs, and returns a 2-tuple containing both output arguments. """ + + def self_func(args, kwargs): + ret = self.func( + *list(args[: len(self.func_args)]), + **{k: v for k, v in kwargs.items() if k in self.func_kwargs}, + ) + if not self.concatenated: + return (ret,) + return ret + return FuncFlow( func=lambda *args, **kwargs: ( - self.func( - *list(args[: len(self.func_args)]), - **{k: v for k, v in kwargs.items() if k in self.func_kwargs}, - ), + *self_func(args, kwargs), other.func( *list(args[len(self.func_args) :]), **{k: v for k, v in kwargs.items() if k in other.func_kwargs}, @@ -478,4 +487,5 @@ class FuncFlow: func_args=self.func_args + other.func_args, func_kwargs=self.func_kwargs | other.func_kwargs, supports_jax=self.supports_jax and other.supports_jax, + concatenated=True, ) diff --git a/src/blender_maxwell/node_trees/maxwell_sim_nodes/nodes/bounds/bound_cond_nodes/absorbing_bound_cond.py b/src/blender_maxwell/node_trees/maxwell_sim_nodes/nodes/bounds/bound_cond_nodes/absorbing_bound_cond.py index 35b9a7a..2f8f276 100644 --- a/src/blender_maxwell/node_trees/maxwell_sim_nodes/nodes/bounds/bound_cond_nodes/absorbing_bound_cond.py +++ b/src/blender_maxwell/node_trees/maxwell_sim_nodes/nodes/bounds/bound_cond_nodes/absorbing_bound_cond.py @@ -218,11 +218,11 @@ class AdiabAbsorbBoundCondNode(base.MaxwellSimNode): if has_sig_order and has_sig_range: return (layers | sig_order | sig_range).compose_within( enclosing_func=lambda els: td.Absorber( - num_layers=els[0][0], + num_layers=els[0], parameters=td.AbsorberParams( - sigma_order=els[0][1], - sigma_min=els[1][0], - sigma_max=els[1][1], + sigma_order=els[1], + sigma_min=els[2][0], + sigma_max=els[2][1], ), ), supports_jax=False, diff --git a/src/blender_maxwell/node_trees/maxwell_sim_nodes/nodes/bounds/bound_cond_nodes/pml_bound_cond.py b/src/blender_maxwell/node_trees/maxwell_sim_nodes/nodes/bounds/bound_cond_nodes/pml_bound_cond.py index 9609a7e..91faec1 100644 --- a/src/blender_maxwell/node_trees/maxwell_sim_nodes/nodes/bounds/bound_cond_nodes/pml_bound_cond.py +++ b/src/blender_maxwell/node_trees/maxwell_sim_nodes/nodes/bounds/bound_cond_nodes/pml_bound_cond.py @@ -140,10 +140,12 @@ class PMLBoundCondNode(base.MaxwellSimNode): col.label(text='2ε₀/Δt') #################### - # - Output + # - FlowKind.Value #################### @events.computes_output_socket( 'BC', + kind=ct.FlowKind.Value, + # Loaded props={'active_socket_set'}, input_sockets={ 'Layers', @@ -162,39 +164,252 @@ class PMLBoundCondNode(base.MaxwellSimNode): 'α Order': True, 'α Range': True, }, + output_sockets={'BC'}, + output_socket_kinds={'BC': ct.FlowKind.Params}, ) - def compute_pml_boundary_cond(self, props, input_sockets) -> td.PML: + def compute_pml_value(self, props, input_sockets, output_sockets) -> td.PML: r"""Computes the PML boundary condition based on the active socket set. - **Simple**: Use `tidy3d`'s default parameters for defining the PML conductor (apart from number of layers). - **Full**: Use the user-defined $\sigma$, $\kappa$, and $\alpha$ parameters, specifically polynomial order and sim-relative min/max conductivity values. """ - log.debug( - '%s: Computing "%s" PML Boundary Condition (Input Sockets = %s)', - self.sim_node_name, - props['active_socket_set'], - input_sockets, - ) + output_params = output_sockets['BC'] + layers = input_sockets['Layers'] - # Simple PML - if props['active_socket_set'] == 'Simple': - return td.PML(num_layers=input_sockets['Layers']) + has_layers = not ct.FlowSignal.check(layers) + has_output_params = not ct.FlowSignal.check(output_params) - # Full PML - return td.PML( - num_layers=input_sockets['Layers'], - parameters=td.PMLParams( - sigma_order=input_sockets['σ Order'], - sigma_min=input_sockets['σ Range'][0], - sigma_max=input_sockets['σ Range'][1], - kappa_order=input_sockets['κ Order'], - kappa_min=input_sockets['κ Range'][0], - kappa_max=input_sockets['κ Range'][1], - alpha_order=input_sockets['α Order'], - alpha_min=input_sockets['α Range'][0], - alpha_max=input_sockets['α Range'][1], - ), - ) + if has_output_params and has_layers and not output_params.symbols: + active_socket_set = props['active_socket_set'] + match active_socket_set: + case 'Simple': + return td.PML(num_layers=layers) + + case 'Full': + sigma_order = input_sockets['σ Order'] + sigma_range = input_sockets['σ Range'] + kappa_order = input_sockets['κ Order'] + kappa_range = input_sockets['κ Range'] + alpha_order = input_sockets['α Order'] + alpha_range = input_sockets['α Range'] + + has_sigma_order = not ct.FlowSignal.check(sigma_order) + has_sigma_range = not ct.FlowSignal.check(sigma_range) + has_kappa_order = not ct.FlowSignal.check(kappa_order) + has_kappa_range = not ct.FlowSignal.check(kappa_range) + has_alpha_order = not ct.FlowSignal.check(alpha_order) + has_alpha_range = not ct.FlowSignal.check(alpha_range) + + if ( + has_sigma_order + and has_sigma_range + and has_kappa_order + and has_kappa_range + and has_alpha_order + and has_alpha_range + ): + return td.PML( + num_layers=layers, + parameters=td.PMLParams( + sigma_order=sigma_order, + sigma_min=sigma_range[0], + sigma_max=sigma_range[1], + kappa_order=kappa_order, + kappa_min=kappa_range[0], + kappa_max=kappa_range[1], + alpha_order=alpha_order, + alpha_min=alpha_range[0], + alpha_max=alpha_range[1], + ), + ) + + return ct.FlowSignal.FlowPending + + #################### + # - FlowKind.Func + #################### + @events.computes_output_socket( + 'BC', + kind=ct.FlowKind.Func, + # Loaded + props={'active_socket_set'}, + input_sockets={ + 'Layers', + 'σ Order', + 'σ Range', + 'κ Order', + 'κ Range', + 'α Order', + 'α Range', + }, + input_socket_kinds={ + 'Layers': ct.FlowKind.Func, + 'σ Order': ct.FlowKind.Func, + 'σ Range': ct.FlowKind.Func, + 'κ Order': ct.FlowKind.Func, + 'κ Range': ct.FlowKind.Func, + 'α Order': ct.FlowKind.Func, + 'α Range': ct.FlowKind.Func, + }, + input_sockets_optional={ + 'σ Order': True, + 'σ Range': True, + 'κ Order': True, + 'κ Range': True, + 'α Order': True, + 'α Range': True, + }, + output_sockets={'BC'}, + output_socket_kinds={'BC': ct.FlowKind.Params}, + ) + def compute_pml_func(self, props, input_sockets, output_sockets) -> td.PML: + output_params = output_sockets['BC'] + layers = input_sockets['Layers'] + + has_output_params = not ct.FlowSignal.check(output_params) + has_layers = not ct.FlowSignal.check(layers) + + if has_output_params and has_layers: + active_socket_set = props['active_socket_set'] + match active_socket_set: + case 'Simple': + return layers.compose_within( + enclosing_func=lambda layers: td.PML(num_layers=layers), + supports_jax=False, + ) + + case 'Full': + sigma_order = input_sockets['σ Order'] + sigma_range = input_sockets['σ Range'] + kappa_order = input_sockets['κ Order'] + kappa_range = input_sockets['κ Range'] + alpha_order = input_sockets['α Order'] + alpha_range = input_sockets['α Range'] + + has_sigma_order = not ct.FlowSignal.check(sigma_order) + has_sigma_range = not ct.FlowSignal.check(sigma_range) + has_kappa_order = not ct.FlowSignal.check(kappa_order) + has_kappa_range = not ct.FlowSignal.check(kappa_range) + has_alpha_order = not ct.FlowSignal.check(alpha_order) + has_alpha_range = not ct.FlowSignal.check(alpha_range) + + if ( + has_sigma_order + and has_sigma_range + and has_kappa_order + and has_kappa_range + and has_alpha_order + and has_alpha_range + ): + return ( + sigma_order + | sigma_range + | kappa_order + | kappa_range + | alpha_order + | alpha_range + ).compose_within( + enclosing_func=lambda els: td.PML( + num_layers=layers, + parameters=td.PMLParams( + sigma_order=els[0], + sigma_min=els[1][0], + sigma_max=els[1][1], + kappa_order=els[2], + kappa_min=els[3][0], + kappa_max=els[3][1], + alpha_order=els[4][1], + alpha_min=els[5][0], + alpha_max=els[5][1], + ), + ) + ) + + return ct.FlowSignal.FlowPending + + #################### + # - FlowKind.Params + #################### + @events.computes_output_socket( + 'BC', + kind=ct.FlowKind.Params, + # Loaded + props={'active_socket_set'}, + input_sockets={ + 'Layers', + 'σ Order', + 'σ Range', + 'κ Order', + 'κ Range', + 'α Order', + 'α Range', + }, + input_socket_kinds={ + 'Layers': ct.FlowKind.Params, + 'σ Order': ct.FlowKind.Params, + 'σ Range': ct.FlowKind.Params, + 'κ Order': ct.FlowKind.Params, + 'κ Range': ct.FlowKind.Params, + 'α Order': ct.FlowKind.Params, + 'α Range': ct.FlowKind.Params, + }, + input_sockets_optional={ + 'σ Order': True, + 'σ Range': True, + 'κ Order': True, + 'κ Range': True, + 'α Order': True, + 'α Range': True, + }, + ) + def compute_pml_params(self, props, input_sockets) -> td.PML: + r"""Computes the PML boundary condition based on the active socket set. + + - **Simple**: Use `tidy3d`'s default parameters for defining the PML conductor (apart from number of layers). + - **Full**: Use the user-defined $\sigma$, $\kappa$, and $\alpha$ parameters, specifically polynomial order and sim-relative min/max conductivity values. + """ + layers = input_sockets['Layers'] + has_layers = not ct.FlowSignal.check(layers) + + if has_layers: + active_socket_set = props['active_socket_set'] + match active_socket_set: + case 'Simple': + return layers + + case 'Full': + sigma_order = input_sockets['σ Order'] + sigma_range = input_sockets['σ Range'] + kappa_order = input_sockets['σ Order'] + kappa_range = input_sockets['σ Range'] + alpha_order = input_sockets['σ Order'] + alpha_range = input_sockets['σ Range'] + + has_sigma_order = not ct.FlowSignal.check(sigma_order) + has_sigma_range = not ct.FlowSignal.check(sigma_range) + has_kappa_order = not ct.FlowSignal.check(kappa_order) + has_kappa_range = not ct.FlowSignal.check(kappa_range) + has_alpha_order = not ct.FlowSignal.check(alpha_order) + has_alpha_range = not ct.FlowSignal.check(alpha_range) + + if ( + has_sigma_order + and has_sigma_range + and has_kappa_order + and has_kappa_range + and has_alpha_order + and has_alpha_range + ): + return ( + sigma_order + | sigma_range + | kappa_order + | kappa_range + | alpha_order + | alpha_range + ) + + return ct.FlowSignal.FlowPending #################### diff --git a/src/blender_maxwell/node_trees/maxwell_sim_nodes/nodes/structures/primitives/box_structure.py b/src/blender_maxwell/node_trees/maxwell_sim_nodes/nodes/structures/primitives/box_structure.py index b2facb0..6dbf856 100644 --- a/src/blender_maxwell/node_trees/maxwell_sim_nodes/nodes/structures/primitives/box_structure.py +++ b/src/blender_maxwell/node_trees/maxwell_sim_nodes/nodes/structures/primitives/box_structure.py @@ -161,20 +161,20 @@ class BoxStructureNode(base.MaxwellSimNode): return (center | size | medium).compose_within( enclosing_func=lambda els: tdadj.JaxStructure( geometry=tdadj.JaxBox( - center=tuple(els[0][0].flatten()), - size=tuple(els[0][1].flatten()), + center=tuple(els[0].flatten()), + size=tuple(els[1].flatten()), ), - medium=els[1], + medium=els[2], ), supports_jax=True, ) return (center | size | medium).compose_within( enclosing_func=lambda els: td.Structure( geometry=td.Box( - center=tuple(els[0][0].flatten()), - size=tuple(els[0][1].flatten()), + center=tuple(els[0].flatten()), + size=tuple(els[1].flatten()), ), - medium=els[1], + medium=els[2], ), supports_jax=False, )