{
"cells": [
{
"cell_type": "code",
"execution_count": 1,
"id": "b780b7af-aee0-4d92-b830-64e9b9d3c2f7",
"metadata": {},
"outputs": [],
"source": [
"%matplotlib ipympl\n",
"import numpy as np\n",
"from numba import jit, vectorize, float64, types, int64\n",
"import matplotlib.pyplot as plt\n",
"from tqdm.notebook import trange, tqdm"
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "6fbf192f-aebc-4bce-ae38-339b5610fad1",
"metadata": {},
"outputs": [],
"source": [
"# Coupling formula\n",
"alpha = 1\n",
"\n",
"# Lattive parameters:\n",
"lattice_x = np.array([1, 0, 0])\n",
"lattice_y = np.array([0, 1, 0])\n",
"lattice_z = np.array([0, 0, 1])\n",
"\n",
"lattice_s = np.array([\n",
" [0,0,0],\n",
" [0.2,0,0],\n",
" [-0.2,0,0],\n",
" [0,0.2,0],\n",
" [0,-0.2,0],\n",
" [0,0,0.2],\n",
" [0,0,-0.2],\n",
"])\n",
"\n",
"site_nb = lattice_s.shape[0]\n",
"\n",
"@vectorize([float64(int64, int64, int64, int64, int64)])\n",
"def compute_coupling(x, y, z, a, b):\n",
" vec = x*lattice_x + y*lattice_y + z*lattice_z + lattice_s[a] - lattice_s[b]\n",
" norm = np.linalg.norm(vec)\n",
" return alpha/norm**3*np.abs(3*vec[2]**2/norm**2 - 1)\n"
]
},
{
"cell_type": "markdown",
"id": "8719ddbb-3348-48d2-b949-3221ca298b09",
"metadata": {},
"source": [
"# State generation"
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "062e190f-1561-4f9f-8ffc-aca8e43db8f4",
"metadata": {},
"outputs": [],
"source": [
"def generate_lattice_sites(n, cell_number = 5):\n",
" # TODO: Improve using random picking\n",
" r = np.unique(np.random.randint([[cell_number, cell_number, cell_number, len(lattice_s)]], size=(n, 4)), axis = 0)\n",
" np.random.shuffle(r)\n",
" return r\n",
"\n",
"def add_noise(couplings, sigma):\n",
" return np.max(0, couplings + np.random.normal(scale = sigma, size = couplings.shape))\n",
"\n",
"@jit\n",
"def couplings(sites):\n",
" n = sites.shape[0]\n",
" couplings = np.full((n , n), np.nan)\n",
" for i in range(sites.shape[0]-1):\n",
" for j in range(i+1, sites.shape[0]):\n",
" x = sites[i][0] - sites[j][0]\n",
" y = sites[i][1] - sites[j][1]\n",
" z = sites[i][2] - sites[j][2]\n",
" couplings[i, j] = compute_coupling(x, y, z, sites[i][3], sites[j][3])\n",
" return couplings"
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "3d1d36c3-1fc4-42e4-af71-8eda4fb34d86",
"metadata": {},
"outputs": [],
"source": [
"test_sites = generate_lattice_sites(10, cell_number=5)"
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "ae4ecda6-14db-4e50-bfd2-affcd6153630",
"metadata": {},
"outputs": [],
"source": [
"test_couplings = couplings(test_sites)"
]
},
{
"cell_type": "markdown",
"id": "311c7c64-84ed-45c2-9f49-5993c67bc21e",
"metadata": {},
"source": [
"# State reconstruction"
]
},