climemu 0.1.6__tar.gz

climemu-0.1.6/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Shahine Bouabid
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

climemu-0.1.6/PKG-INFO

@@ -0,0 +1,41 @@
+Metadata-Version: 2.4
+Name: climemu
+Version: 0.1.6
+Summary: Score-based generative emulation of impact-relevant earth system model outputs
+Author-email: Shahine Bouabid <shahineb@mit.edu>
+License: MIT
+Requires-Python: >=3.12
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: dask>=2025.7.0
+Requires-Dist: diffrax>=0.7.0
+Requires-Dist: einops>=0.8.1
+Requires-Dist: equinox>=0.13.0
+Requires-Dist: huggingface-hub>=0.35.0
+Requires-Dist: jax>=0.7.1
+Requires-Dist: netcdf4>=1.7.2
+Requires-Dist: numpy>=2.3.2
+Requires-Dist: optax>=0.2.5
+Requires-Dist: scikit-learn>=1.7.1
+Requires-Dist: scipy>=1.16.1
+Requires-Dist: tqdm>=4.67.1
+Requires-Dist: xarray>=2025.9.0
+Provides-Extra: gpu
+Requires-Dist: jax[cuda12]>=0.7.1; extra == "gpu"
+Provides-Extra: train
+Requires-Dist: torch>=2.8.0; extra == "train"
+Requires-Dist: healpy>=1.18.1; extra == "train"
+Requires-Dist: wandb>=0.21.3; extra == "train"
+Provides-Extra: plots
+Requires-Dist: matplotlib>=3.10.6; extra == "plots"
+Requires-Dist: seaborn>=0.13.2; extra == "plots"
+Requires-Dist: cartopy>=0.25.0; extra == "plots"
+Requires-Dist: regionmask>=0.13.0; extra == "plots"
+Requires-Dist: pandas>=2.3.2; extra == "plots"
+Requires-Dist: pyshtools>=4.13.1; extra == "plots"
+Provides-Extra: test
+Requires-Dist: pytest>=8.0.0; extra == "test"
+Requires-Dist: pytest-cov>=4.0.0; extra == "test"
+Requires-Dist: pytest-mock>=3.10.0; extra == "test"
+Requires-Dist: jaxlib>=0.4.0; extra == "test"
+Dynamic: license-file

climemu-0.1.6/pyproject.toml

@@ -0,0 +1,80 @@
+[project]
+name = "climemu"
+version = "0.1.6"
+description = "Score-based generative emulation of impact-relevant earth system model outputs"
+readme = "README.md"
+authors = [
+    { name = "Shahine Bouabid", email = "shahineb@mit.edu" },
+]
+requires-python = ">=3.12"
+license = { text = "MIT" }
+dependencies = [
+    "dask>=2025.7.0",
+    "diffrax>=0.7.0",
+    "einops>=0.8.1",
+    "equinox>=0.13.0",
+    "huggingface-hub>=0.35.0",
+    "jax>=0.7.1",
+    "netcdf4>=1.7.2",
+    "numpy>=2.3.2",
+    "optax>=0.2.5",
+    "scikit-learn>=1.7.1",
+    "scipy>=1.16.1",
+    "tqdm>=4.67.1",
+    "xarray>=2025.9.0",
+]
+
+[project.optional-dependencies]
+gpu = [
+    "jax[cuda12]>=0.7.1"
+]
+train = [
+    "torch>=2.8.0",
+    "healpy>=1.18.1",
+    "wandb>=0.21.3",
+]
+plots = [
+    "matplotlib>=3.10.6",
+    "seaborn>=0.13.2",
+    "cartopy>=0.25.0",
+    "regionmask>=0.13.0",
+    "pandas>=2.3.2",
+    "pyshtools>=4.13.1",
+]
+test = [
+    "pytest>=8.0.0",
+    "pytest-cov>=4.0.0",
+    "pytest-mock>=3.10.0",
+    "jaxlib>=0.4.0",
+]
+
+
+
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[tool.setuptools]
+package-dir = {"" = "src"}
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.pytest.ini_options]
+filterwarnings = [
+    "ignore::DeprecationWarning:equinox.internal._noinline",
+    "ignore::DeprecationWarning:jax.interpreters.batching",
+]
+
+[dependency-groups]
+dev = [
+    "pytest>=8.4.2",
+    "pytest-cov>=7.0.0",
+    "pytest-mock>=3.15.1",
+]
+
+[[tool.uv.index]]
+name = "testpypi"
+url = "https://test.pypi.org/simple/"
+publish-url = "https://test.pypi.org/legacy/"
+explicit = true

