Fix RDFs for ase 3.28 (#691)

Co-authored-by: Elliott Kasoar <45317199+ElliottKasoar@users.noreply.github.com>

Author: harveydevereux

janus_core/calculations/md.py

@@ -14,7 +14,6 @@
 from warnings import warn
 
 from ase import Atoms
-from ase.geometry.analysis import Analysis
 from ase.io import read
 from ase.md.bussi import Bussi
 from ase.md.langevin import Langevin
@@ -1293,8 +1292,6 @@ def _post_process(self) -> None:
 
         data = read(self.traj_file, index=":")
 
-        ana = Analysis(data)
-
         if self.post_process_kwargs.get("rdf_compute", False):
             rdf_args = {
                 name: self.post_process_kwargs.get(key, default)
@@ -1312,7 +1309,7 @@ def _post_process(self) -> None:
             )
             rdf_args["index"] = slice_
 
-            compute_rdf(data, ana, filenames=self._rdf_files, **rdf_args)
+            compute_rdf(data, filenames=self._rdf_files, **rdf_args)
 
         if self.post_process_kwargs.get("vaf_compute", False):
             use_vel = self.post_process_kwargs.get("vaf_velocities", False)

janus_core/processing/post_process.py

@@ -4,9 +4,11 @@
 
 from collections.abc import Sequence
 from itertools import combinations_with_replacement
+from warnings import warn
 
 from ase import Atoms
 from ase.geometry.analysis import Analysis
+from ase.geometry.rdf import get_rdf
 import numpy as np
 from numpy import float64
 from numpy.typing import NDArray
@@ -40,7 +42,7 @@ def compute_rdf(
     data
         Dataset to compute RDF of.
     ana
-        ASE Analysis object for data reuse.
+        Deprecated. Please do not use. ASE Analysis object for data reuse.
     filenames
         Filenames to output data to. Must match number of RDFs computed.
     by_elements
@@ -69,6 +71,18 @@ def compute_rdf(
         If `by_elements` is true returns a `dict` of RDF by element pairs.
         Otherwise returns RDF of total system filtered by elements.
     """
+    if ana is not None:
+        warn(
+            "ana has been deprecated.",
+            FutureWarning,
+            stacklevel=2,
+        )
+        if by_elements:
+            raise ValueError(
+                "Analysis.get_rdf has known bugs with by_elements."
+                "Call without ana to use ase.geometry.rdf.get_rdf directly."
+            )
+
     index = slicelike_to_startstopstep(index)
 
     if not isinstance(data, Sequence):
@@ -82,32 +96,34 @@ def compute_rdf(
     ):
         volume = (2 * rmax) ** 3
 
-    if ana is None:
-        ana = Analysis(data)
-
     if by_elements:
         elements = (
             tuple(sorted(set(data[0].get_chemical_symbols())))
             if elements is None
             else elements
         )
 
-        rdf = {
-            element: ana.get_rdf(
-                rmax=rmax,
-                nbins=nbins,
-                elements=element,
-                imageIdx=slice(*index),
-                return_dists=True,
-                volume=volume,
-            )
+        rdfs = {
+            element: [
+                get_rdf(
+                    atoms,
+                    rmax,
+                    nbins,
+                    elements=element,
+                    volume=volume,
+                )
+                for atoms in data
+            ]
             for element in combinations_with_replacement(elements, 2)
         }
 
         # Compute RDF average
         rdf = {
-            element: (rdf[0][1], np.average([rdf_i[0] for rdf_i in rdf], axis=0))
-            for element, rdf in rdf.items()
+            element: (
+                element_rdfs[0][1],
+                np.average([rdf_i[0] for rdf_i in element_rdfs], axis=0),
+            )
+            for element, element_rdfs in rdfs.items()
         }
 
         if filenames is not None:
@@ -129,14 +145,24 @@ def compute_rdf(
                         print(dist, rdf_i, file=out_file)
 
     else:
-        rdf = ana.get_rdf(
-            rmax=rmax,
-            nbins=nbins,
-            elements=elements,
-            imageIdx=slice(*index),
-            return_dists=True,
-            volume=volume,
-        )
+        if ana is not None:
+            rdf = ana.get_rdf(
+                rmax=rmax,
+                nbins=nbins,
+                imageIdx=slice(*index),
+                return_dists=True,
+                volume=volume,
+            )
+        else:
+            rdf = [
+                get_rdf(
+                    atoms,
+                    rmax,
+                    nbins,
+                    volume=volume,
+                )
+                for atoms in data
+            ]
 
         assert isinstance(rdf, list)
 

pyproject.toml

@@ -26,7 +26,7 @@ classifiers = [
 readme = "README.md"
 
 dependencies = [
-    "ase<4.0,==3.27",
+    "ase<4.0,>=3.28",
     "click<9,>=8.2.1",
     "codecarbon<4.0.0,>=3.0.7",
     "numpy<3.0.0,>=1.26.4",
@@ -73,7 +73,7 @@ visualise = [
 
 # MLIPs with updated e3nn
 mattersim = [
-    "mattersim == 1.2.0; sys_platform != 'win32'",
+    "mattersim == 1.2.2; sys_platform != 'win32'",
 ]
 nequip = [
     "nequip == 0.14.0",

tests/test_post_process.py

@@ -4,6 +4,7 @@
 
 from pathlib import Path
 
+from ase.geometry.analysis import Analysis
 from ase.io import read
 import numpy as np
 import pytest
@@ -132,6 +133,51 @@ def test_rdf():
     assert (np.isclose(expected_peaks, rdf[0][peaks])).all()
 
 
+def test_rdf_with_analysis_deprecation():
+    """Test computation of RDF and ana deprecation."""
+    data = read(DATA_PATH / "benzene.xyz")
+    ana = Analysis(data)
+
+    with pytest.warns(FutureWarning, match="ana has been deprecated."):
+        rdf = post_process.compute_rdf(data, ana, index=0, rmax=5.0, nbins=100)
+
+    assert isinstance(rdf, tuple)
+    assert isinstance(rdf[0], np.ndarray)
+
+    expected_peaks = np.asarray(
+        (
+            1.075,
+            1.375,
+            2.175,
+            2.425,
+            2.475,
+            2.775,
+            3.425,
+            3.875,
+            4.275,
+            4.975,
+        )
+    )
+    peaks = np.where(rdf[1] > 0.0)
+    assert (np.isclose(expected_peaks, rdf[0][peaks])).all()
+
+
+def test_rdf_by_elements_analysis_error():
+    """Test the by_elements method with Analysis raises ValueError."""
+    data = read(DATA_PATH / "benzene.xyz")
+    ana = Analysis(data)
+
+    with pytest.raises(ValueError, match="Analysis.get_rdf has known bugs"):
+        post_process.compute_rdf(
+            data,
+            ana,
+            index=0,
+            rmax=5.0,
+            nbins=100,
+            by_elements=True,
+        )
+
+
 def test_rdf_by_elements():
     """Test the by_elements method of compute rdf."""
     data = read(DATA_PATH / "benzene.xyz")
@@ -150,18 +196,7 @@ def test_rdf_by_elements():
 
     expected_peaks = {
         ("C", "C"): (1.375, 2.425, 2.775),
-        ("C", "H"): (
-            1.075,
-            1.375,
-            2.175,
-            2.425,
-            2.475,
-            2.775,
-            3.425,
-            3.875,
-            4.275,
-            4.975,
-        ),
+        ("C", "H"): (1.075, 2.175, 3.425, 3.875),
         ("H", "H"): (2.475, 4.275, 4.975),
     }