Pin PET-MAD versions and did a bit of maintainance. (#164)

* Remove unnecessary mpl tweak

This was a side-effect of a bug in STK brough in through chemiscope (!)

* Start fixing the pet-mad-latest issue

* Add copy button to the code boxes

* Working version of the pet-mad snapshots

* Update flashmd

* flashmd update

* Stray 1.0.1 in nc

* Import thumbnail images directly rather than through a convoluted mpimg workaround

* Linter

* Use a recent pet-mad version for nc

Author: ceriottm

docs/requirements.txt

@@ -2,6 +2,7 @@ sphinx
 sphinx-sitemap
 sphinx-design
 sphinx-gallery
+sphinx-copybutton
 sphinx-toggleprompt
 docutils
 furo

docs/src/conf.py

@@ -7,6 +7,7 @@
 extensions = [
     "sphinx_sitemap",
     "sphinx_design",
+    "sphinx_copybutton",
     "sphinx.ext.viewcode",
     "sphinx.ext.intersphinx",
     "sphinx_gallery.load_style",

examples/flashmd/data/input-al110-base.xml

@@ -10,7 +10,7 @@
    </prng>
    <ffdirect name='driver' pbc="false">
       <pes>metatomic</pes>
-      <parameters>{template:data/al110.xyz,model:pet-mad-latest.pt,device:cpu} </parameters>
+      <parameters>{template:data/al110.xyz,model:pet-mad-v1.1.0.pt,device:cpu} </parameters>
    </ffdirect>
    <system>
       <forces> 

examples/flashmd/data/input-ala2-base.xml

@@ -10,7 +10,7 @@
    </prng>
    <ffdirect name='driver' pbc="false">
       <pes>metatomic</pes>
-      <parameters>{template:data/ala2.xyz,model:pet-mad-latest.pt,device:cpu} </parameters>
+      <parameters>{template:data/ala2.xyz,model:pet-mad-v1.1.0.pt,device:cpu} </parameters>
    </ffdirect>
    <system>
       <forces> 

examples/flashmd/environment.yml

@@ -7,7 +7,7 @@ dependencies:
   - pip:
     - ase>=3.23
     - ipi==3.1.5
-    - pet-mad>=1.2,<1.3
+    - pet-mad>=1.3.1,<1.4
     - vesin[torch]>=0.3.2,<0.4
     - flashmd>=0.1.2
     - metatrain  

examples/flashmd/flashmd-demo.py

@@ -13,6 +13,8 @@
 `this preprint <http://arxiv.org/abs/2505.19350>`_.
 """
 
+# sphinx_gallery_thumbnail_path = '../../examples/flashmd/flashmd-scheme.png'
+
 # %%
 #
 # Start by importing the required libraries. You will need the
@@ -26,8 +28,6 @@
 #
 
 import chemiscope
-import matplotlib.image as mpimg
-import matplotlib.pyplot as plt
 from flashmd import get_universal_model
 from flashmd.ipi import get_npt_stepper, get_nvt_stepper
 from ipi.utils.parsing import read_output, read_trajectory
@@ -40,14 +40,9 @@
 # Each model is trained for a specific stride length, aiming to
 # reproduce the trajectories obtained with a traditional velocity
 # Verlet integrator.
-
-# use matplotlib to display the image so it also display as a thumbnail
-fig, ax = plt.subplots(figsize=(5728 / 300, 2598 / 300), dpi=300)
-img = mpimg.imread("flashmd-scheme.png")
-ax.imshow(img)
-ax.axis("off")
-fig.tight_layout()
-plt.show()
+#
+# .. image:: flashmd-scheme.png
+#    :align: center
 
 
 # %%
@@ -63,8 +58,8 @@
 flashmd_model_64 = get_universal_model(64)
 flashmd_model_64.to(device)
 
-calculator = PETMADCalculator(version="latest", device=device)
-calculator._model.save("pet-mad-latest.pt")
+calculator = PETMADCalculator(version="1.1.0", device=device)
+calculator._model.save("pet-mad-v1.1.0.pt")
 
 
 # %%

examples/learn-tensors-with-mcov/learn-tensors-with-mcov.py

@@ -17,6 +17,8 @@
 and can then be used in an ASE calculator.
 """
 
+# sphinx_gallery_thumbnail_path = '../../examples/learn-tensors-with-mcov/architecture.png' # noqa: E501
+
 # %%
 
 # Core packages
@@ -27,7 +29,6 @@
 
 # Simulation and visualization tools
 import chemiscope
-import matplotlib.image as mpimg
 import matplotlib.pyplot as plt
 import metatensor as mts
 import numpy as np
@@ -173,16 +174,11 @@
 # Here is a schematic representation of the :math:`\lambda`-MCoV model which, in a
 # nutshell, allows us to learn a tensorial property of a system from a set of scalar
 # features used as linear expansion coefficients of a minimal set of basis tensors.
+#
+# .. image:: architecture.png
+#    :align: center
 
 
-# visualize a scheme of the architecture that is also used as a thumbnail for the recipe
-fig, ax = plt.subplots(figsize=(5728 / 300, 2598 / 300), dpi=300)
-img = mpimg.imread("architecture.png")
-ax.imshow(img)
-ax.axis("off")
-fig.tight_layout()
-plt.show()
-
 # %%
 #
 # We parametrize spherical tensors of order :math:`\lambda` as linear combinations

examples/metatomic-plumed/metatomic-plumed.py

@@ -53,9 +53,6 @@
 import torch
 
 
-if hasattr(__import__("builtins"), "get_ipython"):
-    get_ipython().run_line_magic("matplotlib", "inline")  # noqa: F821
-
 # %%
 #
 # Target structures

examples/pet-mad-nc/data/input-nc-nve-mts.xml

@@ -16,11 +16,11 @@
 
 <ffdirect  name='cons' pbc="false">
     <pes>metatomic</pes>
-    <parameters>{template:data/bmimcl.xyz,model:pet-mad-latest.pt,device:cpu,non_conservative:False} </parameters>
+    <parameters>{template:data/bmimcl.xyz,model:pet-mad-v1.1.0.pt,device:cpu,non_conservative:False} </parameters>
 </ffdirect>
 <ffdirect  name='nocons' pbc="false">
     <pes>metatomic</pes>
-    <parameters>{template:data/bmimcl.xyz,model:pet-mad-latest.pt,device:cpu,non_conservative:True} </parameters>
+    <parameters>{template:data/bmimcl.xyz,model:pet-mad-v1.1.0.pt,device:cpu,non_conservative:True} </parameters>
 </ffdirect>
 <system>
     <initialize nbeads="1">

examples/pet-mad-nc/data/input-nc-nve.xml

@@ -14,7 +14,7 @@
 <prng><seed>12345</seed></prng>
 <ffdirect  name='nocons' pbc="false">
     <pes>metatomic</pes>
-    <parameters>{template:data/bmimcl.xyz,model:pet-mad-latest.pt,device:cpu,non_conservative:True} </parameters>
+    <parameters>{template:data/bmimcl.xyz,model:pet-mad-v1.1.0.pt,device:cpu,non_conservative:True} </parameters>
 </ffdirect>
 <system>
     <initialize nbeads="1">