feat: lazy pml BC and better FuncFlow __or__

main
Sofus Albert Høgsbro Rose 2024-05-27 19:46:11 +02:00
parent d0615d0372
commit 624914f8cb
Signed by: so-rose
GPG Key ID: AD901CB0F3701434
4 changed files with 265 additions and 40 deletions

View File

@ -254,6 +254,8 @@ class FuncFlow:
) )
supports_jax: bool = False supports_jax: bool = False
concatenated: bool = False
#################### ####################
# - Functions # - Functions
#################### ####################
@ -464,12 +466,19 @@ class FuncFlow:
Returns: Returns:
A lazy function that takes all arguments of both inputs, and returns a 2-tuple containing both output arguments. A lazy function that takes all arguments of both inputs, and returns a 2-tuple containing both output arguments.
""" """
return FuncFlow(
func=lambda *args, **kwargs: ( def self_func(args, kwargs):
self.func( ret = self.func(
*list(args[: len(self.func_args)]), *list(args[: len(self.func_args)]),
**{k: v for k, v in kwargs.items() if k in self.func_kwargs}, **{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(args, kwargs),
other.func( other.func(
*list(args[len(self.func_args) :]), *list(args[len(self.func_args) :]),
**{k: v for k, v in kwargs.items() if k in other.func_kwargs}, **{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_args=self.func_args + other.func_args,
func_kwargs=self.func_kwargs | other.func_kwargs, func_kwargs=self.func_kwargs | other.func_kwargs,
supports_jax=self.supports_jax and other.supports_jax, supports_jax=self.supports_jax and other.supports_jax,
concatenated=True,
) )

View File

@ -218,11 +218,11 @@ class AdiabAbsorbBoundCondNode(base.MaxwellSimNode):
if has_sig_order and has_sig_range: if has_sig_order and has_sig_range:
return (layers | sig_order | sig_range).compose_within( return (layers | sig_order | sig_range).compose_within(
enclosing_func=lambda els: td.Absorber( enclosing_func=lambda els: td.Absorber(
num_layers=els[0][0], num_layers=els[0],
parameters=td.AbsorberParams( parameters=td.AbsorberParams(
sigma_order=els[0][1], sigma_order=els[1],
sigma_min=els[1][0], sigma_min=els[2][0],
sigma_max=els[1][1], sigma_max=els[2][1],
), ),
), ),
supports_jax=False, supports_jax=False,

View File

@ -140,10 +140,12 @@ class PMLBoundCondNode(base.MaxwellSimNode):
col.label(text='2ε₀/Δt') col.label(text='2ε₀/Δt')
#################### ####################
# - Output # - FlowKind.Value
#################### ####################
@events.computes_output_socket( @events.computes_output_socket(
'BC', 'BC',
kind=ct.FlowKind.Value,
# Loaded
props={'active_socket_set'}, props={'active_socket_set'},
input_sockets={ input_sockets={
'Layers', 'Layers',
@ -162,40 +164,253 @@ class PMLBoundCondNode(base.MaxwellSimNode):
'α Order': True, 'α Order': True,
'α Range': 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. 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). - **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. - **Full**: Use the user-defined $\sigma$, $\kappa$, and $\alpha$ parameters, specifically polynomial order and sim-relative min/max conductivity values.
""" """
log.debug( output_params = output_sockets['BC']
'%s: Computing "%s" PML Boundary Condition (Input Sockets = %s)', layers = input_sockets['Layers']
self.sim_node_name,
props['active_socket_set'],
input_sockets,
)
# Simple PML has_layers = not ct.FlowSignal.check(layers)
if props['active_socket_set'] == 'Simple': has_output_params = not ct.FlowSignal.check(output_params)
return td.PML(num_layers=input_sockets['Layers'])
# Full PML 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( return td.PML(
num_layers=input_sockets['Layers'], num_layers=layers,
parameters=td.PMLParams( parameters=td.PMLParams(
sigma_order=input_sockets['σ Order'], sigma_order=sigma_order,
sigma_min=input_sockets['σ Range'][0], sigma_min=sigma_range[0],
sigma_max=input_sockets['σ Range'][1], sigma_max=sigma_range[1],
kappa_order=input_sockets['κ Order'], kappa_order=kappa_order,
kappa_min=input_sockets['κ Range'][0], kappa_min=kappa_range[0],
kappa_max=input_sockets['κ Range'][1], kappa_max=kappa_range[1],
alpha_order=input_sockets['α Order'], alpha_order=alpha_order,
alpha_min=input_sockets['α Range'][0], alpha_min=alpha_range[0],
alpha_max=input_sockets['α Range'][1], 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
#################### ####################
# - Blender Registration # - Blender Registration

View File

@ -161,20 +161,20 @@ class BoxStructureNode(base.MaxwellSimNode):
return (center | size | medium).compose_within( return (center | size | medium).compose_within(
enclosing_func=lambda els: tdadj.JaxStructure( enclosing_func=lambda els: tdadj.JaxStructure(
geometry=tdadj.JaxBox( geometry=tdadj.JaxBox(
center=tuple(els[0][0].flatten()), center=tuple(els[0].flatten()),
size=tuple(els[0][1].flatten()), size=tuple(els[1].flatten()),
), ),
medium=els[1], medium=els[2],
), ),
supports_jax=True, supports_jax=True,
) )
return (center | size | medium).compose_within( return (center | size | medium).compose_within(
enclosing_func=lambda els: td.Structure( enclosing_func=lambda els: td.Structure(
geometry=td.Box( geometry=td.Box(
center=tuple(els[0][0].flatten()), center=tuple(els[0].flatten()),
size=tuple(els[0][1].flatten()), size=tuple(els[1].flatten()),
), ),
medium=els[1], medium=els[2],
), ),
supports_jax=False, supports_jax=False,
) )