release: create release-0.0.1a2 branch

Author: github-actions[bot]

README.md

@@ -33,7 +33,7 @@ bring MLIP to large-scale industrial application.
 See the [Installation](#-installation) section for details on how to install
 MLIP-JAX and the example Google Colab notebooks linked below for a quick way
 to get started. For detailed instructions, visit our extensive
-[code documentation](https://instadeep.github.io/mlip/).
+[code documentation](https://instadeepai.github.io/mlip/).
 
 This repository currently supports implementations of:
 - [MACE](https://arxiv.org/abs/2206.07697)
@@ -80,7 +80,7 @@ pip install git+https://github.com/jax-md/jax-md.git
 ## ⚡ Examples
 
 In addition to the in-depth tutorials provided as part of our documentation
-[here](https://instadeep.github.io/mlip/user_guide/index.html#deep-dive-tutorials),
+[here](https://instadeepai.github.io/mlip/user_guide/index.html#deep-dive-tutorials),
 we also provide example Jupyter notebooks that can be used as
 simple templates to build your own MLIP pipelines:
 

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "mlip"
-version = "0.0.1a1"
+version = "0.0.1a2"
 description = ""
 license = "LICENSE"
 authors = [

tutorials/model_addition_tutorial.ipynb

@@ -22,7 +22,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "As a first step, we will run the installation of the *mlip* library directly from pip. We also install the appropriate Jax CUDA backend to run on GPU (comment it out to run on CPU). In this notebook, we will not run any simulation and therefore do not install Jax-MD, for details on how to do so, please refer to our *simulation* tutorial. Note that if you have ran another tutorial in the same environment, this installation is not required. Please refer to [our installation page](https://mlip-jax-dot-int-research-tpu.uc.r.appspot.com/installation/index.html) for more information."
+    "As a first step, we will run the installation of the *mlip* library directly from pip. We also install the appropriate Jax CUDA backend to run on GPU (comment it out to run on CPU). In this notebook, we will not run any simulation and therefore do not install Jax-MD, for details on how to do so, please refer to our *simulation* tutorial. Note that if you have ran another tutorial in the same environment, this installation is not required. Please refer to [our installation page](https://instadeepai.github.io/mlip/installation/index.html) for more information."
    ]
   },
   {
@@ -74,20 +74,20 @@
     "- [`MLIPNetwork`][MLIPNetwork] is a base class for GNNs that **computes node-wise energy** summands from edge vectors, node species, and graph edges passed as `senders` and `receivers` index arrays.\n",
     "- [`ForceFieldPredictor`][ForceFieldPredictor] is a generic wrapper around any [`MLIPNetwork`][MLIPNetwork].\n",
     "\n",
-    "  It gathers **total energy, forces (and, if required, stress)** in the [`Prediction`](https://mlip-jax-dot-int-research-tpu.uc.r.appspot.com/api_reference/models/prediction.html) dataclass, by summing the node energies obtained from [`MLIPNetwork`][MLIPNetwork] on a [`jraph.GraphsTuple`](https://jraph.readthedocs.io/en/latest/api.html) object, and differentiating with respect to positions (and unit cell).\n",
+    "  It gathers **total energy, forces (and, if required, stress)** in the [`Prediction`](https://instadeepai.github.io/mlip/api_reference/models/prediction.html) dataclass, by summing the node energies obtained from [`MLIPNetwork`][MLIPNetwork] on a [`jraph.GraphsTuple`](https://jraph.readthedocs.io/en/latest/api.html) object, and differentiating with respect to positions (and unit cell).\n",
     "\n",
     "\n",
-    "For convenience, our training loop and simulation engines finally work with [`ForceField`](https://mlip-jax-dot-int-research-tpu.uc.r.appspot.com/api_reference/models/force_field.html) objects that **wrap a force field predictor and its learnable parameters within a frozen dataclass object**.\n",
+    "For convenience, our training loop and simulation engines finally work with [`ForceField`](https://instadeepai.github.io/mlip/api_reference/models/force_field.html) objects that **wrap a force field predictor and its learnable parameters within a frozen dataclass object**.\n",
     "\n",
     "For illustration, in this notebook we will\n",
     "\n",
     "2. Define a very simple model that returns constant energies,\n",
     "3. Define a more involved GNN model without equivariance constraints.\n",
     "\n",
-    "[MLIPNetwork]: https://mlip-jax-dot-int-research-tpu.uc.r.appspot.com/api_reference/models/mlip_network.html\n",
-    "[ForceFieldPredictor]: https://mlip-jax-dot-int-research-tpu.uc.r.appspot.com/api_reference/models/predictor.html\n",
-    "[ForceField]: https://mlip-jax-dot-int-research-tpu.uc.r.appspot.com/api_reference/models/force_field.html\n",
-    "[Prediction]: https://mlip-jax-dot-int-research-tpu.uc.r.appspot.com/api_reference/models/prediction.html"
+    "[MLIPNetwork]: https://instadeepai.github.io/mlip/api_reference/models/mlip_network.html\n",
+    "[ForceFieldPredictor]: https://instadeepai.github.io/mlip/api_reference/models/predictor.html\n",
+    "[ForceField]: https://instadeepai.github.io/mlip/api_reference/models/force_field.html\n",
+    "[Prediction]: https://instadeepai.github.io/mlip/api_reference/models/prediction.html"
    ]
   },
   {
@@ -102,16 +102,16 @@
     "\n",
     "### a. *Config and DatasetInfo*\n",
     "\n",
-    "To facilitate model loading and saving, our [`MLIPNetwork`](https://mlip-jax-dot-int-research-tpu.uc.r.appspot.com/api_reference/models/mlip_network.html) class **gathers (almost) all of their hyperparameters within a `pydantic.BaseModel` subclass**. Their class attribute `.Config` points to this configuration class. Only exceptions consist of hyperparameters that are data dependent, and might\n",
+    "To facilitate model loading and saving, our [`MLIPNetwork`](https://instadeepai.github.io/mlip/api_reference/models/mlip_network.html) class **gathers (almost) all of their hyperparameters within a `pydantic.BaseModel` subclass**. Their class attribute `.Config` points to this configuration class. Only exceptions consist of hyperparameters that are data dependent, and might\n",
     "conflict with the data processing pipeline.\n",
     "\n",
-    "This is why [`MLIPNetwork`](https://mlip-jax-dot-int-research-tpu.uc.r.appspot.com/api_reference/models/mlip_network.html) **also accept a [`DatasetInfo`](https://mlip-jax-dot-int-research-tpu.uc.r.appspot.com/api_reference/data/dataset_info.html) object** upon initialization, that notably stores:\n",
+    "This is why [`MLIPNetwork`](https://instadeepai.github.io/mlip/api_reference/models/mlip_network.html) **also accept a [`DatasetInfo`](https://instadeepai.github.io/mlip/api_reference/data/dataset_info.html) object** upon initialization, that notably stores:\n",
     "- `cutoff_distance_angstrom : float`\n",
     "- `atomic_energies_map : dict[int, float]`\n",
     "- `avg_num_neighbours : float`\n",
     "- and some other data computed when processing the dataset.\n",
     "\n",
-    "This way, we are sure that our models can only be used in the context they were trained for, and will not be evaluated e.g. on atomic numbers they have never seen. We create a dummy [`DatasetInfo`](https://mlip-jax-dot-int-research-tpu.uc.r.appspot.com/api_reference/data/dataset_info.html) for the purpose of this example:"
+    "This way, we are sure that our models can only be used in the context they were trained for, and will not be evaluated e.g. on atomic numbers they have never seen. We create a dummy [`DatasetInfo`](https://instadeepai.github.io/mlip/api_reference/data/dataset_info.html) for the purpose of this example:"
    ]
   },
   {
@@ -197,7 +197,7 @@
    "source": [
     "### c. *Constant force field*\n",
     "\n",
-    "Now that we have defined this simple `ConstantMLIP` subclass, we can already define a state-holding [`ForceField`](https://mlip-jax-dot-int-research-tpu.uc.r.appspot.com/api_reference/models/force_field.html) object. The quickest (but slightly opaque) way is to use the helper classmethod `ForceField.from_mlip_network()`:"
+    "Now that we have defined this simple `ConstantMLIP` subclass, we can already define a state-holding [`ForceField`](https://instadeepai.github.io/mlip/api_reference/models/force_field.html) object. The quickest (but slightly opaque) way is to use the helper classmethod `ForceField.from_mlip_network()`:"
    ]
   },
   {
@@ -237,9 +237,9 @@
    "source": [
     "For the sake of transparency, let us detail what is actually being done here.\n",
     "\n",
-    "First, a [`ForceFieldPredictor`](https://mlip-jax-dot-int-research-tpu.uc.r.appspot.com/api_reference/models/predictor.html) instance is created on top of the `constant_mlip` model.\n",
+    "First, a [`ForceFieldPredictor`](https://instadeepai.github.io/mlip/api_reference/models/predictor.html) instance is created on top of the `constant_mlip` model.\n",
     "\n",
-    "Then, random parameters are initialized by calling the predictor's `.init()` method on a random seed and a dummy graph. These two objects (the predictor and its parameter dict) are wrapped for convenience inside the [`ForceField`](https://mlip-jax-dot-int-research-tpu.uc.r.appspot.com/api_reference/models/force_field.html) dataclass. The following is thus equivalent:"
+    "Then, random parameters are initialized by calling the predictor's `.init()` method on a random seed and a dummy graph. These two objects (the predictor and its parameter dict) are wrapped for convenience inside the [`ForceField`](https://instadeepai.github.io/mlip/api_reference/models/force_field.html) dataclass. The following is thus equivalent:"
    ]
   },
   {
@@ -265,9 +265,9 @@
     "id": "Po9NjwTo3tvX"
    },
    "source": [
-    "We'll see below how to manually initialize parameters, and call the [`ForceField`](https://mlip-jax-dot-int-research-tpu.uc.r.appspot.com/api_reference/models/force_field.html) default constructor : this only requires an input graph.\n",
+    "We'll see below how to manually initialize parameters, and call the [`ForceField`](https://instadeepai.github.io/mlip/api_reference/models/force_field.html) default constructor : this only requires an input graph.\n",
     "\n",
-    "**N.B.** The [`ForceField`](https://mlip-jax-dot-int-research-tpu.uc.r.appspot.com/api_reference/models/force_field.html) dataclass is frozen: this is to prevent any stateful operations to be performed on the parameters, which would be incompatible with JAX compilation and tracing mechanisms. You can think of [`ForceField`](https://mlip-jax-dot-int-research-tpu.uc.r.appspot.com/api_reference/models/force_field.html) as holding the _state_ of a learnable [`ForceFieldPredictor`](https://mlip-jax-dot-int-research-tpu.uc.r.appspot.com/api_reference/models/predictor.html), although _it remains immutable_."
+    "**N.B.** The [`ForceField`](https://instadeepai.github.io/mlip/api_reference/models/force_field.html) dataclass is frozen: this is to prevent any stateful operations to be performed on the parameters, which would be incompatible with JAX compilation and tracing mechanisms. You can think of [`ForceField`](https://instadeepai.github.io/mlip/api_reference/models/force_field.html) as holding the _state_ of a learnable [`ForceFieldPredictor`](https://instadeepai.github.io/mlip/api_reference/models/predictor.html), although _it remains immutable_."
    ]
   },
   {
@@ -360,9 +360,9 @@
    "source": [
     "### e. *Wrapping the model state in ForceField*\n",
     "\n",
-    "In order to hide the `flax` logic for downstream applications, our `TrainingLoop` class takes in and returns a [`ForceField`](https://mlip-jax-dot-int-research-tpu.uc.r.appspot.com/api_reference/models/force_field.html) object that simply wraps the predictor with its initial and final parameters respectively.\n",
+    "In order to hide the `flax` logic for downstream applications, our `TrainingLoop` class takes in and returns a [`ForceField`](https://instadeepai.github.io/mlip/api_reference/models/force_field.html) object that simply wraps the predictor with its initial and final parameters respectively.\n",
     "\n",
-    "This frozen dataclass can then be easily passed to the [`SimulationEngine`](https://mlip-jax-dot-int-research-tpu.uc.r.appspot.com/api_reference/simulation/simulation_engine.html), or just saved for later (by JSON-serializing the MLIPNetwork's `.config` and `.dataset_info`, and dumping the flattened parameter dict as `.npz`)."
+    "This frozen dataclass can then be easily passed to the [`SimulationEngine`](https://instadeepai.github.io/mlip/api_reference/simulation/simulation_engine.html), or just saved for later (by JSON-serializing the MLIPNetwork's `.config` and `.dataset_info`, and dumping the flattened parameter dict as `.npz`)."
    ]
   },
   {
@@ -393,7 +393,7 @@
     "id": "m0lNIXkOovaI"
    },
    "source": [
-    "Note that [`ForceField`](https://mlip-jax-dot-int-research-tpu.uc.r.appspot.com/api_reference/models/force_field.html) instances are also callable, and morally equivalent to `functools.partial(predictor.apply, params)`.\n",
+    "Note that [`ForceField`](https://instadeepai.github.io/mlip/api_reference/models/force_field.html) instances are also callable, and morally equivalent to `functools.partial(predictor.apply, params)`.\n",
     "\n",
     "This means they can be directly evaluated on a graph by forgetting about the (frozen) learnable parameters, as done during simulation."
    ]
@@ -420,7 +420,7 @@
     "id": "_AnG7n_vszX5"
    },
    "source": [
-    "In theory, the [`ForceField`](https://mlip-jax-dot-int-research-tpu.uc.r.appspot.com/api_reference/models/force_field.html) class is duck-typed for the [`SimulationEngine`](https://mlip-jax-dot-int-research-tpu.uc.r.appspot.com/api_reference/simulation/simulation_engine.html), and you could provide any other object with the following methods and properties (e.g. to wrap models defined in another JAX framework):\n",
+    "In theory, the [`ForceField`](https://instadeepai.github.io/mlip/api_reference/models/force_field.html) class is duck-typed for the [`SimulationEngine`](https://instadeepai.github.io/mlip/api_reference/simulation/simulation_engine.html), and you could provide any other object with the following methods and properties (e.g. to wrap models defined in another JAX framework):\n",
     "- `.__call__(graph: GraphsTuple) -> Prediction`\n",
     "- `.cutoff_distance: float`\n",
     "- `.allowed_atomic_numbers: set[int]`\n",
@@ -472,9 +472,9 @@
     "id": "0MdUJqneXoay"
    },
    "source": [
-    "Having defined our config, we can now create our MLIP model class. Our custom model must inherits the [`MLIPNetwork`](https://mlip-jax-dot-int-research-tpu.uc.r.appspot.com/api_reference/models/mlip_network.html) class, which is itself a `flax.linen.Module` object. As such, we can easily define our network using flax `@nn.compact` decorator, see [the flax docs](https://flax-linen.readthedocs.io/en/latest/quick_start.html) for more information.\n",
+    "Having defined our config, we can now create our MLIP model class. Our custom model must inherits the [`MLIPNetwork`](https://instadeepai.github.io/mlip/api_reference/models/mlip_network.html) class, which is itself a `flax.linen.Module` object. As such, we can easily define our network using flax `@nn.compact` decorator, see [the flax docs](https://flax-linen.readthedocs.io/en/latest/quick_start.html) for more information.\n",
     "\n",
-    "Our model must also have a dataset_info attribute of type [`DatasetInfo`](https://mlip-jax-dot-int-research-tpu.uc.r.appspot.com/api_reference/data/dataset_info.html). This object encapsulates the relevant informations about the dataset at hand that can be used to create the model. For instance, this attribute contains the average number of neighbors per atom in the dataset, which is used in models like [MACE](https://arxiv.org/pdf/2206.07697) to normalize the messages passed to each nodes.\n",
+    "Our model must also have a dataset_info attribute of type [`DatasetInfo`](https://instadeepai.github.io/mlip/api_reference/data/dataset_info.html). This object encapsulates the relevant informations about the dataset at hand that can be used to create the model. For instance, this attribute contains the average number of neighbors per atom in the dataset, which is used in models like [MACE](https://arxiv.org/pdf/2206.07697) to normalize the messages passed to each nodes.\n",
     "\n",
     "We provide a very simple example of MPNN below, which computes messages through an `MLP` encoding of sender and receiver features with edge distances."
    ]
@@ -583,7 +583,7 @@
     "id": "MgYcKGGXZZmO"
    },
    "source": [
-    "Having defined both our model and its associated config classes, we can now instantiate our model and turn it into a [`ForceField`](https://mlip-jax-dot-int-research-tpu.uc.r.appspot.com/api_reference/models/force_field.html) object that can be used for training and simulations."
+    "Having defined both our model and its associated config classes, we can now instantiate our model and turn it into a [`ForceField`](https://instadeepai.github.io/mlip/api_reference/models/force_field.html) object that can be used for training and simulations."
    ]
   },
   {