c5-optimization.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Optimization of instrument designs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from definitions import *\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import os \n",
    "from util import *\n",
    "from instrument import *\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "target_interval = (10e-9, 5000e-9)\n",
    "C_log = 0.01\n",
    "C_linear = 0.00001\n",
    "\n",
    "# L_max = 20.0 # m\n",
    "d_1 = 0.3 # m\n",
    "d_2 = 0.3 # m\n",
    "epsilon = 0.01 # m\n",
    "def normalize_solution(instance, L_s_min):    \n",
    "    # Swap if the precession component lengths are the wrong way around\n",
    "    if instance[2] < instance[3]:\n",
    "        t = instance[2]\n",
    "        instance[2] = instance[3]\n",
    "        instance[3] = t\n",
    "    min_dist = (d_1 + d_2) / 2 + epsilon\n",
    "    prec_dist = instance[2] - instance[3]\n",
    "    if prec_dist <= min_dist:\n",
    "        delta = min_dist - prec_dist\n",
    "        instance[2] += delta/2\n",
    "        instance[3] -= delta/2\n",
    "    # Limit the sample position to be straight after the second precession device with some margin given by epsilon\n",
    "\n",
    "    # First push L_s as far as possible as it is constrained\n",
    "    if instance[1] >= instance[3] - d_2 / 2 - epsilon:\n",
    "        instance[1] = max(L_s_min, instance[3] - d_2 / 2 - epsilon)\n",
    "    # instance[1] = min(instance[1], instance[3] - d_2 / 2 - epsilon)\n",
    "\n",
    "    # Next, push away L_2 as far as still needed\n",
    "    instance[3] = max(instance[3], instance[1] + d_2 / 2 + epsilon)\n",
    "    for i in range(len(instance)):\n",
    "        instance[i] = round(instance[i], 4)\n",
    "\n",
    "    return instance\n",
    "\n",
    "def optimize_instrument(type='wsp', PG=True, L_max = 10, detector_size = detector_size, theta_a_max = 15e-3, N=100000):\n",
    "    L_min = round(detector_size / (2 * np.tan(theta_a_max)),4)\n",
    "    # If a PG is used, change permitted L0 range\n",
    "    if PG:\n",
    "        quality = 0.01\n",
    "        # Parameters are L0, L_s, L_1, L_2\n",
    "        param_space = [{'low': 3.0, 'high': 4.4}, {'low': L_min, 'high': L_max - (d_1 + d_2)}, {'low': L_min + d_1 / 2 + d_2, 'high': L_max - d_1/2}, {'low': L_min + d_2/2, 'high': L_max - (d_1 + d_2/2)}]\n",
    "        monochrom_name = 'PG'\n",
    "    else:\n",
    "        quality = 0.1\n",
    "        param_space = [{'low': 8.0, 'high': 12.0}, {'low': L_min, 'high': L_max - (d_1 + d_2)}, {'low': L_min + d_1 / 2 + d_2, 'high': L_max - d_1/2}, {'low': L_min + d_2/2, 'high': L_max - (d_1 + d_2/2)}]\n",
    "        monochrom_name = 'VS'\n",
    "\n",
    "    def instrument_from_solution(solution, type='wsp'):\n",
    "        L0 = solution[0] * 1e-10\n",
    "        L_s = solution[1]\n",
    "        L_1 = solution[2]\n",
    "        L_2 = solution[3]\n",
    "        match type:\n",
    "            case 'wsp':\n",
    "                theta_0 = np.deg2rad(45)\n",
    "                By_min = 0.1e-3\n",
    "                By_max = 63e-3\n",
    "            case 'iso':\n",
    "                theta_0 = np.deg2rad(45)\n",
    "                By_min = 0.1e-3\n",
    "                By_max = 15e-3\n",
    "            case 'foil':\n",
    "                theta_0 = tune_foil(L0)\n",
    "                By_min = 0.3e-3\n",
    "                By_max = 30e-3\n",
    "            case 'foil2':\n",
    "                theta_0 = tune_foil(L0)\n",
    "                By_min = 0.3e-3\n",
    "                By_max = 150e-3\n",
    "        instr = Instrument('', '', type, L0, quality * L0 / FWHM_factor, theta_0, By_min, By_max, L_s, L_1, L_2, detector_size)\n",
    "        return instr\n",
    "\n",
    "    def fitness_func(solution):\n",
    "        instr = instrument_from_solution(solution, type)\n",
    "        delta_range = instr.delta_range()\n",
    "        fitness = log_overlap_percentage(delta_range, target_interval) * C_log + overlap_percentage(delta_range, target_interval) * C_linear\n",
    "        return fitness\n",
    "    \n",
    "    def generate_normalized_solutions(N):\n",
    "        N_params = 4\n",
    "        population = np.zeros((N, N_params))\n",
    "        fitnesses = np.zeros(N)\n",
    "        for i, param in enumerate(param_space):\n",
    "            population[:, i] = np.random.uniform(low=param['low'], high=param['high'], size=N)\n",
    "        for j in range(0,N):\n",
    "            population[j,:] = normalize_solution(population[j,:], L_min)\n",
    "            fitnesses[j] = fitness_func(population[j,:])\n",
    "        return population, fitnesses\n",
    "    \n",
    "    def choose_best(population, fitnesses):\n",
    "        best_id = np.argmax(fitnesses)\n",
    "        best_sol = population[best_id, :]\n",
    "        best_instr = instrument_from_solution(best_sol, type)\n",
    "        fitness = fitnesses[best_id]\n",
    "        return best_instr, fitness\n",
    "\n",
    "    def brute_force(N = N):\n",
    "        population, fitnesses = generate_normalized_solutions(N)\n",
    "        best_instr, fitness = choose_best(population, fitnesses)\n",
    "        return best_instr, fitness\n",
    "\n",
    "    def boltzmann_P(E,T):\n",
    "        if E > 0:\n",
    "            return True\n",
    "        else:\n",
    "            return np.random.rand() < np.exp(E/T)\n",
    "    \n",
    "    def annealing(N = 10):    \n",
    "        population, fitnesses = generate_normalized_solutions(N)\n",
    "        # Until this point this is just like the brute force method\n",
    "        T_max = 100\n",
    "        T_min = 0.0000001\n",
    "        cool_factor = 0.999\n",
    "        T = T_max\n",
    "        while T > T_min:\n",
    "            candidate_population, candidate_fitnesses = generate_normalized_solutions(N)\n",
    "            for j in range(0,N):\n",
    "                delta_fitness = candidate_fitnesses[j] - fitnesses[j]\n",
    "                if boltzmann_P(delta_fitness, T):\n",
    "                    population[j] = candidate_population[j]\n",
    "                    fitnesses[j] = candidate_fitnesses[j]\n",
    "            T *= cool_factor\n",
    "        best_instr, fitness = choose_best(population, fitnesses)\n",
    "        return best_instr, fitness\n",
    "    \n",
    "    best_instr, fitness = annealing()\n",
    "    best_instr.name = f'{type.upper()} {monochrom_name}'\n",
    "    return best_instr, fitness"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def optimize_all_designs(detector_size=10e-3, theta_a_max = 15e-3):\n",
    "    instrs = [] \n",
    "    for type in ['foil', 'wsp', 'iso','foil2',]:\n",
    "        print(f\"==========Best {type} instrument compatible with pyroletic graphite monochromator==========\")\n",
    "        pg_instr, fitness = optimize_instrument(type=type, PG = True, L_max=5.0, detector_size=detector_size, theta_a_max = theta_a_max)\n",
    "        print(fitness)\n",
    "        print(str(pg_instr))\n",
    "        instrs.append((pg_instr, fitness))\n",
    "    for type in ['foil', 'wsp', 'iso','foil2']:\n",
    "        print(f\"==========Best {type} instrument compatible with velocity selector monochromator==========\")\n",
    "        vs_instr, fitness = optimize_instrument(type=type, PG = False, L_max=5.0, detector_size=detector_size, theta_a_max = theta_a_max)\n",
    "        print(fitness)\n",
    "        print(str(vs_instr))\n",
    "        instrs.append((vs_instr,fitness))\n",
    "    return instrs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "instrs = optimize_all_designs()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "detector_size = 30e-3\n",
    "instrs_2 = optimize_all_designs(detector_size)\n",
    "for instr in instrs_2:\n",
    "    instr[0].name+= ' B'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import pickle\n",
    "# with open('objects/optimized_instruments_1.pkl', 'wb') as file:\n",
    "#     pickle.dump(instrs, file)\n",
    "# with open('objects/optimized_instruments_2.pkl', 'wb') as file:\n",
    "#     pickle.dump(instrs_2, file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Pretty print $\\LaTeX$ tables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def print_optimized_design_tables(instrs):\n",
    "    for i in [0,1,2]:\n",
    "        print(f\"====TABLE {i}======\")\n",
    "        for (instr, fitness) in instrs:\n",
    "            # print(instr)\n",
    "            delta_max_field = instr.delta_max_B_field()\n",
    "            delta_max_env = instr.delta_max_envelope()\n",
    "            delta_max_ten_samples = instr.delta_max_sampling()\n",
    "            # print(F\"Max delta ideal sampling (10 samples per period) (f_0 = {round(f_ten_samples*1e-3)}mm^-1: {round(delta_max_ten_samples * 1e9,2)}nm\")\n",
    "            maxes = [(delta_max_ten_samples, 'sampling'), (delta_max_env, 'envelope'), (delta_max_field, 'precession devices') ]\n",
    "            (delta_max, max_name) = min(maxes)\n",
    "            delta_min_field = instr.delta_min_B_field()\n",
    "            delta_min_single_period = instr.delta_min_detector()\n",
    "            mins = [(delta_min_single_period, 'detector size'), (delta_min_field, 'precession devices')]\n",
    "            (delta_min, min_name) = max(mins)\n",
    "            # print(min_max)\n",
    "            Q_max = instr.Q_max()\n",
    "            L_s_min, L_s_max = instr.L_s_range()\n",
    "            delta_min_var, delta_max_var = instr.delta_range_var_L_s()\n",
    "\n",
    "            # print(f\"{a}\")\n",
    "            num = lambda x: r\"\\num{\" + f\"{x:#.3g}\" + \"}\"\n",
    "            bold = lambda x: r\"\\textbf{\" + f\"{x:#.3g}\" + \"}\"\n",
    "\n",
    "            r = lambda x: num(round(x * 1e9,1)) \n",
    "            r_bold = lambda x: bold(round(x * 1e9,1)) \n",
    "            r2 = lambda x: num(round(x * 1e6,2)) \n",
    "            r2_bold = lambda x: bold(round(x * 1e6,2)) \n",
    "            def r_dynamic(x):\n",
    "                if round(x * 1e9,1) == round(delta_min * 1e9,1):\n",
    "                    return r_bold(x)\n",
    "                else:\n",
    "                    return r(x)\n",
    "            def r2_dynamic(x):\n",
    "                if round(x * 1e6,2) == round(delta_max * 1e6,2):\n",
    "                    return r2_bold(x)\n",
    "                else:\n",
    "                    return r2(x)\n",
    "            if i==0:\n",
    "                # label, L0, DL, L1, L2, LS\n",
    "                print(f\"{instr.name} & {num(round(fitness, 3))} & {num(round(instr.L0 * 1e10,2))} & {num(round(instr.DL * FWHM_factor * 1e10,3))} & {num(round(instr.L_1, 3))} & {num(round(instr.L_2, 3))} & {num(round(instr.L_s, 3))} & {num(round(L_s_min, 3))} & {num(round(L_s_max, 3))} \\\\\\\\\")\n",
    "            elif i==1:\n",
    "                print(f\"{instr.name} & {num(round(Q_max * 1e-10, 5))} & {num(round(instr.theta_a() * 1e3,2))} & {r(delta_min)} & {r2(delta_max)} & {r(delta_min_var)} & {r2(delta_max_var)} \\\\\\\\\")\n",
    "            else:\n",
    "                print(f\"{instr.name} & {r_dynamic(delta_min_single_period)} & {r_dynamic(delta_min_field)} & {r2_dynamic(delta_max_ten_samples)} & {r2_dynamic(delta_max_env)} & {r2_dynamic(delta_max_field)} \\\\\\\\\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def seperate_FOIL2(instrs):\n",
    "    has_foil2 = lambda x: \"FOIL2\" in x[0].name\n",
    "    return list(filter(has_foil2, instrs)), list(filter(lambda x: not has_foil2(x), instrs))\n",
    "instr_1_foil2, instr_1_normal = seperate_FOIL2(instrs)\n",
    "print_optimized_design_tables(instr_1_normal)\n",
    "print_optimized_design_tables(instr_1_foil2)\n",
    "# print(seperate_FOIL2(instrs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "instr_2_foil2, instr_2_normal = seperate_FOIL2(instrs_2)\n",
    "print_optimized_design_tables(instr_2_normal)\n",
    "print_optimized_design_tables(instr_2_foil2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def print_latex_tables(instrs):\n",
    "    for i in [0,1]:\n",
    "        print(f\"====TABLE {i}======\")\n",
    "        for instr in instrs[3:-6]:\n",
    "            delta_max_field = instr.delta_max_B_field()\n",
    "            delta_max_env = instr.delta_max_envelope()\n",
    "            delta_max_ten_samples = instr.delta_max_sampling()\n",
    "            # print(F\"Max delta ideal sampling (10 samples per period) (f_0 = {round(f_ten_samples*1e-3)}mm^-1: {round(delta_max_ten_samples * 1e9,2)}nm\")\n",
    "            maxes = [(delta_max_ten_samples, 'sampling'), (delta_max_env, 'envelope'), (delta_max_field, 'precession devices') ]\n",
    "            (delta_max, max_name) = min(maxes)\n",
    "            delta_min_field = instr.delta_min_B_field()\n",
    "            delta_min_single_period = instr.delta_min_detector()\n",
    "            mins = [(delta_min_single_period, 'detector size'), (delta_min_field, 'precession devices')]\n",
    "            (delta_min, min_name) = max(mins)\n",
    "            # print(min_max)\n",
    "            Q_max = instr.Q_max()\n",
    "            L_s_min, L_s_max = instr.L_s_range()\n",
    "            delta_min_var, delta_max_var = instr.delta_range_var_L_s()\n",
    "\n",
    "            # print(f\"{a}\")\n",
    "            num = lambda x: r\"\\num{\" + f\"{x:#.3g}\" + \"}\"\n",
    "            bold = lambda x: r\"\\textbf{\" + f\"{x:#.3g}\" + \"}\"\n",
    "\n",
    "            r = lambda x: num(round(x * 1e9,1)) \n",
    "            r_bold = lambda x: bold(round(x * 1e9,1)) \n",
    "            r2 = lambda x: num(round(x * 1e6,2)) \n",
    "            r2_bold = lambda x: bold(round(x * 1e6,2))\n",
    "            def r_dynamic(x):\n",
    "                if round(x * 1e9,1) == round(delta_min * 1e9,1):\n",
    "                    return r_bold(x)\n",
    "                else:\n",
    "                    return r(x)\n",
    "            def r2_dynamic(x):\n",
    "                if round(x * 1e6,2) == round(delta_max * 1e6,2):\n",
    "                    return r2_bold(x)\n",
    "                else:\n",
    "                    return r2(x)\n",
    "            if i==0:\n",
    "                print(f\"{instr.name} & {r_dynamic(delta_min_single_period)} & {r_dynamic(delta_min_field)} & {r2_dynamic(delta_max_ten_samples)} & {r2_dynamic(delta_max_env)} & {r2_dynamic(delta_max_field)} \\\\\\\\\")\n",
    "            elif i==1:\n",
    "                print(f\"{instr.name} & {num(round(Q_max * 1e-10, 5))} & {num(round(instr.theta_a() * 1e3,2))} & {r(delta_min)} & {r2(delta_max)} & {r(delta_min_var)} & {r2(delta_max_var)} \\\\\\\\\")\n",
    "if __name__ == '__main__':\n",
    "    import util\n",
    "    instrs_sheet = util.load_instruments('instruments.csv')\n",
    "    for instr in instrs_sheet:\n",
    "        print(instr)\n",
    "    print_latex_tables(instrs_sheet)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Optimization convergence statistics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fitness_statistics(precession_type = 'wsp', PG = False, N_lim = 10000):\n",
    "    fitness_sample = np.zeros((16,50))\n",
    "    N_range = np.round(np.logspace(1, np.log10(N_lim), 50))\n",
    "    for (i,N) in enumerate(N_range):\n",
    "        N = int(N)\n",
    "        for j in range(16):\n",
    "            theta_a_max = 15e-3\n",
    "            _, fitness = optimize_instrument(type=precession_type, PG = PG, L_max=5.0, detector_size=detector_size, theta_a_max = theta_a_max, N = N)\n",
    "            fitness_sample[j,i] = fitness\n",
    "    fitness_mean = np.mean(fitness_sample, axis=0)\n",
    "    fitness_std = np.std(fitness_sample, axis=0, ddof=1)\n",
    "    return N_range, fitness_mean, fitness_std"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fit_stat_map = {}\n",
    "for prec_type in ['foil', 'wsp', 'iso','foil2',]:\n",
    "    fit_stat_map[(prec_type, 'PG')] = fitness_statistics(prec_type,True, 100000)\n",
    "    fit_stat_map[(prec_type, 'VS')] = fitness_statistics(prec_type,False, 100000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import pickle\n",
    "# with open('optimized_instruments_1.pkl', 'wb') as file:\n",
    "#     pickle.dump(fit_stat_map, file)\n",
    "# with open('optimized_instruments_2.pkl', 'wb') as file:\n",
    "#     pickle.dump(instrs_2, file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_with_error_shaded(x,y,y_std,label):\n",
    "    plt.plot(x, y, '.', label=label)\n",
    "    plt.fill_between(x, y-y_std, y+y_std, alpha = 0.5)\n",
    "\n",
    "for mc_type in ['PG', 'VS']:\n",
    "    for prec_type in ['iso','foil', 'wsp','foil2',]:\n",
    "        N_range, fitness_mean, fitness_std = fit_stat_map[(prec_type,mc_type)]\n",
    "        N_start_ix = 25\n",
    "        plot_with_error_shaded(N_range[N_start_ix:], fitness_mean[N_start_ix:], fitness_std[N_start_ix:], prec_type)\n",
    "    plt.legend()\n",
    "    plt.xlabel('N')\n",
    "    plt.ylabel(r'$\\bar{g}_3$')\n",
    "    plt.title(r'$\\bar{g}_3$ for' + f' {mc_type} instruments, average over 16 runs')\n",
    "    plt.grid()\n",
    "    plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "optimal_min = []\n",
    "optimal_max = []\n",
    "optimal_fit = []\n",
    "\n",
    "num_samples = 10  # Number of samples to average\n",
    "L_maxes = np.linspace(3.0, 9.0, 3)\n",
    "for L_max in L_maxes:\n",
    "    delta_ranges = []\n",
    "    fitnesses = []\n",
    "    for _ in range(num_samples):\n",
    "        pg_instr, fitness = optimize_instrument(PG=False, L_max=L_max, N=1000)\n",
    "        delta_range = np.array(pg_instr.delta_range())\n",
    "        delta_ranges.append(delta_range)\n",
    "        fitnesses.append(fitness)\n",
    "    \n",
    "    # Calculate the average delta_range over all samples\n",
    "    avg = np.average(delta_ranges, axis=0)\n",
    "    avg_fit = np.average(fitnesses, axis=0)\n",
    "    \n",
    "    optimal_min.append(avg[0])\n",
    "    optimal_max.append(avg[1])\n",
    "    optimal_fit.append(avg_fit)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.title(r'$\\delta_{max}$ as function of instrument length')\n",
    "plt.plot(L_maxes, np.array(optimal_fit), '.', label='Wollaston prism')\n",
    "plt.xlabel(r'$L_{1,max}$ [m]')\n",
    "plt.ylabel(r'$\\delta_{max}$ [nm]')\n",
    "plt.legend()\n",
    "plt.grid()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "delta_range = pg_instr.delta_range()\n",
    "\n",
    "fitness = log_overlap_percentage(delta_range, target_interval) \n",
    "np.log(delta_range), np.log(target_interval)\n",
    "print(fitness)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "bep",
   "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": 2
}