fastatomstruct 2.0.0__cp39-abi3-win_amd64.whl

fastatomstruct/__init__.py

@@ -0,0 +1,143 @@
+from __future__ import annotations
+
+from .fastatomstruct import *
+
+# ── units: Quantity type + unit registry, built from the Rust unit table ────
+# ``fs.units`` becomes the unit namespace: ``fs.units.fs`` is the unit, so
+# ``2 * fs.units.fs`` is a ``Quantity``; the ASE-style float constants move
+# to ``fs.units.ase``, the registry is ``fs.units.u``. Everything is
+# built from ``fs.units.units_table()``, which mirrors
+# ``fastatomstruct_core::units::quantities::UNITS`` — the single source of
+# truth for symbols, factors and dimensions.
+from ._quantity import attach_to_units_module as _attach_units  # noqa: E402
+
+_attach_units(globals().get("units"))
+
+# ── Convenience re-exports of the most common entry points. ────────────────
+# The canonical location of each name is still the submodule; the aliases
+# below are only there so quick scripts and notebooks stay compact. New
+# symbols should not be added here lightly — submodules are the default.
+#
+# Every alias is guarded so a wheel that was built without the relevant
+# optional Cargo feature (``mace``, ``cp2k``, ``io``, ``md``, …) still
+# imports cleanly — the missing name simply doesn't appear at the top level.
+
+
+def _alias(submodule_name: str, *names: str) -> None:
+    """Bind ``fastatomstruct.<name>`` for every ``name`` in *names*, picking
+    each one from ``fastatomstruct.<submodule_name>``. Anything missing
+    (submodule or attribute) is silently skipped so the import never blows
+    up because of an optional Cargo feature."""
+    submodule = globals().get(submodule_name)
+    if submodule is None:
+        return
+    for name in names:
+        attr = getattr(submodule, name, None)
+        if attr is not None:
+            globals()[name] = attr
+
+
+def _get_exported_names() -> list[str]:
+    """Return a list of public names exported at the top level."""
+    return [n for n in globals() if not n.startswith("_")]
+
+
+def _is_available(name: str) -> bool:
+    """Check whether *name* is exported at the top level."""
+    return name in globals()
+
+
+# Verbs.
+_alias("viewer", "view", "view_many", "view_interactive", "view_image")
+_alias("io", "open", "read", "write")
+_alias("md", "run_molecular_dynamics")
+_alias("optimizers", "run_optimization")
+
+# Calculators.
+_alias(
+    "calc",
+    "LennardJones",
+    "EmtCalculator",
+    "ACECalculator",
+    "MACECalculator",
+    "CP2KCalculator",
+)
+
+# MD building blocks — integrator + standard thermostats + standard barostats.
+_alias(
+    "md",
+    "VelocityVerlet",
+    "LangevinIntegrator",
+    "NoThermostat",
+    "BerendsenThermostat",
+    "BussiThermostat",
+    "NoseHooverThermostat",
+    "AndersenThermostat",
+    "NoBarostat",
+    "BussiBarostat",
+    "SemiIsotropicBussiBarostat",
+)
+
+# Optimizer building blocks.
+_alias(
+    "optimizers",
+    "FireOptimizer",
+    "LbfgsbOptimizer",
+    "BfgsOptimizer",
+    "NagOptimizer",
+    "ConvergenceCriteria",
+    "CellConstraint",
+)
+
+# Phonon — only the three most frequently typed entry points stay at top
+# level. Everything else lives at ``fs.phonon.*``.
+_alias("phonon", "calculate_force_constants", "BandPath", "ForceConstants")
+
+# NEGF electron-transport workflows (CP2K + SMEAGOL via cp2k_rs).
+# Pure Python; imported as a submodule so `fs.transport` works. cp2k_rs is
+# only required when a workflow actually runs, not at import time.
+from . import transport  # noqa: F401
+
+# Plotting helpers (matplotlib-backed)
+# These are pure-Python utilities shipped with the package.
+# Every plotting helper is re-exported at the top level, because that is how
+# they are used everywhere: docs, examples and tests write `fs.plot_dos(...)`,
+# never `fs.plotting.plot_dos(...)`. Only three of the ten used to be listed
+# here, so `fs.plot_cohp` — which docs/examples/cohp_coop.rst and
+# docs/functions/bonding/index.rst both call — raised AttributeError.
+from .plotting import (
+    plot_absorption_reflectivity,
+    plot_band_structure,
+    plot_band_structure_and_dos,
+    plot_cohp,
+    plot_dos,
+    plot_epsilon_complex,
+    plot_icohp_bars,
+    plot_n_kappa,
+    plot_optical_spectrum,
+    plot_sigma_complex,
+)
+
+try:
+    from .active_learning import run_active_learning_md
+except ImportError:
+    import logging as _logging
+
+    _logging.getLogger(__name__).warning(
+        "active_learning requires numpy, ase, and mace-torch; "
+        "run_active_learning_md will not be available. "
+        "Install with: pip install fastatomstruct[mace-torch]",
+    )
+
+# Interactive curve editor for temperature/pressure profiles (matplotlib).
+# Pure-Python; attaches itself to the ``md`` submodule if available.
+try:
+    from .curve_editor import curve_editor as _curve_editor_func
+
+    _md_mod = globals().get("md")
+    if _md_mod is not None:
+        setattr(_md_mod, "curve_editor", _curve_editor_func)
+    # The top-level name ``curve_editor`` is always the subpackage (as the stub says); without
+    # the ``md`` feature use ``from fastatomstruct.curve_editor import curve_editor``.
+except ImportError:
+    pass  # numpy or submodule not available

