Merge pull request #27 from laurinbrunner/fix_example_5_and_6

Fix example 5 and 6
markusschmitt · Jun 8, 2022 · 6a26405 · 6a26405
2 parents 55dd431 + 267a647
commit 6a26405
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 6 deletions.
diff --git a/examples/ex5_dissipative_Lindblad.py b/examples/ex5_dissipative_Lindblad.py
@@ -54,10 +54,10 @@ def norm_fun(v, df=lambda x: x):
 prob_dist = jVMC.operator.povm.get_1_particle_distributions("y_up", Lindbladian.povm)
 prob_dist /= prob_dist[0]
 biases = jnp.log(prob_dist[1:])
-params = copy_dict(psi._param_unflatten_cpx(psi.get_parameters()))
+params = copy_dict(psi._param_unflatten(psi.get_parameters()))
 
-params["params"]["outputDense"]["bias"] = biases
-params["params"]["outputDense"]["kernel"] = 1e-15 * params["params"]["outputDense"]["kernel"]
+params["outputDense"]["bias"] = biases
+params["outputDense"]["kernel"] = 1e-15 * params["outputDense"]["kernel"]
 params = jnp.concatenate([p.ravel()
                           for p in jax.tree_util.tree_flatten(params)[0]])
 psi.set_parameters(params)

diff --git a/examples/ex6_dissipative_Lindblad_2D.py b/examples/ex6_dissipative_Lindblad_2D.py
@@ -58,10 +58,10 @@ def xy_to_id(x, y, L):
 prob_dist = jVMC.operator.povm.get_1_particle_distributions("z_up", Lindbladian.povm)
 prob_dist /= prob_dist[0]
 biases = jnp.log(prob_dist[1:])
-params = copy_dict(psi._param_unflatten_cpx(psi.get_parameters()))
+params = copy_dict(psi._param_unflatten(psi.get_parameters()))
 
-params["params"]["outputDense"]["bias"] = biases
-params["params"]["outputDense"]["kernel"] = 1e-15 * params["params"]["outputDense"]["kernel"]
+params["outputDense"]["bias"] = biases
+params["outputDense"]["kernel"] = 1e-15 * params["outputDense"]["kernel"]
 params = jnp.concatenate([p.ravel()
                           for p in jax.tree_util.tree_flatten(params)[0]])
 psi.set_parameters(params)