Link to JAX examples

Author: benmwebb

doc/.template.coding.ipynb

@@ -273,7 +273,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "> You can also add arbitrary Python code to your module. This is added to the `swig.i-in` file using the SWIG `%pythoncode` directive. See the [PMI module](https://github.com/salilab/pmi/blob/develop/pyext/swig.i-in) for an example.\n",
+    "> You can also add arbitrary Python code to your module. This is added to the `swig.i-in` file using the SWIG `pythoncode` directive. See the [PMI module](https://github.com/salilab/pmi/blob/develop/pyext/swig.i-in) for an example.\n",
     ">\n",
     "> You can also add entire Python submodules by adding Python files to the `pyext/src` subdirectory. For example the file `pyext/src/my_python.py` can be imported in Python using `import IMP.foo.my_python`. This is also [used in the PMI module](https://github.com/salilab/pmi/tree/develop/pyext/src)."
    ]
@@ -624,10 +624,13 @@
     "}\n",
     "```\n",
     "\n",
-    "Our `_get_jax` method returns a new function which, given the current state of the model `X`, calculates and returns the same score as our original C++ implementation. `X` is provided by IMP and is a simple Python dict containing model data. For example `X['xyz']` contains all of the Cartesian\n",
-    "coordinates. The additional parameters used by our `jax_restraint` function, namely the force constant `k` and the particle index `pi`, are extracted from the original C++ object using the corresponding getter methods and baked in to the function using `functools.partial`.\n",
+    "Our `_get_jax` method returns a new function which, given the current state of the model `X`, calculates and returns the same score as our original C++ implementation. `X` is provided by IMP and is a simple Python dict containing model data. ",
+    "For example `X['xyz']` contains all of the Cartesian\n",
+    "coordinates as an N x 3 array. The additional parameters used by our `jax_restraint` function, namely the force constant `k` and the particle index `pi`, are extracted from the original C++ object using the corresponding getter methods and baked in to the function using ``functools.partial``.\n",
     "\n",
-    "Note that one key feature of the JAX library is automatic differentiation by which analytic first derivatives of JAX functions are determined automatically by the JAX library. Thus, it is not necessary for us to calculate them explicitly (as is done in the C++ implementation)."
+    "Note that one key feature of the JAX library is automatic differentiation by which analytic first derivatives of JAX functions are determined automatically by the JAX library. Thus, it is not necessary for us to calculate them explicitly (as is done in the C++ implementation).\n",
+    "\n",
+    "Other key IMP classes, such as ``IMP::Constraint``, ``IMP::PairScore``, and ``IMP::UnaryFunction`` can also be provided with a JAX implementation in the same fashion. See the [IMP.example module](https://github.com/salilab/imp/blob/develop/modules/example/pyext/IMP_example.jax.i) for some examples."
    ]
   }
  ],

doc/coding.ipynb

@@ -289,7 +289,7 @@
       "cell_type": "markdown",
       "metadata": {},
       "source": [
-        "> You can also add arbitrary Python code to your module. This is added to the `swig.i-in` file using the SWIG `%pythoncode` directive. See the [PMI module](https://github.com/salilab/pmi/blob/develop/pyext/swig.i-in) for an example.\n",
+        "> You can also add arbitrary Python code to your module. This is added to the `swig.i-in` file using the SWIG `pythoncode` directive. See the [PMI module](https://github.com/salilab/pmi/blob/develop/pyext/swig.i-in) for an example.\n",
         ">\n",
         "> You can also add entire Python submodules by adding Python files to the `pyext/src` subdirectory. For example the file `pyext/src/my_python.py` can be imported in Python using `import IMP.foo.my_python`. This is also [used in the PMI module](https://github.com/salilab/pmi/tree/develop/pyext/src)."
       ]
@@ -640,10 +640,13 @@
         "}\n",
         "```\n",
         "\n",
-        "Our `_get_jax` method returns a new function which, given the current state of the model `X`, calculates and returns the same score as our original C++ implementation. `X` is provided by IMP and is a simple Python dict containing model data. For example `X['xyz']` contains all of the Cartesian\n",
-        "coordinates. The additional parameters used by our `jax_restraint` function, namely the force constant `k` and the particle index `pi`, are extracted from the original C++ object using the corresponding getter methods and baked in to the function using `functools.partial`.\n",
+        "Our `_get_jax` method returns a new function which, given the current state of the model `X`, calculates and returns the same score as our original C++ implementation. `X` is provided by IMP and is a simple Python dict containing model data. ",
+        "For example `X['xyz']` contains all of the Cartesian\n",
+        "coordinates as an N x 3 array. The additional parameters used by our `jax_restraint` function, namely the force constant `k` and the particle index `pi`, are extracted from the original C++ object using the corresponding getter methods and baked in to the function using [functools.partial](https://docs.python.org/3/library/functools.html#functools.partial).\n",
         "\n",
-        "Note that one key feature of the JAX library is automatic differentiation by which analytic first derivatives of JAX functions are determined automatically by the JAX library. Thus, it is not necessary for us to calculate them explicitly (as is done in the C++ implementation)."
+        "Note that one key feature of the JAX library is automatic differentiation by which analytic first derivatives of JAX functions are determined automatically by the JAX library. Thus, it is not necessary for us to calculate them explicitly (as is done in the C++ implementation).\n",
+        "\n",
+        "Other key IMP classes, such as [IMP::Constraint](https://integrativemodeling.org/2.23.0/doc/ref/classIMP_1_1Constraint.html), [IMP::PairScore](https://integrativemodeling.org/2.23.0/doc/ref/classIMP_1_1PairScore.html), and [IMP::UnaryFunction](https://integrativemodeling.org/2.23.0/doc/ref/classIMP_1_1UnaryFunction.html) can also be provided with a JAX implementation in the same fashion. See the [IMP.example module](https://github.com/salilab/imp/blob/develop/modules/example/pyext/IMP_example.jax.i) for some examples."
       ]
     }
   ],