{
  "cells": [
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "# %%\n",
        "'''\n",
        "__WELCOME__ \n",
        "\n",
        "Welcome to a cosma modeling script Python script, which illustrates how to load a strong lens dataset and analyse it on cosma.\n",
        "\n",
        "This example illustrates how to fit a single dataset with a parallelized Nautilus model-fit. You should\n",
        "only read this example after reading and understanding this example.\n",
        "\n",
        "All aspects of this script which are explained in `example_0.py`, for example setting up the cosma dataset and output\n",
        "directories, are not rexplained in this script. Therefore, if anything does not make sence refer back to `example_0.py`\n",
        "for an examplanation.\n",
        "'''"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "\n",
        "# %%\n",
        "# %%\n",
        "'''\n",
        "__COSMA PATHS SETUP__\n",
        "\n",
        "All of the code below is a repeat of `example_0.py`\n",
        "'''"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "\n",
        "from os import path\n",
        "\n",
        "cosma_path = path.join(path.sep, \"cosma7\", \"data\", \"dp004\", \"cosma_username\")\n",
        "\n",
        "dataset_folder = \"example\"\n",
        "dataset_name = \"simple__no_lens_light\"\n",
        "\n",
        "cosma_dataset_path = path.join(cosma_path, \"dataset\", dataset_folder, dataset_name)\n",
        "\n",
        "cosma_output_path = path.join(cosma_path, \"output\")\n",
        "\n",
        "workspace_path = \"/cosma/home/dp004/cosma_username/autolens_workspace/\"\n",
        "\n",
        "config_path = path.join(workspace_path, \"cosma\", \"config\")\n",
        "\n",
        "from autoconf import conf\n",
        "\n",
        "conf.instance.push(new_path=config_path, output_path=cosma_output_path)"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Cosma submissions require a`batch script`, which tells Cosma the PyAutoLens runners you want it to execute and \n",
        "distributes them to nodes and CPUs. \n",
        "\n",
        "In this previosu example, the batch script ran a multi-program submission which set off many jobs on single CPUs. \n",
        "\n",
        "By inspecting the batch script `autolens_workspace/misc/hpc/cosma/batch/example_1` one can see that only the last line \n",
        "has changed, from: \n",
        "\n",
        "    srun -n 16 --multi-prog conf/example.conf\n",
        "    \n",
        "Too:\n",
        "\n",
        "    python3 /cosma/home/dp004/cosma_username/autolens_workspace/cosma/runners/example.py 1\n",
        "\n",
        "This is straight forward to understand, instead of calling a `.conf` file and passing many `python3` commands to set\n",
        "off multiply jobs we now simply set off a single `python3` command in the batch script. As a result, the batch script\n",
        "`example_1` has no corresponding `example_1.conf` file.\n",
        "    \n",
        "We still pass the integer on the right which is used  to load a specific dataset. This is somewhat optional, but it is\n",
        "benefitial for scripts which perform single-CPU fits or multi-CPU Nautilus fits to use the same code to load\n",
        "data."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "import sys\n",
        "\n",
        "cosma_id = int(sys.argv[1])"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "There is only one more change to the modeling script script that is necessary, which we explain below.\n",
        "\n",
        "All remaining code is repetition of `example_0.py`."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "\n",
        "dataset_type = \"imaging\"\n",
        "pixel_scales = 0.1\n",
        "\n",
        "dataset_name = []\n",
        "dataset_name.append(\"example_image_1\")  # Index 0\n",
        "dataset_name.append(\"example_image_2\")  # Index 1\n",
        "dataset_name.append(\"example_image_3\")  # Index 2\n",
        "dataset_name.append(\"example_image_4\")  # Index 3\n",
        "dataset_name.append(\"example_image_5\")  # Index 4\n",
        "dataset_name.append(\"example_image_6\")  # Index 5\n",
        "dataset_name.append(\"example_image_7\")  # Index 6\n",
        "dataset_name.append(\"example_image_8\")  # Index 7\n",
        "# ...and so on.\n",
        "\n",
        "dataset_name = dataset_name[cosma_id]\n",
        "\n",
        "dataset_path = path.join(cosma_dataset_path, dataset_type, dataset_name)\n",
        "\n",
        "import autofit as af\n",
        "import autolens as al\n",
        "\n",
        "dataset = al.Imaging.from_fits(\n",
        "    data_path=path.join(dataset_path, \"data.fits\"),\n",
        "    psf_path=path.join(dataset_path, \"psf.fits\"),\n",
        "    noise_map_path=path.join(dataset_path, \"noise_map.fits\"),\n",
        "    pixel_scales=0.1,\n",
        ")\n",
        "\n",
        "mask = al.Mask2D.circular(\n",
        "    shape_native=dataset.shape_native, pixel_scales=dataset.pixel_scales, radius=3.0\n",
        ")\n",
        "\n",
        "dataset = dataset.apply_mask(mask=mask)\n",
        "\n",
        "# Lens:\n",
        "\n",
        "mass = af.Model(al.mp.Isothermal)\n",
        "shear = af.Model(al.mp.ExternalShear)\n",
        "\n",
        "lens = af.Model(al.Galaxy, redshift=0.5, mass=mass, shear=shear)\n",
        "\n",
        "# Source:\n",
        "\n",
        "bulge = af.Model(al.lp.Sersic)\n",
        "\n",
        "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n",
        "\n",
        "# Overall Lens Model:\n",
        "\n",
        "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "__Search__\n",
        "\n",
        "Here is where we differ from `example_0.py`. \n",
        "\n",
        "The only change is that the `number_of_cores` input into `Nautilus` is now 16."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "search = af.Nautilus(\n",
        "    path_prefix=path.join(\"cosma_example\"),\n",
        "    name=\"mass[sie]_source[bulge]\",\n",
        "    unique_tag=dataset_name,\n",
        "    n_live=100,\n",
        "    number_of_cores=16,\n",
        ")"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "All code from here is repeated from `example_0.py`."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "analysis = al.AnalysisImaging(dataset=dataset)\n",
        "\n",
        "result = search.fit(model=model, analysis=analysis)"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "__Wrap Up__\n",
        "\n",
        "In this example, we manually coded in the line `number_of_cores=16`.\n",
        "\n",
        "We could make this a command line input of the `python3` command in the batch script. For example, if the last line\n",
        "of the batch script read:\n",
        "\n",
        "    python3 /cosma/home/dp004/cosma_username/autolens_workspace/cosma/runners/example.py 1 16\n",
        "\n",
        "We could pass the `number_of_cores` using the command:\n",
        "\n",
        "    number_of_cores=int(sys.argv[2])\n",
        "\n",
        "\n",
        "It should also be noted that one cannot combine a `.conf` submission script with multi-CPU Nautilus parallelization.\n",
        "\n",
        "Which this should, in principle, be possible, the Python multi-processing library does not seem to happy about it when\n",
        "we do this. So, just don't bother!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [],
      "outputs": [],
      "execution_count": null
    }
  ],
  "metadata": {
    "anaconda-cloud": {},
    "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.6.1"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 4
}