climemu-0.1.6/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

climemu-0.1.6/src/climemu/__init__.py

@@ -0,0 +1,17 @@
+from climemu.utils import Registry
+
+__version__ = "0.1.6"
+
+"""
+Registery of pretrained emulators for usage
+"""
+EMULATORS = Registry()
+
+
+def build_emulator(name):
+    model = EMULATORS[name]()
+    return model
+
+
+from .emulators import Bouabid2025Emulator
+__all__ = ['build_emulator', 'Bouabid2025Emulator']

climemu-0.1.6/src/climemu/emulators/__init__.py

@@ -0,0 +1 @@
+from .bouabid2025 import Bouabid2025Emulator

climemu-0.1.6/src/climemu/emulators/abstractemulator.py

@@ -0,0 +1,44 @@
+from abc import ABC, abstractmethod
+
+
+class AbstractEmulator(ABC):    
+    @abstractmethod
+    def __call__(self, *args, **kwargs):
+        """Generate climate data samples"""
+        pass
+
+    @classmethod
+    def build(cls, **kwargs):
+        return cls(**kwargs)
+
+
+class GriddedEmulator(AbstractEmulator):
+    @property
+    @abstractmethod
+    def lat(self):
+        """Latitude coordinates"""
+        pass
+    
+    @property
+    @abstractmethod
+    def lon(self):
+        """Longitude coordinates"""
+        pass
+    
+    @property
+    @abstractmethod
+    def vars(self):
+        """Variable list"""
+        pass
+
+    @property
+    def nlat(self):
+        return len(self.lat)
+    
+    @property
+    def nlon(self):
+        return len(self.lon)
+    
+    @property
+    def nvar(self):
+        return len(self.vars)

climemu-0.1.6/src/climemu/emulators/bouabid2025.py

