equationGenerator.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Automatic Symbolic Equation Generator\n",
    "\n",
    "*Written by Forough Arabshahi*\n",
    "\n",
    "This notebook presents the code for the automatic dataset generation scheme presented in [this](https://arxiv.org/abs/1801.04342) paper.\n",
    "\n",
    "This equation generation approach is based on generating new mathematical identities by performing local random changes to known identities, starting with a small number of axioms from the domain under study. These changes result in identities of similar or higher complexity (equal or larger depth), which may be correct or incorrect, that are valid expressions within a grammar. The grammar and its rules are defined in detail in [section 2.1 of the paper](https://arxiv.org/abs/1801.04342). In this notebook we choose elementary algebra and trigonometry as the domain under study.\n",
    "\n",
    "### Generating Possible Equations\n",
    "<a id=\"sec:posEq\"></a>\n",
    "In order to generate a new identity which may be correct or incorrect, we select an equation at random from the set of known equations, and make local changes to it. We randomly select a node in the expression tree, followed by randomly selecting one of the following actions to make the local change to the equation at the selected node:\n",
    "* ShrinkNode: Replace the node, if it’s not a leaf, with one of its children, chosen randomly. \n",
    "* ReplaceNode: Replace the symbol at the node (i.e. the terminal or the function) with another\n",
    "compatible one, chosen randomly.\n",
    "* GrowNode: Provide the node as input to another randomly drawn function f , which then\n",
    "replaces the node. If f takes two inputs, the second input will be generated randomly from\n",
    "the set of terminals.\n",
    "* GrowSides: If the selected node is an equality, either add or multiply both sides with a\n",
    "randomly drawn number, or take both sides to the power of a randomly drawn number.\n",
    "\n",
    "This is implemented in a function called *genNegEquation*\n",
    "\n",
    "### Generating Additional Correct Equations\n",
    "In order to generate only correct identities, we follow the same intuition as above, but only replace structure with others that are equal. In particular, we maintain a dictionary of valid statements (mathDictionary) that maps a mathematical statement to another. For example, the dictionary key $x + y$ has value $y + x$. We use this dictionary in our correct equation generation process where we look up patterns from the dictionary. More specifically, we look for keys that match a subtree of the equation then replace that subtree with the pattern of the value of the key. E.g. given input equation $\\sin^2(\\theta) + \\cos^2(\\theta) = 1$, this subtree matching might produce equality $\\cos^2(\\theta) + \\sin^2(\\theta) = 1$ by finding key-value pair $x + y : y + x$.\n",
    "\n",
    "This is implemented in a function called *genPosEquation*\n",
    "\n",
    "***\n",
    "\n",
    "## Implementation Details\n",
    "\n",
    "The dependency for this notebook if the [latex2sympy](https://github.com/augustt198/latex2sympy) package which converts a latex equation to sympy. Set the path to latex2sympy in file eqGen.py available in the root folder. Our equations are represented using the *EquationTree* object. The attributes of this object are: \n",
    "1. func: string that stores the function name\n",
    "2. args: list of EquationTree objects that are the children of this node\n",
    "3. varname: string that stores the variable name if the node is a leaf, otherwise this is ''\n",
    "4. number: integer that indicates the node's unique number\n",
    "4. depth: integer that indicates the depth of the subtree that is attached to the node.\n",
    "\n",
    "below is an example of some equations\n",
    "\n",
    "$\\sin^2(\\theta) + \\cos^2(\\theta) = 1$ | $\\sin(-2.5) = -0.6$ | Decimal expression tree for $2.5$\n",
    "- | - | -\n",
    "<img src=\"figs/eTree.png\", width=\"300\", height=\"300\"/>  | <img src=\"figs/numTree.png\", width=\"300\", height=\"300\"/> | <img src=\"figs/num_tree.png\", width=\"300\", height=\"300\"/>\n",
    "\n",
    "### 1. Importing Modules\n",
    "<a id=\"sec:import\"></a>\n",
    "\n",
    "Most of the functions are implemented in *eqGen.py*. *readJson* reads an input json in latex format and uses latex2sympy to convert it to a sympy equation. Finally, *buildEq* converts a sympy equation to an EquationTree object. *readAxioms* reads an input txt file and uses the compiler package to convert the equations to a compiler object. Function *parseEquation* then converts this to an EquationTree object. *writeJson* saves the generated EquationTree objects in a json file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/Forough/anaconda2/lib/python2.7/site-packages/ipykernel/__main__.py:7: DeprecationWarning: The compiler package is deprecated and removed in Python 3.x.\n"
     ]
    }
   ],
   "source": [
    "import json\n",
    "from sympy import *\n",
    "import pprint\n",
    "import re\n",
    "import copy\n",
    "import random\n",
    "import compiler\n",
    "import mxnet as mx\n",
    "import numpy as np\n",
    "import sys\n",
    "from eqGen import readJson, readAxioms, parseEquation, buildEq, buildAxiom, genPosEquation,\\\n",
    "                  genNegEquation, isCorrect, writeJson\n",
    "from neuralAlgonometry import catJsons"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. Parsing input axioms\n",
    "<a id=\"sec:axioms\"></a>\n",
    "\n",
    "In this section we a small set of input axioms from trigonometry and elementary algebra and convert them to EquationTree objects. trigonometry equations are in file *\"axioms/trigonometryAxioms.json\"* which are collected from the wikipedia [List of Trigonometric Identities](https://en.wikipedia.org/wiki/List_of_trigonometric_identities) page. Elementary algebra equations are in file *\"axioms/algebraAxioms.txt\"* and are hand generated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# path to trigonometry equations collected from wiki\n",
    "inputPath = \"axioms/trigonometryAxioms.json\"\n",
    "# path to some axioms from elementary algebra\n",
    "axiomPath = \"axioms/algebraAxioms.txt\"\n",
    "jsonAtts = [\"equation\", \"range\", \"variables\",\"labels\"]\n",
    "\n",
    "labels = []\n",
    "inputRanges = []\n",
    "inputEquations = [] \n",
    "eqVariables = []\n",
    "parsedEquations = []\n",
    "parsedRanges = []\n",
    "ranges = []\n",
    "inputAxioms = []\n",
    "axiomVariables = []\n",
    "axioms = []\n",
    "axLabels = []\n",
    "\n",
    "random.seed(1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Function *readJson* parses latex equations. An example of the resulting output is given below"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'\\\\sin\\\\theta = \\\\frac{1}{\\\\csc\\\\theta}'"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# parses latex equations from file:\n",
    "readJson(inputPath, inputEquations, inputRanges, jsonAtts)\n",
    "inputEquations[1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We use function *parseEquation* to convert these equations to a sympy equation using the latex2sympy package. An example of the output equation follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Eq(sin(theta), 1/csc(theta))"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Converts latex equations to sympy equations using process_latex\n",
    "parseEquation(inputEquations, parsedEquations, eqVariables)\n",
    "parsedEquations[1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Function *buildEq* converts each sympy equation to an EquationTree object as shown below. The pre order traversal of the parsed equation is also shown as an example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "currEq: Eq(sin(var_0),Pow(csc(var_0),-1))\n",
      "pre order traversal\n",
      "functionality: Equality\n",
      "assigned node number: 0\n",
      "subtree depth: 3\n",
      "functionality: sin\n",
      "assigned node number: 1\n",
      "subtree depth: 2\n",
      "variable name: var_0\n",
      "functionality: Symbol\n",
      "assigned node number: 3\n",
      "subtree depth: 1\n",
      "functionality: Pow\n",
      "assigned node number: 2\n",
      "subtree depth: 2\n",
      "functionality: csc\n",
      "assigned node number: 4\n",
      "subtree depth: 1\n",
      "variable name: var_0\n",
      "functionality: Symbol\n",
      "assigned node number: 6\n",
      "subtree depth: 0\n",
      "variable name: -1\n",
      "functionality: NegativeOne\n",
      "assigned node number: 5\n",
      "subtree depth: 1\n"
     ]
    }
   ],
   "source": [
    "# converts equations from sympy format to EquationTree object\n",
    "# equations pretty print as well as pre order traversal follows\n",
    "equations = []\n",
    "for i, eq in enumerate(parsedEquations):\n",
    "    # building EquationTree object\n",
    "    currEq = buildEq(eq, eqVariables[i])\n",
    "    # assigning a unique number to each node in the tree as well as indicating subtree depth at each level\n",
    "    currEq.enumerize_queue()\n",
    "    equations.append(currEq)\n",
    "    \n",
    "    # creating training labels\n",
    "    # the first equation in the input function is incorrect. It has been deliberately added\n",
    "    # to include all possible functionalities in the functionDictionary. \n",
    "    # This is for compatibility with MxNet's bucketingModule.\n",
    "    if i == 0:\n",
    "        labels.append(mx.nd.array([0]))\n",
    "    else:\n",
    "        labels.append(mx.nd.array([1]))\n",
    "    \n",
    "print \"currEq:\", equations[1]\n",
    "print \"pre order traversal\"\n",
    "equations[1].preOrder()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After parsing trigonometry axioms, we start parsing algebra axioms using python's built in compiler and function *readAxioms*. The parsed equation is shown below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Compare(Add((Const(1), Const(1))), [('==', Const(2))])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# parses text equations using the compiler package and returns an equation in the compiler format\n",
    "readAxioms(axiomPath, inputAxioms, axiomVariables)\n",
    "inputAxioms[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once we have the compiler object, we can convert it to an EquationTree object using function *buildAxiom*. An example parsed equation as well its pre order parse is given below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "an axiom: Eq(Add(1,1),2)\n",
      "pre order traversal:\n",
      "functionality: Equality\n",
      "assigned node number: 0\n",
      "subtree depth: 2\n",
      "functionality: Add\n",
      "assigned node number: 1\n",
      "subtree depth: 1\n",
      "variable name: 1\n",
      "functionality: One\n",
      "assigned node number: 3\n",
      "subtree depth: 0\n",
      "variable name: 1\n",
      "functionality: One\n",
      "assigned node number: 4\n",
      "subtree depth: 0\n",
      "variable name: 2\n",
      "functionality: Integer\n",
      "assigned node number: 2\n",
      "subtree depth: 1\n"
     ]
    }
   ],
   "source": [
    "# converting compiler object axioms to EquationTree objects and creating training labels\n",
    "for i, ax in enumerate(inputAxioms):\n",
    "    currAxiom = buildAxiom(ax, axiomVariables[i])\n",
    "    currAxiom.enumerize_queue()\n",
    "    axioms.append(currAxiom)\n",
    "    axLabels.append(mx.nd.array([1]))\n",
    "    \n",
    "print \"an axiom:\", axioms[0]\n",
    "print \"pre order traversal:\"\n",
    "axioms[0].preOrder()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "148\n",
      "148\n",
      "148\n"
     ]
    }
   ],
   "source": [
    "# appending algebra axioms to trigonometry axioms\n",
    "equations.extend(axioms)\n",
    "eqVariables.extend(axiomVariables)\n",
    "labels.extend(axLabels)\n",
    "print len(equations)\n",
    "print len(eqVariables)\n",
    "print len(labels)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The distribution of the depth of the equations, that contains all axioms from trigonometry and algebra, is shown below. The vector shows the number of equations of depth *i* at *i*th vector position."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "distribution of depth of equations\n",
      "[0, 10, 35, 39, 21, 29, 7, 6, 0, 0]\n"
     ]
    }
   ],
   "source": [
    "depthMat = [0 for _ in range(26)]\n",
    "for eq in equations:\n",
    "    depthMat[eq.depth] += 1\n",
    "print \"distribution of depth of equations\"\n",
    "print depthMat[:10]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3. Contructing the mathDictionary\n",
    "<a id=\"sec:dict\"></a>\n",
    "\n",
    "Here, we construct the *mathdictionary* which is used for generating additional correct identities. This dictionary contains (key, value) pairs that are mathematically equivalent. E.g. $(x+y : y+x)$ is a (key,value) pair in the mathDictionary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "147\n",
      "294\n"
     ]
    }
   ],
   "source": [
    "# constructing the mathDictionary whose (key,value) pairs are valid math equalities\n",
    "# e.g. (x+y : y+x) is a (key,value) pair in the mathDictionary\n",
    "# the dictionary will be updated as more correct equations are generated\n",
    "mathDictionary = {}\n",
    "strMathDictionary = {}\n",
    "for i, eq in enumerate(equations):\n",
    "    if i!=0:\n",
    "        eqCopy = copy.deepcopy(eq)\n",
    "        if str(eqCopy) not in strMathDictionary:\n",
    "            strMathDictionary[str(eqCopy)] = 1\n",
    "            mathDictionary[eqCopy.args[0]] = eqCopy.args[1]\n",
    "            mathDictionary[eqCopy.args[1]] = eqCopy.args[0]\n",
    "        else:\n",
    "            strMathDictionary[str(eqCopy)] += 1\n",
    "# for k, v in strMathDictionary.iteritems():\n",
    "#     print k, v\n",
    "print len(strMathDictionary)\n",
    "print len(mathDictionary)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4. Generating correct equations using mathDictionary lookup\n",
    "<a id=\"sec:subtreeMatching\"></a>\n",
    "\n",
    "Function *genPosEquation* generates a candidate correct equation using subtree matching from a dictionary lookup. More specifically, it chooses a random node of the equation and looks for patterns in the dictionary key that match the subtree of the chosen node. The subtree is then replaced by the dictionary key's value. \n",
    "\n",
    "The code snippet below generates about 10 correct equations by making local changes to equations selected at random from the input axioms. If no duplicate equation is generated, this equation will be added to the list of equations. The depth of the generated equation is limited to 7 (*maxDepth*). parameter *maxDepthSoFar* along with *thrsh* control the generated number of equations of a certain depth before moving to the next depth. This is a good control for training since it ensures that a minimum number of equations of each depth are present in the final dataset. It shoul dbe noted that as depth increases, the space grows exponentially, and this does not aim to cover this exponential space. \n",
    "\n",
    "The distribution of the generated equations are shown below. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "distribution of depth of equations\n",
      "[0, 10, 36, 41, 24, 29, 11, 6, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1]\n"
     ]
    }
   ],
   "source": [
    "maxDepth = 7\n",
    "numPosEq = 10\n",
    "numNegEq = 10\n",
    "numNegRepeats = 2\n",
    "thrsh = 5\n",
    "\n",
    "# set maxDepthSoFar to 0 to generate up to thrsh number of \n",
    "# repeated equations before moving to equations of higher depth\n",
    "maxDepthSoFar = 7\n",
    "totDisc = 0\n",
    "for i in range(0, numPosEq):\n",
    "    randInd = random.choice(range(1, len(equations)))\n",
    "    while labels[randInd].asnumpy() == 0:\n",
    "        randInd = random.choice(range(1, len(equations)))\n",
    "    randEq = copy.deepcopy(equations[randInd])\n",
    "    randEqVariable = copy.deepcopy(eqVariables[randInd])\n",
    "\n",
    "    posEq = genPosEquation(randEq, mathDictionary, randEqVariable)\n",
    "    posVars = posEq.extractVars()\n",
    "    posEq.enumerize_queue()\n",
    "\n",
    "    old = 0\n",
    "    disc = 0\n",
    "    tries = 0\n",
    "    # this loop is to make sure there are no repeats and also that enough \n",
    "    # number of equations of a certain depth are generated\n",
    "    while str(posEq) in strMathDictionary or posEq.depth > maxDepthSoFar:\n",
    "        if str(posEq) in strMathDictionary:\n",
    "            strMathDictionary[str(posEq)] += 1\n",
    "            old += 1\n",
    "            totDisc += 1\n",
    "        elif posEq.depth > maxDepthSoFar:\n",
    "            disc += 1\n",
    "            totDisc += 1\n",
    "\n",
    "        if old > thrsh:\n",
    "            old = 0\n",
    "            maxDepthSoFar += 1\n",
    "            print \"new max depth %d\" %(maxDepthSoFar)\n",
    "            if maxDepthSoFar > maxDepth:\n",
    "                print \"reached maximum depth\"\n",
    "                maxDepthSoFar = maxDepth\n",
    "                break\n",
    "\n",
    "        randInd = random.choice(range(1, len(equations)))\n",
    "        randEq = equations[randInd]\n",
    "        randEqVariable = copy.deepcopy(eqVariables[randInd])\n",
    "        posEq = genPosEquation(randEq, mathDictionary, randEqVariable)\n",
    "        posVars = posEq.extractVars()\n",
    "        posEq.enumerize_queue()\n",
    "\n",
    "    if posEq.depth <= maxDepth:\n",
    "        posEqCopy = copy.deepcopy(posEq)\n",
    "\n",
    "        if str(posEqCopy) not in strMathDictionary:\n",
    "            strMathDictionary[str(posEqCopy)] = 1\n",
    "            if posEqCopy.args[0] not in mathDictionary:\n",
    "                mathDictionary[posEqCopy.args[0]] = posEqCopy.args[1]\n",
    "            if posEqCopy.args[1] not in mathDictionary:\n",
    "                mathDictionary[posEqCopy.args[1]] = posEqCopy.args[0]\n",
    "\n",
    "            equations.append(posEq)\n",
    "            eqVariables.append(posVars)\n",
    "            labels.append(mx.nd.array([1]))\n",
    "    else:\n",
    "        totDisc += 1\n",
    "        print \"discarded pos equation of depth greater than %d: %s\" %(maxDepth, str(posEq))\n",
    "\n",
    "depthMat = [0 for _ in range(26)]\n",
    "for eq in equations:\n",
    "    depthMat[eq.depth] += 1\n",
    "print \"distribution of depth of equations\"\n",
    "print depthMat"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5. Generating correct and incorrect equations using local changes\n",
    "<a id=\"sec:valid\"></a>\n",
    "\n",
    "Function *genNegEquation* generated a candidate correct or incorrect mathematical equation. This function takes as input an EquationTree object. It selects a node chosen at random from that equations and performs one of the operations explained in [this section](#sec:posEq) to the node depending on the type of node. We check the correctness or incorrectness of this generated equation using function *isCorrect* that uses sympy. The depth of the final equations are shown below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "correct negative Eq: Eq(sec(Mul(1,var_0)),sec(var_0))\n",
      "correct negative Eq: Eq(tan(Add(Mul(pi,Pow(2,-1)),Mul(-1,Mul(var_0,1)))),cot(var_0))\n",
      "correct negative Eq: Eq(cot(Add(var_0,Add(pi,pi))),cot(var_0))\n",
      "correct negative Eq: Eq(Mul(0,-1),0)\n",
      "correct negative Eq: Eq(Pow(0,2),0)\n",
      "correct negative Eq: Eq(Mul(0,2),Mul(0,2/5))\n",
      "correct negative Eq: Eq(Mul(0,3),Mul(0,pi))\n",
      "correct negative Eq: Eq(Mul(Mul(0,0),1),0)\n",
      "correct negative Eq: Eq(Mul(0.7,0),0)\n",
      "correct negative Eq: Eq(Pow(1,1/2),Add(1,0))\n",
      "correct negative Eq: Eq(Pow(1,Mul(1/2,var_0)),1)\n",
      "correct negative Eq: Eq(1,1)\n",
      "correct negative Eq: Eq(1,1)\n",
      "correct negative Eq: Eq(Pow(pi,0),1)\n",
      "correct negative Eq: Eq(Pow(4,tanh(0)),1)\n",
      "correct negative Eq: Eq(tan(Add(0,0)),0)\n",
      "correct negative Eq: Eq(0,0)\n",
      "correct negative Eq: Eq(Add(Mul(0,var_0),var_0),var_0)\n",
      "correct negative Eq: Eq(asin(Mul(0,var_0)),0)\n",
      "correct negative Eq: Eq(Mul(Mul(Pow(var_1,1),var_2),var_0),Mul(var_1,Mul(var_0,var_2)))\n",
      "correct negative Eq: Eq(Mul(Add(-1,var_0),0),0)\n",
      "correct negative Eq: Eq(0,0)\n",
      "correct negative Eq: Eq(Add(var_0,var_0),Add(var_0,var_0))\n",
      "distribution of depth of neg equations\n",
      "[0, 5, 20, 45, 25, 31, 12, 8, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]\n"
     ]
    }
   ],
   "source": [
    "# generating negative equations\n",
    "negLabels= [[] for _ in range(numNegRepeats)]\n",
    "negEquations = [[] for _ in range(numNegRepeats)]\n",
    "negEqVariables = [[] for _ in range(numNegRepeats)]\n",
    "negStrMathDictionary = {}\n",
    "corrNegs = 0\n",
    "totDiscNeg = 0\n",
    "ii = len(equations)\n",
    "for i in range(1, len(equations)): \n",
    "    for rep in range(numNegRepeats):\n",
    "        randInd = i\n",
    "        randEq = copy.deepcopy(equations[i])\n",
    "        randEqVariable = copy.deepcopy(eqVariables[randInd])\n",
    "\n",
    "        negEq = genNegEquation(randEq, randEqVariable)\n",
    "        negVars = negEq.extractVars()\n",
    "        negEq.enumerize_queue()\n",
    "        disc = 0\n",
    "        tries = 0\n",
    "        old = 0\n",
    "        while str(negEq) in negStrMathDictionary or negEq.depth > maxDepth:\n",
    "            if str(negEq) in negStrMathDictionary:\n",
    "                negStrMathDictionary[str(negEq)] += 1\n",
    "                old += 1\n",
    "                totDiscNeg += 1\n",
    "                # print \"repeated neg equation\"\n",
    "            elif negEq > maxDepth:\n",
    "                # print \"equation larger than depth\"\n",
    "                disc += 1\n",
    "                totDiscNeg += 1\n",
    "\n",
    "            if old > thrsh:\n",
    "                old = 0\n",
    "                break\n",
    "\n",
    "            negEq = genNegEquation(randEq, randEqVariable)\n",
    "            negVars = negEq.extractVars()\n",
    "            negEq.enumerize_queue()\n",
    "\n",
    "        if negEq.depth <= maxDepth:\n",
    "            \n",
    "            negEqCopy = copy.deepcopy(negEq)\n",
    "            try:\n",
    "                isCorrect(negEq)\n",
    "\n",
    "                if isCorrect(negEq):\n",
    "                    corrNegs += 1\n",
    "\n",
    "                    print \"correct negative Eq:\", negEq\n",
    "\n",
    "                    if str(negEq) not in strMathDictionary:\n",
    "\n",
    "                        strMathDictionary[str(negEqCopy)] = 1\n",
    "                        if negEqCopy.args[0] not in mathDictionary:\n",
    "                            mathDictionary[negEqCopy.args[0]] = negEqCopy.args[1]\n",
    "                        if negEqCopy.args[1] not in mathDictionary:\n",
    "                            mathDictionary[negEqCopy.args[1]] = negEqCopy.args[0]\n",
    "\n",
    "                        labels.append(mx.nd.array([1]))\n",
    "                        equations.append(negEq)\n",
    "                        eqVariables.append(negVars)\n",
    "\n",
    "                elif str(negEqCopy) not in negStrMathDictionary:\n",
    "                        negStrMathDictionary[str(negEqCopy)] = 1\n",
    "\n",
    "                        negLabels[rep].append(mx.nd.array([0]))\n",
    "                        negEquations[rep].append(negEq)\n",
    "                        negEqVariables[rep].append(negVars)\n",
    "                else:\n",
    "                    totDiscNeg += 1\n",
    "\n",
    "            except:\n",
    "\n",
    "                if str(negEqCopy) not in negStrMathDictionary:\n",
    "                    negStrMathDictionary[str(negEqCopy)] = 1\n",
    "\n",
    "                    negLabels[rep].append(mx.nd.array([0]))\n",
    "                    negEquations[rep].append(negEq)\n",
    "                    negEqVariables[rep].append(negVars)\n",
    "                else:\n",
    "                    totDiscNeg += 1\n",
    "\n",
    "        else:\n",
    "            totDiscNeg += 1\n",
    "            print \"discarded neg equation of depth greater than %d: %s\" %(maxDepth, str(negEq))\n",
    "\n",
    "depthMat = [0 for _ in range(26)]\n",
    "for eq in negEquations[0]:\n",
    "    depthMat[eq.depth] += 1\n",
    "print \"distribution of depth of neg equations\"\n",
    "print depthMat"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 6. Saving the generated dataset\n",
    "<a id=\"sec:save\"></a>\n",
    "\n",
    "Finally, we would like to save the resulting dataset which will be used to train neural network models for mathematical equalities. Function *writeJson* writes this resulting dataset to a json. In order to save an EquationTree object, we use the pre-order traversal of each equation. E.g. In order to save equation $0=0$, we construct:\n",
    "\n",
    "\"equation\": {\n",
    "               * \"vars\": \",0,#,#,0,#,#\", \n",
    "               * \"numNodes\": \"3\", \n",
    "               * \"variables\": {}, \n",
    "               * \"depth\": \"1,0,#,#,0,#,#\", \n",
    "               * \"nodeNum\": \"0,1,#,#,2,#,#\", \n",
    "               * \"func\": \"Equality,Integer,#,#,Integer,#,#\"\n",
    "            },\n",
    "            \n",
    "Where # indicates a null pointer. It should be noted that the trees are binary. \n",
    "\n",
    "Finally function *catJsons* concatenates all the created jsons into a single file containing correct and incorrect equations that can be loaded for training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "data/data10_pos.json : [1, 10, 42, 49, 26, 30, 11, 6]\n",
      "data/data10_neg_rep0.json : [0, 5, 20, 45, 25, 31, 12, 8]\n",
      "data/data10_neg_rep1.json : [0, 6, 26, 40, 25, 30, 11, 7]\n"
     ]
    }
   ],
   "source": [
    "# writing equations to file\n",
    "writeJson(\"data/data%d_pos.json\"%(numPosEq), equations, ranges, eqVariables, labels, maxDepth)\n",
    "for rep in range(numNegRepeats):\n",
    "    writeJson(\"data/data%d_neg_rep%d.json\"%(numNegEq,rep), negEquations[rep], ranges, negEqVariables[rep], negLabels[rep], maxDepth)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "catJsons(['data/data%d_pos.json'%(numPosEq), 'data/data%d_neg_rep%d.json'%(numNegEq,0), 'data/data%d_neg_rep%d.json'%(numNegEq,1)],\n",
    "          'data/data%d_final.json'%(numPosEq), maxDepth=maxDepth)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}