Merge pull request #76 from markusschmitt/operator_interface

Operator interface
markusschmitt · Sep 13, 2024 · 129ea53 · 129ea53
2 parents 1a4c10e + 6c35f19
commit 129ea53
Show file tree

Hide file tree

Showing 17 changed files with 521 additions and 98 deletions.
diff --git a/examples/ex0_ground_state_search.ipynb b/examples/ex0_ground_state_search.ipynb
@@ -14,8 +14,7 @@
     "sys.path.append(sys.path[0] + \"/..\")\n",
     "\n",
     "import jax\n",
-    "from jax.config import config\n",
-    "config.update(\"jax_enable_x64\", True)\n",
+    "jax.config.update(\"jax_enable_x64\", True)\n",
     "\n",
     "import jax.random as random\n",
     "import jax.numpy as jnp\n",

diff --git a/examples/ex0_ground_state_search.py b/examples/ex0_ground_state_search.py
@@ -2,8 +2,7 @@
 # coding: utf-8
 
 import jax
-from jax.config import config
-config.update("jax_enable_x64", True)
+jax.config.update("jax_enable_x64", True)
 
 import jax.random as random
 import jax.numpy as jnp

diff --git a/examples/ex2_unitary_time_evolution.py b/examples/ex2_unitary_time_evolution.py
@@ -1,12 +1,7 @@
 import os
 
 import jax
-from jax.config import config
-config.update("jax_enable_x64", True)
-
-import jax.random as random
-import flax
-import jax.numpy as jnp
+jax.config.update("jax_enable_x64", True)
 
 import numpy as np
 

diff --git a/examples/ex3_custom_net.py b/examples/ex3_custom_net.py
@@ -1,7 +1,6 @@
 import jax
 import flax
 import jVMC
-import numpy as np
 
 # This class defines the network structure of a complex RBM
 

diff --git a/examples/ex4_benchmarking.py b/examples/ex4_benchmarking.py
@@ -1,8 +1,7 @@
 import os
 
 import jax
-from jax.config import config
-config.update("jax_enable_x64", True)
+jax.config.update("jax_enable_x64", True)
 
 import jax.random as random
 import jax.numpy as jnp

diff --git a/examples/ex5_dissipative_Lindblad.py b/examples/ex5_dissipative_Lindblad.py
@@ -1,13 +1,9 @@
 import matplotlib.pyplot as plt
-import numpy as np
-import flax
-import jax.random as random
-from jax.config import config
 import jax.numpy as jnp
 import jax
 import jVMC
 from functools import partial
-config.update("jax_enable_x64", True)
+jax.config.update("jax_enable_x64", True)
 
 
 def copy_dict(a):

diff --git a/examples/ex6_dissipative_Lindblad_2D.py b/examples/ex6_dissipative_Lindblad_2D.py
@@ -1,12 +1,8 @@
 import matplotlib.pyplot as plt
-import numpy as np
-import flax
-import jax.random as random
-from jax.config import config
 import jax.numpy as jnp
 import jax
 import jVMC
-config.update("jax_enable_x64", True)
+jax.config.update("jax_enable_x64", True)
 from functools import partial
 
 

diff --git a/examples/ex7_fermions.ipynb b/examples/ex7_fermions.ipynb
@@ -27,8 +27,7 @@
     "\n",
     "# jax\n",
     "import jax\n",
-    "from jax.config import config\n",
-    "config.update(\"jax_enable_x64\", True)\n",
+    "jax.config.update(\"jax_enable_x64\", True)\n",
     "import jax.numpy as jnp\n",
     "import flax.linen as nn\n",
     "\n",
@@ -138,11 +137,11 @@
     "    if i == flavourL//flavour-1:\n",
     "        continue\n",
     "    # up chain hopping\n",
-    "    hamiltonian.add(op.scal_opstr( t, ( annihilation(site1UP + i) ,creation(site1UP + i + 1)  ) ) )\n",
-    "    hamiltonian.add(op.scal_opstr( t, ( annihilation(site1UP + i + 1) ,creation(site1UP + i)  ) ) )\n",
+    "    hamiltonian.add(op.scal_opstr( t, ( creation(site1UP + i + 1), annihilation(site1UP + i) ,  ) ) )\n",
+    "    hamiltonian.add(op.scal_opstr( t, ( creation(site1UP + i), annihilation(site1UP + i + 1) ,  ) ) )\n",
     "    # down chain hopping\n",
-    "    hamiltonian.add(op.scal_opstr( t, ( annihilation(site1DO - i) ,creation(site1DO - i - 1)  ) ) )\n",
-    "    hamiltonian.add(op.scal_opstr( t, ( annihilation(site1DO - i - 1) ,creation(site1DO - i)  ) ) )"
+    "    hamiltonian.add(op.scal_opstr( t, ( creation(site1DO - i - 1), annihilation(site1DO - i) ,  ) ) )\n",
+    "    hamiltonian.add(op.scal_opstr( t, ( creation(site1DO - i), annihilation(site1DO - i - 1) ,  ) ) )"
    ]
   },
   {