@@ -0,0 +1,200 @@
+import os
+import yaml
+from functools import partial
+import xarray as xr
+import numpy as np
+import jax.numpy as jnp
+import jax.random as jr
+import equinox as eqx
+from huggingface_hub import hf_hub_download
+from diffusion import HealPIXUNet, ContinuousVESchedule, ContinuousHeunSampler
+from .abstractemulator import GriddedEmulator
+from .. import EMULATORS
+
+
+class Bouabid2025Emulator(GriddedEmulator):
+    def __init__(self, esm_name: str):
+        self.esm = esm_name
+        self.repo_id = "shahineb/climemu"
+
+    def load(self, which: str = "default"):
+        # Set files directory in hugging face repo
+        self.files_dir = os.path.join(self.esm, which)
+
+        # Load pattern scaling coefficients
+        self.β = self._load_pattern_scaling()
+
+        # Load climatology data
+        self.climatology = self._load_climatology()
+    
+        # Load the generative model precursor
+        self.precursor = self._load_precursor()
+
+    def compile(self, n_samples, n_steps=30):
+        # Fix number of samples and steps for generation
+        self.generative_model = partial(self.precursor,
+                                        n_samples=n_samples,
+                                        n_steps=n_steps)
+
+        # Perform a dry run to compile the JAX functions (important for performance)
+        dummy_pattern = jnp.zeros((self.nlat, self.nlon))
+        _ = self.generative_model(pattern=dummy_pattern, key=jr.PRNGKey(0))
+
+    def __call__(self, gmst, month, seed=None, xarray=False):
+        # Apply pattern scaling: pattern = β₀ + β₁ * ΔT
+        pattern = self.β[month - 1, :, 1] * gmst + self.β[month - 1, :, 0]
+        pattern = pattern.reshape((self.nlat, self.nlon))
+        
+        # Generate samples using the diffusion model
+        key = jr.PRNGKey(seed) if seed else jr.PRNGKey(np.random.randint(0, 1000000))
+        samples = self.generative_model(pattern=pattern, key=key)
+
+        # Convert to xarray Dataset
+        if xarray:
+            samples = xr.Dataset(
+                {
+                    var: (("member", "lat", "lon"), samples[:, i, :, :])
+                    for i, var in enumerate(self.vars)
+                },
+                coords={
+                    "member": jnp.arange(len(samples)) + 1,
+                    "lat": self.lat,
+                    "lon": self.lon,
+                },
+            )
+        return samples
+
+    def _load_precursor(self):
+        # Build the neural network and noise schedule
+        config = self._load_config()
+        nn = self._load_nn(config)
+        schedule = self._load_schedule(config)
+
+        # Load normalization statistics used during training (needed to denormalize the generated samples)
+        μ, σ = self._load_normalization()
+
+        # Define the output size for the generated samples 
+        output_size = (config['out_channels'], config['input_size'][1], config['input_size'][2])
+
+        # Create an precursor for the generative model
+        precursor = partial(draw_samples_single,
+                            nn=nn,
+                            schedule=schedule,
+                            output_size=output_size,
+                            μ=μ, σ=σ)
+        return precursor
+
+    def _load_config(self):
+        # Load configuration from YAML file
+        config_path = hf_hub_download(self.repo_id, f"{self.files_dir}/config.yaml")
+        with open(config_path, 'r') as file:
+            config = yaml.safe_load(file)
+        return config
+
+    def _load_normalization(self):
+        # Load normalization statistics used during training
+        norm_stats_path = hf_hub_download(self.repo_id, f"{self.files_dir}/μ_σ.npz")
+        stats = jnp.load(norm_stats_path)
+        μ, σ = stats['μ'], stats['σ']
+        return μ, σ
+
+    def _load_pattern_scaling(self):
+        # Load pattern scaling coefficients
+        pattern_scaling_path = hf_hub_download(self.repo_id, f"{self.files_dir}/β.npy")
+        β = jnp.load(pattern_scaling_path)
+        return β
+    
+    def _load_climatology(self):
+        # Load climatology data
+        climatology_path = hf_hub_download(self.repo_id, f"{self.esm}/piControl_climatology.nc")
+        climatology = xr.open_dataset(climatology_path)
+        return climatology
+
+    def _load_nn(self, config):
+        # Load graph edges for HEALPix to lat-lon connectivity
+        edges_path = hf_hub_download(self.repo_id, f"{self.files_dir}/edges.npz")
+        edges_data = jnp.load(edges_path)
+        to_healpix = jnp.array(edges_data['to_healpix']).astype(jnp.int32)
+        to_latlon = jnp.array(edges_data['to_latlon']).astype(jnp.int32)
+
+        # Initialize the neural network
+        nn = HealPIXUNet(input_size=config['input_size'],
+                         nside=config['nside'],
+                         enc_filters=config['enc_filters'],
+                         dec_filters=config['dec_filters'],
+                         out_channels=config['out_channels'],
+                         temb_dim=config['temb_dim'],
+                         healpix_emb_dim=config['healpix_emb_dim'],
+                         edges_to_healpix=to_healpix,
+                         edges_to_latlon=to_latlon)
+
+        # Load the pre-trained weights from the saved model file
+        weights_path = hf_hub_download(self.repo_id, f"{self.files_dir}/weights.eqx")
+        nn = eqx.tree_deserialise_leaves(weights_path, nn)
+        return nn
+    
+    def _load_schedule(self, config):
+        # Load the maximum noise level as used in training
+        sigma_max_path = hf_hub_download(self.repo_id, f"{self.files_dir}/σmax.npy")
+        σmax = jnp.load(sigma_max_path)
+
+        # Create the variance exploding schedule
+        schedule = ContinuousVESchedule(config['sigma_min'], σmax)
+        return schedule
+    
+    @property
+    def lat(self):
+        return self.climatology['lat'].values
+    
+    @property
+    def lon(self):
+        return self.climatology['lon'].values
+
+    @property
+    def vars(self):
+        return list(self.climatology.data_vars)
+
+
+@eqx.filter_jit
+def normalize(x, μ, σ):
+    """Normalize data using mean and standard deviation."""
+    return (x - μ) / σ
+
+@eqx.filter_jit
+def denormalize(x, μ, σ):
+    """Denormalize data using mean and standard deviation."""
+    return σ * x + μ
+
+
+def create_sampler(nn, schedule, pattern, μ, σ, output_size):
+    """Create a sampler for a given pattern."""
+    context = normalize(pattern, μ[-1], σ[-1])[None, ...]
+    def nn_with_context(x, t):
+        x = jnp.concatenate((x, context), axis=0)
+        return nn(x, t)
+    return ContinuousHeunSampler(schedule, nn_with_context, output_size)
+
+@eqx.filter_jit
+def draw_samples_single(nn, schedule, pattern, n_samples, n_steps, μ, σ, output_size, key=jr.PRNGKey(0)):
+    """Draw samples for a given pattern."""
+    sampler = create_sampler(nn, schedule, pattern, μ, σ, output_size)
+    samples = sampler.sample(n_samples, steps=n_steps, key=key)
+    return denormalize(samples, μ[:-1], σ[:-1])
+
+
+@EMULATORS.register("MPI-ESM1-2-LR")
+class MPIEmulator(Bouabid2025Emulator):
+    def __init__(self, which="default"):
+        super().__init__(esm_name="MPI-ESM1-2-LR")
+
+
+@EMULATORS.register("MIROC6")
+class MIROCEmulator(Bouabid2025Emulator):
+    def __init__(self, which="default"):
+        super().__init__(esm_name="MIROC6")
+
+
+@EMULATORS.register("ACCESS-ESM1-5")
+class ACCESSEmulator(Bouabid2025Emulator):
+    def __init__(self, which="default"):
+        super().__init__(esm_name="ACCESS-ESM1-5")

