torch_to_jax_parameters.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## To get a matching ones need to modify the ordering in `MessagePassingConvolution`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"4\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ase.neighborlist\n",
    "import e3nn as e3nn_torch\n",
    "import e3nn_jax as e3nn\n",
    "import haiku as hk\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import torch\n",
    "from mace import modules as torch_modules\n",
    "\n",
    "from mace_jax.modules import MACE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/mgeiger/git/mace/mace/modules/blocks.py:77: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  torch.tensor(atomic_energies, dtype=torch.get_default_dtype()),\n"
     ]
    }
   ],
   "source": [
    "@hk.without_apply_rng\n",
    "@hk.transform\n",
    "def jax_model(\n",
    "    vectors: jnp.ndarray,  # [n_edges, 3]\n",
    "    node_specie: jnp.ndarray,  # [n_nodes, #scalar_features]\n",
    "    senders: jnp.ndarray,  # [n_edges]\n",
    "    receivers: jnp.ndarray,  # [n_edges]\n",
    "):\n",
    "    e3nn.config(\"path_normalization\", \"path\")\n",
    "    e3nn.config(\"gradient_normalization\", \"path\")\n",
    "    return MACE(\n",
    "        r_max=2.0,\n",
    "        radial_basis=lambda r, r_max: e3nn.bessel(r, 8, r_max),\n",
    "        radial_envelope=lambda r, r_max: e3nn.poly_envelope(5 - 1, 2, r_max)(r),\n",
    "        max_ell=3,\n",
    "        num_interactions=2,\n",
    "        num_species=1,\n",
    "        hidden_irreps=\"11x0e+11x1o\",\n",
    "        readout_mlp_irreps=\"16x0e\",\n",
    "        avg_num_neighbors=3.0,\n",
    "        correlation=2,\n",
    "        output_irreps=\"0e\",\n",
    "        symmetric_tensor_product_basis=False,\n",
    "    )(vectors, node_specie, senders, receivers).array[:, :, 0]\n",
    "\n",
    "\n",
    "torch_model = torch_modules.MACE(\n",
    "    r_max=2.0,\n",
    "    num_bessel=8,\n",
    "    num_polynomial_cutoff=5,\n",
    "    max_ell=3,\n",
    "    interaction_cls_first=torch_modules.RealAgnosticInteractionBlock,\n",
    "    interaction_cls=torch_modules.RealAgnosticResidualInteractionBlock,\n",
    "    num_interactions=2,\n",
    "    num_elements=1,\n",
    "    hidden_irreps=e3nn_torch.o3.Irreps(\"11x0e+11x1o\"),\n",
    "    MLP_irreps=\"16x0e\",\n",
    "    avg_num_neighbors=3.0,\n",
    "    correlation=2,\n",
    "    atomic_energies=torch.zeros(1),\n",
    "    atomic_numbers=[],\n",
    "    gate=torch.nn.SiLU(),\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def linear_torch_to_jax(linear):\n",
    "    return {\n",
    "        f\"w[{ins.i_in},{ins.i_out}] {linear.irreps_in[ins.i_in]},{linear.irreps_out[ins.i_out]}\": jnp.asarray(\n",
    "            w.data\n",
    "        )\n",
    "        for i, ins, w in linear.weight_views(yield_instruction=True)\n",
    "    }\n",
    "\n",
    "\n",
    "def skip_tp_torch_to_jax(tp):\n",
    "    return {\n",
    "        f\"w[{ins.i_in1},{ins.i_out}] {tp.irreps_in1[ins.i_in1]},{tp.irreps_out[ins.i_out]}\": jnp.moveaxis(\n",
    "            jnp.asarray(w.data), 1, 0\n",
    "        )\n",
    "        for i, ins, w in tp.weight_views(yield_instruction=True)\n",
    "    }\n",
    "\n",
    "\n",
    "w = {\n",
    "    \"general_mace/~/linear_node_embedding_block\": {\n",
    "        \"embeddings\": (\n",
    "            torch_model.node_embedding.linear.weight.detach()\n",
    "            .numpy()\n",
    "            .reshape((1, -1))\n",
    "        )\n",
    "    },\n",
    "    \"general_mace/layer_0/skip_tp_first\": skip_tp_torch_to_jax(\n",
    "        torch_model.interactions[0].skip_tp\n",
    "    ),\n",
    "    \"general_mace/layer_1/skip_tp\": skip_tp_torch_to_jax(\n",
    "        torch_model.interactions[1].skip_tp\n",
    "    ),\n",
    "    \"general_mace/layer_0/interaction_block/linear_up\": linear_torch_to_jax(\n",
    "        torch_model.interactions[0].linear_up\n",
    "    ),\n",
    "    \"general_mace/layer_0/interaction_block/linear_down\": linear_torch_to_jax(\n",
    "        torch_model.interactions[0].linear\n",
    "    ),\n",
    "    \"general_mace/layer_0/interaction_block/message_passing_convolution/multi_layer_perceptron/linear_0\": {\n",
    "        \"w\": (\n",
    "            torch_model.interactions[0].conv_tp_weights.layer0.weight.detach().numpy()\n",
    "        )\n",
    "    },\n",
    "    \"general_mace/layer_0/interaction_block/message_passing_convolution/multi_layer_perceptron/linear_1\": {\n",
    "        \"w\": (\n",
    "            torch_model.interactions[0].conv_tp_weights.layer1.weight.detach().numpy()\n",
    "        )\n",
    "    },\n",
    "    \"general_mace/layer_0/interaction_block/message_passing_convolution/multi_layer_perceptron/linear_2\": {\n",
    "        \"w\": (\n",
    "            torch_model.interactions[0].conv_tp_weights.layer2.weight.detach().numpy()\n",
    "        )\n",
    "    },\n",
    "    \"general_mace/layer_0/interaction_block/message_passing_convolution/multi_layer_perceptron/linear_3\": {\n",
    "        \"w\": (\n",
    "            torch_model.interactions[0].conv_tp_weights.layer3.weight.detach().numpy()\n",
    "        )\n",
    "    },\n",
    "    \"general_mace/layer_1/interaction_block/linear_up\": linear_torch_to_jax(\n",
    "        torch_model.interactions[1].linear_up\n",
    "    ),\n",
    "    \"general_mace/layer_1/interaction_block/linear_down\": linear_torch_to_jax(\n",
    "        torch_model.interactions[1].linear\n",
    "    ),\n",
    "    \"general_mace/layer_1/interaction_block/message_passing_convolution/multi_layer_perceptron/linear_0\": {\n",
    "        \"w\": (\n",
    "            torch_model.interactions[1].conv_tp_weights.layer0.weight.detach().numpy()\n",
    "        )\n",
    "    },\n",
    "    \"general_mace/layer_1/interaction_block/message_passing_convolution/multi_layer_perceptron/linear_1\": {\n",
    "        \"w\": (\n",
    "            torch_model.interactions[1].conv_tp_weights.layer1.weight.detach().numpy()\n",
    "        )\n",
    "    },\n",
    "    \"general_mace/layer_1/interaction_block/message_passing_convolution/multi_layer_perceptron/linear_2\": {\n",
    "        \"w\": (\n",
    "            torch_model.interactions[1].conv_tp_weights.layer2.weight.detach().numpy()\n",
    "        )\n",
    "    },\n",
    "    \"general_mace/layer_1/interaction_block/message_passing_convolution/multi_layer_perceptron/linear_3\": {\n",
    "        \"w\": (\n",
    "            torch_model.interactions[1].conv_tp_weights.layer3.weight.detach().numpy()\n",
    "        )\n",
    "    },\n",
    "    \"general_mace/layer_0/equivariant_product_basis_block/~/symmetric_contraction\": {\n",
    "        \"w2_0e\": jnp.array(\n",
    "            torch_model.products[0]\n",
    "            .symmetric_contractions.contractions[0]\n",
    "            .weights_max.detach()\n",
    "            .numpy()\n",
    "        ),\n",
    "        \"w2_1o\": jnp.array(\n",
    "            torch_model.products[0]\n",
    "            .symmetric_contractions.contractions[1]\n",
    "            .weights_max.detach()\n",
    "            .numpy()\n",
    "        ),\n",
    "        \"w1_0e\": jnp.array(\n",
    "            torch_model.products[0]\n",
    "            .symmetric_contractions.contractions[0]\n",
    "            .weights[0]\n",
    "            .detach()\n",
    "            .numpy()\n",
    "        ),\n",
    "        \"w1_1o\": jnp.array(\n",
    "            torch_model.products[0]\n",
    "            .symmetric_contractions.contractions[1]\n",
    "            .weights[0]\n",
    "            .detach()\n",
    "            .numpy()\n",
    "        ),\n",
    "    },\n",
    "    \"general_mace/layer_0/equivariant_product_basis_block/linear\": linear_torch_to_jax(\n",
    "        torch_model.products[0].linear\n",
    "    ),\n",
    "    \"general_mace/layer_1/equivariant_product_basis_block/~/symmetric_contraction\": {\n",
    "        \"w2_0e\": jnp.array(\n",
    "            torch_model.products[1]\n",
    "            .symmetric_contractions.contractions[0]\n",
    "            .weights_max.detach()\n",
    "            .numpy()\n",
    "        ),\n",
    "        \"w1_0e\": jnp.array(\n",
    "            torch_model.products[1]\n",
    "            .symmetric_contractions.contractions[0]\n",
    "            .weights[0]\n",
    "            .detach()\n",
    "            .numpy()\n",
    "        ),\n",
    "    },\n",
    "    \"general_mace/layer_1/equivariant_product_basis_block/linear\": linear_torch_to_jax(\n",
    "        torch_model.products[1].linear\n",
    "    ),\n",
    "    \"general_mace/layer_0/linear_readout_block/linear\": linear_torch_to_jax(\n",
    "        torch_model.readouts[0].linear\n",
    "    ),\n",
    "    \"general_mace/layer_1/non_linear_readout_block/linear\": linear_torch_to_jax(\n",
    "        torch_model.readouts[1].linear_1\n",
    "    ),\n",
    "    \"general_mace/layer_1/non_linear_readout_block/linear_1\": linear_torch_to_jax(\n",
    "        torch_model.readouts[1].linear_2\n",
    "    ),\n",
    "}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "n_nodes: 4\n",
      "n_edges: 172\n"
     ]
    }
   ],
   "source": [
    "positions = np.array(\n",
    "    [\n",
    "        [0.0, 0.0, 0.0],\n",
    "        [0.5, 0.0, 0.0],\n",
    "        [0.0, 0.4, 0.0],\n",
    "        [0.0, 0.3, 0.3],\n",
    "    ]\n",
    ")\n",
    "node_specie = np.arange(4) % 1\n",
    "cell = np.identity(3)\n",
    "\n",
    "senders, receivers, receivers_unit_shifts = ase.neighborlist.primitive_neighbor_list(\n",
    "    quantities=\"ijS\",\n",
    "    pbc=(True, True, False),\n",
    "    cell=cell,\n",
    "    positions=positions,\n",
    "    cutoff=2.0,\n",
    ")\n",
    "\n",
    "print(f\"n_nodes: {len(positions)}\")\n",
    "print(f\"n_edges: {len(senders)}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "vectors = (positions[receivers] + receivers_unit_shifts @ cell) - positions[senders]\n",
    "w0 = jax.jit(jax_model.init)(jax.random.PRNGKey(0), vectors, node_specie, senders, receivers)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jax.tree_util.tree_structure(w0) == jax.tree_util.tree_structure(w)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "t_out = torch_model(\n",
    "    {\n",
    "        \"positions\": torch.tensor(positions, dtype=torch.float32),\n",
    "        \"edge_index\": torch.tensor(np.stack([senders, receivers]), dtype=torch.long),\n",
    "        \"shifts\": torch.tensor(receivers_unit_shifts, dtype=torch.float32),\n",
    "        \"node_attrs\": torch.eye(1)[node_specie],\n",
    "        \"ptr\": torch.tensor([0, len(positions)], dtype=torch.long),\n",
    "        \"batch\": torch.zeros(len(positions), dtype=torch.long),\n",
    "        \"cell\": torch.tensor(cell, dtype=torch.float32),\n",
    "    }\n",
    ")\n",
    "\n",
    "t_out = t_out[\"contributions\"][0, 1:].detach().numpy()\n",
    "\n",
    "\n",
    "vectors = (positions[receivers] + receivers_unit_shifts @ cell) - positions[senders]\n",
    "j_out = jnp.sum(jax.jit(jax_model.apply)(w, vectors, node_specie, senders, receivers), axis=0)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([-1.9673237e-03, -2.2330176e-05], dtype=float32)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "d = t_out - j_out\n",
    "\n",
    "d / np.abs(t_out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 0.55468345, 46.466164  ], dtype=float32)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "t_out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([ 0.5557747, 46.4672   ], dtype=float32)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "j_out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.595529907587252"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def moment(f, n, dtype=None, device=None):\n",
    "    gen = torch.Generator(device=\"cpu\").manual_seed(0)\n",
    "    z = torch.randn(1_000_000, generator=gen, dtype=torch.float64).to(dtype=dtype, device=device)\n",
    "    return f(z).pow(n).mean()\n",
    "\n",
    "moment(torch.nn.functional.silu, 2, dtype=torch.float64, device='cpu').pow(0.5).item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.59717857837677"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "k = jax.random.PRNGKey(0)\n",
    "x = jax.random.normal(k, (1_000_000,))\n",
    "(jnp.mean(jax.nn.silu(x) ** 2) ** 0.5).item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([1.], dtype=float32)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "e3nn.sh(\"0e\", jnp.array([1.0, 0.0, 0.0]), True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1x1o+1x2e+1x3o"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "e3nn.Irreps.spherical_harmonics(3)[1:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1x0e+1x1o [1. 1. 1. 1.]"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "e3nn.tensor_product(e3nn.IrrepsArray.ones(\"0e\", ()), e3nn.IrrepsArray.ones(\"0e + 1o\", ()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1x0e+2x1o+1x1e+3x2e+1x2o+2x3o+1x3e+1x4e"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "e3nn.tensor_product(\"1o + 2e + 3o\", \"0e + 1o\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.6 64-bit",
   "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.9.6"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}