
{
"cell_type": "code",
"execution_count": 9,
"id": "daaa4792-069c-44fd-bdbd-90af6f51858f",
"metadata": {},
"outputs": [],
"source": [
"@jit\n",
"def index_to_coord(index, max_distance, site_nb):\n",
" center = max_distance // 2\n",
" return (\n",
" index // (max_distance**2 * site_nb**2) - center,\n",
" index // (max_distance * site_nb**2) % max_distance - center,\n",
" index // site_nb**2 % max_distance - center,\n",
" index // site_nb % site_nb,\n",
" index % site_nb\n",
" ) \n",
"\n",
"@jit\n",
"def coord_to_index(vec, max_distance, site_nb):\n",
" center = max_distance // 2\n",
" return site_nb**2 * ((vec[0] + center)*max_distance**2 + (vec[1] + center) * max_distance + (vec[2] + center)) + vec[3]*site_nb + vec[4]\n",
"\n",
"@jit #(locals={'current_set': types.Set(types.UniTuple(float64, 3))})\n",
"def vector_couplings(max_distance, site_nb, tolerance):\n",
" print(\"Computing vector -> couplings lookup\")\n",
" couplings = np.empty(max_distance**3*site_nb**2)\n",
" for i in range(max_distance**3*site_nb**2):\n",
" couplings[i] = compute_coupling(\n",
" *index_to_coord(i, max_distance, site_nb)\n",
" )\n",
" # Revert the map\n",
" print(\"Reverting the map\")\n",
" coupling_values = [] # OK to append to that list because O(number_spins)\n",
" current_set = set(((0,0,0,0,0),))\n",
" current_set.pop() # HACK for type inference\n",
" vectors = [] # OK to append because O(number_spins)\n",
" sorted_couplings = np.argsort(couplings)\n",
" # i_max = coupling index st coupling + tolerance < current cursor\n",
" # i_min = coupling index st coupling - tolerance < current cursor\n",
" i_min, i_max = 0,0\n",
" while not np.isnan(couplings[sorted_couplings[i_max]]):\n",
" if np.isnan(couplings[sorted_couplings[i_min]]) or couplings[sorted_couplings[i_min]] - tolerance > couplings[sorted_couplings[i_max]] + tolerance:\n",
" # Depile du côté des maxima (ie on retire un vecteur au set)\n",
" vec = index_to_coord(sorted_couplings[i_max], max_distance, site_nb)\n",
" coupling_cursor = couplings[sorted_couplings[i_max]] + tolerance\n",
" if len(coupling_values) == 0 or not np.isclose(coupling_values[-1], coupling_cursor, rtol = 1e-9, atol = tolerance * 1e-9):\n",
" vectors.append(current_set.copy())\n",
" coupling_values.append(coupling_cursor)\n",
" current_set.remove(vec)\n",
" i_max += 1\n",
" else:\n",
" # Empile du côté des minima (ie on ajoute un vecteur au set)\n",
" vec = index_to_coord(sorted_couplings[i_min], max_distance, site_nb)\n",
" coupling_cursor = couplings[sorted_couplings[i_min]] - tolerance\n",
" if len(coupling_values) == 0 or not np.isclose(coupling_values[-1], coupling_cursor, rtol = 1e-9, atol = tolerance * 1e-9):\n",
" vectors.append(current_set.copy())\n",
" coupling_values.append(coupling_cursor)\n",
" current_set.add(vec)\n",
" i_min += 1\n",
" vectors.append(current_set)\n",
" print(\"Successfully reverted the map\")\n",
" return np.array(coupling_values), vectors"
]
},