climemu-0.1.6/src/climemu/utils/__init__.py

@@ -0,0 +1,3 @@
+from .registry import Registry
+
+__all__ = ['Registry']

climemu-0.1.6/src/climemu/utils/registry.py

@@ -0,0 +1,71 @@
+import types
+
+
+class Registry(dict):
+    """
+    A helper class for managing access to builders, it extends a dictionary
+    and provides a registering functions than can be used as a decorator
+
+    Creating a registry:
+        MODULES = Registry()
+
+    There two types of builder callable you can register:
+
+    (1) : Functions which can be registered by either a simple call
+        ```
+        def build_bar():
+            return True
+        MODULES.register('bar', build_bar)
+        ```
+        or using a decorator at function definition
+        ```
+        @MODULES.register('bar')
+        def build_bar():
+            return True
+        ```
+
+    (2) : Class constructor method cls.build which again can be registered
+    with a call if class has a cls.build method
+        ```
+        foo = Foo()
+        MODULES.register('Foo', foo)
+        ```
+        of using a decorator at class definition
+        ```
+        @MODULES.register('foo')
+        class Foo:
+            @classmethod
+            def build(cls, *args, **kwargs):
+                # build a class instance foo
+                return foo
+        ```
+
+    Access of module is just like using a dictionary, eg:
+        build_bar = MODULES['bar']
+        build_foo = MODULES['foo']
+    """
+
+    def __init__(self, *args, **kwargs):
+        super(Registry, self).__init__(*args, **kwargs)
+
+    def register(self, name, module_object=None):
+        # Used as a decorator
+        if module_object is None:
+            def register_func(module_object):
+                self.register(name=name, module_object=module_object)
+                return module_object
+            return register_func
+
+        # Used as a function call
+        else:
+            if isinstance(module_object, type):
+                self._register_generic(module_dict=self, name=name, builder=module_object.build)
+            elif isinstance(module_object, types.FunctionType):
+                self._register_generic(module_dict=self, name=name, builder=module_object)
+            else:
+                raise TypeError("Trying to register unknown data type")
+
+    @staticmethod
+    def _register_generic(module_dict, name, builder):
+        assert name not in module_dict
+        module_dict[name] = builder

climemu-0.1.6/src/climemu.egg-info/PKG-INFO

@@ -0,0 +1,41 @@
+Metadata-Version: 2.4
+Name: climemu
+Version: 0.1.6
+Summary: Score-based generative emulation of impact-relevant earth system model outputs
+Author-email: Shahine Bouabid <shahineb@mit.edu>
+License: MIT
+Requires-Python: >=3.12
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: dask>=2025.7.0
+Requires-Dist: diffrax>=0.7.0
+Requires-Dist: einops>=0.8.1
+Requires-Dist: equinox>=0.13.0
+Requires-Dist: huggingface-hub>=0.35.0
+Requires-Dist: jax>=0.7.1
+Requires-Dist: netcdf4>=1.7.2
+Requires-Dist: numpy>=2.3.2
+Requires-Dist: optax>=0.2.5
+Requires-Dist: scikit-learn>=1.7.1
+Requires-Dist: scipy>=1.16.1
+Requires-Dist: tqdm>=4.67.1
+Requires-Dist: xarray>=2025.9.0
+Provides-Extra: gpu
+Requires-Dist: jax[cuda12]>=0.7.1; extra == "gpu"
+Provides-Extra: train
+Requires-Dist: torch>=2.8.0; extra == "train"
+Requires-Dist: healpy>=1.18.1; extra == "train"
+Requires-Dist: wandb>=0.21.3; extra == "train"
+Provides-Extra: plots
+Requires-Dist: matplotlib>=3.10.6; extra == "plots"
+Requires-Dist: seaborn>=0.13.2; extra == "plots"
+Requires-Dist: cartopy>=0.25.0; extra == "plots"
+Requires-Dist: regionmask>=0.13.0; extra == "plots"
+Requires-Dist: pandas>=2.3.2; extra == "plots"
+Requires-Dist: pyshtools>=4.13.1; extra == "plots"
+Provides-Extra: test
+Requires-Dist: pytest>=8.0.0; extra == "test"
+Requires-Dist: pytest-cov>=4.0.0; extra == "test"
+Requires-Dist: pytest-mock>=3.10.0; extra == "test"
+Requires-Dist: jaxlib>=0.4.0; extra == "test"
+Dynamic: license-file

