diff --git a/.github/workflows/build-wheels.yml b/.github/workflows/build-wheels.yml index 1bb28931a..4118604da 100644 --- a/.github/workflows/build-wheels.yml +++ b/.github/workflows/build-wheels.yml @@ -111,11 +111,11 @@ jobs: matrix: torch-version: ['1.12', '1.13', '2.0', '2.1', '2.2', '2.3'] arch: ['arm64', 'x86_64'] - os: ['ubuntu-22.04', 'macos-11', 'macos-14', 'windows-2019'] + os: ['ubuntu-22.04', 'macos-13', 'macos-14', 'windows-2019'] exclude: # remove mismatched arch for macOS - {os: macos-14, arch: x86_64} - - {os: macos-11, arch: arm64} + - {os: macos-13, arch: arm64} # no arm64-windows build - {os: windows-2019, arch: arm64} # https://github.com/pytorch/pytorch/issues/125109 @@ -124,7 +124,7 @@ jobs: - {os: macos-14, arch: arm64, torch-version: '1.12'} - {os: macos-14, arch: arm64, torch-version: '1.13'} # x86_64-macos is only supported for torch <2.3 - - {os: macos-11, arch: x86_64, torch-version: '2.3'} + - {os: macos-13, arch: x86_64, torch-version: '2.3'} include: # add `cibw-arch` and `rust-target` to the different configurations - name: x86_64 Linux @@ -138,7 +138,7 @@ jobs: rust-target: aarch64-unknown-linux-gnu cibw-arch: aarch64 - name: x86_64 macOS - os: macos-11 + os: macos-13 arch: x86_64 rust-target: x86_64-apple-darwin cibw-arch: x86_64 @@ -233,7 +233,7 @@ jobs: os: ubuntu-22.04 arch: arm64 - name: x86_64 macOS - os: macos-11 + os: macos-13 arch: x86_64 - name: arm64 macOS os: macos-14 diff --git a/docs/generate_examples/conf.py b/docs/generate_examples/conf.py index 2df9bed58..334454ee4 100644 --- a/docs/generate_examples/conf.py +++ b/docs/generate_examples/conf.py @@ -2,8 +2,7 @@ import os -from matplotlib.animation import Animation, FFMpegWriter, ImageMagickWriter -from sphinx_gallery.scrapers import figure_rst, matplotlib_scraper +from chemiscope.sphinx import ChemiscopeScraper from sphinx_gallery.sorting import FileNameSortKey @@ -15,56 +14,6 @@ ROOT = os.path.realpath(os.path.join(HERE, "..", "..")) -class AnimationScrapper: - """ - Alternative image scrapper for sphinx-gallery, rendering the animation to GIF - instead of a base64-encoded HTML video. - - This decreases the size of the documentation and the time taken to render the - documentation (sphinx-gallery with ``matplotlib_animations=True``) renders the GIF - and then re-render to create the HTML video. - """ - - def __init__(self): - # only process animations once by storing their `id()` here - self.seen = set() - - def __call__(self, block, block_vars, gallery_conf): - animations = [] - for variable in block_vars["example_globals"].values(): - if isinstance(variable, Animation) and id(variable) not in self.seen: - self.seen.add(id(variable)) - animations.append(variable) - - if len(animations) == 0: - # just do the default scrapping - return matplotlib_scraper(block, block_vars, gallery_conf) - else: - # process matplotlib static preview of the animation, and ignore it - matplotlib_scraper(block, block_vars, gallery_conf) - - image_names = [] - for anim, image_path in zip(animations, block_vars["image_path_iterator"]): - image_path = str(image_path)[:-4] + ".gif" - image_names.append(image_path) - - # this is strongly inspired by the code in sphinx-gallery - fig_size = anim._fig.get_size_inches() - thumb_size = gallery_conf["thumbnail_size"] - use_dpi = round( - min(t_s / f_s for t_s, f_s in zip(thumb_size, fig_size)) - ) - if FFMpegWriter.isAvailable(): - writer = "ffmpeg" - elif ImageMagickWriter.isAvailable(): - writer = "imagemagick" - else: - writer = None - anim.save(str(image_path), writer=writer, dpi=use_dpi) - - return figure_rst(image_names, gallery_conf["src_dir"]) - - sphinx_gallery_conf = { "filename_pattern": ".*", "copyfile_regex": r".*\.(example|npz)", @@ -79,7 +28,7 @@ def __call__(self, block, block_vars, gallery_conf): os.path.join(ROOT, "docs", "src", "examples", "atomistic"), ], "matplotlib_animations": False, - "image_scrapers": (AnimationScrapper()), + "image_scrapers": ("matplotlib", ChemiscopeScraper()), "remove_config_comments": True, "within_subsection_order": FileNameSortKey, "default_thumb_file": os.path.join( diff --git a/docs/logo/metatensor-512.png b/docs/logo/metatensor-512.png new file mode 100644 index 000000000..ac14d505d Binary files /dev/null and b/docs/logo/metatensor-512.png differ diff --git a/docs/logo/metatensor-64.png b/docs/logo/metatensor-64.png new file mode 100644 index 000000000..d68605f0a Binary files /dev/null and b/docs/logo/metatensor-64.png differ diff --git a/docs/logo/metatensor-horizontal-dark.svg b/docs/logo/metatensor-horizontal-dark.svg new file mode 100644 index 000000000..abbd8b1ef --- /dev/null +++ b/docs/logo/metatensor-horizontal-dark.svg @@ -0,0 +1,144 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + etatensor + + diff --git a/docs/logo/metatensor-horizontal.svg b/docs/logo/metatensor-horizontal.svg new file mode 100644 index 000000000..d575cc71e --- /dev/null +++ b/docs/logo/metatensor-horizontal.svg @@ -0,0 +1,144 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + etatensor + + diff --git a/docs/logo/metatensor-vertical-dark.svg b/docs/logo/metatensor-vertical-dark.svg new file mode 100644 index 000000000..2d65a6dd0 --- /dev/null +++ b/docs/logo/metatensor-vertical-dark.svg @@ -0,0 +1,144 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + metatensor + + diff --git a/docs/logo/metatensor-vertical.svg b/docs/logo/metatensor-vertical.svg new file mode 100644 index 000000000..6e40e3f7f --- /dev/null +++ b/docs/logo/metatensor-vertical.svg @@ -0,0 +1,144 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + metatensor + + diff --git a/docs/logo/metatensor.svg b/docs/logo/metatensor.svg new file mode 100644 index 000000000..f860f5427 --- /dev/null +++ b/docs/logo/metatensor.svg @@ -0,0 +1,139 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/docs/requirements.txt b/docs/requirements.txt index b7fbed197..ec22d3151 100644 --- a/docs/requirements.txt +++ b/docs/requirements.txt @@ -14,3 +14,4 @@ torch # required for examples ase +chemiscope >= 0.7.0 diff --git a/docs/src/conf.py b/docs/src/conf.py index 1e899c2f4..517312674 100644 --- a/docs/src/conf.py +++ b/docs/src/conf.py @@ -161,15 +161,19 @@ def setup(app): # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom # ones. extensions = [ + # official extensions "sphinx.ext.viewcode", "sphinx.ext.autodoc", "sphinx.ext.intersphinx", - "sphinx_toggleprompt", - "sphinx_gallery.gen_gallery", + # third party extensions "sphinxcontrib.details.directive", + "sphinx_gallery.gen_gallery", + "sphinx_toggleprompt", "breathe", "myst_parser", "sphinx_design", + "chemiscope.sphinx", + # local extensions "versions_list", ] @@ -211,12 +215,17 @@ def setup(app): # -- Options for HTML output ------------------------------------------------- -# The theme to use for HTML and HTML Help pages. See the documentation for -# a list of builtin themes. -# +# The theme to use for HTML and HTML Help pages. html_theme = "furo" html_title = "Metatensor" +html_favicon = "../logo/metatensor-64.png" + +html_theme_options = { + "light_logo": "images/metatensor-horizontal.png", + "dark_logo": "images/metatensor-horizontal-dark.png", + "sidebar_hide_name": True, +} # Add any paths that contain custom static files (such as style sheets) here, # relative to this directory. They are copied after the builtin static files, diff --git a/docs/src/index.rst b/docs/src/index.rst index 7fde442ed..dddd2d4ab 100644 --- a/docs/src/index.rst +++ b/docs/src/index.rst @@ -1,5 +1,11 @@ -Metatensor -========== +.. image:: /../static/images/metatensor-horizontal.png + :class: only-light sd-mb-4 + :width: 600px + +.. image:: /../static/images/metatensor-horizontal-dark.png + :class: only-dark sd-mb-4 + :width: 600px + ``metatensor`` is a library for defining, manipulating, storing, and sharing arrays with many, potentially sparse, indices. Think numpy's ``ndarray`` or diff --git a/docs/src/operations/reference/index.rst b/docs/src/operations/reference/index.rst index c400844b1..5b3fad42b 100644 --- a/docs/src/operations/reference/index.rst +++ b/docs/src/operations/reference/index.rst @@ -12,6 +12,10 @@ API reference .. grid:: :margin: 0 0 0 0 + .. grid-item-version:: 0.2.2 + :tag-prefix: metatensor-operations-v + :url-suffix: operations/reference/index.html + .. grid-item-version:: 0.2.1 :tag-prefix: metatensor-operations-v :url-suffix: operations/reference/index.html diff --git a/docs/static/images/metatensor-horizontal-dark.png b/docs/static/images/metatensor-horizontal-dark.png new file mode 100644 index 000000000..d41b8657c Binary files /dev/null and b/docs/static/images/metatensor-horizontal-dark.png differ diff --git a/docs/static/images/metatensor-horizontal.png b/docs/static/images/metatensor-horizontal.png new file mode 100644 index 000000000..b35f11a1e Binary files /dev/null and b/docs/static/images/metatensor-horizontal.png differ diff --git a/python/examples/atomistic/2-running-ase-md.py b/python/examples/atomistic/2-running-ase-md.py index 2e7228d6c..af2c21c49 100644 --- a/python/examples/atomistic/2-running-ase-md.py +++ b/python/examples/atomistic/2-running-ase-md.py @@ -13,7 +13,7 @@ .. TODO: once we have more MD engine, add link to their doc here """ -# sphinx_gallery_thumbnail_number = 4 +# sphinx_gallery_thumbnail_number = 3 from typing import Dict, List, Optional @@ -25,6 +25,7 @@ # tools for visualization import ase.visualize.plot +import chemiscope import matplotlib.pyplot as plt # the usuals suspects @@ -127,7 +128,7 @@ def forward( # ------------------------ # # Now that we have a model for the energy of our system, let's create some initial -# simulation state. We'll build a 3x3x3 super cell of diamond carbon. In pactise, you +# simulation state. We'll build a 3x3x3 super cell of diamond carbon. In practice, you # could also read the initial state from a file. primitive = ase.build.bulk(name="C", crystalstructure="diamond", a=3.567) @@ -234,12 +235,9 @@ def forward( # %% # -# We can use ASE and matplotlib to visualize the trajectory - -animation = ase.visualize.plot.animate( - trajectory, radii=0.5, interval=100, save_count=None -) -plt.show() +# We can use `chemiscope `_ to visualize the trajectory +viewer_settings = {"bonds": False, "playbackDelay": 70} +chemiscope.show(trajectory, mode="structure", settings={"structure": [viewer_settings]}) # %% # diff --git a/python/metatensor-operations/CHANGELOG.md b/python/metatensor-operations/CHANGELOG.md index 692f07ffc..aa72a50b6 100644 --- a/python/metatensor-operations/CHANGELOG.md +++ b/python/metatensor-operations/CHANGELOG.md @@ -17,6 +17,15 @@ follows [Semantic Versioning](https://semver.org/spec/v2.0.0.html). ### Removed --> +## [Version 0.2.2](https://github.com/lab-cosmo/metatensor/releases/tag/metatensor-operations-v0.2.2) - 2024-06-19 + +### Fixed + +- Fixed a bug in `metatensor.torch.sort` where the labels where not properly + sorted (#647) +- Fix `metatensor.abs` when used with complex values (#553) + + ## [Version 0.2.1](https://github.com/lab-cosmo/metatensor/releases/tag/metatensor-operations-v0.2.1) - 2024-03-01 ### Changed diff --git a/python/metatensor-operations/setup.py b/python/metatensor-operations/setup.py index 8f74c1663..778cfd9dd 100644 --- a/python/metatensor-operations/setup.py +++ b/python/metatensor-operations/setup.py @@ -12,7 +12,7 @@ ROOT = os.path.realpath(os.path.dirname(__file__)) METATENSOR_CORE = os.path.realpath(os.path.join(ROOT, "..", "metatensor-core")) -METATENSOR_OPERATIONS_VERSION = "0.2.1" +METATENSOR_OPERATIONS_VERSION = "0.2.2" class bdist_egg_disabled(bdist_egg): diff --git a/python/metatensor-torch/metatensor/torch/atomistic/ase_calculator.py b/python/metatensor-torch/metatensor/torch/atomistic/ase_calculator.py index d7fd18677..aa395eabe 100644 --- a/python/metatensor-torch/metatensor/torch/atomistic/ase_calculator.py +++ b/python/metatensor-torch/metatensor/torch/atomistic/ase_calculator.py @@ -511,7 +511,6 @@ def _ase_to_torch_data(atoms, dtype, device): def _full_3x3_to_voigt_6_stress(stress): - ase.stress.full_3x3_to_voigt_6_stress """ Re-implementation of ``ase.stress.full_3x3_to_voigt_6_stress`` which does not do the stress symmetrization correctly (they do ``(stress[1, 2] + stress[1, 2]) / 2.0``) diff --git a/python/metatensor-torch/metatensor/torch/atomistic/openmm_interface.py b/python/metatensor-torch/metatensor/torch/atomistic/openmm_interface.py new file mode 100644 index 000000000..7162309a4 --- /dev/null +++ b/python/metatensor-torch/metatensor/torch/atomistic/openmm_interface.py @@ -0,0 +1,138 @@ +import torch +from typing import Iterable, Optional +from metatensor.torch.atomistic import load_atomistic_model, System, ModelOutput, ModelEvaluationOptions +from metatensor.torch import Labels +from typing import List + +try: + import openmm + import openmmtorch + from openmmml.mlpotential import MLPotential, MLPotentialImpl, MLPotentialImplFactory + HAS_OPENMM = True +except ImportError as e: + class MLPotential: + pass + class MLPotentialImpl: + pass + class MLPotentialImplFactory: + pass + HAS_OPENMM = False + + +class MetatensorPotentialImplFactory(MLPotentialImplFactory): + + def createImpl( + name: str, **args + ) -> MLPotentialImpl: + # TODO: extensions_directory + return MetatensorPotentialImpl(name, **args) + + +class MetatensorPotentialImpl(MLPotentialImpl): + + def __init__(self, name: str, path: str) -> None: + self.path = path + + def addForces( + self, + topology: openmm.app.Topology, + system: openmm.System, + atoms: Optional[Iterable[int]], + forceGroup: int, + **args, + ) -> None: + + if not HAS_OPENMM: + raise ImportError( + "Could not import openmm. If you want to use metatensor with " + "openmm, please install openmm-ml with conda." + ) + + model = load_atomistic_model( + self.path # TODO: extensions_directory + ) + + # Get the atomic numbers of the ML region. + all_atoms = list(topology.atoms()) + atomic_numbers = [atom.element.atomic_number for atom in all_atoms] + + # TODO: Set up selected_atoms as a Labels object + if atoms is None: + selected_atoms = None + else: + selected_atoms = Labels( + names=["system", "atom"], + values=torch.tensor( + [[0, selected_atom] for selected_atom in atoms], + dtype=torch.int32, + ), + ) + + class MetatensorForce(torch.nn.Module): + + def __init__( + self, + model: torch.jit._script.RecursiveScriptModule, + atomic_numbers: List[int], + selected_atoms: Optional[Labels], + ) -> None: + super(MetatensorForce, self).__init__() + + self.model = model + self.register_buffer("atomic_numbers", torch.tensor(atomic_numbers, dtype=torch.int32)) + self.evaluation_options = ModelEvaluationOptions( + length_unit='nm', + outputs={ + "energy": ModelOutput( + quantity="energy", + unit="kJ/mol", + per_atom=False, + ), + }, + selected_atoms=selected_atoms, + ) + + + def forward( + self, positions: torch.Tensor, cell: Optional[torch.Tensor] = None + ) -> torch.Tensor: + # move labels if necessary + selected_atoms = self.evaluation_options.selected_atoms + if selected_atoms is not None: + if selected_atoms.device != positions.device: + self.evaluation_options.selected_atoms = selected_atoms.to(positions.device) + + if cell is None: + cell = torch.zeros((3, 3), dtype=positions.dtype, device=positions.device) + + # create System + system = System( + types=self.atomic_numbers, + positions=positions, + cell=cell, + ) + + energy = self.model([system], self.evaluation_options, check_consistency=True)["energy"].block().values.reshape(()) + return energy + + metatensor_force = MetatensorForce( + model, + atomic_numbers, + selected_atoms, + ) + + # torchscript everything + module = torch.jit.script(metatensor_force) + + # create the OpenMM force + force = openmmtorch.TorchForce(module) + isPeriodic = ( + topology.getPeriodicBoxVectors() is not None + ) or system.usesPeriodicBoundaryConditions() + force.setUsesPeriodicBoundaryConditions(isPeriodic) + force.setForceGroup(forceGroup) + + system.addForce(force) + + +MLPotential.registerImplFactory("metatensor", MetatensorPotentialImplFactory)