
{
"cell_type": "code",
"execution_count": 10,
"id": "4a78bf20-e87b-4aeb-99f0-2bccc9dc4233",
"metadata": {},
"outputs": [],
"source": [
"@jit\n",
"def cost_function(sites, couplings):\n",
" couplings_theory = couplings(sites)\n",
" cost = 0\n",
" for i in range(sites.shape[0]-1):\n",
" for j in range(i+1, sites.shape[0]):\n",
" if np.isnan(couplings[i,j]):\n",
" continue\n",
" x = sites[i][0] - sites[j][0]\n",
" y = sites[i][1] - sites[j][1]\n",
" z = sites[i][2] - sites[j][2]\n",
" cost += (couplings[i,j] - compute_coupling(x, y, z, sites[i][3], sites[j][3]))**2\n",
" return cost\n",
"\n",
"@jit\n",
"def exchange_columns(couplings, permutation, a, b):\n",
" a, b = min(a, b), max(a, b)\n",
" print(f\"Exchange {a}-{b}\")\n",
" permutation[a], permutation[b] = permutation[b], permutation[a]\n",
" for i in range(a):\n",
" couplings[i, a], couplings[i, b] = couplings[i, b], couplings[i, a]\n",
" for i in range(a+1, b):\n",
" couplings[a, i], couplings[i, b] = couplings[i, b], couplings[a, i]\n",
" for i in range(b+1, couplings.shape[0]): \n",
" couplings[a, i], couplings[b, i] = couplings[b, i], couplings[a, i]\n",
"\n",
"@jit\n",
"def set_placing_order(couplings):\n",
" print(couplings.shape)\n",
" n_tot = couplings.shape[0]\n",
" permutation = np.arange(n_tot)\n",
" m = np.nanmax(couplings[0,:])\n",
" first = 0\n",
" for i in range(1, n_tot):\n",
" if np.nanmax(couplings[i,:])> m:\n",
" m = np.nanmax(couplings[i,:])\n",
" first = i\n",
" exchange_columns(couplings, permutation, 0, first)\n",
" for i in range(1, n_tot):\n",
" next_index = np.argmax(couplings[:i,i:])%(n_tot-i) + i\n",
" if next_index != i:\n",
" exchange_columns(couplings, permutation, i, next_index)\n",
" return couplings, permutation\n",
"\n",
"#@jit\n",
"def compute_sites(couplings, lattice_size, site_nb, tolerance):\n",
" print(\"Begin\")\n",
" n_placed = 1\n",
" n_tot = couplings.shape[0]\n",
" couplings, permutation = set_placing_order(couplings)\n",
" print(\"ordered spins\")\n",
" couplings_vectors_tup = vector_couplings(tolerance = tolerance, max_distance = lattice_size*2, site_nb = site_nb)\n",
" print(\"computed_lookup table\")\n",
" possible_configurations = []\n",
" for i in range(site_nb):\n",
" c = np.zeros((n_tot, 4), dtype = np.int64)\n",
" c[0][3] = i\n",
" possible_configurations.append(c)\n",
" \n",
" current_couplings = np.empty(n_tot) # stores current coupling values to other spins\n",
"\n",
" print(\"Initialization successful\")\n",
" while n_placed < n_tot:\n",
" print(f\"Placing {n_placed}. {len(possible_configurations)} cases to process\")\n",
" # First place them in order\n",
" n_to_place = n_placed\n",
"\n",
" for i in range(n_tot):\n",
" current_couplings[i] = np.nanmax([couplings[n_to_place, i], couplings[i, n_to_place]])\n",
" # Be careful, position relative to n_to_place-1\n",
" candidates = couplings_vectors_tup[1][np.searchsorted(couplings_vectors_tup[0], current_couplings[n_to_place-1])]\n",
" new_possible_configurations = []\n",
" for config in tqdm(possible_configurations):\n",
" for candidate in candidates:\n",
" good_candidate = True\n",
" # First check the origin lattice site\n",
" if candidate[4] != config[n_to_place - 1][3]:\n",
" continue\n",
" for i in range(n_placed-1):\n",
" if np.isnan(current_couplings[i]):\n",
" # no data for this already placed spin\n",
" continue\n",
" x = candidate[0] + config[n_to_place-1][0] - config[i][0]\n",
" y = candidate[1] + config[n_to_place-1][1] - config[i][1]\n",
" z = candidate[2] + config[n_to_place-1][2] - config[i][2]\n",
" a = candidate[3]\n",
" b = config[i][3]\n",
" if (x, y, z) == (0,0,0) and a == b:\n",
" good_candidate = False\n",
" break\n",
" coupling_candidate = compute_coupling(\n",
" x, y, z, a, b\n",
" )\n",
" if np.isnan(coupling_candidate) or np.abs(current_couplings[i] - coupling_candidate) > tolerance:\n",
" good_candidate = False\n",
" break\n",
" if good_candidate:\n",
" new_config = config.copy()\n",
" new_config[n_to_place] = np.array((candidate[0] + config[n_to_place-1][0], candidate[1] + config[n_to_place-1][1], candidate[2] + config[n_to_place-1][2], candidate[3]))\n",
" new_possible_configurations.append(new_config)\n",
" possible_configurations = new_possible_configurations\n",
" n_placed+=1\n",
" return [i[permutation] for i in possible_configurations], permutation\n",
"\n"
]
},
{
"cell_type": "code",
"execution_count": 11,
"id": "9d186b15-36cf-4aff-9287-381e31424afb",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Begin\n",
"(10, 10)\n",
"Exchange 0-0\n",
"ordered spins\n",
"Computing vector -> couplings lookup\n",
"Reverting the map\n",
"Successfully reverted the map\n",
"computed_lookup table\n",
"Initialization successful\n",
"Placing 1. 7 cases to process\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"/tmp/ipykernel_44029/1413072332.py:69: RuntimeWarning: All-NaN axis encountered\n",
" current_couplings[i] = np.nanmax([couplings[n_to_place, i], couplings[i, n_to_place]])\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "fb4ed34663df4985ac8a6a31dc42a50d",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
" 0%| | 0/7 [00:00\n",
" \n",
" Figure\n",
"
\n",
" 
\n",
" \n",
" "
],
"text/plain": [
"Canvas(toolbar=Toolbar(toolitems=[('Home', 'Reset original view', 'home', 'home'), ('Back', 'Back to previous …"
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"fig = plt.figure()\n",
"ax = fig.add_subplot(projection='3d')\n",
"\n",
"positions = lattice_x[None,:]*test_sites[:,0][:,None] + lattice_y[None,:] * test_sites[:,1][:,None] + lattice_z[None,:] * test_sites[:,2][:,None] + lattice_s[test_sites[:,3]]\n",
"\n",
"ax.scatter(positions[:,0], positions[:,1], positions[:,2])\n",
"ax.plot(positions[permutation,0], positions[permutation,1], positions[permutation,2])\n",
"\n",
"ax.set_xlabel('X Label')\n",
"ax.set_ylabel('Y Label')\n",
"ax.set_zlabel('Z Label')\n",
"\n",
"for i in range(test_sites.shape[0]):\n",
" ax.text(positions[i,0], positions[i,1], positions[i,2], str(i))"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "7a904af3-06ae-4e1c-b282-9526dea314bd",
"metadata": {
"scrolled": true
},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"id": "7537e392-4c17-490c-8ec4-c1d147fa3e24",
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.12.9"
}
},
"nbformat": 4,
"nbformat_minor": 5
}