394 lines
10 KiB
Python
394 lines
10 KiB
Python
|
import json
|
||
|
|
import tempfile
|
||
|
|
import functools
|
||
|
|
import typing as typ
|
||
|
|
import json
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
import bpy
|
||
|
|
import sympy as sp
|
||
|
|
import pydantic as pyd
|
||
|
|
import tidy3d as td
|
||
|
|
import tidy3d.web as _td_web
|
||
|
|
|
||
|
|
from ......utils.auth_td_web import g_td_web, is_td_web_authed
|
||
|
|
from .... import contracts as ct
|
||
|
|
from .... import sockets
|
||
|
|
from ... import base
|
||
|
|
|
||
|
|
####################
|
||
|
|
# - Task Getters
|
||
|
|
####################
|
||
|
|
## TODO: We should probably refactor this setup.
|
||
|
|
@functools.cache
|
||
|
|
def estimated_task_cost(task_id: str):
|
||
|
|
return _td_web.api.webapi.estimate_cost(task_id)
|
||
|
|
|
||
|
|
@functools.cache
|
||
|
|
def billed_task_cost(task_id: str):
|
||
|
|
return _td_web.api.webapi.real_cost(task_id)
|
||
|
|
|
||
|
|
@functools.cache
|
||
|
|
def task_status(task_id: str):
|
||
|
|
task = _td_web.api.webapi.get_info(task_id)
|
||
|
|
return task.status
|
||
|
|
|
||
|
|
####################
|
||
|
|
# - Progress Timer
|
||
|
|
####################
|
||
|
|
## TODO: We should probably refactor this too.
|
||
|
|
class Tidy3DTaskStatusModalOperator(bpy.types.Operator):
|
||
|
|
bl_idname = "blender_maxwell.tidy_3d_task_status_modal_operator"
|
||
|
|
bl_label = "Tidy3D Task Status Modal Operator"
|
||
|
|
|
||
|
|
_timer = None
|
||
|
|
_task_id = None
|
||
|
|
_node = None
|
||
|
|
_status = None
|
||
|
|
_reported_done = False
|
||
|
|
|
||
|
|
def modal(self, context, event):
|
||
|
|
# Retrieve New Status
|
||
|
|
task_status.cache_clear()
|
||
|
|
new_status = task_status(self._task_id)
|
||
|
|
if new_status != self._status:
|
||
|
|
task_status.cache_clear()
|
||
|
|
self._status = new_status
|
||
|
|
|
||
|
|
# Check Done Status
|
||
|
|
if self._status in {"success", "error"}:
|
||
|
|
# Report Done
|
||
|
|
if not self._reported_done:
|
||
|
|
self._node.trigger_action("value_changed")
|
||
|
|
self._reported_done = True
|
||
|
|
|
||
|
|
# Finish when Billing is Known
|
||
|
|
if not billed_task_cost(self._task_id):
|
||
|
|
billed_task_cost.cache_clear()
|
||
|
|
else:
|
||
|
|
return {'FINISHED'}
|
||
|
|
|
||
|
|
return {'PASS_THROUGH'}
|
||
|
|
|
||
|
|
def execute(self, context):
|
||
|
|
node = context.node
|
||
|
|
wm = context.window_manager
|
||
|
|
|
||
|
|
self._timer = wm.event_timer_add(0.25, window=context.window)
|
||
|
|
self._task_id = node.uploaded_task_id
|
||
|
|
self._node = node
|
||
|
|
self._status = task_status(self._task_id)
|
||
|
|
|
||
|
|
wm.modal_handler_add(self)
|
||
|
|
return {'RUNNING_MODAL'}
|
||
|
|
|
||
|
|
####################
|
||
|
|
# - Web Uploader / Loader / Runner / Releaser
|
||
|
|
####################
|
||
|
|
## TODO: We should probably refactor this too.
|
||
|
|
class Tidy3DWebUploadOperator(bpy.types.Operator):
|
||
|
|
bl_idname = "blender_maxwell.tidy_3d_web_upload_operator"
|
||
|
|
bl_label = "Tidy3D Web Upload Operator"
|
||
|
|
bl_description = "Upload the attached (locked) simulation, such that it is ready to run on the Tidy3D cloud"
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def poll(cls, context):
|
||
|
|
space = context.space_data
|
||
|
|
return (
|
||
|
|
space.type == 'NODE_EDITOR'
|
||
|
|
and space.node_tree is not None
|
||
|
|
and space.node_tree.bl_idname == "MaxwellSimTreeType"
|
||
|
|
and is_td_web_authed()
|
||
|
|
and hasattr(context, "node")
|
||
|
|
and context.node.lock_tree
|
||
|
|
)
|
||
|
|
|
||
|
|
def execute(self, context):
|
||
|
|
node = context.node
|
||
|
|
node.web_upload()
|
||
|
|
return {'FINISHED'}
|
||
|
|
|
||
|
|
class Tidy3DLoadUploadedOperator(bpy.types.Operator):
|
||
|
|
bl_idname = "blender_maxwell.tidy_3d_load_uploaded_operator"
|
||
|
|
bl_label = "Tidy3D Load Uploaded Operator"
|
||
|
|
bl_description = "Load an already-uploaded simulation, as selected in the dropdown of the 'Cloud Task' socket"
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def poll(cls, context):
|
||
|
|
space = context.space_data
|
||
|
|
return (
|
||
|
|
space.type == 'NODE_EDITOR'
|
||
|
|
and space.node_tree is not None
|
||
|
|
and space.node_tree.bl_idname == "MaxwellSimTreeType"
|
||
|
|
and is_td_web_authed()
|
||
|
|
and hasattr(context, "node")
|
||
|
|
and context.node.lock_tree
|
||
|
|
)
|
||
|
|
|
||
|
|
def execute(self, context):
|
||
|
|
node = context.node
|
||
|
|
node.load_uploaded_task()
|
||
|
|
|
||
|
|
# Load Simulation to Compare
|
||
|
|
## Load Local Sim
|
||
|
|
local_sim = node._compute_input("FDTD Sim")
|
||
|
|
|
||
|
|
## Load Cloud Sim
|
||
|
|
task_id = node.compute_output("Cloud Task")
|
||
|
|
with tempfile.NamedTemporaryFile(delete=False) as f:
|
||
|
|
_path_tmp = Path(f.name)
|
||
|
|
_path_tmp.rename(f.name + ".json")
|
||
|
|
path_tmp = Path(f.name + ".json")
|
||
|
|
cloud_sim = _td_web.api.webapi.load_simulation(task_id, path=str(path_tmp))
|
||
|
|
|
||
|
|
Path(path_tmp).unlink()
|
||
|
|
|
||
|
|
## Compare
|
||
|
|
if local_sim != cloud_sim:
|
||
|
|
node.release_uploaded_task()
|
||
|
|
msg = "Loaded simulation doesn't match input simulation"
|
||
|
|
raise ValueError(msg)
|
||
|
|
|
||
|
|
return {'FINISHED'}
|
||
|
|
|
||
|
|
class RunUploadedTidy3DSim(bpy.types.Operator):
|
||
|
|
bl_idname = "blender_maxwell.run_uploaded_tidy_3d_sim"
|
||
|
|
bl_label = "Run Uploaded Tidy3D Sim"
|
||
|
|
bl_description = "Run the currently uploaded (and loaded) simulation"
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def poll(cls, context):
|
||
|
|
space = context.space_data
|
||
|
|
return (
|
||
|
|
space.type == 'NODE_EDITOR'
|
||
|
|
and space.node_tree is not None
|
||
|
|
and space.node_tree.bl_idname == "MaxwellSimTreeType"
|
||
|
|
and is_td_web_authed()
|
||
|
|
and hasattr(context, "node")
|
||
|
|
and context.node.lock_tree
|
||
|
|
and context.node.uploaded_task_id
|
||
|
|
and task_status(context.node.uploaded_task_id) == "draft"
|
||
|
|
)
|
||
|
|
|
||
|
|
def execute(self, context):
|
||
|
|
node = context.node
|
||
|
|
node.run_uploaded_task()
|
||
|
|
bpy.ops.blender_maxwell.tidy_3d_task_status_modal_operator()
|
||
|
|
return {'FINISHED'}
|
||
|
|
|
||
|
|
class ReleaseTidy3DExportOperator(bpy.types.Operator):
|
||
|
|
bl_idname = "blender_maxwell.release_tidy_3d_export_operator"
|
||
|
|
bl_label = "Release Tidy3D Export Operator"
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def poll(cls, context):
|
||
|
|
space = context.space_data
|
||
|
|
return (
|
||
|
|
space.type == 'NODE_EDITOR'
|
||
|
|
and space.node_tree is not None
|
||
|
|
and space.node_tree.bl_idname == "MaxwellSimTreeType"
|
||
|
|
and is_td_web_authed()
|
||
|
|
and hasattr(context, "node")
|
||
|
|
and context.node.lock_tree
|
||
|
|
and context.node.uploaded_task_id
|
||
|
|
)
|
||
|
|
|
||
|
|
def execute(self, context):
|
||
|
|
node = context.node
|
||
|
|
node.release_uploaded_task()
|
||
|
|
return {'FINISHED'}
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
####################
|
||
|
|
# - Web Exporter Node
|
||
|
|
####################
|
||
|
|
class Tidy3DWebExporterNode(base.MaxwellSimNode):
|
||
|
|
node_type = ct.NodeType.Tidy3DWebExporter
|
||
|
|
bl_label = "Tidy3DWebExporter"
|
||
|
|
|
||
|
|
input_sockets = {
|
||
|
|
"FDTD Sim": sockets.MaxwellFDTDSimSocketDef(),
|
||
|
|
"Cloud Task": sockets.Tidy3DCloudTaskSocketDef(
|
||
|
|
task_exists=False,
|
||
|
|
),
|
||
|
|
}
|
||
|
|
output_sockets = {
|
||
|
|
"Cloud Task": sockets.Tidy3DCloudTaskSocketDef(
|
||
|
|
task_exists=True,
|
||
|
|
),
|
||
|
|
}
|
||
|
|
|
||
|
|
lock_tree: bpy.props.BoolProperty(
|
||
|
|
name="Whether to lock the attached tree",
|
||
|
|
description="Whether or not to lock the attached tree",
|
||
|
|
default=False,
|
||
|
|
update=(lambda self, context: self.sync_lock_tree(context)),
|
||
|
|
)
|
||
|
|
uploaded_task_id: bpy.props.StringProperty(
|
||
|
|
name="Uploaded Task ID",
|
||
|
|
description="The uploaded task ID",
|
||
|
|
default="",
|
||
|
|
)
|
||
|
|
|
||
|
|
####################
|
||
|
|
# - Sync Methods
|
||
|
|
####################
|
||
|
|
def sync_lock_tree(self, context):
|
||
|
|
node_tree = self.id_data
|
||
|
|
|
||
|
|
if self.lock_tree:
|
||
|
|
self.trigger_action("enable_lock")
|
||
|
|
self.locked = False
|
||
|
|
for bl_socket in self.inputs:
|
||
|
|
if bl_socket.name == "FDTD Sim": continue
|
||
|
|
bl_socket.locked = False
|
||
|
|
|
||
|
|
else:
|
||
|
|
self.trigger_action("disable_lock")
|
||
|
|
|
||
|
|
####################
|
||
|
|
# - Output Socket Callbacks
|
||
|
|
####################
|
||
|
|
def web_upload(self):
|
||
|
|
if not (sim := self._compute_input("FDTD Sim")):
|
||
|
|
raise ValueError("Must attach simulation")
|
||
|
|
|
||
|
|
if not (new_task_dict := self._compute_input("Cloud Task")):
|
||
|
|
raise ValueError("No valid cloud task defined")
|
||
|
|
|
||
|
|
td_web = g_td_web(None) ## Presume already auth'ed
|
||
|
|
|
||
|
|
self.uploaded_task_id = td_web.api.webapi.upload(
|
||
|
|
sim,
|
||
|
|
**new_task_dict,
|
||
|
|
verbose=True,
|
||
|
|
)
|
||
|
|
|
||
|
|
self.inputs["Cloud Task"].sync_task_loaded(self.uploaded_task_id)
|
||
|
|
|
||
|
|
def load_uploaded_task(self):
|
||
|
|
self.inputs["Cloud Task"].sync_task_loaded(None)
|
||
|
|
self.uploaded_task_id = self._compute_input("Cloud Task")
|
||
|
|
|
||
|
|
self.trigger_action("value_changed")
|
||
|
|
|
||
|
|
def run_uploaded_task(self):
|
||
|
|
td_web = g_td_web(None) ## Presume already auth'ed
|
||
|
|
td_web.api.webapi.start(self.uploaded_task_id)
|
||
|
|
|
||
|
|
self.trigger_action("value_changed")
|
||
|
|
|
||
|
|
def release_uploaded_task(self):
|
||
|
|
self.uploaded_task_id = ""
|
||
|
|
self.inputs["Cloud Task"].sync_task_released(specify_new_task=True)
|
||
|
|
|
||
|
|
self.trigger_action("value_changed")
|
||
|
|
|
||
|
|
####################
|
||
|
|
# - UI
|
||
|
|
####################
|
||
|
|
def draw_operators(self, context, layout):
|
||
|
|
is_authed = is_td_web_authed()
|
||
|
|
has_uploaded_task_id = bool(self.uploaded_task_id)
|
||
|
|
|
||
|
|
# Row: Run Simulation
|
||
|
|
row = layout.row(align=True)
|
||
|
|
if has_uploaded_task_id: row.enabled = False
|
||
|
|
row.operator(
|
||
|
|
Tidy3DWebUploadOperator.bl_idname,
|
||
|
|
text="Upload Sim",
|
||
|
|
)
|
||
|
|
tree_lock_icon = "LOCKED" if self.lock_tree else "UNLOCKED"
|
||
|
|
row.prop(self, "lock_tree", toggle=True, icon=tree_lock_icon, text="")
|
||
|
|
|
||
|
|
# Row: Run Simulation
|
||
|
|
row = layout.row(align=True)
|
||
|
|
if is_authed and has_uploaded_task_id:
|
||
|
|
run_sim_text = f"Run Sim (~{estimated_task_cost(self.uploaded_task_id):.3f} credits)"
|
||
|
|
else:
|
||
|
|
run_sim_text = f"Run Sim"
|
||
|
|
|
||
|
|
row.operator(
|
||
|
|
RunUploadedTidy3DSim.bl_idname,
|
||
|
|
text=run_sim_text,
|
||
|
|
)
|
||
|
|
if has_uploaded_task_id:
|
||
|
|
tree_lock_icon = "LOOP_BACK"
|
||
|
|
row.operator(
|
||
|
|
ReleaseTidy3DExportOperator.bl_idname,
|
||
|
|
icon="LOOP_BACK",
|
||
|
|
text="",
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
row.operator(
|
||
|
|
Tidy3DLoadUploadedOperator.bl_idname,
|
||
|
|
icon="TRIA_UP_BAR",
|
||
|
|
text="",
|
||
|
|
)
|
||
|
|
|
||
|
|
# Row: Simulation Progress
|
||
|
|
if is_authed and has_uploaded_task_id:
|
||
|
|
progress = {
|
||
|
|
"draft": (0.0, "Waiting to Run..."),
|
||
|
|
"initialized": (0.0, "Initializing..."),
|
||
|
|
"queued": (0.0, "Queued..."),
|
||
|
|
"preprocessing": (0.05, "Pre-processing..."),
|
||
|
|
"running": (0.2, "Running..."),
|
||
|
|
"postprocessing": (0.85, "Post-processing..."),
|
||
|
|
"success": (1.0, f"Success (={billed_task_cost(self.uploaded_task_id)} credits)"),
|
||
|
|
"error": (1.0, f"Error (={billed_task_cost(self.uploaded_task_id)} credits)"),
|
||
|
|
}[task_status(self.uploaded_task_id)]
|
||
|
|
|
||
|
|
layout.separator()
|
||
|
|
row = layout.row(align=True)
|
||
|
|
row.progress(
|
||
|
|
factor=progress[0],
|
||
|
|
type="BAR",
|
||
|
|
text=progress[1],
|
||
|
|
)
|
||
|
|
|
||
|
|
####################
|
||
|
|
# - Output Methods
|
||
|
|
####################
|
||
|
|
@base.computes_output_socket(
|
||
|
|
"Cloud Task",
|
||
|
|
input_sockets={"Cloud Task"},
|
||
|
|
)
|
||
|
|
def compute_cloud_task(self, input_sockets: dict) -> str | None:
|
||
|
|
if self.uploaded_task_id: return self.uploaded_task_id
|
||
|
|
return None
|
||
|
|
|
||
|
|
####################
|
||
|
|
# - Update
|
||
|
|
####################
|
||
|
|
@base.on_value_changed(socket_name="FDTD Sim")
|
||
|
|
def on_value_changed__fdtd_sim(self):
|
||
|
|
estimated_task_cost.cache_clear()
|
||
|
|
task_status.cache_clear()
|
||
|
|
billed_task_cost.cache_clear()
|
||
|
|
|
||
|
|
@base.on_value_changed(socket_name="Cloud Task")
|
||
|
|
def on_value_changed__cloud_task(self):
|
||
|
|
estimated_task_cost.cache_clear()
|
||
|
|
task_status.cache_clear()
|
||
|
|
billed_task_cost.cache_clear()
|
||
|
|
|
||
|
|
|
||
|
|
####################
|
||
|
|
# - Blender Registration
|
||
|
|
####################
|
||
|
|
BL_REGISTER = [
|
||
|
|
Tidy3DWebUploadOperator,
|
||
|
|
Tidy3DTaskStatusModalOperator,
|
||
|
|
RunUploadedTidy3DSim,
|
||
|
|
Tidy3DLoadUploadedOperator,
|
||
|
|
ReleaseTidy3DExportOperator,
|
||
|
|
Tidy3DWebExporterNode,
|
||
|
|
]
|
||
|
|
BL_NODES = {
|
||
|
|
ct.NodeType.Tidy3DWebExporter: (
|
||
|
|
ct.NodeCategory.MAXWELLSIM_OUTPUTS_EXPORTERS
|
||
|
|
)
|
||
|
|
}
|