{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import underworld as uw\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "mesh = uw.mesh.FeMesh_Cartesian()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "swarm = uw.swarm.Swarm(mesh)\n",
    "def addvars(swarm_inst):\n",
    "    return (\n",
    "            ( swarm_inst.add_variable('char'  ,1), np.iinfo(np.int8).max    , 'char'   ),\n",
    "            ( swarm_inst.add_variable('short' ,1), np.iinfo(np.int16).max   , 'short'  ),\n",
    "            ( swarm_inst.add_variable('int'   ,1), np.iinfo(np.int32).max   , 'int'    ),\n",
    "            ( swarm_inst.add_variable('long'  ,1), np.iinfo(np.int64).max   , 'long'   ),\n",
    "            ( swarm_inst.add_variable('float' ,1), np.finfo(np.float32).max , 'float'  ),\n",
    "            ( swarm_inst.add_variable('double',1), np.finfo(np.float64).max , 'double' ) )\n",
    "vars = addvars(swarm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# add particles\n",
    "swarm.populate_using_layout(uw.swarm.layouts.PerCellSpaceFillerLayout(swarm,20))\n",
    "# initialise\n",
    "for var in vars:\n",
    "    var[0].data[:] = var[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test evaluation\n",
    "for var in vars:\n",
    "    res = var[0].evaluate( (0.5,0.5) )\n",
    "    if not (np.allclose( res, var[1])):\n",
    "        strguy  = \"There was an error evaluating a swarm variable of type '{}'.\\n\".format(var[2])\n",
    "        strguy += \"Expected value was '{}', but encountered value was '{}'\".format(res,var[1])\n",
    "        raise RuntimeError(strguy)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test checkpointing\n",
    "# first dump data to disk\n",
    "swarm.save(\"test_swarm.h5\")\n",
    "for var in vars:\n",
    "    var[0].save(\"test_swarvar_\" + var[2] + \".h5\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# now create swarm which we will use to clone the original swarm\n",
    "swarm2 = uw.swarm.Swarm(mesh)\n",
    "vars2 = addvars(swarm2)\n",
    "# load saved data onto new swarm\n",
    "swarm2.load(\"test_swarm.h5\")\n",
    "for var in vars2:\n",
    "    var[0].load(\"test_swarvar_\" + var[2] + \".h5\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# now check data on new swarm is identical to old\n",
    "for var in vars2:\n",
    "    if not np.allclose(var[0].data, var[1] ):\n",
    "        raise RuntimeError(\"Something went wrong with swarm variable object.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# cleanup\n",
    "if uw.mpi.rank == 0: \n",
    "    import os\n",
    "    os.remove( \"test_swarm.h5\" )\n",
    "    for var in vars:\n",
    "        os.remove( \"test_swarvar_\" + var[2] + \".h5\" ) "
   ]
  }
 ],
 "metadata": {
  "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.5.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}