JaxPM/notebooks/03-MultiHost_PM.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "22803ddc",
      "metadata": {
        "colab_type": "text",
        "id": "view-in-github"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/DifferentiableUniverseInitiative/JaxPM/blob/main/notebooks/Introduction.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "id": "9Jy5BL1XiK1s",
      "metadata": {
        "id": "9Jy5BL1XiK1s"
      },
      "outputs": [],
      "source": [
        "!pip install --quiet git+https://github.com/DifferentiableUniverseInitiative/JaxPM.git\n",
        "!pip install diffrax"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "id": "c5f42bbe",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "c5f42bbe",
        "outputId": "a7841b28-5f20-4856-bd1d-f8a3572095b5"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "%pylab is deprecated, use %matplotlib inline and import the required libraries.\n",
            "Populating the interactive namespace from numpy and matplotlib\n"
          ]
        }
      ],
      "source": [
        "!salloc --account=tkc@a100 -C a100 --gres=gpu:8 --ntasks-per-node=8 --time=00:30:00  --cpus-per-task=8 --hint=nomultithread --qos=qos_gpu-dev --nodes=2"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "7ebdfc00",
      "metadata": {},
      "outputs": [],
      "source": [
        "!squeue -u $USER"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "c014316c",
      "metadata": {},
      "outputs": [],
      "source": [
        "export JOB_ID=123456"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "b7eabac5",
      "metadata": {},
      "outputs": [],
      "source": [
        "!srun --jobid=$JOB_ID -n 16 python 03-MultiHost_PM.py"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "472dd4bf",
      "metadata": {},
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "\n",
        "data = np.load(\"multihost_pm.npz\")\n",
        "initial_conditions = data['initial_conditions']\n",
        "lpt_displacements = data['lpt_displacements']\n",
        "ode_solutions = data['ode_solutions']\n",
        "solver_stats = data['solver_stats']\n",
        "print(f\"Solver stats: {solver_stats}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "4e012ce8",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 323
        },
        "id": "4e012ce8",
        "outputId": "75390318-8072-481f-ffb9-ec09cd71cb1d"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "shape of grid_mesh: (256, 256, 256)\n"
          ]
        }
      ],
      "source": [
        "from visualize import plot_fields\n",
        "\n",
        "fields = {\"Initial Conditions\" : initial_conditions , \"LPT Field\" : cic_paint(jnp.zeros(mesh_shape) ,lpt_particles)}\n",
        "for i , field in enumerate(ode_particles):\n",
        "    fields[f\"field_{i}\"] = cic_paint(jnp.zeros(mesh_shape) , field)\n",
        "plot_fields(fields)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "include_colab_link": true,
      "name": "Introduction.ipynb",
      "provenance": []
    },
    "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.10.4"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}