mnt: fixed TS_04 example

zerothi · Apr 28, 2021 · 0cdb1cd · 0cdb1cd
1 parent e3995f3
commit 0cdb1cd
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 51 deletions.
diff --git a/TS_04/RUN.fdf b/TS_04/RUN.fdf
@@ -13,6 +13,9 @@ TBT.DOS.A.All
 TS.BTD.Pivot PCG
 
 TS.Voltage 0 eV
+
+TS.Hartree.Fix -A+A
+
 # Although we have 4 electrodes we will only
 # use two different chemical potentials
 # One for each chain.

diff --git a/TS_04/run.ipynb b/TS_04/run.ipynb
@@ -190,65 +190,39 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "G = sisl.Grid\n",
     "# Define the boundary conditions in the unit-cell\n",
-    "bc = [[G.PERIODIC, G.PERIODIC],\n",
-    "      [G.PERIODIC, G.PERIODIC], # even though left/right meet, we can still use periodic because of the fixed boundaries of the electrodes\n",
-    "      [G.NEUMANN, G.NEUMANN]]\n",
+    "bc = [['dirichlet', 'dirichlet'],\n",
+    "      ['dirichlet', 'dirichlet'],\n",
+    "      ['neumann', 'neumann']]\n",
     "\n",
-    "def solve(A, b, tol=1e-10):\n",
-    "    \"\"\" Don't worry too much about this, if you don't have pyamg installed it will fail\"\"\"\n",
-    "    import pyamg\n",
-    "    from scipy.sparse.linalg import cg\n",
-    "    print('Initiating solver:')\n",
-    "    ml = pyamg.aggregation.smoothed_aggregation_solver(A, max_levels=1000)\n",
-    "    M = ml.aspreconditioner(cycle='W') # pre-conditioner\n",
-    "    residuals = []\n",
-    "    def callback(x):\n",
-    "        residuals.append(np.linalg.norm(b-A*x))\n",
-    "        print('    itt {} with residual {:.2e} / {:.2e}'.format(len(residuals), residuals[-1], residuals[-1] / residuals[0]))\n",
-    "    x, info = cg(A, b, tol=tol, M=M, callback=callback)\n",
-    "    if info != 0:\n",
-    "        raise ValueError('Could not find solution')\n",
-    "    print('Solved the Poisson problem!!!')\n",
-    "    return x\n",
+    "# Import the required machinery for solving the boundary conditions\n",
+    "# There is also a command-line utility to do this from a siesta.TBT.nc file with\n",
+    "# some easier to use command-line options. It isn't fully automated but almost.\n",
+    "from sisl_toolbox.transiesta.poisson.fftpoisson_fix import solve_poisson\n",
     "\n",
-    "# Atomic indices for the 4 electrodes\n",
-    "elecs = []\n",
+    "device_name = device.copy()\n",
+    "\n",
+    "# define the electrodes in the device together with their potential\n",
+    "elecs = {}\n",
     "# X-left\n",
-    "elecs.append([1., elec_x, np.arange(elec_x.na)])\n",
+    "device_name['elec-x-1'] = np.arange(elec_x.na)\n",
+    "elecs['elec-x-1'] = 1.\n",
     "# X-right\n",
-    "elecs.append([1., elec_x, np.arange(chain_x.na - elec_x.na, chain_x.na)])\n",
+    "device_name['elec-x-2'] = np.arange(chain_x.na - elec_x.na, chain_x.na)\n",
+    "elecs['elec-x-2'] = 1.\n",
     "# Y-left\n",
-    "elecs.append([-1., elec_y, np.arange(chain_x.na, chain_x.na + elec_y.na)])\n",
+    "device_name['elec-y-1'] = np.arange(chain_x.na, chain_x.na + elec_y.na)\n",
+    "elecs['elec-y-1'] = -1.\n",
     "# Y-right\n",
-    "elecs.append([-1., elec_y, np.arange(device.na - elec_y.na, device.na)])\n",
-    "\n",
-    "# We use the extra dimension because of idx % shape further down below, i.e. broad-casting\n",
-    "shape = np.array([grid0.shape])\n",
-    "grid = sisl.Grid(shape, bc=bc, geometry=device)\n",
-    "\n",
-    "# Now create the pyamg matrix\n",
-    "# Because we already have the \"correct\" boundary conditions as provided above\n",
-    "# this call will also setup the correct boundary conditions.\n",
-    "A, b = grid.topyamg()\n",
-    "\n",
-    "# Find placement\n",
-    "for i, (V, elec, elec_idx) in enumerate(elecs):\n",
-    "    print('Finding indices for electrode {}'.format(i + 1))\n",
-    "    # Note this will also work for *any* skewed electrode cell.\n",
-    "    cube = sisl.Cuboid(elec.sc.cell, center=np.average(device.xyz[elec_idx, :], 0))\n",
-    "    idx = grid.index(cube)\n",
-    "    # fold indices into primary cell\n",
-    "    idx = np.unique(idx % shape, axis=0)\n",
-    "    indices = grid.pyamg_index(idx)\n",
-    "    grid.pyamg_fix(A, b, indices, V)\n",
+    "device_name['elec-y-2'] = np.arange(device.na - elec_y.na, device.na)\n",
+    "elecs['elec-y-2'] = -1.\n",
     "\n",
-    "# Ensure we have memory enough (for really large systems\n",
-    "# the sparse matrix may be *VERY* big).\n",
-    "# small trick to not have the solution array twice.\n",
-    "del grid.grid\n",
-    "grid.grid = solve(A, b).reshape(shape.ravel())\n",
+    "# Now solve\n",
+    "print(\"Starting solution... Please hold... This takes time! :)\")\n",
+    "grid = solve_poisson(device_name, grid0.shape,\n",
+    "                     dtype=np.float32, tolerance=2e-6,\n",
+    "                     boundary=bc, radius=1.7, **elecs)\n",
+    "print(\"Done... storing boundary conditions to disk\")\n",
     "grid.write('V.TSV.nc')"
    ]
   },