scattering_predictor MJOLNIR.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "6c956100-a914-45b2-beba-3cdfbd26ca25",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from MJOLNIR.TasUBlibDEG import calTwoTheta, calculateBMatrix, calcCell\n",
    "from MJOLNIR.Geometry.Instrument import prediction, converterToA3A4, converterToQxQy\n",
    "from MJOLNIR.Data import Sample\n",
    "import MJOLNIR\n",
    "\n",
    "class PredictionToolManager:\n",
    "    def __init__(self):\n",
    "        self.predictionAX = None\n",
    "        self.sample = None\n",
    "        self.Efs = None\n",
    "        self.A4Instrument = None\n",
    "        self.loadEnergies()\n",
    "\n",
    "    def updateSample(self, cell, alignment1, alignment2):\n",
    "        self.sample = Sample.Sample(*cell, projectionVector1=alignment1, projectionVector2=alignment2)\n",
    "        self.calcualteHKLtoA3A4()\n",
    "\n",
    "    def getAlignment(self, alignment):\n",
    "        H, K, L, A3, A4, Ei, Ef = alignment\n",
    "        return [H, K, L, A3, A4, 0.0, 0.0, Ei, Ef]  # H, K, L, A3, A4, phi, chi, Ei, Ef\n",
    "\n",
    "    def setAlignment(self, R, alignment):\n",
    "        H, K, L, A3, A4, Ei, Ef = R\n",
    "        alignment[0] = H\n",
    "        alignment[1] = K\n",
    "        alignment[2] = L\n",
    "        alignment[3] = A3\n",
    "        alignment[4] = A4\n",
    "        alignment[5] = Ei\n",
    "        alignment[6] = Ef\n",
    "\n",
    "    def calculateA4(self, alignment):\n",
    "        cell = self.getCell()\n",
    "        r = self.getAlignment(alignment=alignment)\n",
    "        Cell = calcCell(cell)\n",
    "        B = calculateBMatrix(Cell)\n",
    "        qe = np.concatenate([r[:3], r[-2:]])\n",
    "        A4 = calTwoTheta(B=B, qe=qe, ss=-1)\n",
    "        return A4\n",
    "\n",
    "    def loadEnergies(self, instrument='CAMEA'):\n",
    "        if instrument == 'CAMEA':\n",
    "            calibPath = MJOLNIR.__CAMEANormalizationBinning1__\n",
    "            calib = np.loadtxt(calibPath, delimiter=',', skiprows=3)\n",
    "            self.A4Instrument = calib[:, -1].reshape(104, 8)\n",
    "            self.EfInstrument = calib[:, 4].reshape(104, 8)\n",
    "            self.EfInstrument[np.isclose(self.EfInstrument, 0.0)] = np.nan\n",
    "        else:\n",
    "            raise NotImplementedError('Currently only CAMEA is supported.')\n",
    "        self.Efs = np.nanmean(self.EfInstrument, axis=0)[::-1]\n",
    "\n",
    "    def validateR(self, r):\n",
    "        H, K, L, _, A4, _, _, Ei, Ef = r\n",
    "        if np.isclose(np.linalg.norm([H, K, L]), 0.0):\n",
    "            return False\n",
    "        if np.isclose(A4, 0.0):\n",
    "            return False\n",
    "        if np.any([np.isclose(E, 0.0) for E in [Ei, Ef]]):\n",
    "            return False\n",
    "        return True\n",
    "\n",
    "    def generatePrediction(self, scan, alignment1, alignment2, cell, instrument='CAMEA'):\n",
    "        A3Start, A3Stop, A3Steps, Ei, A4, points, Monitor = scan\n",
    "        r1 = np.array(self.getAlignment(alignment=alignment1))\n",
    "        r2 = np.array(self.getAlignment(alignment=alignment2))\n",
    "        sample = Sample.calculateSample(cell=cell, HKL1=r1[:3], HKL2=r2[:3], A3R1=r1[3], A3R2=r2[3], Ei=r1[-2], Ef=r1[-1])\n",
    "        ax = prediction(A3Start=A3Start, A3Stop=A3Stop, A3Steps=A3Steps, A4Positions=A4, Ei=Ei, sample=sample, points=points, instrument=instrument)\n",
    "        self.predictionAx = ax\n",
    "\n",
    "        for a in ax[:-1]:\n",
    "            a.set_ylim(*ax[-1].get_ylim())\n",
    "            a.set_xlim(*ax[-1].get_xlim())\n",
    "\n",
    "        plt.show()\n",
    "\n",
    "    def getScan(self, A3Start, A3Stop, A3Steps, Ei, A4String, Monitor):\n",
    "        A4 = self.formatA4String(A4String)\n",
    "        points = True\n",
    "        return A3Start, A3Stop, A3Steps, Ei, A4, points, Monitor\n",
    "\n",
    "    def formatA4String(self, A4String):\n",
    "        return [float(x) for x in A4String.split(',') if x]\n",
    "\n",
    "    def getCell(self, a, b, c, alpha, beta, gamma):\n",
    "        return a, b, c, alpha, beta, gamma\n",
    "\n",
    "    def calcualteHKLtoA3A4(self, Ei, Ef, H, K, L):\n",
    "        \"\"\"\n",
    "        Calculate A3 and A4 angles from HKL values and energies.\n",
    "\n",
    "        Parameters:\n",
    "        - Ei: Incident energy\n",
    "        - Ef: Final energy\n",
    "        - H, K, L: Miller indices\n",
    "\n",
    "        Returns:\n",
    "        - A3, A4: Calculated angles\n",
    "        \"\"\"\n",
    "        UB = self.sample.UB  # Get the UB matrix from the sample\n",
    "        Qx, Qy, _ = np.dot(UB, [H, K, L])  # Transform HKL to Q-space\n",
    "        A3, A4 = converterToA3A4(Qx, Qy, Ei, Ef, A3Off=0.0, A4Sign=np.sign(self.sample.plane_vector1[4]))  # Convert to angles\n",
    "        return A3, A4\n",
    "\n",
    "\n",
    "    def calcualteA3A4toHKL(self, A3, A4, Ei, Ef):\n",
    "        Qx, Qy = converterToQxQy(A3, A4, Ei, Ef)\n",
    "        H, K, L = self.sample.calculateQxQyToHKL(Qx, Qy)\n",
    "        return H, K, L\n",
    "\n",
    "    def getCalculation(self, Ei, EfIndex, H, K, L, A3, A4):\n",
    "        if EfIndex == len(self.Efs):\n",
    "            Ef = Ei\n",
    "        else:\n",
    "            Ef = self.Efs[EfIndex]\n",
    "        return Ei, Ef, H, K, L, A3, A4\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "f9fd839e",
   "metadata": {},
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "PredictionToolManager.calcualteHKLtoA3A4() missing 5 required positional arguments: 'Ei', 'Ef', 'H', 'K', and 'L'",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[12], line 43\u001b[0m\n\u001b[0;32m     40\u001b[0m     \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mPrediction generation complete.\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m     42\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;18m__name__\u001b[39m \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m__main__\u001b[39m\u001b[38;5;124m\"\u001b[39m:\n\u001b[1;32m---> 43\u001b[0m     \u001b[43mrun_prediction_tool\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
      "Cell \u001b[1;32mIn[12], line 17\u001b[0m, in \u001b[0;36mrun_prediction_tool\u001b[1;34m()\u001b[0m\n\u001b[0;32m     14\u001b[0m alignment2 \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39marray([\u001b[38;5;241m0\u001b[39m, \u001b[38;5;241m1\u001b[39m, \u001b[38;5;241m0\u001b[39m, \u001b[38;5;241m15.0\u001b[39m, \u001b[38;5;241m25.0\u001b[39m, \u001b[38;5;241m0.0\u001b[39m, \u001b[38;5;241m0.0\u001b[39m, \u001b[38;5;241m14.7\u001b[39m, \u001b[38;5;241m14.7\u001b[39m])  \u001b[38;5;66;03m# Alignment vector 2 as numpy array\u001b[39;00m\n\u001b[0;32m     16\u001b[0m \u001b[38;5;66;03m# Update the sample configuration in the manager\u001b[39;00m\n\u001b[1;32m---> 17\u001b[0m \u001b[43mmanager\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mupdateSample\u001b[49m\u001b[43m(\u001b[49m\u001b[43mcell\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43malignment1\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43malignment2\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m     19\u001b[0m \u001b[38;5;66;03m# Perform HKL to A3, A4 calculation using the refactored method\u001b[39;00m\n\u001b[0;32m     20\u001b[0m Ei \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m14.7\u001b[39m\n",
      "Cell \u001b[1;32mIn[11], line 18\u001b[0m, in \u001b[0;36mPredictionToolManager.updateSample\u001b[1;34m(self, cell, alignment1, alignment2)\u001b[0m\n\u001b[0;32m     16\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mupdateSample\u001b[39m(\u001b[38;5;28mself\u001b[39m, cell, alignment1, alignment2):\n\u001b[0;32m     17\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39msample \u001b[38;5;241m=\u001b[39m Sample\u001b[38;5;241m.\u001b[39mSample(\u001b[38;5;241m*\u001b[39mcell, projectionVector1\u001b[38;5;241m=\u001b[39malignment1, projectionVector2\u001b[38;5;241m=\u001b[39malignment2)\n\u001b[1;32m---> 18\u001b[0m     \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcalcualteHKLtoA3A4\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
      "\u001b[1;31mTypeError\u001b[0m: PredictionToolManager.calcualteHKLtoA3A4() missing 5 required positional arguments: 'Ei', 'Ef', 'H', 'K', and 'L'"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "\n",
    "# Assuming the refactored PredictionToolManager class is already defined and imported\n",
    "\n",
    "def run_prediction_tool():\n",
    "    # Initialize the prediction tool manager\n",
    "    manager = PredictionToolManager()\n",
    "    \n",
    "    # Define a sample unit cell (a, b, c, alpha, beta, gamma)\n",
    "    cell = (5.43, 5.43, 5.43, 90.0, 90.0, 90.0)  # Example cubic cell, typical for silicon\n",
    "\n",
    "    # Define two alignments with H, K, L, A3, A4, phi, chi, Ei, Ef values\n",
    "    alignment1 = np.array([1, 0, 0, 10.0, 20.0, 0.0, 0.0, 14.7, 14.7])  # Alignment vector 1 as numpy array\n",
    "    alignment2 = np.array([0, 1, 0, 15.0, 25.0, 0.0, 0.0, 14.7, 14.7])  # Alignment vector 2 as numpy array\n",
    "\n",
    "    # Update the sample configuration in the manager\n",
    "    manager.updateSample(cell, alignment1, alignment2)\n",
    "\n",
    "    # Perform HKL to A3, A4 calculation using the refactored method\n",
    "    Ei = 14.7\n",
    "    Ef = 14.7  # Assuming Ei = Ef for this example\n",
    "    H, K, L = 1, 0, 0  # Example HKL values\n",
    "    A3, A4 = manager.calcualteHKLtoA3A4(Ei, Ef, H, K, L)\n",
    "    print(f\"Calculated A3: {A3}, A4: {A4}\")\n",
    "\n",
    "    # Define scan parameters: A3 start/stop, number of steps, incident energy, A4 positions, and monitor value\n",
    "    A3Start = 0.0  # Starting angle for A3\n",
    "    A3Stop = 60.0  # Ending angle for A3\n",
    "    A3Steps = 10   # Number of steps in the scan\n",
    "    Ei = 14.7      # Incident energy in meV\n",
    "    A4String = \"20.0,30.0,40.0\"  # Example A4 positions\n",
    "    Monitor = 1000  # Monitor counts\n",
    "\n",
    "    # Get the formatted scan parameters tuple\n",
    "    scan = manager.getScan(A3Start, A3Stop, A3Steps, Ei, A4String, Monitor)\n",
    "    \n",
    "    # Generate the prediction using the provided scan parameters and sample configuration\n",
    "    manager.generatePrediction(scan, alignment1, alignment2, cell, instrument='CAMEA')\n",
    "\n",
    "    print(\"Prediction generation complete.\")\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    run_prediction_tool()\n"
   ]
  }
 ],
 "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}