climemu-0.1.6/src/climemu.egg-info/SOURCES.txt

@@ -0,0 +1,46 @@
+LICENSE
+README.md
+pyproject.toml
+src/climemu/__init__.py
+src/climemu.egg-info/PKG-INFO
+src/climemu.egg-info/SOURCES.txt
+src/climemu.egg-info/dependency_links.txt
+src/climemu.egg-info/requires.txt
+src/climemu.egg-info/top_level.txt
+src/climemu/emulators/__init__.py
+src/climemu/emulators/abstractemulator.py
+src/climemu/emulators/bouabid2025.py
+src/climemu/utils/__init__.py
+src/climemu/utils/registry.py
+src/datasets/__init__.py
+src/datasets/cmip6.py
+src/datasets/constants.py
+src/datasets/pattern_to_cmip6.py
+src/diffusion/__init__.py
+src/diffusion/losses/__init__.py
+src/diffusion/losses/denoising_score_matching.py
+src/diffusion/nn/__init__.py
+src/diffusion/nn/healpixunet.py
+src/diffusion/nn/backbones/__init__.py
+src/diffusion/nn/backbones/convnet.py
+src/diffusion/nn/modules/__init__.py
+src/diffusion/nn/modules/remap.py
+src/diffusion/nn/modules/healpix/__init__.py
+src/diffusion/nn/modules/healpix/conv.py
+src/diffusion/nn/modules/healpix/padding.py
+src/diffusion/nn/timeencoder/__init__.py
+src/diffusion/nn/timeencoder/gaussianfourier.py
+src/diffusion/samplers/__init__.py
+src/diffusion/samplers/continuous_ode_sampler.py
+src/diffusion/schedules/__init__.py
+src/diffusion/schedules/variance_exploding.py
+src/utils/__init__.py
+src/utils/collate.py
+src/utils/arrays/__init__.py
+src/utils/arrays/aggregate.py
+src/utils/arrays/convert.py
+src/utils/arrays/filter.py
+src/utils/arrays/reshape.py
+src/utils/arrays/smooth.py
+src/utils/graphs/__init__.py
+src/utils/graphs/haversine.py

climemu-0.1.6/src/climemu.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

climemu-0.1.6/src/climemu.egg-info/requires.txt

@@ -0,0 +1,35 @@
+dask>=2025.7.0
+diffrax>=0.7.0
+einops>=0.8.1
+equinox>=0.13.0
+huggingface-hub>=0.35.0
+jax>=0.7.1
+netcdf4>=1.7.2
+numpy>=2.3.2
+optax>=0.2.5
+scikit-learn>=1.7.1
+scipy>=1.16.1
+tqdm>=4.67.1
+xarray>=2025.9.0
+
+[gpu]
+jax[cuda12]>=0.7.1
+
+[plots]
+matplotlib>=3.10.6
+seaborn>=0.13.2
+cartopy>=0.25.0
+regionmask>=0.13.0
+pandas>=2.3.2
+pyshtools>=4.13.1
+
+[test]
+pytest>=8.0.0
+pytest-cov>=4.0.0
+pytest-mock>=3.10.0
+jaxlib>=0.4.0
+
+[train]
+torch>=2.8.0
+healpy>=1.18.1
+wandb>=0.21.3

climemu-0.1.6/src/climemu.egg-info/top_level.txt

@@ -0,0 +1,4 @@
+climemu
+datasets
+diffusion
+utils

climemu-0.1.6/src/datasets/__init__.py

@@ -0,0 +1,12 @@
+from .cmip6 import (
+    CMIP6Data
+)
+
+from .pattern_to_cmip6 import (
+    PatternToCMIP6Dataset
+)
+
+__all__ = [
+    "CMIP6Data",
+    "PatternToCMIP6Dataset"
+]