fastatomstruct/__init__.pyi

@@ -0,0 +1,133 @@
+# This file is generated by gen_stub.py -- do not edit by hand.
+
+"""
+This type stub file was generated by pyright.
+"""
+
+from .fastatomstruct import *
+from .fastatomstruct import (
+    AoMatrixCsrResult as AoMatrixCsrResult,
+    AoMatrixResult as AoMatrixResult,
+    Atom as Atom,
+    Atoms as Atoms,
+    AtomsIterator as AtomsIterator,
+    BandStructureResult as BandStructureResult,
+    BuildOptions as BuildOptions,
+    CutoffMode as CutoffMode,
+    CylindricalAxis as CylindricalAxis,
+    DomainDecomposition as DomainDecomposition,
+    ElectronDensityResult as ElectronDensityResult,
+    Element as Element,
+    EnergyComponents as EnergyComponents,
+    FilterAnd as FilterAnd,
+    FilterElement as FilterElement,
+    FilterIndices as FilterIndices,
+    FilterNot as FilterNot,
+    FilterOr as FilterOr,
+    FilterPosition as FilterPosition,
+    FilterTag as FilterTag,
+    FixAtoms as FixAtoms,
+    GenericParticle as GenericParticle,
+    GridInfo as GridInfo,
+    HomoLumoResult as HomoLumoResult,
+    LatticeParameters as LatticeParameters,
+    LatticeType as LatticeType,
+    MoCoefficientsResult as MoCoefficientsResult,
+    NeighbourMode as NeighbourMode,
+    NeighbourSelectionCriteria as NeighbourSelectionCriteria,
+    QuantumATKDataType as QuantumATKDataType,
+    ReadIndex as ReadIndex,
+    ScfInfo as ScfInfo,
+    TimeAxis as TimeAxis,
+    Trajectory as Trajectory,
+    TrajectoryChunks as TrajectoryChunks,
+    TrajectoryIterator as TrajectoryIterator,
+    VDOSMethod as VDOSMethod,
+    VacfMethod as VacfMethod,
+    analysis as analysis,
+    bonding as bonding,
+    build as build,
+    calc as calc,
+    convert_ase as convert_ase,
+    disable_telemetry as disable_telemetry,
+    elph as elph,
+    enable_telemetry as enable_telemetry,
+    fetch_cp2k_worker as fetch_cp2k_worker,
+    filter_indices as filter_indices,
+    finalize_mpi as finalize_mpi,
+    get_mpi_status as get_mpi_status,
+    get_num_threads as get_num_threads,
+    init_mpi as init_mpi,
+    io as io,
+    is_log_enabled as is_log_enabled,
+    kubo_greenwood_transport as kubo_greenwood_transport,
+    md as md,
+    nald as nald,
+    optical_spectrum as optical_spectrum,
+    optimizers as optimizers,
+    phonon as phonon,
+    queue as queue,
+    read_obs_log as read_obs_log,
+    run_active_learning_md as run_active_learning_md,
+    run_dd_md as run_dd_md,
+    set_log_enabled as set_log_enabled,
+    set_num_threads as set_num_threads,
+    suppress_telemetry_notice as suppress_telemetry_notice,
+    units as units,
+    viewer as viewer,
+)
+from .viewer import (
+    view as view,
+    view_image as view_image,
+    view_interactive as view_interactive,
+    view_many as view_many,
+)
+from .io import open as open, read as read, write as write
+from .md import (
+    AndersenThermostat as AndersenThermostat,
+    BerendsenThermostat as BerendsenThermostat,
+    BussiBarostat as BussiBarostat,
+    BussiThermostat as BussiThermostat,
+    LangevinIntegrator as LangevinIntegrator,
+    NoBarostat as NoBarostat,
+    NoThermostat as NoThermostat,
+    NoseHooverThermostat as NoseHooverThermostat,
+    SemiIsotropicBussiBarostat as SemiIsotropicBussiBarostat,
+    VelocityVerlet as VelocityVerlet,
+    run_molecular_dynamics as run_molecular_dynamics,
+)
+from .optimizers import (
+    BfgsOptimizer as BfgsOptimizer,
+    CellConstraint as CellConstraint,
+    ConvergenceCriteria as ConvergenceCriteria,
+    FireOptimizer as FireOptimizer,
+    LbfgsbOptimizer as LbfgsbOptimizer,
+    NagOptimizer as NagOptimizer,
+    run_optimization as run_optimization,
+)
+from .calc import (
+    ACECalculator as ACECalculator,
+    CP2KCalculator as CP2KCalculator,
+    EmtCalculator as EmtCalculator,
+    LennardJones as LennardJones,
+    MACECalculator as MACECalculator,
+)
+from .phonon import (
+    BandPath as BandPath,
+    ForceConstants as ForceConstants,
+    calculate_force_constants as calculate_force_constants,
+)
+from .plotting import (
+    plot_absorption_reflectivity as plot_absorption_reflectivity,
+    plot_band_structure as plot_band_structure,
+    plot_band_structure_and_dos as plot_band_structure_and_dos,
+    plot_cohp as plot_cohp,
+    plot_dos as plot_dos,
+    plot_epsilon_complex as plot_epsilon_complex,
+    plot_icohp_bars as plot_icohp_bars,
+    plot_n_kappa as plot_n_kappa,
+    plot_optical_spectrum as plot_optical_spectrum,
+    plot_sigma_complex as plot_sigma_complex,
+)
+from . import curve_editor, transport
+from .active_learning import run_active_learning_md

