{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ae971fd",
   "metadata": {
    "tags": [
     "remove-input"
    ]
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from numpy.linalg import norm\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "\n",
    "plt.rcParams['animation.frame_format'] = \"svg\"\n",
    "%config InlineBackend.figure_formats = ['svg']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99dd2b41",
   "metadata": {},
   "source": [
    "## Square ASI robustness to dilution defects\n",
    "\n",
    "In this example we generate geometries of dilute square ASI and relax them by a rotating, global field.\n",
    "The results analyzed in terms of vertex type population and their robustness to increasing dilution fraction.\n",
    "\n",
    "We first create the dilute ensembles from a square ASI.\n",
    "The ensembles are relaxed with a rotating field protocol in flatspin.\n",
    "Finally, we analyze the results with respect to vertex type populations."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4174ec32",
   "metadata": {},
   "source": [
    "### Creating dilute ensembles\n",
    "First, we generate the example system: a complete square ASI.\n",
    "We will then iteratively remove magnets and save the resulting systems."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92eb6561",
   "metadata": {},
   "outputs": [],
   "source": [
    "from flatspin.model import SquareSpinIceClosed, CustomSpinIce\n",
    "from flatspin.plotting import plot_vector_image\n",
    "\n",
    "size = (50, 50)\n",
    "sq_asi = SquareSpinIceClosed(size = size, lattice_spacing = 1)\n",
    "\n",
    "plt.figure(figsize=(10,6))\n",
    "plt.title(f'Full original ensemble, size = {size}')\n",
    "sq_asi.plot();"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c7b21628",
   "metadata": {},
   "source": [
    "Next, we define the dilution fractions and the corresponding number of magnets to be removed at each dilution fraction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "687e3bab",
   "metadata": {},
   "outputs": [],
   "source": [
    "asi_size_n = sq_asi.spin_count\n",
    "\n",
    "dilutions = np.array([0.001, 0.002, 0.005, 0.01, 0.02, 0.05, 0.1, 0.15, 0.2, 0.3, 0.4])\n",
    "num_removed = [round(d*asi_size_n) for d in dilutions]\n",
    "\n",
    "print(f\"For {asi_size_n} total magnets\")\n",
    "df = pd.DataFrame({'dilution fraction': dilutions, 'magnets removed': num_removed })\n",
    "display(df)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88ea9e96",
   "metadata": {},
   "source": [
    "We now define the dilute geometries, using several random seeds for the stochastic removal of magnets.\n",
    "The geometries are saved as files where the positions and angles of each spin are stored, to be used by the {class}`CustomSpinIce <flatspin.model.CustomSpinIce>` model class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b56cd20d",
   "metadata": {},
   "outputs": [],
   "source": [
    "root_dir = 'dilution'\n",
    "if not os.path.exists(root_dir):\n",
    "    os.makedirs(root_dir)\n",
    "    print(f\"Creating directory {root_dir}\")\n",
    "\n",
    "square_geom = (sq_asi.pos, sq_asi.angle)\n",
    "\n",
    "random_iterations = 10\n",
    "for rand_seed in range(random_iterations):\n",
    "    # Order the ids of the magnets. The order in which they are removed\n",
    "    np.random.seed(rand_seed)\n",
    "    removed_id = np.random.choice(asi_size_n, size=max(num_removed), replace=False)\n",
    "\n",
    "    for i, n in enumerate(num_removed):\n",
    "        square_geom_removed = (np.delete(square_geom[0], removed_id[:n], axis = 0),\n",
    "                           np.delete(square_geom[1], removed_id[:n], axis = 0))\n",
    "        # Save coordinates and angles\n",
    "        np.savetxt(f'{root_dir}/coords-dilution-{str(i).zfill(2)}-seed-{str(rand_seed).zfill(2)}.csv', (square_geom_removed[0]))\n",
    "        np.savetxt(f'{root_dir}/angles-dilution-{str(i).zfill(2)}-seed-{str(rand_seed).zfill(2)}.csv', (square_geom_removed[1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "734a9010",
   "metadata": {},
   "source": [
    "We can instantiate a CustomSpinIce example model from one of these generated geometries:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00011cae",
   "metadata": {},
   "outputs": [],
   "source": [
    "dilute_example = CustomSpinIce(magnet_coords = f'{root_dir}/coords-dilution-06-seed-00.csv',\n",
    "                               magnet_angles = f'{root_dir}/angles-dilution-06-seed-00.csv',\n",
    "                               radians = True)\n",
    "\n",
    "plt.figure(figsize=(10,6))\n",
    "plt.title(f'Dilute ensemble, size = {size}, dilution = {1-dilute_example.spin_count/asi_size_n:.2f}')\n",
    "dilute_example.plot();"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35ee152f",
   "metadata": {},
   "source": [
    "### Baseline field annealing\n",
    "\n",
    "We test a field annealing approach for the undiluted system here, before we generate the runs for all diluted systems. See the [field annealing documentation](anneal-field) for an example of how to set this up.\n",
    "\n",
    "First, we define some physical parameters of the simulation and a model instance:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "debbdd58",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Size of magnets\n",
    "w, l, t = 220e-9, 80e-9, 20e-9 # [m]\n",
    "volume = w*l*t # [m^3]\n",
    "M_sat = 860e3 # [A/m]\n",
    "M_i = M_sat * volume\n",
    "\n",
    "a = 1e-9 # lattice_spacing unit [m]\n",
    "lattice_spacing = 300 # [a]\n",
    "\n",
    "mu_0 = 1.25663706e-6 # m kg /s^2 A^2\n",
    "alpha = mu_0 * M_i / (4 * np.pi * a**3)\n",
    "\n",
    "# Astroid parameters\n",
    "sw_b = 0.38\n",
    "sw_c = 1\n",
    "sw_beta = 1.3\n",
    "sw_gamma = 3.6\n",
    "\n",
    "\n",
    "rand_seed = 0\n",
    "disorder = 0.05\n",
    "neighbor_distance = 3\n",
    "\n",
    "test_size = (10,10)\n",
    "\n",
    "model = SquareSpinIceClosed(size=test_size, lattice_spacing=lattice_spacing, alpha=alpha,\n",
    "                            disorder=disorder, neighbor_distance=neighbor_distance,\n",
    "                            sw_b=sw_b, sw_c=sw_c, sw_beta=sw_beta, sw_gamma=sw_gamma,\n",
    "                            random_seed=rand_seed, use_opencl=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c89a65e",
   "metadata": {},
   "source": [
    "Next, we use the `Rotate` encoder and find some fitting field values that will relax the system with a rotating field.\n",
    "We can adjust the parameter values to see the effect.\n",
    "The current parameters are the ones used for the experiment in the paper, which takes a while to run, even for the smaller test system."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f51239e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from flatspin.encoder import Rotate\n",
    "\n",
    "# We define the decreasing field strength\n",
    "H_hi, H_low = 0.080, 0.072\n",
    "input_values = np.arange(1,0,-0.001)\n",
    "timesteps = 100 # (per input value)\n",
    "\n",
    "# Define the external field values based on the Rotate encoder\n",
    "encoder = Rotate(H = H_hi, H0 = H_low, timesteps=timesteps)\n",
    "h_ext = encoder(input_values)\n",
    "\n",
    "# Start in polarized state\n",
    "model.polarize()\n",
    "\n",
    "# Record spins and number of spin flips over time\n",
    "spins = []\n",
    "flips = []\n",
    "for i, h in enumerate(h_ext):\n",
    "    model.set_h_ext(h)\n",
    "    s = model.relax()\n",
    "    # Record spin state at the end of each rotation\n",
    "    if (i+1) % timesteps == 0:\n",
    "        spins.append(model.spin.copy())\n",
    "    flips.append(s)\n",
    "\n",
    "print(f\"Completed {sum(flips)} flips\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ce565be",
   "metadata": {},
   "outputs": [],
   "source": [
    "H = norm(h_ext, axis=-1).round(10)\n",
    "df = pd.DataFrame({'H': H, 'flips': flips})\n",
    "df.groupby('H', sort=False).sum().plot(figsize=(10,6))\n",
    "plt.gca().invert_xaxis()\n",
    "plt.ylabel(\"Spin flips\")\n",
    "plt.xlabel(\"Field strength [T]\");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e76fc94b",
   "metadata": {},
   "source": [
    "For a too strong (saturating) field strength, all the magnets will be flipped by the field.\n",
    "The maximum number of flips in one rotation of the field, i.e., at one field strength, is twice the total number of magnets (1 flip + 1 flip back).\n",
    "The number of timesteps and the field strength parameters should be tuned so that we start out flipping almost all the magnets (i.e., a strong *enough* field to change the system state) and go until we don't flip any magnets as we don't need to simulate weaker fields than that which results in no activity.\n",
    "\n",
    "\n",
    "Below, we animate over the states of the relaxation process (skipping quite a few frames)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98ed4811",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "from matplotlib.animation import FuncAnimation\n",
    "from IPython.display import HTML\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "\n",
    "skip_frames = 20\n",
    "def animate_spin(i):\n",
    "    i = i*skip_frames\n",
    "    H = norm(h_ext[(i+1) * timesteps - 1])\n",
    "    ax.set_title(f\"H={H:.3f} [T]\")\n",
    "    model.set_spin(spins[i])\n",
    "    model.plot_vertex_mag(ax=ax, replace=True)\n",
    "\n",
    "anim = FuncAnimation(fig, animate_spin, frames=len(spins[::skip_frames]), interval=200, blit=False)\n",
    "plt.close() # Only show the animation\n",
    "HTML(anim.to_jshtml())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2bc1ab44",
   "metadata": {},
   "source": [
    "### Create the runs to generate the data.\n",
    "\n",
    "We can now generate the data for all the dilute systems and different random seeds.\n",
    "We generate a command that can schedule a whole sweep of jobs based on the generated dilute geometries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9d8ba75",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate flatspin-run-sweep command\n",
    "print(\"\\nGENERATED COMMAND: \\n\")\n",
    "print(f\"flatspin-run-sweep -r dist -m CustomSpinIce -e Rotate \", end='\\n')\n",
    "print(f\"-p sw_b={sw_b} -p sw_c={sw_c} -p sw_beta={sw_beta} -p sw_gamma={sw_gamma} \", end='\\n')\n",
    "print(f\"-p alpha={alpha:.2f} -p lattice_spacing={lattice_spacing} -p disorder={disorder} \", end='\\n')\n",
    "print(f\"-p neighbor_distance={neighbor_distance} \", end='\\n')\n",
    "print(f\"-p 'input=np.arange(1,0,-0.001)' \", end='\\n')\n",
    "print(f\"-p timesteps={timesteps} -p H={H_hi} -p H0={H_low} \", end='\\n')\n",
    "print(f\"-s 'num=range(10)' -s 'rs=range(10)' \", end='\\n')\n",
    "print(f\"-s 'magnet_coords=[\\\"coords-dilution-\\\"+str(num).zfill(2)+\\\"-seed-\\\"+str(rs).zfill(2)+\\\".csv\\\"]' \", end='\\n')\n",
    "print(f\"-s 'magnet_angles=[\\\"angles-dilution-\\\"+str(num).zfill(2)+\\\"-seed-\\\"+str(rs).zfill(2)+\\\".csv\\\"]' \", end='\\n')\n",
    "\n",
    "print(\"-o CHANGE_TO_OUTPUT_FOLDER\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62df3add",
   "metadata": {},
   "source": [
    "### Analyze generated data\n",
    "\n",
    "The dataset can also be {download}`downloaded here </data/flatspin/dilution-example-run.zip>`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57dfd0fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "from flatspin.data import Dataset, read_table, read_geometry, vector_grid\n",
    "from flatspin.grid import Grid\n",
    "from flatspin.vertices import find_vertices, vertex_type\n",
    "\n",
    "data_path = '/data/flatspin/dilution-example-run'\n",
    "\n",
    "D = Dataset.read(data_path)\n",
    "display(D.index)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2cda8a1f",
   "metadata": {},
   "source": [
    "We now count all the vertices and their types for all the different runs.\n",
    "This might take some time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6942c65",
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import defaultdict\n",
    "from joblib import Parallel, delayed\n",
    "\n",
    "def count_types(ds):\n",
    "    dilution_num = int(ds.index['num'])\n",
    "    rand_seed = int(ds.index['rs'])\n",
    "    spins = read_table(ds.tablefile('spin'), index_col='t')\n",
    "    \n",
    "    # Choose the final spin state at t=-1\n",
    "    spins_final = spins.iloc[-1]\n",
    "\n",
    "    geom = read_geometry(ds.tablefile('geometry'))\n",
    "    pos, angles = geom\n",
    "    grid = Grid(pos)\n",
    "\n",
    "    # Identify vertices (only complete vertices)\n",
    "    _, _, vertices = find_vertices(grid, pos, angles, 3)\n",
    "\n",
    "    # Sort vertices by type\n",
    "    types = [vertex_type([spins_final[s] for s in v],\n",
    "                         [pos[s] for s in v],\n",
    "                         [angles[s] for s in v])\n",
    "             for v in vertices]\n",
    "\n",
    "    # Count the vertex types\n",
    "    vertex_types, counts = np.unique(types, return_counts=True)\n",
    "    typecount = dict(zip(vertex_types, counts))\n",
    "    \n",
    "    return dilution_num, rand_seed, counts\n",
    "\n",
    "queue = D[0::10]\n",
    "types_by_dilution = Parallel(n_jobs=10, verbose=0)(delayed(count_types)(ds) for ds in queue)\n",
    "\n",
    "df_types = pd.DataFrame(types_by_dilution, columns={'dilution': 'num', 'rand_seed': 'rs', 'type': 'type'})\n",
    "label_types = ['Type I', 'Type II', 'Type III', 'Type IV']\n",
    "\n",
    "df_types[label_types[:-1]] = pd.DataFrame(df_types.type.tolist(), index= df_types.index)\n",
    "df_types_avg = df_types.groupby('dilution').mean()\n",
    "sum_magnets = df_types_avg[label_types[:-1]].sum(axis=1)\n",
    "df_types_avg[[l+'_frac' for l in label_types[:-1]]] = (df_types_avg[label_types[:-1]].T/sum_magnets).T\n",
    "df_types_avg"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80785f56",
   "metadata": {},
   "source": [
    "### Plotting the results\n",
    "\n",
    "First, let's look at each dilution fraction for one selected random seed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a95cbd53",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "def read_data_vertices_grid(dataset, t=[0]):\n",
    "    ''' Return the positions and magnetization of the dataset's vertices '''\n",
    "    grid_size = None\n",
    "    crop_width = ((1,1),(1,1))\n",
    "    win_shape, win_step = ((3,3), (2, 2))\n",
    "\n",
    "    df = read_table(dataset.tablefile('mag'), index_col='t')\n",
    "    UV = np.array(df)\n",
    "    UV = UV.reshape((UV.shape[0], -1, 2))\n",
    "    UV = UV[t]\n",
    "\n",
    "    pos, angle = read_geometry(dataset.tablefile('geometry'))\n",
    "\n",
    "    XY, UV = vector_grid(pos, UV, grid_size, crop_width, win_shape, win_step, normalize=False)\n",
    "    return XY, UV/2\n",
    "\n",
    "def plot_vertices(XY, UVi, ax):\n",
    "    ax.axis('off')\n",
    "    plot_vector_image(XY, UVi, ax=ax, replace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73053760",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load final states (might take some time)\n",
    "final_dilution_states = dict()\n",
    "for num, ds in D[D.index['rs']==0].groupby('num'):\n",
    "    final_dilution_states[dilutions[num]] = read_data_vertices_grid(ds, t=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8cb9525",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot final states\n",
    "fig, axs = plt.subplots(2, 5, figsize=(10,4), dpi=100)\n",
    "for ax, (dilution, vstate) in zip(axs.flat, final_dilution_states.items()):\n",
    "    ax.set_title(f'dilution = {dilution}')\n",
    "    XY, UVi = vstate\n",
    "    mag = np.amax(np.linalg.norm(UVi, axis=-1))\n",
    "    plot_vertices(XY, UVi/mag, ax)\n",
    "fig.suptitle('Final timestep', fontsize='x-large');"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad2f14d1",
   "metadata": {},
   "source": [
    "Finally, we plot the vertex type as a function of dilution fraction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49f13476",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_types_avg['Type IV_frac'] = 0.0\n",
    "\n",
    "plt.figure(figsize=(10,6))\n",
    "for label, ts in zip(label_types, df_types_avg[[l+'_frac' for l in label_types]]):\n",
    "    plt.plot(dilutions[:-1], df_types_avg[ts], marker = 'x', label=label)\n",
    "    plt.xscale('log')\n",
    "plt.ylabel('Average vertex fraction')\n",
    "plt.xlabel('Dilution fraction')\n",
    "plt.legend();"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "formats": "ipynb,md:myst",
   "text_representation": {
    "extension": ".md",
    "format_name": "myst",
    "format_version": 0.13,
    "jupytext_version": "1.11.5"
   }
  },
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.10"
  },
  "source_map": [
   8,
   19,
   30,
   36,
   46,
   50,
   59,
   64,
   84,
   88,
   96,
   104,
   134,
   140,
   169,
   176,
   185,
   204,
   211,
   225,
   231,
   240,
   245,
   287,
   293,
   317,
   324,
   333,
   337
  ]
 },
 "nbformat": 4,
 "nbformat_minor": 5
}