{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import hfda\n",
    "import numpy as np\n",
    "import scipy.integrate as si\n",
    "import scipy.interpolate as so\n",
    "import scipy.constants as sc\n",
    "import scipy.stats as ss\n",
    "import matplotlib.pyplot as plt\n",
    "import sympy as sm\n",
    "import plotly.graph_objs as go\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "import pandas as pd\n",
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def ring(radius, location):\n",
    "    alpha, beta, gamma, x0, y0, z0 = location\n",
    "    phi = np.linspace(0, 2*np.pi, 100)\n",
    "    xring = np.cos(alpha)*np.cos(beta)*np.cos(phi) + (np.cos(alpha)*np.sin(beta)*np.sin(gamma)-np.sin(alpha)*np.cos(gamma))*np.sin(phi)\n",
    "    yring = np.sin(alpha)*np.cos(beta)*np.cos(phi) + (np.cos(alpha)*np.cos(gamma)+np.sin(alpha)*np.sin(beta)*np.sin(gamma))*np.sin(phi)\n",
    "    zring = -1*np.sin(beta)*np.cos(phi) + np.cos(beta)*np.sin(gamma)*np.sin(phi)\n",
    "    xring *= radius\n",
    "    yring *= radius\n",
    "    zring *= radius\n",
    "    return (xring+x0, yring+y0, zring+z0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rotation(angles):\n",
    "    alpha, beta, gamma, x0, y0, z0 = angles\n",
    "    matrix = np.asarray([[np.cos(alpha)*np.cos(beta), np.cos(alpha)*np.sin(beta)*np.sin(gamma)-np.sin(alpha)*np.cos(gamma), np.sin(alpha)*np.sin(gamma)+np.cos(alpha)*np.sin(beta)*np.cos(gamma)],\n",
    "                      [np.sin(alpha)*np.cos(beta), np.cos(alpha)*np.cos(gamma)+np.sin(alpha)*np.sin(beta)*np.sin(gamma), np.sin(alpha)*np.sin(beta)*np.cos(gamma)-np.cos(alpha)*np.sin(gamma)],\n",
    "                      [-1*np.sin(beta), np.cos(beta)*np.sin(gamma), np.cos(beta)*np.cos(gamma)]])\n",
    "    return matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "#magnetic field, no interpolation\n",
    "def rkmagfield(radius, current):\n",
    "    t, x, y, z = sm.symbols('t, x, y, z')\n",
    "    l = sm.Matrix([radius*sm.cos(t), radius*sm.sin(t), 0])\n",
    "    r = sm.Matrix([x,y,z])\n",
    "    rprime = r-l\n",
    "    rphat = rprime/rprime.norm()\n",
    "    inside = sm.diff(l,t).cross(rphat)/(rprime.norm()**2)\n",
    "    intx = sm.lambdify([t,x,y,z], inside[0])\n",
    "    inty = sm.lambdify([t,x,y,z], inside[1])\n",
    "    intz = sm.lambdify([t,x,y,z], inside[2])\n",
    "    field = lambda point: ((sc.mu_0*current)/(4*np.pi)) * np.array([si.quad(intx,0,2*np.pi,args=(point[0], point[1], point[2]))[0], si.quad(inty,0,2*np.pi,args=(point[0], point[1], point[2]))[0], si.quad(intz,0,2*np.pi,args=(point[0], point[1], point[2]))[0]])\n",
    "    return field"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "#interpolation evaluator for rk step\n",
    "def newnetdir(func, rings, point, back):\n",
    "    total = np.asarray(back)\n",
    "    for p in rings:\n",
    "        orgpoint = point - np.asarray(p)[3:6]\n",
    "        R = rotation(p)\n",
    "        orgpoint = np.matmul(R.T, orgpoint.T).T\n",
    "        field = np.asarray(func(orgpoint))\n",
    "        if np.any(np.isnan(field)) == True:\n",
    "            return False, False\n",
    "        field = np.matmul(R, field).flatten()\n",
    "        total = total + field\n",
    "    norm = np.sqrt(total[0]**2 + total[1]**2 + total[2]**2)\n",
    "    direction = total/norm\n",
    "    return direction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "#runge-kutta 4 step\n",
    "def steprk(func, rings, point, h, back):\n",
    "    original = point\n",
    "    dir1 = newnetdir(func, rings, point, back)\n",
    "    if type(dir1) == bool:\n",
    "        return False\n",
    "    k1 = h*dir1\n",
    "    point = original + k1/2\n",
    "    dir2 = newnetdir(func, rings, point, back)\n",
    "    if type(dir2) == bool:\n",
    "        return False\n",
    "    k2 = h*dir2\n",
    "    point = original + k2/2\n",
    "    dir3 = newnetdir(func, rings, point, back)\n",
    "    if type(dir3) == bool:\n",
    "        return False\n",
    "    k3 = h*dir3\n",
    "    point = original + k3\n",
    "    dir4 = newnetdir(func, rings, point, back)\n",
    "    if type(dir4) == bool:\n",
    "        return False\n",
    "    k4 = h*dir4\n",
    "    nextpoint = original + k1/6 + k2/3 + k3/3 + k4/6\n",
    "    return nextpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "#RK step closed field line\n",
    "def rkfieldline(function, rings, startpoint, stepsize, limit=2e5):\n",
    "    nowpoint = startpoint\n",
    "    newpoint = np.asarray([0,0,-1])\n",
    "    line = startpoint\n",
    "    iteration = 0\n",
    "    difference = 1\n",
    "    dif2 = -1\n",
    "    dif1 = -1\n",
    "    #while iteration < limit and difference >= stepsize:\n",
    "    closeapp = False\n",
    "    while closeapp == False and iteration < limit:\n",
    "        if (iteration % 10000) == 0:\n",
    "            print(iteration)\n",
    "        newpoint = steprk(function, rings, nowpoint, stepsize, [0,0,0])\n",
    "        if type(newpoint) == bool:\n",
    "            if newpoint == False:\n",
    "                return line, nowpoint, -1, iteration\n",
    "        line = np.vstack([line, newpoint])\n",
    "        lastpoint = nowpoint\n",
    "        nowpoint = newpoint\n",
    "        dif2 = dif1\n",
    "        dif1 = difference\n",
    "        difference = np.sqrt((line[0,0]-nowpoint[0])**2 + (line[0,1]-nowpoint[1])**2 + (line[0,2]-nowpoint[2])**2)\n",
    "        if ((dif2 > dif1) and (difference > dif1) and (iteration > 2)):\n",
    "            nowpoint = lastpoint\n",
    "            difference = dif1\n",
    "            closeapp = True\n",
    "            break\n",
    "        iteration += 1\n",
    "    diff = nowpoint - line[0, ::]\n",
    "    return line, diff, difference, iteration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "#RK step field line full\n",
    "def rktrace(function, rings, startpoint, stepsize, limit=2e5, background = [0,0,0]):\n",
    "    nowpoint = startpoint\n",
    "    newpoint = np.asarray([0,0,-1])\n",
    "    line = startpoint\n",
    "    iteration = 0\n",
    "    difference = 1\n",
    "    while iteration < limit:\n",
    "        if (iteration % 1000) == 0:\n",
    "            print(iteration)\n",
    "        newpoint = steprk(function, rings, nowpoint, stepsize, background)\n",
    "        if type(newpoint) == bool:\n",
    "            if newpoint == False:\n",
    "                return line, nowpoint, -1, iteration\n",
    "        line = np.vstack([line, newpoint])\n",
    "        lastpoint = nowpoint\n",
    "        nowpoint = newpoint\n",
    "        difference = np.sqrt((line[0,0]-nowpoint[0])**2 + (line[0,1]-nowpoint[1])**2 + (line[0,2]-nowpoint[2])**2)\n",
    "        iteration += 1\n",
    "    diff = nowpoint - line[0, ::]\n",
    "    return line, diff, difference, iteration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "R = 1\n",
    "I = 10\n",
    "p1 = [0,0,0,0,0,0]\n",
    "p2 = [0,0,np.pi/2,3,0,0]\n",
    "p3 = [0,0,np.pi/2,1,0,0]\n",
    "loop1 = ring(R, p1)\n",
    "loop2 = ring(R, p2)\n",
    "loop3 = ring(R, p3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def lyapunov(traj1, traj2):\n",
    "    x1 = traj1[::,0]\n",
    "    y1 = traj1[::,1]\n",
    "    z1 = traj1[::,2]\n",
    "    x2 = traj2[::,0]\n",
    "    y2 = traj2[::,1]\n",
    "    z2 = traj2[::,2]\n",
    "    diffx, diffy, diffz = x1-x2, y1-y2, z1-z2\n",
    "    sep = np.sqrt(np.power(diffx,2) + np.power(diffy,2) + np.power(diffz,2))\n",
    "    isep = sep[0]\n",
    "    lsep = np.log(sep)\n",
    "    iters = np.arange(len(sep)) + 1\n",
    "    lia = np.log(sep/isep)/iters\n",
    "    return sep, iters, lia, lsep"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def quants(s, N=1e4):\n",
    "    hfd = []\n",
    "    kmaxmax = int(np.floor(N/2))\n",
    "    for k in [2,kmaxmax]:\n",
    "        h = hfda.measure(s,k)\n",
    "        hfd.append(h)\n",
    "        print(k, h)\n",
    "    hfdev = h - hfd[0]\n",
    "    temp=0\n",
    "    var = s[1]-s[0]\n",
    "    count=0\n",
    "    for i in range(1,len(s)):\n",
    "        prev = var\n",
    "        var = s[i]-s[i-1]\n",
    "        psign = prev/np.abs(prev)\n",
    "        sign = var/np.abs(var)\n",
    "        if psign != sign:\n",
    "            count += 1\n",
    "        temp = temp+np.abs(var)\n",
    "    temp = temp/(np.amax(s)-np.amin(s))\n",
    "    print(temp,count)\n",
    "    return hfdev, temp, count"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def strength(func, rings, point, back):\n",
    "    total = np.asarray(back)\n",
    "    for p in rings:\n",
    "        orgpoint = point - np.asarray(p)[3:6]\n",
    "        R = rotation(p)\n",
    "        orgpoint = np.matmul(R.T, orgpoint.T).T\n",
    "        field = np.asarray(func(orgpoint))\n",
    "        if np.any(np.isnan(field)) == True:\n",
    "            return False, False\n",
    "        field = np.matmul(R, field).flatten()\n",
    "        total = total + field\n",
    "    norm = np.sqrt(total[0]**2 + total[1]**2 + total[2]**2)\n",
    "    return norm"
   ]
  }
 ],
 "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.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}