I use the “resume” in the "opt_disc" stage, it will raise KeyError: 'discrete_scaling', how should I solve this problem?

[xuanlunhuang]

I edited the "wlens" for the broadband target. I used the google colab to run the code but it was interrupted during the "opt_disc" transformation. So I used "resume" to continue the optimization but the following error occurred:

[2021-11-30 03:54:12,347][INFO][solver][run_plan] Setting up workspace.
[2021-11-30 03:54:12,646][INFO][solver][restore_workspace] Restoring from checkpoint 20211125_C1550_40nm_60etch/cont_to_disc.chkpt.pkl
Traceback (most recent call last):
  File "grating.py", line 655, in <module>
    resume_opt(args.save_folder)
  File "grating.py", line 585, in resume_opt
    problem_graph.run_plan(plan, ".", save_folder=save_folder, resume=True)
  File "/content/drive/MyDrive/spins-b/examples/invdes/grating_coupler/spins/invdes/problem_graph/solver.py", line 55, in run_plan
    console_logger)
  File "/content/drive/MyDrive/spins-b/examples/invdes/grating_coupler/spins/invdes/problem_graph/solver.py", line 152, in restore_workspace
    work.get_object(param).set_parameter_value(param_value)
  File "/content/drive/MyDrive/spins-b/examples/invdes/grating_coupler/spins/invdes/problem_graph/workspace.py", line 172, in get_object
    node = self._nodes[name_or_node]
KeyError: 'discrete_scaling'

Then, I checked the "cont_to_disc.chkpt.pkl" file using pickle and found that there exist a "discrete_scaling" key in the "parameters" dicts:

{'time': '2021-11-27 09:38:36.049518', 'parameters': {'param_obj_final_val': 1.0, 'param_discrete_penalty_val': 1.0, 'discrete_scaling': 5.0}

but in the "restore_workspace" function (File "/content/drive/MyDrive/spins-b/examples/invdes/grating_coupler/spins/invdes/problem_graph/solver.py", line 152, in restore_workspace), the nodes which are to add do not contain "discrete_scaling":

        # Iterate through all the previous transformations, restoring any
        # parametrizations and parameters along the way. It is not strictly
        # necessary to restore parametrizations/parameters along the way,
        # but it was done out of implementation convenience.
        for transform in plan.transformations[:transform_index]:
            # Restore any parametrizations.
            work.get_object(transform.parametrization).deserialize(
                chkpt_data["parametrizations"][transform.parametrization.name])

            # Add any parameter descriptions. Actual values restored below.
            if transform.parameter_list:
                for set_param in transform.parameter_list:
                    work._add_node(set_param.parameter)

This is confirmed by print(work._nodes) after the above work._add_node(set_param.parameter), because in the transform.parameter_list I only find 'param_obj_final_val' and 'param_discrete_penalty_val' but no 'discrete_scaling' (in "create_transformations" fuction of grating.py):

        trans_list.append(
            optplan.Transformation(
                name="opt_cont_disc",
                parameter_list=[
                    optplan.SetParam(
                        parameter=obj_val_param,
                        function=obj,
                        parametrization=cont_param),
                    optplan.SetParam(
                        parameter=discrete_penalty_val,
                        function=optplan.DiscretePenalty(),
                        parametrization=cont_param)
                ],

This might be the reason why this KeyError occurs.

Then I checked how "chkpt.pkl" file is saved. I went to work.logger.write_checkpoint(transformation_param.name) in the "run_plan" function in solver.py ( File "/content/drive/MyDrive/spins-b/examples/invdes/grating_coupler/spins/invdes/problem_graph/solver.py", line 63) and run the following code, which is copied from "write_checkpoint", by creating a new work object.

# Get workspace parameters.
parameter_data = {}
parameter_list = work.get_objects_by_type(optplan.Parameter)
for param_name, param_obj in parameter_list.items():
    parameter_data[param_name] = param_obj.value

print(parameter_list)
print(parameter_data)

And for the printed parameter_data, 'discrete_scaling' does not exist, which means the 'discrete_scaling' won't be saved in "chkpt.pkl" file

{'param_obj_final_val': <spins.invdes.problem.objective.Parameter object at 0x000001613C52EA90>, 'param_discrete_penalty_val': <spins.invdes.problem.objective.Parameter object at 0x000001613C52E438>}
{'param_obj_final_val': 1.0, 'param_discrete_penalty_val': 1.0}

BUT why there exists a 'discrete_scaling' in "chkpt.pkl" file, making the KeyError occur.