uv python package manager

The python package manager ecosystem has become more refined and improved over the last 14 years since I started programming in python. When I first started with python pypi and conda were the de facto package managers, with conda being the more popular one because it did a better job with managing dependencies and system libraries. As time has gone by, pypi became better at managing dependencies and conda became ideal for very complex environments -- although containerization and virtual environments have made this less an issue.

One of the problems with just relying on pypi or conda is that a lot of the mainstream data science and ML/AI libraries are large and heavy so the setup time can be grueling in terms of time (and disk space!). And because the focus has shifted to virtual environments and containerization for python scripts, projects, and packages, the package managers really haven't kept up with the times.

A recent project called uv has focused on trying to make the python package manager experience more streamlined and faster. The biggest concern by the community is that uv is written in Rust and thus breaks the convention of python developers being able to read and maintain the package manager codebase. But since I'm not a core python developer, I'm not going to worry about that.

Poetry

For a while I used poetry as my package manager. I liked it a lot, but uv seems to be much faster from start to finish for setting up a new project. However, I wouldn't write poetry off yet and to be honest just vanilla python -m venv is still a good solution.

How to use uv

First things first is installing uv itself. I'm going to do this from a linux perspective, so bear that in mind. I typically like to install using the install script approach but you can also use pip if you prefer to:

curl -fsSL https://get.uv.dev | sh

and from pypi:

pip install uv

you can also use cargo or configure with a uv docker image, but I don't think most will do that.

Python versions

With uv you can easily install different versions of python by running uv python install <version> which will make the python version available to use with a script, project, or package. To list the python versions available to uv you can just do uv python list.

Running scripts

The overall structure of uv is to assume a pyproject.toml file is present in the current working directory and that this can be used to create a virtual environment and install the dependencies. The beauty is uv aims to manage all this under the hood so you don't have to worry about it. So you just need to have your pyproject.toml file and use:

uv run <script_name>.py

where the run command is doing all the coordinating work for you to run your script. However if you have a script that doesn't have any dependencies, i.e. using vanilla python, you can just run the script you would just do uv --run --no-project <script_name>.py. However this is rarely the case for most python scenarios, so it's not very useful. Additionally, in many cases you're just using numpy, scipy, matplotlib, etc. But it might not make sense to create a pyproject.toml file for just a single script. The uv project has a very elegant and clean solution for this. You just specify the dependencies in the script itself with specific markup. Here is what it looks like:

# /// script
# requires-python = ">=3.10"
# dependencies = [
#   "numpy",
#   "scipy",
#   "matplotlib",
# ]
# ///

import numpy as np
import scipy as sp
import matplotlib.pyplot as plt

Now when you run this script with uv run <script_name>.py it will install the dependencies and run the script.

Script Dependencies

I don't know exactly where or how uv configures the script dependencies environment, but essentially it's created some cached virtual environment in the .uv (or in .cache?) directory for your script that it can reference. It seems smart enough to know when the environment for the script already exists and doesn't need to be recreated.

In my opinion this is one of the best features of uv because you don't have to manage the venv directly in the folder or elsewhere.

Creating venv

If you're working on a project or have forked some repo, you can create a virtual environment using uv with the following command:

uv venv .venv

There are a bunch of options for setting things like python version etc., just type uv venv --help to see the options.

Then as with any python virtual environment you activate it in your shell using source .venv/bin/activate and treat it like a normal python virtual environment.

The one thing that is a bit strange is how you add python packages to your .venv, you can't just use pip install <package> because that will install the package globally, rather you have to use uv pip install <package>.

If you're contributing to an existing python project and need to install it in development mode you can do the normal pip but you have to use uv pip install -e .. Some different options are highlighted here.

Creating a project

The last kind of configuration setup you can use uv for is to create a project. What this does is create the minimal pyproject.toml file and a the virtual environment. You can do this with the following command:

uv init --name <project_name>

If you plan on making this a python package, app, or lib you can add the flags --package, --app, or --lib. I think the main difference is the folder structure so it's typically like this:

lib_project/ ├── utilities/ │ ├── __init__.py │ └── function.py ├── tests/ │ ├── __init__.py │ └── test_function.py

app_project/ ├── main.py ├── tests/ │ ├── __init__.py │ └── test_main.py

package_project/ ├── src/ │ ├── __init__.py │ └── main.py ├── tests/ ├── docs/ ├── pyproject.toml

uv folder Structure

The above is my guess but don't quote me on it. Could be very different from what uv actually creates.

requirements.txt

If you have a requirements.txt file, you can use uv pip install -r requirements.txt to install the dependencies. I believe if you do this though you'll need to manage dependencies exclusively with uv pip and can't do things like uv add <package> or uv remove <package>.

It would be nice if uv did what pyscaffold or cookiecutter do and create a more comprehensive project setup. I like the way those tools handle repo details, pre-commit hooks, docs, etc. I can never remember all the configuration options I prefer, so having a basic but complete setup would be great.

This is uv for you! I've been using it most of the time for my python related stuff. The only issue that I've run into is sometimes it fails to install a package because of some compatibility issue with my system, but it's hard to figure out what the issue is. I know it's a uv issue because if I just use my system python pip it works fine. I would say if you like poetry then you'll probably like uv as well and find value in the speed-up you gain in setting this up. It would be nice if VSCode had some kind of extension for uv to make it easier to setup things, but not exactly sure how much value that would actually add to be honest.

Reuse and Attribution

Orb v3 MLIP

