py-neuralgcm: add new package (#43117)

Co-authored-by: Tamara Dahlgren <35777542+tldahlgren@users.noreply.github.com>
This commit is contained in:
Adam J. Stewart 2024-03-12 04:02:03 +01:00 committed by GitHub
parent 001af62585
commit e12a8a69c7
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 112 additions and 0 deletions

View File

@ -0,0 +1,33 @@
# Copyright 2013-2024 Lawrence Livermore National Security, LLC and other
# Spack Project Developers. See the top-level COPYRIGHT file for details.
#
# SPDX-License-Identifier: (Apache-2.0 OR MIT)
from spack.package import *
class PyDinosaur(PythonPackage):
"""Dinosaur: differentiable dynamics for global atmospheric modeling."""
homepage = "https://github.com/google-research/dinosaur"
git = "https://github.com/google-research/dinosaur.git"
license("Apache-2.0")
version("main", branch="main")
depends_on("py-setuptools", type="build")
with default_args(type=("build", "run")):
depends_on("python@3.10:")
depends_on("py-fsspec")
depends_on("py-jax")
depends_on("py-jaxlib")
depends_on("py-numpy")
depends_on("py-pandas")
depends_on("py-pint")
depends_on("py-scipy")
depends_on("py-scikit-learn")
depends_on("py-tree-math")
depends_on("py-xarray")
depends_on("py-xarray-tensorstore")

View File

@ -0,0 +1,32 @@
# Copyright 2013-2024 Lawrence Livermore National Security, LLC and other
# Spack Project Developers. See the top-level COPYRIGHT file for details.
#
# SPDX-License-Identifier: (Apache-2.0 OR MIT)
from spack.package import *
class PyNeuralgcm(PythonPackage):
"""NeuralGCM: Hybrid ML + Physics model of Earth's atmosphere."""
homepage = "https://github.com/google-research/neuralgcm"
git = "https://github.com/google-research/neuralgcm.git"
license("Apache-2.0")
version("main", branch="main")
depends_on("py-setuptools", type="build")
with default_args(type=("build", "run")):
depends_on("python@3.10:")
depends_on("py-dinosaur")
depends_on("py-dm-haiku")
depends_on("py-gin-config")
depends_on("py-jax")
depends_on("py-jaxlib")
depends_on("py-numpy")
depends_on("py-optax")
depends_on("py-pandas")
depends_on("py-tensorflow-probability")
depends_on("py-xarray")

View File

@ -0,0 +1,20 @@
# Copyright 2013-2024 Lawrence Livermore National Security, LLC and other
# Spack Project Developers. See the top-level COPYRIGHT file for details.
#
# SPDX-License-Identifier: (Apache-2.0 OR MIT)
from spack.package import *
class PyTreeMath(PythonPackage):
"""Mathematical operations for JAX pytrees."""
homepage = "https://github.com/google/tree-math"
pypi = "tree-math/tree-math-0.2.0.tar.gz"
license("Apache-2.0")
version("0.2.0", sha256="fced2b436fa265b4e24ab46b5768d7b03a4a8d0b75de8a5ab110abaeac3b5772")
depends_on("py-setuptools", type="build")
depends_on("py-jax", type=("build", "run"))

View File

@ -0,0 +1,27 @@
# Copyright 2013-2024 Lawrence Livermore National Security, LLC and other
# Spack Project Developers. See the top-level COPYRIGHT file for details.
#
# SPDX-License-Identifier: (Apache-2.0 OR MIT)
from spack.package import *
class PyXarrayTensorstore(PythonPackage):
"""Xarray-TensorStore is a small library that allows opening Zarr arrays into Xarray
via TensorStore, instead of the standard Zarr-Python library.
"""
homepage = "https://github.com/google/xarray-tensorstore"
pypi = "xarray-tensorstore/xarray-tensorstore-0.1.1.tar.gz"
license("Apache-2.0")
version("0.1.1", sha256="2ee6f164c9f1bc43328245b8d06c21863204fcd4e6159ddd6d8867c313c1d9b4")
depends_on("py-setuptools", type="build")
with default_args(type=("build", "run")):
depends_on("py-numpy")
depends_on("py-xarray")
depends_on("py-zarr")
depends_on("py-tensorstore")