conf.py docs fix

.github/workflows/link-check-config.json

@@ -1,7 +1,6 @@
 {
   "ignorePatterns": [
-    {
-      "pattern": "../tutorials/.+.ipynb$"
-    }
+    {"pattern": "../tutorials/.+.ipynb$"},
+    {"pattern": "*github.com/user-attachments*"}
   ]
 }

.pre-commit-config.yaml

@@ -22,7 +22,7 @@ repos:
       - id: check-case-conflict
       - id: check-symlinks
       - id: check-yaml
-        exclude: ".*\/copilot\/.*"
+        exclude: .*\/copilot\/.*
       - id: destroyed-symlinks
       - id: end-of-file-fixer
       - id: forbid-new-submodules
@@ -34,7 +34,7 @@ repos:
     hooks:
       - id: codespell
         stages: [pre-commit, commit-msg]
-        args: ["--ignore-words-list", "statics,crate,annote,atomate,nd,te,titel,coo,slite,fro"]
+        args: [--ignore-words-list, "statics,crate,annote,atomate,nd,te,titel,coo,slite,fro"]
 
   - repo: https://github.com/igorshubovych/markdownlint-cli
     rev: v0.44.0

README.md

@@ -9,10 +9,12 @@
 [zenodo]: https://zenodo.org/records/15127004
 
 <!-- help docs find start of prose in readme, DO NOT REMOVE -->
-torch-sim is a next-generation open-source atomistic simulation engine for the MLIP era. By rewriting the core primitives of atomistic simulation in Pytorch, it allows orders of magnitude acceleration of popular machine learning potentials.
+torch-sim is a next-generation open-source atomistic simulation engine for the MLIP
+era. By rewriting the core primitives of atomistic simulation in Pytorch, it allows
+orders of magnitude acceleration of popular machine learning potentials.
 
 * Automatic batching and GPU memory management allowing significant simulation speedup
-* Support for MACE and Fairchem MLIP models
+* Support for MACE, Fairchem, and SevenNet MLIP models with more in progress
 * Support for classical lennard jones, morse, and soft-sphere potentials
 * Molecular dynamics integration schemes like NVE, NVT Langevin, and NPT Langevin
 * Relaxation of atomic positions and cell with gradient descent and FIRE
@@ -24,7 +26,7 @@ torch-sim is a next-generation open-source atomistic simulation engine for the M
 
 ## Quick Start
 
-Here is a quick demonstration of many of the core features of TorchSim:
+Here is a quick demonstration of many of the core features of torch-sim:
 native support for GPUs, MLIP models, ASE integration, simple API,
 autobatching, and trajectory reporting, all in under 40 lines of code.
 
