py-neuralgcm: add new package (#43117)

Co-authored-by: Tamara Dahlgren <35777542+tldahlgren@users.noreply.github.com>

4 files changed, 112 insertions, 0 deletions

diff --git a/var/spack/repos/builtin/packages/py-dinosaur/package.py b/var/spack/repos/builtin/packages/py-dinosaur/package.py
new file mode 100644
index 0000000000..a5e37777b5
--- /dev/null
+++ b/var/spack/repos/builtin/packages/py-dinosaur/package.py
@@ -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")
diff --git a/var/spack/repos/builtin/packages/py-neuralgcm/package.py b/var/spack/repos/builtin/packages/py-neuralgcm/package.py
new file mode 100644
index 0000000000..af57eb8ee2
--- /dev/null
+++ b/var/spack/repos/builtin/packages/py-neuralgcm/package.py
@@ -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")
diff --git a/var/spack/repos/builtin/packages/py-tree-math/package.py b/var/spack/repos/builtin/packages/py-tree-math/package.py
new file mode 100644
index 0000000000..4012457d95
--- /dev/null
+++ b/var/spack/repos/builtin/packages/py-tree-math/package.py
@@ -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"))
diff --git a/var/spack/repos/builtin/packages/py-xarray-tensorstore/package.py b/var/spack/repos/builtin/packages/py-xarray-tensorstore/package.py
new file mode 100644
index 0000000000..ea315c8d0e
--- /dev/null
+++ b/var/spack/repos/builtin/packages/py-xarray-tensorstore/package.py
@@ -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")