Merge branch 'dev' into squash-merge-release

README.md

@@ -3,12 +3,16 @@
 
 <img src="ACE-logo.png" alt="Logo for the ACE Project" style="width: auto; height: 50px;">
 
+<<<<<<< HEAD
+Development documentation can be found on [github-pages](https://ai2cm.github.io/full-model/index.html).
+=======
 # Ai2 Climate Emulator
 This repo contains code accompanying four papers describing ACE models:
 - "ACE: A fast, skillful learned global atmospheric model for climate prediction" ([link](https://arxiv.org/abs/2310.02074))
 - "Application of the Ai2 Climate Emulator to E3SMv2's global atmosphere model, with a focus on precipitation fidelity" ([link](https://agupubs.onlinelibrary.wiley.com/doi/full/10.1029/2024JH000136))
 - "ACE2: Accurately learning subseasonal to decadal atmospheric variability and forced responses" ([link](https://arxiv.org/abs/2411.11268))
 - "ACE2-SOM: Coupling to a slab ocean and learning the sensitivity of climate to changes in CO2" ([link](https://arxiv.org/abs/2412.04418))
+>>>>>>> a7734af7f1435d88bb71651c10b1ac2031ddbe17
 
 ## Installation
 

fme/fme/ace/__init__.py

@@ -1,6 +1,5 @@
 import sys
 
-from fme.ace.data_loading.config import TimeSlice, XarrayDataConfig
 from fme.ace.data_loading.inference import (
     ExplicitIndices,
     InferenceInitialConditionIndices,
@@ -41,6 +40,7 @@
 from fme.ace.registry.sfno import SFNO_V0_1_0, SphericalFourierNeuralOperatorBuilder
 from fme.core.corrector.corrector import CorrectorConfig
 from fme.core.corrector.ocean import OceanCorrectorConfig
+from fme.core.dataset.config import TimeSlice, XarrayDataConfig
 from fme.core.gridded_ops import GriddedOperations
 from fme.core.loss import WeightedMappingLossConfig
 from fme.core.normalizer import NormalizationConfig

fme/fme/ace/aggregator/inference/annual.py

@@ -7,7 +7,7 @@
 import xarray as xr
 from matplotlib.figure import Figure
 
-from fme.ace.data_loading.data_typing import VariableMetadata
+from fme.core.dataset.data_typing import VariableMetadata
 from fme.core.device import get_device
 from fme.core.distributed import Distributed
 from fme.core.gridded_ops import GriddedOperations

fme/fme/ace/aggregator/inference/enso/enso.py

@@ -8,7 +8,7 @@
 import xarray as xr
 
 from fme.ace.aggregator.plotting import get_cmap_limits, plot_imshow, plot_paneled_data
-from fme.ace.data_loading.data_typing import VariableMetadata
+from fme.core.dataset.data_typing import VariableMetadata
 from fme.core.device import get_device
 from fme.core.distributed import Distributed
 from fme.core.gridded_ops import GriddedOperations

fme/fme/ace/aggregator/inference/main.py

@@ -17,12 +17,12 @@
 import xarray as xr
 
 from fme.ace.data_loading.batch_data import BatchData, PairedData, PrognosticState
-from fme.ace.data_loading.data_typing import VariableMetadata
 from fme.core.coordinates import (
     HorizontalCoordinates,
     HybridSigmaPressureCoordinate,
     LatLonCoordinates,
 )
+from fme.core.dataset.data_typing import VariableMetadata
 from fme.core.generics.aggregator import (
     InferenceAggregatorABC,
     InferenceLog,

fme/fme/ace/aggregator/inference/reduced.py

@@ -6,7 +6,7 @@
 import torch
 import xarray as xr
 
-from fme.ace.data_loading.data_typing import VariableMetadata
+from fme.core.dataset.data_typing import VariableMetadata
 from fme.core.device import get_device
 from fme.core.distributed import Distributed
 from fme.core.gridded_ops import GriddedOperations

fme/fme/ace/aggregator/inference/seasonal.py

@@ -6,7 +6,7 @@
 import xarray as xr
 
 from fme.ace.aggregator.plotting import plot_paneled_data
-from fme.ace.data_loading.data_typing import VariableMetadata
+from fme.core.dataset.data_typing import VariableMetadata
 from fme.core.device import get_device
 from fme.core.distributed import Distributed
 from fme.core.gridded_ops import GriddedOperations

fme/fme/ace/aggregator/inference/time_mean.py

@@ -5,7 +5,7 @@
 import torch
 import xarray as xr
 
-from fme.ace.data_loading.data_typing import VariableMetadata
+from fme.core.dataset.data_typing import VariableMetadata
 from fme.core.distributed import Distributed
 from fme.core.gridded_ops import GriddedOperations
 from fme.core.typing_ import TensorDict, TensorMapping

fme/fme/ace/aggregator/inference/video.py

@@ -5,7 +5,7 @@
 import torch
 import xarray as xr
 
-from fme.ace.data_loading.data_typing import VariableMetadata
+from fme.core.dataset.data_typing import VariableMetadata
 from fme.core.distributed import Distributed
 from fme.core.typing_ import TensorDict, TensorMapping
 from fme.core.wandb import WandB

fme/fme/ace/aggregator/inference/zonal_mean.py

@@ -6,7 +6,7 @@
 import torch
 import xarray as xr
 
-from fme.ace.data_loading.data_typing import VariableMetadata
+from fme.core.dataset.data_typing import VariableMetadata
 from fme.core.device import get_device
 from fme.core.distributed import Distributed
 from fme.core.typing_ import TensorDict, TensorMapping

fme/fme/ace/aggregator/one_step/main.py

@@ -3,8 +3,8 @@
 import numpy as np
 import torch
 
-from fme.ace.data_loading.data_typing import VariableMetadata
 from fme.ace.stepper import TrainOutput
+from fme.core.dataset.data_typing import VariableMetadata
 from fme.core.generics.aggregator import AggregatorABC
 from fme.core.gridded_ops import GriddedOperations
 from fme.core.typing_ import TensorMapping

fme/fme/ace/aggregator/one_step/map.py

@@ -2,7 +2,7 @@
 
 import torch
 
-from fme.ace.data_loading.data_typing import VariableMetadata
+from fme.core.dataset.data_typing import VariableMetadata
 from fme.core.distributed import Distributed
 from fme.core.typing_ import TensorDict, TensorMapping
 from fme.core.wandb import Image

fme/fme/ace/aggregator/one_step/snapshot.py

@@ -2,7 +2,7 @@
 
 import torch
 
-from fme.ace.data_loading.data_typing import VariableMetadata
+from fme.core.dataset.data_typing import VariableMetadata
 from fme.core.typing_ import TensorMapping
 from fme.core.wandb import Image
 

fme/fme/ace/data_loading/batch_data.py

@@ -25,10 +25,10 @@
 import xarray as xr
 from torch.utils.data import default_collate
 
-from fme.ace.data_loading._xarray import DatasetProperties
-from fme.ace.data_loading.data_typing import VariableMetadata
-from fme.ace.data_loading.requirements import PrognosticStateDataRequirements
+from fme.ace.requirements import PrognosticStateDataRequirements
 from fme.core.coordinates import HorizontalCoordinates, HybridSigmaPressureCoordinate
+from fme.core.dataset.data_typing import VariableMetadata
+from fme.core.dataset.xarray import DatasetProperties
 from fme.core.device import get_device
 from fme.core.generics.data import DataLoader, GriddedDataABC, InferenceDataABC
 from fme.core.gridded_ops import GriddedOperations