Last week, Orbital Materials released their state-of-the-art (SOTA) MLIP model v3. Based on the pareto front performance plot of select models in Figure 1, seems to be the best-performing MLIP model to date [1]. The one thing that stood out to me was the introduction of a new loss regularization technique they call equigrad. I think this is a really clever/efficient way to learn the rotational equivariance without having to use any SO(3) equivariant layers, which are computationally intensive and architecturally complex¹.

Figure 1. Performance of orb-v3 models vs others (fig. 1 from ref. [1])

Equigrad

The orb-v3 class of models introduce a loss-based approach to rotational symmetry, which enforces the physical constraint that the total potential energy of a system should remain invariant under global rigid-body rotation. Rather than embedding rotational equivariance of the node states in the network architecture itself, as is done with MACE or SevenNet models, with equigrad one introduces a differentiable penalty on the model's energy prediction when atomic positions and the cell matrix are infinitesimally rotated.

The authors define a global rotation as a matrix exponential of a skew-symmetric matrix given as:

$ \begin{equation} \mathbf{R} = e^{\mathbf{G} - \mathbf{G}^T} \label{eq:rot_matrix} \end{equation} $

where $\mathbf{G} \in \mathbb{R}^{3 \times 3}$. The rotational gradient of the predicted energy, evaluated at the identity rotation ($\mathbf{G}=\mathbf{0}$), is then given as:

$ \begin{equation} \Delta_{\text{rot}} = \left. \frac{\partial E(\mathbf{r}^T \mathbf{R}, \mathbf{hR})}{\partial \mathbf{G}} \right|_{\mathbf{G}=\mathbf{0}} \label{eq:rot_grad} \end{equation} $

where $\mathbf{r}$ is the atomic positions and $\mathbf{h}$ the cell matrix. What is nice is this measures the sensitivity of the energy to infinitesimally small global rotations. If the energy is correctly invariant to rotation, then $\Delta_{\text{rot}}$ is zero and thus any deviation penalizes the model as:

$ \begin{equation} \mathcal{L}_{\text{equigrad}} = \lambda \left| \Delta_{\text{rot}} \right|_2^2 \end{equation} $

Computational Efficiency

I think the reason this remains efficient is because if we expand $\mathbf{R}$ via a Taylor series:

$ \begin{equation} \mathbf{R} = \mathbf{I} + (\mathbf{G} - \mathbf{G}^T) + \frac{1}{2!}(\mathbf{G} - \mathbf{G}^T)^2 + \cdots. \end{equation} $

and then because one evaluates the derivative w.r.t $\mathbf{G} = \mathbf{0}$, all the higher-order terms vanish due to the presence of $\mathbf{G}$ in the higher-order terms. Therefore one gets:

$ \begin{equation} \left. \frac{d\mathbf{R}}{d\mathbf{G}} \right|_{\mathbf{G}=\mathbf{0}} = \mathbf{G} - \mathbf{G}^T \end{equation} $

and $\Delta_{\text{rot}}$ is now a first-order variation. The linearization of $\mathbf{R}$ avoids evaluating the full matrix exponential in eq. $\ref{eq:rot_matrix}$, which saves compute time. Furthermore, for conservative models where forces and stress are computed via autograd, the evaluation of $\Delta_{\text{rot}}$ will use the same gradient computations and therefore the loss term adds negligible overhead [1]. At least this is what makes sense to me after reading the preprint [1].

Why This Works

This approach works because the force is defined as the negative gradient of the scalar potential energy:

$ \begin{equation} \mathbf{F}_i = -\nabla_{\mathbf{r}_i} E \end{equation} $

and for a conserved force field where $E$ is invariant to global rotation, then the forces transform equivariantly under rotation as:

$ \begin{equation} \mathbf{r}'_i = \mathbf{R} \mathbf{r}_i \quad \Rightarrow \quad \mathbf{F}'_i = \mathbf{R} \mathbf{F}_i \end{equation} $

What results is that by the enforcement of rotational invariance of the energy via equigrad, it implicitly leads to learning rotational equivariance with respect to the forces. The paper shows that equigrad significantly improves rotational invariance (by ~5x according to their tests, see Fig 3 in [1]) and also leads to improved energy RMSD (see Figure 2). The key is that this is achieved without including in the architecture any SO(3)-equivariant layers and without having to evaluate the full matrix exponential in eq. $\eqref{eq:rot_matrix}$, thereby making it very computationally efficient.

Figure 2. Energy RMSD with and without equigrad regularization (fig. 5 from ref. [1])

Summarizing things

The equigrad loss regularization introduced by B. Rhodes et al. [1] is clean and principled (I think), and seems to be a computationally efficient alternative to including equivariant layers in MLIP models. It yields:

• Enforces energy invariance and learns force equivariance via the loss function.
• Avoids the equivariant layer tensor field operations and spherical harmonics.
• Adds negligible overhead, particularly for conservative models, due to linearization of $\mathbf{R}$ at $\mathbf{G} = \mathbf{0}$.

Footnotes

---

1. Adding SO(3) equivariant layers to the network architecture can be computationally expensive and architecturally complex compared to standard message passing GNN layers. For example, SO(3)-equivariant networks must compute spherical harmonics for each interaction, perform tensor product operations then contract these representations through Clebsch–Gordan decompositions, and maintain irreducible representation (irrep) channels at multiple orders. ↩

References

[1] B. Rhodes, S. Vandenhaute, V. Šimkus, J. Gin, J. Godwin, T. Duignan, M. Neumann, Orb-v3: atomistic simulation at scale, (2025). DOI.

Reuse and Attribution