@@ -85,6 +87,19 @@ relaxed_state = ts.optimize(
 print(relaxed_state.energy)
 ```
 
+## Speedup
+
+torch-sim achieves up to 100x speedup compared to ASE with popular MLIPs.
+
+![Speedup comparison](https://github.com/user-attachments/assets/2ad1d8b0-a7aa-467b-9260-acb76a1ed591)
+
+This figure compares the time per atom of ASE and torch_sim. Time per atom is defined
+as the number of atoms / total time. While ASE can only run a single system of n_atoms
+(on the x axis), torch_sim can run as many systems as will fit in memory. On an H100,
+the max atoms that could fit in memory was 8000 for gemnet, 10000 for MACE, and 2500
+for SevenNet. This metric describes model performance by capturing speed and memory
+usage simultaneously.
+
 ## Installation
 
 ### PyPI Installation
@@ -101,19 +116,6 @@ cd torch-sim
 pip install .
 ```
 
-## Speedup
-
-torch-sim achieves up to 100x speedup compared to ASE with popular MLIPs.
-
-![Speedup comparison](docs/_static/speedup_plot.png)
-
-This figure compares the time per atom of ASE and torch_sim. Time per atom is defined
-as the number of atoms / total time. While ASE can only run a single system of n_atoms
-(on the x axis), torch_sim can run as many systems as will fit in memory. On an H100,
-the max atoms that could fit in memory was 8000 for gemnet, 10000 for MACE, and 2500
-for SevenNet. This metric gives a clean comparison of how different models make use of
-the available hardware.
-
 ## Examples
 
 To understand how torch-sim works, start with the [comprehensive tutorials](https://radical-ai.github.io/torch-sim/user/overview.html) in the documentation.
@@ -130,4 +132,4 @@ torch-sim's structure is summarized in the [API reference](https://radical-ai.gi
 
 ## Citation
 
-A manuscript is in preparation. Meanwhile, if you use TorchSim in your research, please [cite the Zenodo archive][zenodo].
+A manuscript is in preparation. Meanwhile, if you use torch-sim in your research, please [cite the Zenodo archive][zenodo].

docs/conf.py

@@ -17,14 +17,14 @@
 
 # -- Project information -----------------------------------------------------
 
-project = "torch-sim"
+project = "torch-sim-atomistic"
 copyright = "2025, Radical AI"  # noqa: A001
 author = "Abhijeet Gangan, Orion Cohen, Janosh Riebesell"
 
 # The short X.Y version
-version = importlib.metadata.version("torch-sim")
+version = importlib.metadata.version(project)
 # The full version, including alpha/beta/rc tags
-release = importlib.metadata.version("torch-sim")
+release = importlib.metadata.version(project)
 
 # -- General configuration ---------------------------------------------------
 

docs/dev/dev_install.md

@@ -81,6 +81,7 @@ python -m http.server -d docs_build
 To locally generate the tutorials, they must be copied to the docs folder,
 converted to `.ipynb` files, and executed. Then the .py files and any generated
 trajectory files must be cleaned up.
+
 ```bash
 cp -r examples/tutorials docs/ && \
 jupytext --set-formats "py:percent,ipynb" docs/tutorials/*.py && \

docs/user/overview.md

@@ -29,4 +29,4 @@ Learn more in [Understanding Reporting](../tutorials/reporting_tutorial.ipynb)
 
 Under the hood, `torch-sim` takes a modular functional approach to atomistic simulation. Each integrator or optimizer function, such as `nvt_langevin,` takes in a model and parameters and returns `init` and `update` functions that act on a unique `State.` The state inherits from `SimState` and tracks the fixed and fluctuating parameters of the simulation, such as the `momenta` for NVT or the timestep for FIRE. The runner functions take this basic structure and wrap it in a convenient interface with autobatching and reporting.
 
-Learn more in [Fundamentals of `torch-sim`](../tutorials/low_level_tutorial.ipynb) and [Hybrid Swap Tutorial](../tutorials/hybrid_swap_tutorial.ipynb)
+Learn more in [Fundamentals of `torch-sim`](../tutorials/low_level_tutorial.ipynb) and [Implementing New Methods](../tutorials/hybrid_swap_tutorial.ipynb)

pyproject.toml

@@ -87,6 +87,7 @@ ignore = [
     "ERA001",  # Found commented-out code
     "FIX002",  # Line contains TODO, consider resolving the issue
     "G003",    # logging-string-concat
+    "G004",    # logging uses f-string
     "INP001",  # implicit-namespace-package
     "ISC001",  # avoid conflicts with the formatter
     "N803",    # Variable name should be lowercase
@@ -118,7 +119,6 @@ pep8-naming.ignore-names = ["get_kT", "kT"]
 "examples/**/*" = ["B018"]
 "examples/tutorials/**/*" = ["ALL"]
 
-
 [tool.ruff.format]
 docstring-code-format = true
 
@@ -128,3 +128,8 @@ check-filenames = true
 [tool.pytest]
 addopts = ["--cov-report=term-missing", "--cov=torch_sim", "-v"]
 testpaths = ["tests"]
+
+[tool.coverage.run]
+omit = [
+    "torch_sim/unbatched/*",
+]

tests/test_trajectory.py

@@ -354,12 +354,12 @@ def test_invalid_dtype_handling(test_file: Path) -> None:
     complex_data = {
         "complex": np.random.default_rng(seed=0).random((10, 3)).astype(np.float16)
     }
-    with pytest.raises(ValueError, match="Unsupported dtype"):
+    with pytest.raises(ValueError, match="Unsupported array.dtype="):
         traj.write_arrays(complex_data, steps=0)
 
     # Test string data
     string_data = {"strings": np.array([["a", "b", "c"]] * 10)}
-    with pytest.raises(ValueError, match="Unsupported dtype"):
+    with pytest.raises(ValueError, match="Unsupported array.dtype="):
         traj.write_arrays(string_data, steps=0)
 
     traj.close()

torch_sim/autobatching.py

@@ -86,36 +86,36 @@ def _argmax_bins(lst: list[float]) -> int:
     def _revargsort_bins(lst: list[float]) -> list[int]:
         return sorted(range(len(lst)), key=lambda i: -lst[i])
 
-    isdict = isinstance(items, dict)
+    is_dict = isinstance(items, dict)
 
     if not hasattr(items, "__len__"):
         raise TypeError("d must be iterable")
 
-    if not isdict and hasattr(items[0], "__len__"):
+    if not is_dict and hasattr(items[0], "__len__"):
         if weight_pos is not None:
             key = lambda x: x[weight_pos]  # noqa: E731
         if key is None:
             raise ValueError("Must provide weight_pos or key for tuple list")
 
-    if not isdict and key:
+    if not is_dict and key:
         new_dict = dict(enumerate(items))
         items = {i: key(val) for i, val in enumerate(items)}
-        isdict = True
+        is_dict = True
         is_tuple_list = True
     else:
         is_tuple_list = False
 
-    if isdict:
+    if is_dict:
         # get keys and values (weights)
         keys_vals = items.items()
         keys = [k for k, v in keys_vals]
         vals = [v for k, v in keys_vals]
 
         # sort weights decreasingly
-        ndcs = _revargsort_bins(vals)
+        n_dcs = _revargsort_bins(vals)
 
-        weights = _get_bins(vals, ndcs)
-        keys = _get_bins(keys, ndcs)
+        weights = _get_bins(vals, n_dcs)
+        keys = _get_bins(keys, n_dcs)
 
         bins = [{}]
     else:
@@ -140,15 +140,15 @@ def _revargsort_bins(lst: list[float]) -> list[int]:
 
     weights = _get_bins(weights, valid_ndcs)
 
-    if isdict:
+    if is_dict:
         keys = _get_bins(keys, valid_ndcs)
 
     # prepare array containing the current weight of the bins
     weight_sum = [0.0]
 
     # iterate through the weight list, starting with heaviest
     for item, weight in enumerate(weights):
-        if isdict:
+        if is_dict:
             key = keys[item]
 
         # find candidate bins where the weight might fit
@@ -170,7 +170,7 @@ def _revargsort_bins(lst: list[float]) -> list[int]:
             # open a new bin
             b = len(weight_sum)
             weight_sum.append(0.0)
-            if isdict:
+            if is_dict:
                 bins.append({})
             else:
                 bins.append([])
@@ -180,7 +180,7 @@ def _revargsort_bins(lst: list[float]) -> list[int]:
             b = 0
 
         # put it in
-        if isdict:
+        if is_dict:
             bins[b][key] = weight
         else:
             bins[b].append(weight)
@@ -234,9 +234,7 @@ def measure_model_memory_forward(state: SimState, model: ModelInterface) -> floa
 
     logging.info(  # noqa: LOG015
         "Model Memory Estimation: Running forward pass on state with "
-        "%s atoms and %s batches.",
-        state.n_atoms,
-        state.n_batches,
+        f"{state.n_atoms} atoms and {state.n_batches} batches.",
     )
     # Clear GPU memory
     torch.cuda.synchronize()
@@ -324,7 +322,7 @@ def calculate_memory_scaler(
     Args:
         state (SimState): State to calculate metric for, with shape information
             specific to the SimState instance.
-        memory_scales_with (Literal["n_atoms_x_density", "n_atoms"]): Type of metric
+        memory_scales_with ("n_atoms_x_density" |s "n_atoms"): Type of metric
             to use. "n_atoms" uses only atom count and is suitable for models that
             have a fixed number of neighbors. "n_atoms_x_density" uses atom count
             multiplied by number density and is better for models with radial cutoffs
@@ -404,12 +402,9 @@ def estimate_max_memory_scaler(
 
     logging.info(  # noqa: LOG015
         "Model Memory Estimation: Estimating memory from worst case of "
-        "largest and smallest system. Largest system has %s atoms and %s batches, "
-        "and smallest system has %s atoms and %s batches.",
-        max_state.n_atoms,
-        max_state.n_batches,
-        min_state.n_atoms,
-        min_state.n_batches,
+        f"largest and smallest system. Largest system has {max_state.n_atoms} atoms "
+        f"and {max_state.n_batches} batches, and smallest system has "
+        f"{min_state.n_atoms} atoms and {min_state.n_batches} batches.",
     )
     min_state_max_batches = determine_max_batch_size(min_state, model, **kwargs)
     max_state_max_batches = determine_max_batch_size(max_state, model, **kwargs)
@@ -474,7 +469,7 @@ def __init__(
         Args:
             model (ModelInterface): Model to batch for, used to estimate memory
                 requirements.
-            memory_scales_with (Literal["n_atoms", "n_atoms_x_density"]): Metric to use
+            memory_scales_with ("n_atoms" | "n_atoms_x_density"): Metric to use
                 for estimating memory requirements:
                 - "n_atoms": Uses only atom count
                 - "n_atoms_x_density": Uses atom count multiplied by number density
@@ -767,7 +762,7 @@ def __init__(
         Args:
             model (ModelInterface): Model to batch for, used to estimate memory
                 requirements.
-            memory_scales_with (Literal["n_atoms", "n_atoms_x_density"]): Metric to use
+            memory_scales_with ("n_atoms" | "n_atoms_x_density"): Metric to use
                 for estimating memory requirements:
                 - "n_atoms": Uses only atom count
                 - "n_atoms_x_density": Uses atom count multiplied by number density