fastatomstruct/_pjrt/__init__.py

@@ -0,0 +1,185 @@
+"""XLA runtime libraries (PJRT plugins) for the Rust MACE backend.
+
+The libraries are large and optional:
+
+* CPU: the companion package ``fastatomstruct-xla-cpu``
+  (``pip install "fastatomstruct[xla]"``)
+* CUDA: Google's CUDA 13 plugin ``jax-cuda13-pjrt`` (the XLA build of JAX)
+  and NVIDIA's CUDA 13 pip packages (``pip install "fastatomstruct[cuda13]"``),
+  in the versions verified against the MACE reference tests
+
+:func:`plugin_path` finds whichever is installed, :func:`cueq_libraries` the
+optional cuEquivariance kernels.  For development,
+``scripts/fetch_pjrt.py`` can instead place the pinned builds in this
+directory (wheels built afterwards ship them).  Nothing here imports jax.
+"""
+
+from __future__ import annotations
+
+import ctypes
+import importlib
+import os
+from importlib import metadata
+from pathlib import Path
+
+_HERE = Path(__file__).resolve().parent
+_NAMES = {"cpu": "libpjrt_cpu.so", "cuda": "libpjrt_cuda.so"}
+_ENV = {"cpu": "MACE_PJRT_PLUGIN", "cuda": "MACE_PJRT_CUDA_PLUGIN"}
+_CPU_COMPANION = "fastatomstruct_xla_cpu"
+
+#: The distribution that ships the CUDA plugin, and the plugin inside it.
+CUDA_PLUGIN_DISTRIBUTION = "jax-cuda13-pjrt"
+_CUDA_PLUGIN_FILE = "jax_plugins/xla_cuda13/xla_cuda_plugin.so"
+# CUDA libraries the plugin opens by soname, by NVIDIA package, in loading
+# order (the list of jax_plugins.xla_cuda13). The plugin's RPATH finds them
+# in a standard site-packages layout as well; loading them first also covers
+# the plugin and NVIDIA's packages in different site directories.
+_CUDA_PRELOAD = (
+    ("cu13", ("libcudart.so.13",)),
+    ("cu13", ("libnvrtc.so.13",)),
+    ("cu13", ("libcublas.so.13", "libcublasLt.so.13")),
+    ("nccl", ("libnccl.so.2",)),
+    ("cu13", ("libcupti.so.13",)),
+    ("cu13", ("libcusparse.so.12",)),
+    ("cu13", ("libcusolver.so.12",)),
+    ("cu13", ("libcufft.so.12",)),
+    ("nvshmem", ("libnvshmem_host.so.3",)),
+    ("cudnn", ("libcudnn.so.9",)),
+)
+_CUDA_REQUIRED = ("libcudart.so.13", "libcublas.so.13")
+_cuda_loaded: list[ctypes.CDLL] = []
+
+
+def _cpu_companion() -> str | None:
+    try:
+        module = importlib.import_module(_CPU_COMPANION)
+    except ImportError:
+        return None
+    return module.library_path()
+
+
+def _cuda_plugin() -> str | None:
+    """The plugin of the installed ``jax-cuda13-pjrt``, with NVIDIA's
+    libraries loaded first; ``None`` if it is not installed."""
+    try:
+        dist = metadata.distribution(CUDA_PLUGIN_DISTRIBUTION)
+    except metadata.PackageNotFoundError:
+        return None
+    path = Path(str(dist.locate_file(_CUDA_PLUGIN_FILE))).resolve()
+    if not path.is_file():
+        raise FileNotFoundError(
+            f"{path} is missing from {CUDA_PLUGIN_DISTRIBUTION}; reinstall "
+            '"fastatomstruct[cuda13]"'
+        )
+    _preload_cuda()
+    return str(path)
+
+
+def _nvidia_root() -> Path | None:
+    try:
+        dist = metadata.distribution("nvidia-cuda-runtime")
+    except metadata.PackageNotFoundError:
+        return None
+    return Path(str(dist.locate_file("nvidia"))).resolve()
+
+
+def _preload_cuda() -> None:
+    """Load the CUDA libraries from NVIDIA's pip packages into the process,
+    as ``jax_plugins.xla_cuda13`` does before it registers the plugin.
+
+    Libraries a package does not ship are skipped (the plugin opens optional
+    ones itself or does without); the CUDA runtime and cuBLAS must be there.
+    XLA finds ``nvvm/libdevice`` and ``ptxas`` in ``nvidia/cu13`` on its own.
+    """
+    if _cuda_loaded:
+        return
+    root = _nvidia_root()
+    found = set()
+    for package, names in _CUDA_PRELOAD:
+        for name in names:
+            path = None if root is None else root / package / "lib" / name
+            if path is None or not path.is_file():
+                continue
+            try:
+                _cuda_loaded.append(
+                    ctypes.CDLL(str(path), mode=os.RTLD_NOW | os.RTLD_GLOBAL)
+                )
+            except OSError as exc:
+                raise OSError(f"could not load {path}: {exc}") from None
+            found.add(name)
+    missing = [name for name in _CUDA_REQUIRED if name not in found]
+    if missing:
+        _cuda_loaded.clear()
+        raise FileNotFoundError(
+            f"{missing} not found in NVIDIA's pip packages ({root}); "
+            'reinstall "fastatomstruct[cuda13]"'
+        )
+
+
+def plugin_path(device: str = "cpu") -> str | None:
+    """Return the PJRT plugin for ``device`` (``'cpu'`` or ``'cuda'``).
+
+    Order: the device's environment variable, the installed runtime
+    (``fastatomstruct-xla-cpu`` / ``jax-cuda13-pjrt``), then a library placed
+    in this package.  ``None`` if none exists.
+    """
+    if device not in _NAMES:
+        raise ValueError(f"unknown device {device!r} (expected 'cpu' or 'cuda')")
+    env = os.environ.get(_ENV[device])
+    if env:
+        return env
+    installed = _cpu_companion() if device == "cpu" else _cuda_plugin()
+    if installed is not None:
+        return installed
+    bundled = _HERE / _NAMES[device]
+    return str(bundled) if bundled.is_file() else None
+
+
+def library_dirs(device: str = "cpu") -> list[str]:
+    """Directories with the shared libraries the plugin for ``device`` needs.
+
+    In this process :func:`plugin_path` loads them itself.  A process that is
+    not Python (the domain decomposition workers) gets the directories on
+    ``LD_LIBRARY_PATH``; the CUDA plugin also finds them through its RPATH.
+    """
+    if device not in _NAMES:
+        raise ValueError(f"unknown device {device!r} (expected 'cpu' or 'cuda')")
+    if device == "cpu":
+        return []
+    root = _nvidia_root()
+    if root is None:
+        return []
+    return [str(p) for p in sorted(root.glob("*/lib")) if p.is_dir()]
+
+
+def cueq_libraries(device: str = "cuda") -> list[str]:
+    """The kernel library of NVIDIA cuEquivariance and what it depends on, in
+    loading order, from the installed packages ``cuequivariance-ops-jax`` and
+    ``cuequivariance-ops`` (proprietary; fastatomstruct does not ship them)."""
+    import importlib.util
+
+    if device != "cuda":
+        raise ValueError(
+            "enable_cueq needs device='cuda' (the kernels are CUDA kernels)"
+        )
+    paths = []
+    for package, library in (
+        ("cuequivariance_ops", "libcue_ops.so"),
+        ("cuequivariance_ops_jax", "libcue_ops_jax.so"),
+    ):
+        # (located without importing: the jax package would import jax)
+        spec = importlib.util.find_spec(package)
+        locations = list(getattr(spec, "submodule_search_locations", None) or [])
+        found = [
+            os.path.join(location, "lib", library)
+            for location in locations
+            if os.path.isfile(os.path.join(location, "lib", library))
+        ]
+        if not found:
+            raise ValueError(
+                "enable_cueq needs NVIDIA's cuEquivariance kernels: "
+                "pip install cuequivariance-ops-jax-cu13 "
+                f"({library} of the package {package} was not found)"
+            )
+        paths.append(found[0])
+    return paths